diff --git a/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/SnowflakeMetadataConnector.java b/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/SnowflakeMetadataConnector.java index b404ba00f..5dd1845b6 100644 --- a/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/SnowflakeMetadataConnector.java +++ b/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/SnowflakeMetadataConnector.java @@ -287,67 +287,71 @@ public final void addTasksTo( if (isAssessment) { for (AssessmentQuery item : planner.generateAssessmentQueries()) { - String formatString = - item.needsOverride - ? getOverrideableQuery(arguments, item.formatString, TABLE_STORAGE_METRICS) - : item.formatString; - String whereCondition; - // Check whether the overrides changed anything. - if (formatString.equals(item.formatString)) { - // Overrides either not applied or equal to default values. - whereCondition = - " WHERE deleted = FALSE AND schema_dropped IS NULL AND table_dropped IS NULL"; - } else { - whereCondition = ""; - } - // The condition is always passed to String.format. SHOW queries simply ignore it. - String query = String.format(formatString, ACCOUNT_USAGE_SCHEMA_NAME, whereCondition); + String query = queryForAssessment(item, arguments); Task task = new JdbcSelectTask(item.zipEntryName, query, TaskCategory.REQUIRED, TaskOptions.DEFAULT) .withHeaderTransformer(item.transformer()); out.add(task); } - } else { - if (!arguments.getDatabases().isEmpty()) { - TaskOptions taskOptions = TaskOptions.DEFAULT; - for (String database : arguments.getDatabases()) { - String formatString = - String.format( - "%s IN DATABASE %s", - SnowflakePlanner.SHOW_EXTERNAL_TABLES.formatString, databaseNameQuoted(database)); - addAssessmentQuery( - SnowflakePlanner.SHOW_EXTERNAL_TABLES.withFormatString(formatString), - out, - arguments, - ACCOUNT_USAGE_SCHEMA_NAME, - taskOptions); - // Next tasks will append to the same file. - taskOptions = taskOptions.withWriteMode(WriteMode.APPEND_EXISTING); - } - } else { - addAssessmentQuery( - SnowflakePlanner.SHOW_EXTERNAL_TABLES, - out, - arguments, - ACCOUNT_USAGE_SCHEMA_NAME, - TaskOptions.DEFAULT); - } + return; + } + ImmutableList databases = arguments.getDatabases(); + + if (databases.isEmpty()) { + AssessmentQuery query = SnowflakePlanner.SHOW_EXTERNAL_TABLES; + Task task = convertAssessmentQuery(query, arguments, TaskOptions.DEFAULT); + out.add(task); + return; + } + + TaskOptions taskOptions = TaskOptions.DEFAULT; + + for (String item : databases) { + String quotedName = databaseNameQuoted(item); + AssessmentQuery baseQuery = SnowflakePlanner.SHOW_EXTERNAL_TABLES; + + String formatString = String.format("%s IN DATABASE %s", baseQuery.formatString, quotedName); + AssessmentQuery query = baseQuery.withFormatString(formatString); + Task task = convertAssessmentQuery(query, arguments, taskOptions); + out.add(task); + // Next tasks will append to the same file. + taskOptions = taskOptions.withWriteMode(WriteMode.APPEND_EXISTING); + } + } + + private String queryForAssessment(AssessmentQuery item, ConnectorArguments arguments) { + MetadataView view = TABLE_STORAGE_METRICS; + String schema = ACCOUNT_USAGE_SCHEMA_NAME; + if (!item.needsOverride) { + return item.substitute(schema, ""); + } + + ConnectorProperty propertyQuery = PropertyAction.QUERY.toProperty(view); + String overrideQuery = arguments.getDefinition(propertyQuery); + if (overrideQuery != null) { + return String.format(overrideQuery, schema, ""); } + + ConnectorProperty propertyWhere = PropertyAction.WHERE.toProperty(view); + String overrideWhere = arguments.getDefinition(propertyWhere); + if (overrideWhere != null) { + return item.substitute(schema, overrideWhere); + } + + String whereCondition = + " WHERE deleted = FALSE AND schema_dropped IS NULL AND table_dropped IS NULL"; + return item.substitute(schema, whereCondition); } - private void addAssessmentQuery( + private Task convertAssessmentQuery( @Nonnull AssessmentQuery item, - @Nonnull List> out, @Nonnull ConnectorArguments arguments, - @Nonnull String AU, @Nonnull TaskOptions taskOptions) { String formatString = overrideFormatString(item, arguments); - String query = String.format(formatString, AU, EMPTY_WHERE_CONDITION); + String query = String.format(formatString, ACCOUNT_USAGE_SCHEMA_NAME, EMPTY_WHERE_CONDITION); String zipName = item.zipEntryName; - Task task = - new JdbcSelectTask(zipName, query, TaskCategory.REQUIRED, taskOptions) - .withHeaderTransformer(item.transformer()); - out.add(task); + return new JdbcSelectTask(zipName, query, TaskCategory.REQUIRED, taskOptions) + .withHeaderTransformer(item.transformer()); } private String overrideFormatString(AssessmentQuery query, ConnectorArguments arguments) { diff --git a/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/SnowflakePlanner.java b/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/SnowflakePlanner.java index 63b84246a..ca2bee523 100644 --- a/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/SnowflakePlanner.java +++ b/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/SnowflakePlanner.java @@ -147,6 +147,10 @@ static AssessmentQuery createShow(String view, Format zipFormat, CaseFormat case ResultSetTransformer transformer() { return HeaderTransformerUtil.toCamelCaseFrom(caseFormat); } + + String substitute(String schema, String whereCondition) { + return String.format(formatString, schema, whereCondition); + } } Task eventStateTask() {