🐛 Destination Redshift: Improved discovery for redshift-destination not SUPER streams

airbyte-integrations/connectors/destination-jdbc/src/main/java/io/airbyte/integrations/destination/jdbc/JdbcBufferedConsumerFactory.java

@@ -5,7 +5,6 @@
 package io.airbyte.integrations.destination.jdbc;
 
 import static io.airbyte.integrations.destination.jdbc.constants.GlobalDataSizeConstants.DEFAULT_MAX_BATCH_SIZE_BYTES;
-import static java.util.stream.Collectors.toSet;
 
 import com.fasterxml.jackson.databind.JsonNode;
 import com.google.common.base.Preconditions;
@@ -49,11 +48,11 @@ public class JdbcBufferedConsumerFactory {
   private static final Logger LOGGER = LoggerFactory.getLogger(JdbcBufferedConsumerFactory.class);
 
   public static AirbyteMessageConsumer create(final Consumer<AirbyteMessage> outputRecordCollector,
-                                              final JdbcDatabase database,
-                                              final SqlOperations sqlOperations,
-                                              final NamingConventionTransformer namingResolver,
-                                              final JsonNode config,
-                                              final ConfiguredAirbyteCatalog catalog) {
+      final JdbcDatabase database,
+      final SqlOperations sqlOperations,
+      final NamingConventionTransformer namingResolver,
+      final JsonNode config,
+      final ConfiguredAirbyteCatalog catalog) {
     final List<WriteConfig> writeConfigs = createWriteConfigs(namingResolver, config, catalog, sqlOperations.isSchemaRequired());
 
     return new BufferedStreamConsumer(
@@ -66,9 +65,9 @@ public static AirbyteMessageConsumer create(final Consumer<AirbyteMessage> outpu
   }
 
   private static List<WriteConfig> createWriteConfigs(final NamingConventionTransformer namingResolver,
-                                                      final JsonNode config,
-                                                      final ConfiguredAirbyteCatalog catalog,
-                                                      final boolean schemaRequired) {
+      final JsonNode config,
+      final ConfiguredAirbyteCatalog catalog,
+      final boolean schemaRequired) {
     if (schemaRequired) {
       Preconditions.checkState(config.has("schema"), "jdbc destinations must specify a schema.");
     }
@@ -77,10 +76,10 @@ private static List<WriteConfig> createWriteConfigs(final NamingConventionTransf
   }
 
   private static Function<ConfiguredAirbyteStream, WriteConfig> toWriteConfig(
-                                                                              final NamingConventionTransformer namingResolver,
-                                                                              final JsonNode config,
-                                                                              final Instant now,
-                                                                              final boolean schemaRequired) {
+      final NamingConventionTransformer namingResolver,
+      final JsonNode config,
+      final Instant now,
+      final boolean schemaRequired) {
     return stream -> {
       Preconditions.checkNotNull(stream.getDestinationSyncMode(), "Undefined destination sync mode");
       final AirbyteStream abStream = stream.getStream();
@@ -102,23 +101,21 @@ private static Function<ConfiguredAirbyteStream, WriteConfig> toWriteConfig(
   }
 
   /**
-   * Defer to the {@link AirbyteStream}'s namespace. If this is not set, use the destination's default
-   * schema. This namespace is source-provided, and can be potentially empty.
-   *
-   * The logic here matches the logic in the catalog_process.py for Normalization. Any modifications
-   * need to be reflected there and vice versa.
+   * Defer to the {@link AirbyteStream}'s namespace. If this is not set, use the destination's default schema. This namespace is source-provided, and can be potentially empty.
+   * <p>
+   * The logic here matches the logic in the catalog_process.py for Normalization. Any modifications need to be reflected there and vice versa.
    */
   private static String getOutputSchema(final AirbyteStream stream,
-                                        final String defaultDestSchema,
-                                        final NamingConventionTransformer namingResolver) {
+      final String defaultDestSchema,
+      final NamingConventionTransformer namingResolver) {
     return stream.getNamespace() != null
         ? namingResolver.getNamespace(stream.getNamespace())
         : namingResolver.getNamespace(defaultDestSchema);
   }
 
   private static OnStartFunction onStartFunction(final JdbcDatabase database,
-                                                 final SqlOperations sqlOperations,
-                                                 final List<WriteConfig> writeConfigs) {
+      final SqlOperations sqlOperations,
+      final List<WriteConfig> writeConfigs) {
     return () -> {
       LOGGER.info("Preparing tmp tables in destination started for {} streams", writeConfigs.size());
       for (final WriteConfig writeConfig : writeConfigs) {
@@ -135,9 +132,9 @@ private static OnStartFunction onStartFunction(final JdbcDatabase database,
   }
 
   private static RecordWriter<AirbyteRecordMessage> recordWriterFunction(final JdbcDatabase database,
-                                                                         final SqlOperations sqlOperations,
-                                                                         final List<WriteConfig> writeConfigs,
-                                                                         final ConfiguredAirbyteCatalog catalog) {
+      final SqlOperations sqlOperations,
+      final List<WriteConfig> writeConfigs,
+      final ConfiguredAirbyteCatalog catalog) {
     final Map<AirbyteStreamNameNamespacePair, WriteConfig> pairToWriteConfig = writeConfigs.stream()
         .collect(Collectors.toUnmodifiableMap(JdbcBufferedConsumerFactory::toNameNamespacePair, Function.identity()));
 
@@ -153,13 +150,13 @@ private static RecordWriter<AirbyteRecordMessage> recordWriterFunction(final Jdb
   }
 
   private static OnCloseFunction onCloseFunction(final JdbcDatabase database,
-                                                 final SqlOperations sqlOperations,
-                                                 final List<WriteConfig> writeConfigs) {
+      final SqlOperations sqlOperations,
+      final List<WriteConfig> writeConfigs) {
     return (hasFailed) -> {
       // copy data
       if (!hasFailed) {
         final List<String> queryList = new ArrayList<>();
-        sqlOperations.onDestinationCloseOperations(database, writeConfigs.stream().map(WriteConfig::getOutputSchemaName).collect(toSet()));
+        sqlOperations.onDestinationCloseOperations(database, writeConfigs);
         LOGGER.info("Finalizing tables in destination started for {} streams", writeConfigs.size());
         for (final WriteConfig writeConfig : writeConfigs) {
           final String schemaName = writeConfig.getOutputSchemaName();
@@ -178,24 +175,26 @@ private static OnCloseFunction onCloseFunction(final JdbcDatabase database,
           queryList.add(sqlOperations.copyTableQuery(database, schemaName, srcTableName, dstTableName));
         }
 
-        LOGGER.info("Executing finalization of tables.");
-        sqlOperations.executeTransaction(database, queryList);
-        LOGGER.info("Finalizing tables in destination completed.");
-      }
-      // clean up
-      LOGGER.info("Cleaning tmp tables in destination started for {} streams", writeConfigs.size());
-      for (final WriteConfig writeConfig : writeConfigs) {
-        final String schemaName = writeConfig.getOutputSchemaName();
-        final String tmpTableName = writeConfig.getTmpTableName();
-        LOGGER.info("Cleaning tmp table in destination started for stream {}. schema {}, tmp table name: {}", writeConfig.getStreamName(), schemaName,
-            tmpTableName);
-
-        sqlOperations.dropTableIfExists(database, schemaName, tmpTableName);
-      }
-      LOGGER.info("Cleaning tmp tables in destination completed.");
-    };
+      LOGGER.info("Executing finalization of tables.");
+      sqlOperations.executeTransaction(database, queryList);
+      LOGGER.info("Finalizing tables in destination completed.");
+    }
+    // clean up
+    LOGGER.info("Cleaning tmp tables in destination started for {} streams", writeConfigs.size());
+    for (final WriteConfig writeConfig : writeConfigs) {
+      final String schemaName = writeConfig.getOutputSchemaName();
+      final String tmpTableName = writeConfig.getTmpTableName();
+      LOGGER.info("Cleaning tmp table in destination started for stream {}. schema {}, tmp table name: {}", writeConfig.getStreamName(), schemaName,
+          tmpTableName);
+
+      sqlOperations.dropTableIfExists(database, schemaName, tmpTableName);
+    }
+    LOGGER.info("Cleaning tmp tables in destination completed.");
   }
 
+  ;
+}
+
   private static AirbyteStreamNameNamespacePair toNameNamespacePair(final WriteConfig config) {
     return new AirbyteStreamNameNamespacePair(config.getStreamName(), config.getNamespace());
   }

airbyte-integrations/connectors/destination-jdbc/src/main/java/io/airbyte/integrations/destination/jdbc/SqlOperations.java

@@ -8,7 +8,6 @@
 import io.airbyte.db.jdbc.JdbcDatabase;
 import io.airbyte.protocol.models.AirbyteRecordMessage;
 import java.util.List;
-import java.util.Set;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -130,10 +129,10 @@ default boolean isSchemaExists(final JdbcDatabase database, final String schemaN
    * Redshift destination:
    *
    * @param database - Database that the connector is interacting with
-   * @param schemaNames - schemas will be discovered
+   * @param writeConfigs - schemas and tables (streams) will be discovered
    * @see io.airbyte.integrations.destination.redshift.RedshiftSqlOperations#onDestinationCloseOperations
    */
-  default void onDestinationCloseOperations(JdbcDatabase database, Set<String> schemaNames) {
+  default void onDestinationCloseOperations(final JdbcDatabase database, final List<WriteConfig> writeConfigs) {
     // do nothing
     LOGGER.info("No onDestinationCloseOperations required for this destination.");
   }

airbyte-integrations/connectors/destination-jdbc/src/main/java/io/airbyte/integrations/destination/jdbc/copy/CopyConsumerFactory.java

@@ -162,8 +162,6 @@ private static void closeAsOneTransaction(final Map<AirbyteStreamNameNamespacePa
         }
       }
       if (!hasFailed) {
-        sqlOperations.onDestinationCloseOperations(db,
-            pairToCopier.keySet().stream().map(AirbyteStreamNameNamespacePair::getNamespace).collect(toSet()));
         sqlOperations.executeTransaction(db, queries);
       }
     } finally {

airbyte-integrations/connectors/destination-jdbc/src/main/java/io/airbyte/integrations/destination/staging/StagingConsumerFactory.java

@@ -212,8 +212,7 @@ private OnCloseFunction onCloseFunction(final JdbcDatabase database,
           }
           queryList.add(stagingOperations.copyTableQuery(database, schemaName, srcTableName, dstTableName));
         }
-        stagingOperations.onDestinationCloseOperations(database,
-            writeConfigs.stream().map(WriteConfig::getOutputSchemaName).collect(Collectors.toSet()));
+        stagingOperations.onDestinationCloseOperations(database, writeConfigs);
         LOGGER.info("Executing finalization of tables.");
         stagingOperations.executeTransaction(database, queryList);
         LOGGER.info("Finalizing tables in destination completed.");