Skip to content

Commit e38f04d

Browse files
misoltmisolt
committed
Add more tests for Snowflake
1 parent 79a68a3 commit e38f04d

File tree

2 files changed

+65
-33
lines changed

2 files changed

+65
-33
lines changed

‎dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/SnowflakeLogsConnector.java‎

Lines changed: 27 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -226,7 +226,7 @@ private String createQueryFromAccountUsage(ConnectorArguments arguments)
226226
+ "FROM SNOWFLAKE.ACCOUNT_USAGE.QUERY_HISTORY\n"
227227
+ "WHERE end_time >= to_timestamp_ltz('%s')\n"
228228
+ "AND end_time <= to_timestamp_ltz('%s')\n");
229-
queryBuilder.append(extractEarliestTimestamp(arguments));
229+
queryBuilder.append(earliestTimestamp(arguments));
230230
if (overrideWhere != null) queryBuilder.append(" AND ").append(overrideWhere);
231231
return queryBuilder.toString().replace('\n', ' ');
232232
}
@@ -269,7 +269,7 @@ private String createQueryFromInformationSchema(ConnectorArguments arguments)
269269
// because we always iterate over the full 7 trailing days; maybe it's worth
270270
// preventing that in the future. To do that, we should require getQueryLogEarliestTimestamp()
271271
// to parse and return an ISO instant, not a database-server-specific format.
272-
queryBuilder.append(extractEarliestTimestamp(arguments));
272+
queryBuilder.append(earliestTimestamp(arguments));
273273

274274
String overrideWhere = getOverrideWhere(arguments);
275275
if (overrideWhere != null) {
@@ -336,13 +336,13 @@ private String createExtendedQueryFromAccountUsage(ConnectorArguments arguments)
336336
+ "AND end_time <= to_timestamp_ltz('%s')\n"
337337
+ "AND is_client_generated_statement = FALSE\n");
338338

339-
queryBuilder.append(extractEarliestTimestamp(arguments));
339+
queryBuilder.append(earliestTimestamp(arguments));
340340
if (overrideWhere != null) queryBuilder.append(" AND ").append(overrideWhere);
341341
return queryBuilder.toString().replace('\n', ' ');
342342
}
343343

344344
@Nonnull
345-
static String extractEarliestTimestamp(@Nonnull ConnectorArguments arguments) {
345+
static String earliestTimestamp(@Nonnull ConnectorArguments arguments) {
346346
String timestamp = arguments.getQueryLogEarliestTimestamp();
347347
if (isBlank(timestamp)) {
348348
return "";
@@ -429,17 +429,16 @@ private static MetadataDumperUsageException unsupportedOption(String option) {
429429
return new MetadataDumperUsageException(message);
430430
}
431431

432-
private String getOverrideableQuery(
433-
@Nullable String overrideQuery, @Nonnull String defaultSql, @Nonnull String whereField) {
434-
String sql = overrideQuery != null ? overrideQuery : defaultSql;
435-
return sql
436-
+ "\n"
437-
+ "WHERE "
438-
+ whereField
439-
+ " >= to_timestamp_ltz('%s')\n"
440-
+ "AND "
441-
+ whereField
442-
+ " <= to_timestamp_ltz('%s')";
432+
static String overrideableQuery(
433+
@Nullable String override, @Nonnull String defaultSql, @Nonnull String whereField) {
434+
String start = whereField + " >= to_timestamp_ltz('%s')";
435+
String end = whereField + " <= to_timestamp_ltz('%s')";
436+
437+
if (override != null) {
438+
return String.format("%s\nWHERE %s\nAND %s", override, start, end);
439+
} else {
440+
return String.format("%s\nWHERE %s\nAND %s", defaultSql, start, end);
441+
}
443442
}
444443

445444
private String parseColumnsFromHeader(Class<? extends Enum<?>> headerClass) {
@@ -454,7 +453,7 @@ private List<TaskDescription> createTimeSeriesTasks(ConnectorArguments arguments
454453
return Arrays.asList(
455454
new TaskDescription(
456455
WarehouseEventsHistoryFormat.ZIP_ENTRY_PREFIX,
457-
getOverrideableQuery(
456+
overrideableQuery(
458457
arguments.getDefinition(
459458
SnowflakeLogConnectorProperties.WAREHOUSE_EVENTS_HISTORY_OVERRIDE_QUERY),
460459
String.format(
@@ -465,7 +464,7 @@ private List<TaskDescription> createTimeSeriesTasks(ConnectorArguments arguments
465464
WarehouseEventsHistoryFormat.Header.class),
466465
new TaskDescription(
467466
AutomaticClusteringHistoryFormat.ZIP_ENTRY_PREFIX,
468-
getOverrideableQuery(
467+
overrideableQuery(
469468
arguments.getDefinition(
470469
SnowflakeLogConnectorProperties.AUTOMATIC_CLUSTERING_HISTORY_OVERRIDE_QUERY),
471470
String.format(
@@ -476,7 +475,7 @@ private List<TaskDescription> createTimeSeriesTasks(ConnectorArguments arguments
476475
AutomaticClusteringHistoryFormat.Header.class),
477476
new TaskDescription(
478477
CopyHistoryFormat.ZIP_ENTRY_PREFIX,
479-
getOverrideableQuery(
478+
overrideableQuery(
480479
arguments.getDefinition(
481480
SnowflakeLogConnectorProperties.COPY_HISTORY_OVERRIDE_QUERY),
482481
String.format(
@@ -487,7 +486,7 @@ private List<TaskDescription> createTimeSeriesTasks(ConnectorArguments arguments
487486
CopyHistoryFormat.Header.class),
488487
new TaskDescription(
489488
DatabaseReplicationUsageHistoryFormat.ZIP_ENTRY_PREFIX,
490-
getOverrideableQuery(
489+
overrideableQuery(
491490
arguments.getDefinition(
492491
SnowflakeLogConnectorProperties
493492
.DATABASE_REPLICATION_USAGE_HISTORY_OVERRIDE_QUERY),
@@ -499,7 +498,7 @@ private List<TaskDescription> createTimeSeriesTasks(ConnectorArguments arguments
499498
DatabaseReplicationUsageHistoryFormat.Header.class),
500499
new TaskDescription(
501500
LoginHistoryFormat.ZIP_ENTRY_PREFIX,
502-
getOverrideableQuery(
501+
overrideableQuery(
503502
arguments.getDefinition(
504503
SnowflakeLogConnectorProperties.LOGIN_HISTORY_OVERRIDE_QUERY),
505504
String.format(
@@ -510,7 +509,7 @@ private List<TaskDescription> createTimeSeriesTasks(ConnectorArguments arguments
510509
LoginHistoryFormat.Header.class),
511510
new TaskDescription(
512511
MeteringDailyHistoryFormat.ZIP_ENTRY_PREFIX,
513-
getOverrideableQuery(
512+
overrideableQuery(
514513
arguments.getDefinition(
515514
SnowflakeLogConnectorProperties.METERING_DAILY_HISTORY_OVERRIDE_QUERY),
516515
String.format(
@@ -521,7 +520,7 @@ private List<TaskDescription> createTimeSeriesTasks(ConnectorArguments arguments
521520
MeteringDailyHistoryFormat.Header.class),
522521
new TaskDescription(
523522
PipeUsageHistoryFormat.ZIP_ENTRY_PREFIX,
524-
getOverrideableQuery(
523+
overrideableQuery(
525524
arguments.getDefinition(
526525
SnowflakeLogConnectorProperties.PIPE_USAGE_HISTORY_OVERRIDE_QUERY),
527526
String.format(
@@ -532,7 +531,7 @@ private List<TaskDescription> createTimeSeriesTasks(ConnectorArguments arguments
532531
PipeUsageHistoryFormat.Header.class),
533532
new TaskDescription(
534533
QueryAccelerationHistoryFormat.ZIP_ENTRY_PREFIX,
535-
getOverrideableQuery(
534+
overrideableQuery(
536535
arguments.getDefinition(
537536
SnowflakeLogConnectorProperties.QUERY_ACCELERATION_HISTORY_OVERRIDE_QUERY),
538537
String.format(
@@ -544,7 +543,7 @@ private List<TaskDescription> createTimeSeriesTasks(ConnectorArguments arguments
544543
TaskCategory.OPTIONAL),
545544
new TaskDescription(
546545
ReplicationGroupUsageHistoryFormat.ZIP_ENTRY_PREFIX,
547-
getOverrideableQuery(
546+
overrideableQuery(
548547
arguments.getDefinition(
549548
SnowflakeLogConnectorProperties.REPLICATION_GROUP_USAGE_HISTORY_OVERRIDE_QUERY),
550549
String.format(
@@ -555,7 +554,7 @@ private List<TaskDescription> createTimeSeriesTasks(ConnectorArguments arguments
555554
ReplicationGroupUsageHistoryFormat.Header.class),
556555
new TaskDescription(
557556
ServerlessTaskHistoryFormat.ZIP_ENTRY_PREFIX,
558-
getOverrideableQuery(
557+
overrideableQuery(
559558
arguments.getDefinition(
560559
SnowflakeLogConnectorProperties.SERVERLESS_TASK_HISTORY_OVERRIDE_QUERY),
561560
String.format(
@@ -566,7 +565,7 @@ private List<TaskDescription> createTimeSeriesTasks(ConnectorArguments arguments
566565
ServerlessTaskHistoryFormat.Header.class),
567566
new TaskDescription(
568567
TaskHistoryFormat.ZIP_ENTRY_PREFIX,
569-
getOverrideableQuery(
568+
overrideableQuery(
570569
arguments.getDefinition(
571570
SnowflakeLogConnectorProperties.TASK_HISTORY_OVERRIDE_QUERY),
572571
String.format(
@@ -577,7 +576,7 @@ private List<TaskDescription> createTimeSeriesTasks(ConnectorArguments arguments
577576
TaskHistoryFormat.Header.class),
578577
new TaskDescription(
579578
WarehouseLoadHistoryFormat.ZIP_ENTRY_PREFIX,
580-
getOverrideableQuery(
579+
overrideableQuery(
581580
arguments.getDefinition(
582581
SnowflakeLogConnectorProperties.WAREHOUSE_LOAD_HISTORY_OVERRIDE_QUERY),
583582
String.format(
@@ -588,7 +587,7 @@ private List<TaskDescription> createTimeSeriesTasks(ConnectorArguments arguments
588587
WarehouseLoadHistoryFormat.Header.class),
589588
new TaskDescription(
590589
WarehouseMeteringHistoryFormat.ZIP_ENTRY_PREFIX,
591-
getOverrideableQuery(
590+
overrideableQuery(
592591
arguments.getDefinition(
593592
SnowflakeLogConnectorProperties.WAREHOUSE_METERING_HISTORY_OVERRIDE_QUERY),
594593
String.format(

‎dumper/app/src/test/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/SnowflakeLogsConnectorTest.java‎

Lines changed: 38 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,10 @@
1717
package com.google.edwmigration.dumper.application.dumper.connector.snowflake;
1818

1919
import static com.google.common.base.Predicates.alwaysTrue;
20+
import static com.google.edwmigration.dumper.application.dumper.connector.snowflake.SnowflakeLogsConnector.earliestTimestamp;
21+
import static com.google.edwmigration.dumper.application.dumper.connector.snowflake.SnowflakeLogsConnector.overrideableQuery;
2022
import static org.junit.Assert.assertEquals;
23+
import static org.junit.Assert.assertFalse;
2124
import static org.junit.Assert.assertTrue;
2225

2326
import com.google.edwmigration.dumper.application.dumper.ConnectorArguments;
@@ -37,7 +40,7 @@ public class SnowflakeLogsConnectorTest {
3740
@Test
3841
public void testExecution() throws Exception {
3942
class ExecutionTest extends AbstractSnowflakeConnectorExecutionTest {
40-
void test(File output) {
43+
void test(File output) throws Exception {
4144
if (run(ARGS(new SnowflakeLogsConnector(), output))) {
4245
new ZipValidator()
4346
.withFormat("snowflake.logs.zip")
@@ -52,32 +55,62 @@ void test(File output) {
5255
}
5356

5457
@Test
55-
public void extractEarliestTimestamp_notProvided_emptyResult() throws IOException {
58+
public void earliestTimestamp_notProvided_emptyResult() throws IOException {
5659
ConnectorArguments arguments = new ConnectorArguments("--connector", "snowflake-logs");
5760

58-
String result = SnowflakeLogsConnector.extractEarliestTimestamp(arguments);
61+
String result = earliestTimestamp(arguments);
5962

6063
assertEquals("", result);
6164
}
6265

6366
@Test
64-
public void extractEarliestTimestamp_provided_resultMatches() throws IOException {
67+
public void earliestTimestamp_provided_resultMatches() throws IOException {
6568
ConnectorArguments arguments =
6669
new ConnectorArguments(
6770
"--connector",
6871
"snowflake-logs",
6972
"--" + ConnectorArguments.OPT_QUERY_LOG_EARLIEST_TIMESTAMP,
7073
"2024-03-21");
7174

72-
String result = SnowflakeLogsConnector.extractEarliestTimestamp(arguments);
75+
String result = earliestTimestamp(arguments);
7376

7477
assertTrue(result, result.contains("2024-03-21"));
7578
assertTrue(result, result.contains("start_time"));
7679
assertTrue(result, result.endsWith("\n"));
7780
}
7881

82+
@Test
83+
public void overrideableQuery_overrideAbsent_defaultUsed() throws IOException {
84+
String defaultSql = "SELECT event_name, query_id FROM WAREHOUSE_EVENTS_HISTORY";
85+
86+
String result = overrideableQuery(null, defaultSql, "timestamp");
87+
88+
assertTrue(result, result.contains("event_name"));
89+
}
90+
91+
@Test
92+
public void overrideableQuery_overrideEmpty_resultEmpty() throws IOException {
93+
String defaultSql = "SELECT event_name, query_id FROM WAREHOUSE_EVENTS_HISTORY";
94+
String override = "";
95+
96+
String result = overrideableQuery(override, defaultSql, "timestamp");
97+
98+
assertFalse(result, result.contains("event_name"));
99+
}
100+
101+
@Test
102+
public void overrideableQuery_overridePresent_defaultIgnored() throws IOException {
103+
String defaultSql = "SELECT event_name, query_id FROM WAREHOUSE_EVENTS_HISTORY";
104+
String override = "SELECT query_id FROM WAREHOUSE_EVENTS_HISTORY";
105+
106+
String result = overrideableQuery(override, defaultSql, "timestamp");
107+
108+
assertFalse(result, result.contains("event_name"));
109+
}
110+
79111
@Test
80112
public void validate_unsupportedOption_throwsException() throws IOException {
113+
SnowflakeLogsConnector connector = new SnowflakeLogsConnector();
81114
ConnectorArguments arguments =
82115
new ConnectorArguments(
83116
"--connector",

0 commit comments

Comments
 (0)