cdk-java move the generationId handling to its own class (#43329)

This pull request introduces a new interface, `JdbcGenerationHandler`, to handle the retrieval of the `_airbyte_generation_id` for any row in a table. The `getGenerationIdInTable` method has been moved from `SqlOperations` to this new interface. The necessary changes have been made across various classes to integrate this new interface properly.

Author: stephane-airbyte

airbyte-cdk/java/airbyte-cdk/core/src/main/kotlin/io/airbyte/cdk/integrations/destination/jdbc/JdbcGenerationHandler.kt

@@ -0,0 +1,16 @@
+/*
+ * Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.cdk.integrations.destination.jdbc
+
+import io.airbyte.cdk.db.jdbc.JdbcDatabase
+
+interface JdbcGenerationHandler {
+    /**
+     * get the value of _airbyte_generation_id for any row in table {@code rawNamespace.rawName}
+     *
+     * @returns true if the table exists and contains such a row, false otherwise
+     */
+    fun getGenerationIdInTable(database: JdbcDatabase, namespace: String, name: String): Long?
+}

airbyte-cdk/java/airbyte-cdk/core/src/main/kotlin/io/airbyte/cdk/integrations/destination/jdbc/SqlOperations.kt

@@ -132,13 +132,6 @@ interface SqlOperations {
     /** Check if the data record is valid and ok to be written to destination */
     fun isValidData(data: JsonNode?): Boolean
 
-    /**
-     * get the value of _airbyte_generation_id for any row in table {@code rawNamespace.rawName}
-     *
-     * @returns true if the table exists and contains such a row, false otherwise
-     */
-    fun getGenerationIdInTable(database: JdbcDatabase, namespace: String, name: String): Long?
-
     /** overwrite the raw table with the temporary raw table */
     fun overwriteRawTable(database: JdbcDatabase, rawNamespace: String, rawName: String)
 

airbyte-cdk/java/airbyte-cdk/db-destinations/src/main/kotlin/io/airbyte/cdk/integrations/destination/jdbc/AbstractJdbcDestination.kt

@@ -189,6 +189,8 @@ abstract class AbstractJdbcDestination<DestinationState : MinimumDestinationStat
 
     protected abstract fun getSqlOperations(config: JsonNode): SqlOperations
 
+    protected abstract fun getGenerationHandler(): JdbcGenerationHandler
+
     protected abstract fun getDestinationHandler(
         config: JsonNode,
         databaseName: String,
@@ -289,6 +291,7 @@ abstract class AbstractJdbcDestination<DestinationState : MinimumDestinationStat
             outputRecordCollector,
             database,
             getSqlOperations(config),
+            getGenerationHandler(),
             namingResolver,
             config,
             catalog,

airbyte-cdk/java/airbyte-cdk/db-destinations/src/main/kotlin/io/airbyte/cdk/integrations/destination/jdbc/JdbcBufferedConsumerFactory.kt

@@ -55,6 +55,7 @@ object JdbcBufferedConsumerFactory {
         outputRecordCollector: Consumer<AirbyteMessage>,
         database: JdbcDatabase,
         sqlOperations: SqlOperations,
+        generationIdHandler: JdbcGenerationHandler,
         namingResolver: NamingConventionTransformer,
         config: JsonNode,
         catalog: ConfiguredAirbyteCatalog,
@@ -77,12 +78,19 @@ object JdbcBufferedConsumerFactory {
             onStartFunction(
                 database,
                 sqlOperations,
+                generationIdHandler,
                 writeConfigs,
                 typerDeduper,
                 namingResolver,
                 parsedCatalog
             ),
-            onCloseFunction(database, sqlOperations, parsedCatalog, typerDeduper),
+            onCloseFunction(
+                database,
+                sqlOperations,
+                generationIdHandler,
+                parsedCatalog,
+                typerDeduper
+            ),
             JdbcInsertFlushFunction(
                 defaultNamespace,
                 recordWriterFunction(database, sqlOperations, writeConfigs, catalog),
@@ -98,15 +106,15 @@ object JdbcBufferedConsumerFactory {
 
     private fun createWriteConfigs(
         database: JdbcDatabase,
-        sqlOperations: SqlOperations,
+        generationIdHandler: JdbcGenerationHandler,
         namingResolver: NamingConventionTransformer,
         parsedCatalog: ParsedCatalog,
     ): List<WriteConfig> {
         return parsedCatalog.streams.map {
             val rawSuffix: String =
                 if (
                     it.minimumGenerationId == 0L ||
-                        sqlOperations.getGenerationIdInTable(
+                        generationIdHandler.getGenerationIdInTable(
                             database,
                             it.id.rawNamespace,
                             it.id.rawName
@@ -162,6 +170,7 @@ object JdbcBufferedConsumerFactory {
     private fun onStartFunction(
         database: JdbcDatabase,
         sqlOperations: SqlOperations,
+        generationIdHandler: JdbcGenerationHandler,
         writeConfigs: MutableList<WriteConfig>,
         typerDeduper: TyperDeduper,
         namingResolver: NamingConventionTransformer,
@@ -170,7 +179,7 @@ object JdbcBufferedConsumerFactory {
         return OnStartFunction {
             typerDeduper.prepareSchemasAndRunMigrations()
             writeConfigs.addAll(
-                createWriteConfigs(database, sqlOperations, namingResolver, parsedCatalog)
+                createWriteConfigs(database, generationIdHandler, namingResolver, parsedCatalog)
             )
             LOGGER.info {
                 "Preparing raw tables in destination started for ${writeConfigs.size} streams"
@@ -194,7 +203,7 @@ object JdbcBufferedConsumerFactory {
                     0L -> {}
                     writeConfig.generationId ->
                         if (
-                            sqlOperations.getGenerationIdInTable(
+                            generationIdHandler.getGenerationIdInTable(
                                 database,
                                 schemaName,
                                 dstTableName + writeConfig.rawTableSuffix
@@ -268,6 +277,7 @@ object JdbcBufferedConsumerFactory {
     private fun onCloseFunction(
         database: JdbcDatabase,
         sqlOperations: SqlOperations,
+        generationIdHandler: JdbcGenerationHandler,
         catalog: ParsedCatalog,
         typerDeduper: TyperDeduper
     ): OnCloseFunction {
@@ -278,7 +288,7 @@ object JdbcBufferedConsumerFactory {
                 catalog.streams.forEach {
                     if (
                         it.minimumGenerationId != 0L &&
-                            sqlOperations.getGenerationIdInTable(
+                            generationIdHandler.getGenerationIdInTable(
                                 database,
                                 it.id.rawNamespace,
                                 it.id.rawName

airbyte-cdk/java/airbyte-cdk/db-destinations/src/main/kotlin/io/airbyte/cdk/integrations/destination/jdbc/typing_deduping/JdbcDestinationHandler.kt

@@ -9,7 +9,7 @@ import io.airbyte.cdk.db.jdbc.JdbcDatabase
 import io.airbyte.cdk.integrations.base.JavaBaseConstants
 import io.airbyte.cdk.integrations.base.JavaBaseConstants.DestinationColumns
 import io.airbyte.cdk.integrations.destination.jdbc.ColumnDefinition
-import io.airbyte.cdk.integrations.destination.jdbc.SqlOperations
+import io.airbyte.cdk.integrations.destination.jdbc.JdbcGenerationHandler
 import io.airbyte.cdk.integrations.destination.jdbc.TableDefinition
 import io.airbyte.cdk.integrations.util.ConnectorExceptionUtil.getResultsOrLogAndThrowFirst
 import io.airbyte.commons.concurrency.CompletableFutures
@@ -55,7 +55,7 @@ abstract class JdbcDestinationHandler<DestinationState>(
     protected val rawTableNamespace: String,
     private val dialect: SQLDialect,
     private val columns: DestinationColumns = DestinationColumns.V2_WITH_GENERATION,
-    private val sqlOperations: SqlOperations,
+    protected val generationHandler: JdbcGenerationHandler,
 ) : DestinationHandler<DestinationState> {
     protected val dslContext: DSLContext
         get() = DSL.using(dialect)
@@ -343,13 +343,13 @@ abstract class JdbcDestinationHandler<DestinationState>(
                     isFinalTableEmpty,
                     destinationState,
                     finalTableGenerationId =
-                        sqlOperations.getGenerationIdInTable(
+                        generationHandler.getGenerationIdInTable(
                             jdbcDatabase,
                             streamConfig.id.rawNamespace,
                             streamConfig.id.rawName
                         ),
                     finalTempTableGenerationId =
-                        sqlOperations.getGenerationIdInTable(
+                        generationHandler.getGenerationIdInTable(
                             jdbcDatabase,
                             streamConfig.id.rawNamespace,
                             streamConfig.id.rawName + AbstractStreamOperation.TMP_TABLE_SUFFIX

airbyte-cdk/java/airbyte-cdk/db-destinations/src/main/kotlin/io/airbyte/cdk/integrations/destination/jdbc/typing_deduping/NoOpJdbcDestinationHandler.kt

@@ -6,7 +6,7 @@ package io.airbyte.cdk.integrations.destination.jdbc.typing_deduping
 
 import com.fasterxml.jackson.databind.JsonNode
 import io.airbyte.cdk.db.jdbc.JdbcDatabase
-import io.airbyte.cdk.integrations.destination.jdbc.SqlOperations
+import io.airbyte.cdk.integrations.destination.jdbc.JdbcGenerationHandler
 import io.airbyte.integrations.base.destination.typing_deduping.AirbyteType
 import io.airbyte.integrations.base.destination.typing_deduping.DestinationInitialStatus
 import io.airbyte.integrations.base.destination.typing_deduping.Sql
@@ -19,14 +19,14 @@ class NoOpJdbcDestinationHandler<DestinationState>(
     jdbcDatabase: JdbcDatabase,
     rawTableSchemaName: String,
     sqlDialect: SQLDialect,
-    sqlOperations: SqlOperations,
+    generationHandler: JdbcGenerationHandler,
 ) :
     JdbcDestinationHandler<DestinationState>(
         databaseName,
         jdbcDatabase,
         rawTableSchemaName,
         sqlDialect,
-        sqlOperations = sqlOperations
+        generationHandler = generationHandler,
     ) {
 
     override fun execute(sql: Sql) {

airbyte-cdk/java/airbyte-cdk/db-destinations/src/test/kotlin/io/airbyte/cdk/integrations/destination/jdbc/AbstractJdbcDestinationTest.kt

@@ -153,6 +153,8 @@ class AbstractJdbcDestinationTest {
 
         override fun getSqlOperations(config: JsonNode): SqlOperations = TestJdbcSqlOperations()
 
+        override fun getGenerationHandler(): JdbcGenerationHandler = mock()
+
         override fun getDestinationHandler(
             config: JsonNode,
             databaseName: String,

airbyte-cdk/java/airbyte-cdk/db-destinations/src/test/kotlin/io/airbyte/cdk/integrations/destination/jdbc/TestJdbcSqlOperations.kt

@@ -23,15 +23,6 @@ class TestJdbcSqlOperations : JdbcSqlOperations() {
         // Not required for the testing
     }
 
-    override fun getGenerationIdInTable(
-        database: JdbcDatabase,
-        rawNamespace: String,
-        rawName: String
-    ): Long {
-        return -1L
-        // Not required for the testing
-    }
-
     override fun overwriteRawTable(database: JdbcDatabase, rawNamespace: String, rawName: String) {
         // Not required for the testing
     }