[b/394312229] Dump Snowflake external table metadata without the assessment flag (#828)

Author: Hackerpilot

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

@@ -133,8 +133,7 @@ final ImmutableList<Task<?>> getSqlTasks(
   private void setCurrentDatabase(@Nonnull String databaseName, @Nonnull JdbcTemplate jdbcTemplate)
       throws MetadataDumperUsageException {
     String currentDatabase =
-        jdbcTemplate.queryForObject(
-            String.format("USE DATABASE \"%s\";", databaseName), String.class);
+        jdbcTemplate.queryForObject(String.format("USE DATABASE %s;", databaseName), String.class);
     if (currentDatabase == null) {
       List<String> dbNames =
           jdbcTemplate.query("SHOW DATABASES", (rs, rowNum) -> rs.getString("name"));
@@ -146,8 +145,7 @@ private void setCurrentDatabase(@Nonnull String databaseName, @Nonnull JdbcTempl
     }
   }
 
-  private String sanitizeDatabaseName(@Nonnull String databaseName)
-      throws MetadataDumperUsageException {
+  String sanitizeDatabaseName(@Nonnull String databaseName) throws MetadataDumperUsageException {
     CharMatcher doubleQuoteMatcher = CharMatcher.is('"');
     String trimmedName = doubleQuoteMatcher.trimFrom(databaseName);
     int charLengthWithQuotes = databaseName.length() + 2;

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

@@ -28,7 +28,7 @@
 @ParametersAreNonnullByDefault
 final class HeaderTransformerUtil {
 
-  static ResultSetTransformer<String[]> toCamelCaseFrom(CaseFormat format) {
+  public static ResultSetTransformer<String[]> toCamelCaseFrom(CaseFormat format) {
     return resultSet -> transformToCamelCase(resultSet, format);
   }
 

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

@@ -20,6 +20,7 @@ enum MetadataView {
   DATABASES("databases"),
   SCHEMATA("schemata"),
   TABLES("tables"),
+  EXTERNAL_TABLES("externalTables"),
   COLUMNS("columns"),
   VIEWS("views"),
   FUNCTIONS("functions"),

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

@@ -146,7 +146,18 @@ public final void addTasksTo(
     out.add(new FormatTask(FORMAT_NAME));
 
     boolean INJECT_IS_FAULT = arguments.isTestFlag('A');
-    final String IS = INJECT_IS_FAULT ? "__NONEXISTENT__" : "INFORMATION_SCHEMA";
+    // INFORMATION_SCHEMA queries must be qualified with a database
+    // name or that a "USE DATABASE" command has previously been run
+    // in the same session. Qualify the name to avoid this dependency.
+    final String databaseName = arguments.getDatabaseSingleName();
+    final String IS;
+    if (INJECT_IS_FAULT) {
+      IS = "__NONEXISTENT__";
+    } else if (databaseName == null) {
+      IS = "INFORMATION_SCHEMA";
+    } else {
+      IS = sanitizeDatabaseName(databaseName) + ".INFORMATION_SCHEMA";
+    }
     final String AU = "SNOWFLAKE.ACCOUNT_USAGE";
     final String AU_WHERE = " WHERE DELETED IS NULL";
 
@@ -240,15 +251,25 @@ public final void addTasksTo(
 
     if (isAssessment) {
       for (AssessmentQuery item : planner.generateAssessmentQueries()) {
-        String formatString = overrideFormatString(item, arguments);
-        String query = String.format(formatString, AU, /* an empty WHERE clause */ "");
-        String zipName = item.zipEntryName;
-        Task<?> task = new JdbcSelectTask(zipName, query).withHeaderTransformer(item.transformer());
-        out.add(task);
+        addAssessmentQuery(item, out, arguments, AU);
       }
+    } else {
+      addAssessmentQuery(SnowflakePlanner.SHOW_EXTERNAL_TABLES, out, arguments, AU);
     }
   }
 
+  private void addAssessmentQuery(
+      @Nonnull AssessmentQuery item,
+      @Nonnull List<? super Task<?>> out,
+      @Nonnull ConnectorArguments arguments,
+      @Nonnull String AU) {
+    String formatString = overrideFormatString(item, arguments);
+    String query = String.format(formatString, AU, /* an empty WHERE clause */ "");
+    String zipName = item.zipEntryName;
+    Task<?> task = new JdbcSelectTask(zipName, query).withHeaderTransformer(item.transformer());
+    out.add(task);
+  }
+
   private String overrideFormatString(AssessmentQuery query, ConnectorArguments arguments) {
     if (query.needsOverride) {
       return getOverrideableQuery(arguments, query.formatString, TABLE_STORAGE_METRICS);

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

@@ -42,6 +42,9 @@
 @ParametersAreNonnullByDefault
 final class SnowflakePlanner {
 
+  public static final AssessmentQuery SHOW_EXTERNAL_TABLES =
+      AssessmentQuery.createShow("EXTERNAL TABLES", Format.EXTERNAL_TABLES, LOWER_UNDERSCORE);
+
   private enum Format {
     EXTERNAL_TABLES(ExternalTablesFormat.AU_ZIP_ENTRY_NAME),
     FUNCTION_INFO(FunctionInfoFormat.AU_ZIP_ENTRY_NAME),
@@ -62,7 +65,7 @@ ImmutableList<AssessmentQuery> generateAssessmentQueries() {
     return ImmutableList.of(
         AssessmentQuery.createMetricsSelect(Format.TABLE_STORAGE_METRICS, UPPER_UNDERSCORE),
         AssessmentQuery.createShow("WAREHOUSES", Format.WAREHOUSES, LOWER_UNDERSCORE),
-        AssessmentQuery.createShow("EXTERNAL TABLES", Format.EXTERNAL_TABLES, LOWER_UNDERSCORE),
+        SHOW_EXTERNAL_TABLES,
         AssessmentQuery.createShow("FUNCTIONS", Format.FUNCTION_INFO, LOWER_UNDERSCORE));
   }
 

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

@@ -56,11 +56,13 @@ public void testExecution() throws Exception {
     validator.withEntryValidator(
         SnowflakeMetadataDumpFormat.FunctionsFormat.AU_ZIP_ENTRY_NAME,
         SnowflakeMetadataDumpFormat.FunctionsFormat.Header.class);
+    validator.withEntryValidator(
+        SnowflakeMetadataDumpFormat.ExternalTablesFormat.AU_ZIP_ENTRY_NAME,
+        SnowflakeMetadataDumpFormat.ExternalTablesFormat.Header.class);
 
     validator.withExpectedEntries(
         SnowflakeMetadataDumpFormat.TableStorageMetricsFormat.AU_ZIP_ENTRY_NAME,
         SnowflakeMetadataDumpFormat.FunctionInfoFormat.AU_ZIP_ENTRY_NAME,
-        SnowflakeMetadataDumpFormat.ExternalTablesFormat.AU_ZIP_ENTRY_NAME,
         SnowflakeMetadataDumpFormat.WarehousesFormat.AU_ZIP_ENTRY_NAME);
   }
 }

dumper/lib-dumper-spi/src/main/java/com/google/edwmigration/dumper/plugin/lib/dumper/spi/SnowflakeMetadataDumpFormat.java

@@ -139,5 +139,27 @@ interface WarehousesFormat {
 
   interface ExternalTablesFormat {
     String AU_ZIP_ENTRY_NAME = "external_tables.csv";
+
+    enum Header {
+      CreatedOn,
+      Name,
+      DatabaseName,
+      SchemaName,
+      Invalid,
+      InvalidReason,
+      Owner,
+      Comment,
+      Stage,
+      Location,
+      FileFormatName,
+      FileFormatType,
+      Cloud,
+      Region,
+      NotificationChannel,
+      LastRefreshedOn,
+      TableFormat,
+      LastRefreshDetails,
+      OwnerRoleType
+    }
   }
 }