Move getSqlTasks to SnowflakeInput

Author: misolt

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

@@ -21,7 +21,6 @@
 
 import com.google.common.base.CharMatcher;
 import com.google.common.base.Joiner;
-import com.google.common.collect.ImmutableList;
 import com.google.edwmigration.dumper.application.dumper.ConnectorArguments;
 import com.google.edwmigration.dumper.application.dumper.MetadataDumperUsageException;
 import com.google.edwmigration.dumper.application.dumper.annotations.RespectsArgumentDriver;
@@ -37,9 +36,6 @@
 import com.google.edwmigration.dumper.application.dumper.connector.Connector;
 import com.google.edwmigration.dumper.application.dumper.handle.Handle;
 import com.google.edwmigration.dumper.application.dumper.handle.JdbcHandle;
-import com.google.edwmigration.dumper.application.dumper.task.AbstractJdbcTask;
-import com.google.edwmigration.dumper.application.dumper.task.Summary;
-import com.google.edwmigration.dumper.application.dumper.task.Task;
 import java.sql.Driver;
 import java.sql.SQLException;
 import java.util.ArrayList;
@@ -208,23 +204,6 @@ private String getUrlFromArguments(@Nonnull ConnectorArguments arguments) {
     return buf.toString();
   }
 
-  final ImmutableList<Task<?>> getSqlTasks(
-      @Nonnull SnowflakeInput inputSource,
-      @Nonnull Class<? extends Enum<?>> header,
-      @Nonnull String format,
-      @Nonnull AbstractJdbcTask<Summary> schemaTask,
-      @Nonnull AbstractJdbcTask<Summary> usageTask) {
-    switch (inputSource) {
-      case USAGE_THEN_SCHEMA_SOURCE:
-        return ImmutableList.of(usageTask, schemaTask.onlyIfFailed(usageTask));
-      case SCHEMA_ONLY_SOURCE:
-        return ImmutableList.of(schemaTask);
-      case USAGE_ONLY_SOURCE:
-        return ImmutableList.of(usageTask);
-    }
-    throw new AssertionError();
-  }
-
   private void setCurrentDatabase(@Nonnull String databaseName, @Nonnull JdbcTemplate jdbcTemplate)
       throws MetadataDumperUsageException {
     String currentDatabase =

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

@@ -16,6 +16,12 @@
  */
 package com.google.edwmigration.dumper.application.dumper.connector.snowflake;
 
+import com.google.common.collect.ImmutableList;
+import com.google.edwmigration.dumper.application.dumper.task.AbstractJdbcTask;
+import com.google.edwmigration.dumper.application.dumper.task.Task;
+import javax.annotation.Nonnull;
+import javax.annotation.ParametersAreNonnullByDefault;
+
 /**
  * Represents a strategy of getting Snowflake data.
  *
@@ -29,11 +35,33 @@
  * https://docs.snowflake.net/manuals/user-guide/data-share-consumers.html You must: GRANT IMPORTED
  * PRIVILEGES ON DATABASE snowflake TO ROLE <SOMETHING>;
  */
+@ParametersAreNonnullByDefault
 enum SnowflakeInput {
   /** Get data from ACCOUNT_USAGE contents, with a fallback to INFORMATION_SCHEMA. */
-  USAGE_THEN_SCHEMA_SOURCE,
+  USAGE_THEN_SCHEMA_SOURCE {
+    @Override
+    @Nonnull
+    ImmutableList<Task<?>> sqlTasks(AbstractJdbcTask<?> schemaTask, AbstractJdbcTask<?> usageTask) {
+      return ImmutableList.of(usageTask, schemaTask.onlyIfFailed(usageTask));
+    }
+  },
   /** Get data relying only on the contents of INFORMATION_SCHEMA */
-  SCHEMA_ONLY_SOURCE,
+  SCHEMA_ONLY_SOURCE {
+    @Override
+    @Nonnull
+    ImmutableList<Task<?>> sqlTasks(AbstractJdbcTask<?> schemaTask, AbstractJdbcTask<?> usageTask) {
+      return ImmutableList.of(schemaTask);
+    }
+  },
   /** Get data relying only on the contents of ACCOUNT_USAGE */
-  USAGE_ONLY_SOURCE;
+  USAGE_ONLY_SOURCE {
+    @Override
+    @Nonnull
+    ImmutableList<Task<?>> sqlTasks(AbstractJdbcTask<?> schemaTask, AbstractJdbcTask<?> usageTask) {
+      return ImmutableList.of(usageTask);
+    }
+  };
+
+  @Nonnull
+  abstract ImmutableList<Task<?>> sqlTasks(AbstractJdbcTask<?> schema, AbstractJdbcTask<?> usage);
 }

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

@@ -162,9 +162,7 @@ private void addSqlTasksWithInfoSchemaFallback(
     if (isAssessment) {
       out.add(usageTask);
     } else {
-      ImmutableList<Task<?>> tasks =
-          getSqlTasks(inputSource, header, format, schemaTask, usageTask);
-      out.addAll(tasks);
+      out.addAll(inputSource.sqlTasks(schemaTask, usageTask));
     }
   }