upgrade cdk; fix compiler errors

edgao · edgao · commit 221b105e0184 · 2024-07-12T11:28:51.000-07:00
diff --git a/airbyte-integrations/connectors/destination-databricks/build.gradle b/airbyte-integrations/connectors/destination-databricks/build.gradle
@@ -18,7 +18,7 @@ plugins {
 }
 
 airbyteJavaConnector {
-    cdkVersionRequired = '0.35.6'
+    cdkVersionRequired = '0.38.3'
     features = ['db-destinations', 's3-destinations', 'typing-deduping']
     useLocalCdk = false
 }
diff --git a/airbyte-integrations/connectors/destination-databricks/src/main/kotlin/io/airbyte/integrations/destination/databricks/DatabricksDestination.kt b/airbyte-integrations/connectors/destination-databricks/src/main/kotlin/io/airbyte/integrations/destination/databricks/DatabricksDestination.kt
@@ -15,10 +15,10 @@ import io.airbyte.cdk.integrations.destination.async.AsyncStreamConsumer
 import io.airbyte.cdk.integrations.destination.async.buffers.BufferManager
 import io.airbyte.cdk.integrations.destination.async.deser.AirbyteMessageDeserializer
 import io.airbyte.cdk.integrations.destination.s3.FileUploadFormat
-import io.airbyte.cdk.integrations.util.addDefaultNamespaceToStreams
 import io.airbyte.integrations.base.destination.operation.DefaultFlush
 import io.airbyte.integrations.base.destination.operation.DefaultSyncOperation
 import io.airbyte.integrations.base.destination.typing_deduping.CatalogParser
+import io.airbyte.integrations.base.destination.typing_deduping.StreamConfig
 import io.airbyte.integrations.base.destination.typing_deduping.StreamId
 import io.airbyte.integrations.destination.databricks.jdbc.DatabricksDestinationHandler
 import io.airbyte.integrations.destination.databricks.jdbc.DatabricksNamingTransformer
@@ -60,7 +60,7 @@ class DatabricksDestination : BaseConnector(), Destination {
         val datasource = DatabricksConnectorClientsFactory.createDataSource(connectorConfig)
         val jdbcDatabase = DefaultJdbcDatabase(datasource)
         val destinationHandler =
-            DatabricksDestinationHandler(connectorConfig.database, jdbcDatabase)
+            DatabricksDestinationHandler(sqlGenerator, connectorConfig.database, jdbcDatabase)
         val workspaceClient =
             DatabricksConnectorClientsFactory.createWorkspaceClient(
                 connectorConfig.hostname,
@@ -85,6 +85,17 @@ class DatabricksDestination : BaseConnector(), Destination {
                 dummyNamespace,
                 dummyName
             )
+        val streamConfig =
+            StreamConfig(
+                id = streamId,
+                destinationSyncMode = DestinationSyncMode.OVERWRITE,
+                primaryKey = listOf(),
+                cursor = Optional.empty(),
+                columns = linkedMapOf(),
+                generationId = 0,
+                minimumGenerationId = 0,
+                syncId = 0
+            )
 
         try {
             storageOperations.prepareStage(streamId, DestinationSyncMode.OVERWRITE)
@@ -98,9 +109,14 @@ class DatabricksDestination : BaseConnector(), Destination {
         try {
             val writeBuffer = DatabricksFileBufferFactory.createBuffer(FileUploadFormat.CSV)
             writeBuffer.use {
-                it.accept("{\"airbyte_check\":\"passed\"}", "{}", System.currentTimeMillis())
+                it.accept(
+                    "{\"airbyte_check\":\"passed\"}",
+                    "{}",
+                    generationId = 0,
+                    System.currentTimeMillis()
+                )
                 it.flush()
-                storageOperations.writeToStage(streamId, writeBuffer)
+                storageOperations.writeToStage(streamConfig, writeBuffer)
             }
         } catch (e: Exception) {
             log.error(e) { "Failed to write to stage as part of CHECK" }
@@ -129,9 +145,6 @@ class DatabricksDestination : BaseConnector(), Destination {
 
         // TODO: Deserialization should be taken care by connector runner framework later
         val connectorConfig = DatabricksConnectorConfig.deserialize(config)
-        // TODO: This abomination continues to stay, this call should be implicit in ParsedCatalog
-        //  with defaultNamespace injected
-        addDefaultNamespaceToStreams(catalog, connectorConfig.schema)
 
         val sqlGenerator =
             DatabricksSqlGenerator(DatabricksNamingTransformer(), connectorConfig.database)
@@ -145,7 +158,7 @@ class DatabricksDestination : BaseConnector(), Destination {
         val datasource = DatabricksConnectorClientsFactory.createDataSource(connectorConfig)
         val jdbcDatabase = DefaultJdbcDatabase(datasource)
         val destinationHandler =
-            DatabricksDestinationHandler(connectorConfig.database, jdbcDatabase)
+            DatabricksDestinationHandler(sqlGenerator, connectorConfig.database, jdbcDatabase)
 
         // Minimum surface area for AsyncConsumer's lifecycle functions to call.
         val storageOperations =
@@ -179,9 +192,9 @@ class DatabricksDestination : BaseConnector(), Destination {
             catalog = catalog,
             bufferManager =
                 BufferManager(
+                    defaultNamespace = connectorConfig.schema,
                     (Runtime.getRuntime().maxMemory() * BufferManager.MEMORY_LIMIT_RATIO).toLong(),
                 ),
-            defaultNamespace = Optional.of(connectorConfig.schema),
             airbyteMessageDeserializer = AirbyteMessageDeserializer(),
         )
     }
diff --git a/airbyte-integrations/connectors/destination-databricks/src/main/kotlin/io/airbyte/integrations/destination/databricks/jdbc/DatabricksDestinationHandler.kt b/airbyte-integrations/connectors/destination-databricks/src/main/kotlin/io/airbyte/integrations/destination/databricks/jdbc/DatabricksDestinationHandler.kt
@@ -29,6 +29,7 @@ import java.util.*
 import kotlin.streams.asSequence
 
 class DatabricksDestinationHandler(
+    private val sqlGenerator: DatabricksSqlGenerator,
     private val databaseName: String,
     private val jdbcDatabase: JdbcDatabase,
 ) : DestinationHandler<MinimumDestinationState.Impl> {
@@ -111,6 +112,13 @@ class DatabricksDestinationHandler(
             .toList()
     }
 
+    override fun createNamespaces(schemas: Set<String>) {
+        for (schema in schemas) {
+            // TODO: Optimize by running SHOW SCHEMAS; rather than CREATE SCHEMA if not exists
+            execute(sqlGenerator.createSchema(schema))
+        }
+    }
+
     private fun findExistingTable(
         streamIds: List<StreamId>
     ): Map<String, LinkedHashMap<String, TableDefinition>> {
@@ -182,7 +190,7 @@ class DatabricksDestinationHandler(
             tableDefinition.columns.entries
                 .filter { (it.key != abRawId && it.key != abExtractedAt && it.key != abMeta) }
                 .associate {
-                    it.key!! to if (it.value.type != "DECIMAL") it.value.type else "DECIMAL(38, 10)"
+                    it.key to if (it.value.type != "DECIMAL") it.value.type else "DECIMAL(38, 10)"
                 }
         return actualColumns == expectedColumns
     }
@@ -221,7 +229,7 @@ class DatabricksDestinationHandler(
                 // Handle resultset call in the function which will be closed
                 // after the scope is exited
                 val resultSet =
-                    metadata?.getTables(
+                    metadata.getTables(
                         databaseName,
                         id.rawNamespace,
                         id.rawName,
diff --git a/airbyte-integrations/connectors/destination-databricks/src/main/kotlin/io/airbyte/integrations/destination/databricks/operation/DatabricksStorageOperation.kt b/airbyte-integrations/connectors/destination-databricks/src/main/kotlin/io/airbyte/integrations/destination/databricks/operation/DatabricksStorageOperation.kt
@@ -36,7 +36,8 @@ class DatabricksStorageOperation(
     //  Hoist them to SqlGenerator interface in CDK, until then using concrete instance.
     private val databricksSqlGenerator = sqlGenerator as DatabricksSqlGenerator
 
-    override fun writeToStage(streamId: StreamId, data: SerializableBuffer) {
+    override fun writeToStage(streamConfig: StreamConfig, data: SerializableBuffer) {
+        val streamId = streamConfig.id
         val stagedFile = "${stagingDirectory(streamId, database)}/${data.filename}"
         workspaceClient.files().upload(stagedFile, data.inputStream)
         destinationHandler.execute(
@@ -122,12 +123,6 @@ class DatabricksStorageOperation(
         }
     }
 
-    override fun createFinalNamespace(streamId: StreamId) {
-        val finalSchema = streamId.finalNamespace
-        // TODO: Optimize by running SHOW SCHEMAS; rather than CREATE SCHEMA if not exists
-        destinationHandler.execute(sqlGenerator.createSchema(finalSchema))
-    }
-
     override fun createFinalTable(streamConfig: StreamConfig, suffix: String, replace: Boolean) {
         // The table doesn't exist. Create it. Don't force.
         destinationHandler.execute(
diff --git a/airbyte-integrations/connectors/destination-databricks/src/main/kotlin/io/airbyte/integrations/destination/databricks/operation/DatabricksStreamOperation.kt b/airbyte-integrations/connectors/destination-databricks/src/main/kotlin/io/airbyte/integrations/destination/databricks/operation/DatabricksStreamOperation.kt
@@ -27,10 +27,12 @@ class DatabricksStreamOperation(
     private val storageOperation: StorageOperation<SerializableBuffer>,
     destinationInitialStatus: DestinationInitialStatus<MinimumDestinationState.Impl>,
     private val fileUploadFormat: FileUploadFormat,
+    disableTypeDedupe: Boolean,
 ) :
     AbstractStreamOperation<MinimumDestinationState.Impl, SerializableBuffer>(
         storageOperation,
         destinationInitialStatus,
+        disableTypeDedupe = disableTypeDedupe
     ) {
     private val log = KotlinLogging.logger {}
     override fun writeRecords(streamConfig: StreamConfig, stream: Stream<PartialAirbyteMessage>) {
@@ -40,6 +42,7 @@ class DatabricksStreamOperation(
                 it.accept(
                     record.serialized!!,
                     Jsons.serialize(record.record!!.meta),
+                    streamConfig.generationId,
                     record.record!!.emittedAt,
                 )
             }
@@ -51,7 +54,7 @@ class DatabricksStreamOperation(
                     )
                 }) to staging"
             }
-            storageOperation.writeToStage(streamConfig.id, writeBuffer)
+            storageOperation.writeToStage(streamConfig, writeBuffer)
         }
     }
 }
diff --git a/airbyte-integrations/connectors/destination-databricks/src/main/kotlin/io/airbyte/integrations/destination/databricks/operation/DatabricksStreamOperationFactory.kt b/airbyte-integrations/connectors/destination-databricks/src/main/kotlin/io/airbyte/integrations/destination/databricks/operation/DatabricksStreamOperationFactory.kt
@@ -13,12 +13,14 @@ import io.airbyte.integrations.base.destination.typing_deduping.migrators.Minimu
 class DatabricksStreamOperationFactory(private val storageOperations: DatabricksStorageOperation) :
     StreamOperationFactory<MinimumDestinationState.Impl> {
     override fun createInstance(
-        destinationInitialStatus: DestinationInitialStatus<MinimumDestinationState.Impl>
+        destinationInitialStatus: DestinationInitialStatus<MinimumDestinationState.Impl>,
+        disableTypeDedupe: Boolean,
     ): StreamOperation<MinimumDestinationState.Impl> {
         return DatabricksStreamOperation(
             storageOperations,
             destinationInitialStatus,
-            FileUploadFormat.CSV
+            FileUploadFormat.CSV,
+            disableTypeDedupe,
         )
     }
 }
diff --git a/airbyte-integrations/connectors/destination-databricks/src/main/kotlin/io/airbyte/integrations/destination/databricks/staging/DatabricksFileBufferFactory.kt b/airbyte-integrations/connectors/destination-databricks/src/main/kotlin/io/airbyte/integrations/destination/databricks/staging/DatabricksFileBufferFactory.kt
@@ -55,7 +55,9 @@ object DatabricksFileBufferFactory {
                             id: UUID,
                             formattedString: String,
                             emittedAt: Long,
-                            formattedAirbyteMetaString: String
+                            formattedAirbyteMetaString: String,
+                            // TODO use this value
+                            generationId: Long,
                         ): List<Any> {
                             return listOf(
                                 id,
diff --git a/airbyte-integrations/connectors/destination-databricks/src/test-integration/kotlin/io/airbyte/integrations/destination/databricks/typededupe/DatabricksSqlGeneratorIntegrationTest.kt b/airbyte-integrations/connectors/destination-databricks/src/test-integration/kotlin/io/airbyte/integrations/destination/databricks/typededupe/DatabricksSqlGeneratorIntegrationTest.kt
@@ -61,23 +61,23 @@ class DatabricksSqlGeneratorIntegrationTest :
         }
     }
 
+    override val sqlGenerator: SqlGenerator
+        get() = DatabricksSqlGenerator(DatabricksNamingTransformer(), connectorConfig.database)
+    private val databricksSqlGenerator = sqlGenerator as DatabricksSqlGenerator
     override val destinationHandler: DestinationHandler<MinimumDestinationState.Impl>
         get() =
             DatabricksDestinationHandler(
+                databricksSqlGenerator,
                 connectorConfig.database,
                 jdbcDatabase,
             )
-    override val sqlGenerator: SqlGenerator
-        get() = DatabricksSqlGenerator(DatabricksNamingTransformer(), connectorConfig.database)
     override val supportsSafeCast: Boolean
         get() = true
 
     override fun createNamespace(namespace: String) {
         destinationHandler.execute(sqlGenerator.createSchema(namespace))
     }
 
-    private val databricksSqlGenerator = sqlGenerator as DatabricksSqlGenerator
-
     override fun createRawTable(streamId: StreamId) {
         destinationHandler.execute(databricksSqlGenerator.createRawTable(streamId))
     }
@@ -136,7 +136,9 @@ class DatabricksSqlGeneratorIntegrationTest :
         includeCdcDeletedAt: Boolean,
         streamId: StreamId,
         suffix: String?,
-        records: List<JsonNode>
+        records: List<JsonNode>,
+        // TODO
+        generationId: Long,
     ) {
         val columnNames =
             if (includeCdcDeletedAt) FINAL_TABLE_COLUMN_NAMES_CDC else FINAL_TABLE_COLUMN_NAMES

Original file line number	Diff line number	Diff line change
`@@ -18,7 +18,7 @@ plugins {`
`18`	`18`	`}`
`19`	`19`
`20`	`20`	`airbyteJavaConnector {`
`21`		`- cdkVersionRequired = '0.35.6'`
	`21`	`+ cdkVersionRequired = '0.38.3'`
`22`	`22`	`features = ['db-destinations', 's3-destinations', 'typing-deduping']`
`23`	`23`	`useLocalCdk = false`
`24`	`24`	`}`
Original file line number	Diff line number	Diff line change
`@@ -27,10 +27,12 @@ class DatabricksStreamOperation(`
`27`	`27`	`private val storageOperation: StorageOperation<SerializableBuffer>,`
`28`	`28`	`destinationInitialStatus: DestinationInitialStatus<MinimumDestinationState.Impl>,`
`29`	`29`	`private val fileUploadFormat: FileUploadFormat,`
	`30`	`+ disableTypeDedupe: Boolean,`
`30`	`31`	`) :`
`31`	`32`	`AbstractStreamOperation<MinimumDestinationState.Impl, SerializableBuffer>(`
`32`	`33`	`storageOperation,`
`33`	`34`	`destinationInitialStatus,`
	`35`	`+ disableTypeDedupe = disableTypeDedupe`
`34`	`36`	`) {`
`35`	`37`	`private val log = KotlinLogging.logger {}`
`36`	`38`	`override fun writeRecords(streamConfig: StreamConfig, stream: Stream<PartialAirbyteMessage>) {`
`@@ -40,6 +42,7 @@ class DatabricksStreamOperation(`
`40`	`42`	`it.accept(`
`41`	`43`	`record.serialized!!,`
`42`	`44`	`Jsons.serialize(record.record!!.meta),`
	`45`	`+ streamConfig.generationId,`
`43`	`46`	`record.record!!.emittedAt,`
`44`	`47`	`)`
`45`	`48`	`}`
`@@ -51,7 +54,7 @@ class DatabricksStreamOperation(`
`51`	`54`	`)`
`52`	`55`	`}) to staging"`
`53`	`56`	`}`
`54`		`- storageOperation.writeToStage(streamConfig.id, writeBuffer)`
	`57`	`+ storageOperation.writeToStage(streamConfig, writeBuffer)`
`55`	`58`	`}`
`56`	`59`	`}`
`57`	`60`	`}`
Original file line number	Diff line number	Diff line change
`@@ -13,12 +13,14 @@ import io.airbyte.integrations.base.destination.typing_deduping.migrators.Minimu`
`13`	`13`	`class DatabricksStreamOperationFactory(private val storageOperations: DatabricksStorageOperation) :`
`14`	`14`	`StreamOperationFactory<MinimumDestinationState.Impl> {`
`15`	`15`	`override fun createInstance(`
`16`		`- destinationInitialStatus: DestinationInitialStatus<MinimumDestinationState.Impl>`
	`16`	`+ destinationInitialStatus: DestinationInitialStatus<MinimumDestinationState.Impl>,`
	`17`	`+ disableTypeDedupe: Boolean,`
`17`	`18`	`): StreamOperation<MinimumDestinationState.Impl> {`
`18`	`19`	`return DatabricksStreamOperation(`
`19`	`20`	`storageOperations,`
`20`	`21`	`destinationInitialStatus,`
`21`		`- FileUploadFormat.CSV`
	`22`	`+ FileUploadFormat.CSV,`
	`23`	`+ disableTypeDedupe,`
`22`	`24`	`)`
`23`	`25`	`}`
`24`	`26`	`}`