Destination bigquery: pull in cdk update for refreshes bugfix (#42504)

Author: edgao

airbyte-integrations/connectors/destination-bigquery/build.gradle

@@ -3,7 +3,7 @@ plugins {
 }
 
 airbyteJavaConnector {
-    cdkVersionRequired = '0.41.4'
+    cdkVersionRequired = '0.44.14'
     features = [
             'db-destinations',
             'datastore-bigquery',

airbyte-integrations/connectors/destination-bigquery/metadata.yaml

@@ -5,7 +5,7 @@ data:
   connectorSubtype: database
   connectorType: destination
   definitionId: 22f6c74f-5699-40ff-833c-4a879ea40133
-  dockerImageTag: 2.8.6
+  dockerImageTag: 2.8.7
   dockerRepository: airbyte/destination-bigquery
   documentationUrl: https://docs.airbyte.com/integrations/destinations/bigquery
   githubIssueLabel: destination-bigquery

airbyte-integrations/connectors/destination-bigquery/src/main/kotlin/io/airbyte/integrations/destination/bigquery/BigQueryDestination.kt

@@ -38,6 +38,7 @@ import io.airbyte.integrations.base.destination.operation.DefaultSyncOperation
 import io.airbyte.integrations.base.destination.operation.StandardStreamOperation
 import io.airbyte.integrations.base.destination.typing_deduping.CatalogParser
 import io.airbyte.integrations.base.destination.typing_deduping.DestinationInitialStatus
+import io.airbyte.integrations.base.destination.typing_deduping.ImportType
 import io.airbyte.integrations.base.destination.typing_deduping.InitialRawTableStatus
 import io.airbyte.integrations.base.destination.typing_deduping.ParsedCatalog
 import io.airbyte.integrations.base.destination.typing_deduping.Sql
@@ -58,7 +59,6 @@ import io.airbyte.protocol.models.v0.AirbyteMessage
 import io.airbyte.protocol.models.v0.AirbyteRecordMessageMeta
 import io.airbyte.protocol.models.v0.AirbyteStreamStatusTraceMessage
 import io.airbyte.protocol.models.v0.ConfiguredAirbyteCatalog
-import io.airbyte.protocol.models.v0.DestinationSyncMode
 import io.github.oshai.kotlinlogging.KotlinLogging
 import java.io.ByteArrayInputStream
 import java.io.IOException
@@ -174,12 +174,12 @@ class BigQueryDestination : BaseConnector(), Destination {
             val streamConfig =
                 StreamConfig(
                     id = streamId,
-                    destinationSyncMode = DestinationSyncMode.OVERWRITE,
+                    postImportAction = ImportType.APPEND,
                     primaryKey = listOf(),
                     cursor = Optional.empty(),
                     columns = linkedMapOf(),
-                    generationId = 0,
-                    minimumGenerationId = 0,
+                    generationId = 1,
+                    minimumGenerationId = 1,
                     syncId = 0
                 )
 
@@ -206,7 +206,9 @@ class BigQueryDestination : BaseConnector(), Destination {
                         ),
                     isSchemaMismatch = true,
                     isFinalTableEmpty = true,
-                    destinationState = BigQueryDestinationState(needsSoftReset = false)
+                    destinationState = BigQueryDestinationState(needsSoftReset = false),
+                    finalTempTableGenerationId = null,
+                    finalTableGenerationId = null,
                 )
 
             // We simulate a mini-sync to see the raw table code path is exercised. and disable T+D

airbyte-integrations/connectors/destination-bigquery/src/main/kotlin/io/airbyte/integrations/destination/bigquery/operation/BigQueryGcsStorageOperation.kt

@@ -65,20 +65,27 @@ class BigQueryGcsStorageOperation(
         suffix: String,
         data: SerializableBuffer
     ) {
-        val stagedFileName: String = uploadRecordsToStage(streamConfig.id, suffix, data)
+        val stagedFileName: String =
+            uploadRecordsToStage(streamConfig.id, suffix, data, streamConfig.generationId)
         copyIntoTableFromStage(streamConfig.id, suffix, stagedFileName)
     }
 
     private fun uploadRecordsToStage(
         streamId: StreamId,
         suffix: String,
-        buffer: SerializableBuffer
+        buffer: SerializableBuffer,
+        generationId: Long,
     ): String {
         val objectPath: String = stagingFullPath(streamId)
         log.info {
             "Uploading records to for ${streamId.rawNamespace}.${streamId.rawName}$suffix to path $objectPath"
         }
-        return gcsStorageOperations.uploadRecordsToBucket(buffer, streamId.rawNamespace, objectPath)
+        return gcsStorageOperations.uploadRecordsToBucket(
+            buffer,
+            streamId.rawNamespace,
+            objectPath,
+            generationId
+        )
     }
 
     private fun copyIntoTableFromStage(streamId: StreamId, suffix: String, stagedFileName: String) {

airbyte-integrations/connectors/destination-bigquery/src/main/kotlin/io/airbyte/integrations/destination/bigquery/operation/BigQueryStorageOperation.kt

@@ -8,6 +8,7 @@ import com.google.cloud.bigquery.BigQuery
 import com.google.cloud.bigquery.QueryJobConfiguration
 import com.google.cloud.bigquery.TableId
 import com.google.cloud.bigquery.TableResult
+import io.airbyte.cdk.integrations.base.JavaBaseConstants
 import io.airbyte.integrations.base.destination.operation.StorageOperation
 import io.airbyte.integrations.base.destination.typing_deduping.StreamConfig
 import io.airbyte.integrations.base.destination.typing_deduping.StreamId
@@ -90,7 +91,12 @@ abstract class BigQueryStorageOperation<Data>(
         if (result.totalRows == 0L) {
             return null
         }
-        return result.iterateAll().first()["_airbyte_generation_id"].longValue
+        val value = result.iterateAll().first().get(JavaBaseConstants.COLUMN_NAME_AB_GENERATION_ID)
+        return if (value == null || value.isNull) {
+            0
+        } else {
+            value.longValue
+        }
     }
 
     private fun createStagingTable(streamId: StreamId, suffix: String) {

airbyte-integrations/connectors/destination-bigquery/src/main/kotlin/io/airbyte/integrations/destination/bigquery/typing_deduping/BigQueryDestinationHandler.kt

@@ -11,6 +11,7 @@ import io.airbyte.cdk.integrations.base.JavaBaseConstants
 import io.airbyte.cdk.integrations.util.ConnectorExceptionUtil
 import io.airbyte.commons.exceptions.ConfigErrorException
 import io.airbyte.integrations.base.destination.operation.AbstractStreamOperation
+import io.airbyte.integrations.base.destination.operation.AbstractStreamOperation.Companion.TMP_TABLE_SUFFIX
 import io.airbyte.integrations.base.destination.typing_deduping.*
 import io.airbyte.integrations.base.destination.typing_deduping.CollectionUtils.containsAllIgnoreCase
 import io.airbyte.integrations.base.destination.typing_deduping.CollectionUtils.containsIgnoreCase
@@ -60,7 +61,7 @@ class BigQueryDestinationHandler(private val bq: BigQuery, private val datasetLo
             SELECT TIMESTAMP_SUB(MIN(_airbyte_extracted_at), INTERVAL 1 MICROSECOND)
             FROM ${'$'}{raw_table}
             WHERE _airbyte_loaded_at IS NULL
-            
+
             """.trimIndent()
                                 )
                         )
@@ -95,7 +96,7 @@ class BigQueryDestinationHandler(private val bq: BigQuery, private val datasetLo
                                     """
             SELECT MAX(_airbyte_extracted_at)
             FROM ${'$'}{raw_table}
-            
+
             """.trimIndent()
                                 )
                         )
@@ -118,6 +119,43 @@ class BigQueryDestinationHandler(private val bq: BigQuery, private val datasetLo
         }
     }
 
+    private fun getFinalTableGeneration(id: StreamId, suffix: String): Long? {
+        val finalTable = bq.getTable(TableId.of(id.finalNamespace, id.finalName + suffix))
+        if (finalTable == null || !finalTable.exists()) {
+            return null
+        }
+
+        val tableDef = finalTable.getDefinition<StandardTableDefinition>()
+        val hasGenerationId: Boolean =
+            tableDef.schema
+                ?.fields
+                // Field doesn't have a hasColumn(String) method >.>
+                ?.any { it.name == JavaBaseConstants.COLUMN_NAME_AB_GENERATION_ID }
+                ?: false
+        if (!hasGenerationId) {
+            return null
+        }
+
+        val result =
+            bq.query(
+                QueryJobConfiguration.of(
+                    """
+                SELECT ${JavaBaseConstants.COLUMN_NAME_AB_GENERATION_ID}
+                FROM ${id.finalTableId(BigQuerySqlGenerator.QUOTE, suffix)}
+                """.trimIndent()
+                )
+            )
+        if (result.totalRows == 0L) {
+            return null
+        }
+        val value = result.iterateAll().first().get(JavaBaseConstants.COLUMN_NAME_AB_GENERATION_ID)
+        return if (value == null || value.isNull) {
+            0
+        } else {
+            value.longValue
+        }
+    }
+
     @Throws(InterruptedException::class)
     override fun execute(sql: Sql) {
         val transactions = sql.asSqlStrings("BEGIN TRANSACTION", "COMMIT TRANSACTION")
@@ -234,7 +272,9 @@ class BigQueryDestinationHandler(private val bq: BigQuery, private val datasetLo
                         isFinalTableEmpty(
                             id
                         ), // Return a default state blob since we don't actually track state.
-                    BigQueryDestinationState(false)
+                    BigQueryDestinationState(false),
+                    finalTableGenerationId = getFinalTableGeneration(id, ""),
+                    finalTempTableGenerationId = getFinalTableGeneration(id, TMP_TABLE_SUFFIX)
                 )
             )
         }

airbyte-integrations/connectors/destination-bigquery/src/main/kotlin/io/airbyte/integrations/destination/bigquery/typing_deduping/BigQuerySqlGenerator.kt

@@ -8,7 +8,6 @@ import com.google.common.annotations.VisibleForTesting
 import io.airbyte.integrations.base.destination.typing_deduping.*
 import io.airbyte.integrations.base.destination.typing_deduping.Array
 import io.airbyte.integrations.destination.bigquery.BigQuerySQLNameTransformer
-import io.airbyte.protocol.models.v0.DestinationSyncMode
 import java.time.Instant
 import java.util.*
 import java.util.function.Function
@@ -296,7 +295,7 @@ class BigQuerySqlGenerator
         useExpensiveSaferCasting: Boolean
     ): Sql {
         val handleNewRecords =
-            if (stream.destinationSyncMode == DestinationSyncMode.APPEND_DEDUP) {
+            if (stream.postImportAction == ImportType.DEDUPE) {
                 upsertNewRecords(stream, finalSuffix, useExpensiveSaferCasting, minRawTimestamp)
             } else {
                 insertNewRecords(stream, finalSuffix, useExpensiveSaferCasting, minRawTimestamp)
@@ -587,7 +586,7 @@ class BigQuerySqlGenerator
                 .collect(Collectors.joining("\n"))
         val extractedAtCondition = buildExtractedAtCondition(minRawTimestamp)
 
-        if (stream.destinationSyncMode == DestinationSyncMode.APPEND_DEDUP) {
+        if (stream.postImportAction == ImportType.DEDUPE) {
             // When deduping, we need to dedup the raw records. Note the row_number() invocation in
             // the SQL
             // statement. Do the same extract+cast CTE + airbyte_meta construction as in non-dedup
@@ -929,7 +928,7 @@ class BigQuerySqlGenerator
 
         fun clusteringColumns(stream: StreamConfig): List<String> {
             val clusterColumns: MutableList<String> = ArrayList()
-            if (stream.destinationSyncMode == DestinationSyncMode.APPEND_DEDUP) {
+            if (stream.postImportAction == ImportType.DEDUPE) {
                 // We're doing de-duping, therefore we have a primary key.
                 // Cluster on the first 3 PK columns since BigQuery only allows up to 4 clustering
                 // columns,

airbyte-integrations/connectors/destination-bigquery/src/test-integration/kotlin/io/airbyte/integrations/destination/bigquery/operation/BigQueryStorageOperationTest.kt

@@ -17,6 +17,7 @@ import io.airbyte.cdk.integrations.destination.async.model.PartialAirbyteRecordM
 import io.airbyte.commons.json.Jsons
 import io.airbyte.commons.string.Strings
 import io.airbyte.integrations.base.destination.operation.AbstractStreamOperation.Companion.TMP_TABLE_SUFFIX
+import io.airbyte.integrations.base.destination.typing_deduping.ImportType
 import io.airbyte.integrations.base.destination.typing_deduping.StreamConfig
 import io.airbyte.integrations.base.destination.typing_deduping.StreamId
 import io.airbyte.integrations.destination.bigquery.BigQueryConsts
@@ -27,7 +28,6 @@ import io.airbyte.integrations.destination.bigquery.typing_deduping.BigQuerySqlG
 import io.airbyte.integrations.destination.bigquery.typing_deduping.BigQuerySqlGeneratorIntegrationTest
 import io.airbyte.protocol.models.v0.AirbyteMessage.Type
 import io.airbyte.protocol.models.v0.AirbyteRecordMessageMeta
-import io.airbyte.protocol.models.v0.DestinationSyncMode
 import java.nio.file.Files
 import java.nio.file.Path
 import java.util.Optional
@@ -62,7 +62,7 @@ class BigQueryDirectLoadingStorageOperationTest {
     private val streamConfig =
         StreamConfig(
             streamId,
-            DestinationSyncMode.APPEND,
+            ImportType.APPEND,
             emptyList(),
             Optional.empty(),
             LinkedHashMap(),

airbyte-integrations/connectors/destination-bigquery/src/test-integration/kotlin/io/airbyte/integrations/destination/bigquery/typing_deduping/AbstractBigQueryTypingDedupingTest.kt

@@ -290,6 +290,65 @@ abstract class AbstractBigQueryTypingDedupingTest : BaseTypingDedupingTest() {
         }
     }
 
+    @Test
+    fun testGenerationIdMigrationForOverwrite() {
+        // First sync
+        val catalog1 =
+            ConfiguredAirbyteCatalog()
+                .withStreams(
+                    listOf(
+                        ConfiguredAirbyteStream()
+                            .withSyncMode(SyncMode.FULL_REFRESH)
+                            .withDestinationSyncMode(DestinationSyncMode.OVERWRITE)
+                            .withSyncId(41L)
+                            .withGenerationId(42L)
+                            .withMinimumGenerationId(0L)
+                            .withStream(
+                                AirbyteStream()
+                                    .withNamespace(streamNamespace)
+                                    .withName(streamName)
+                                    .withJsonSchema(SCHEMA),
+                            ),
+                    ),
+                )
+        val messages1 = readMessages("dat/sync1_messages.jsonl")
+        runSync(
+            catalog1,
+            messages1,
+            "airbyte/destination-bigquery:2.4.20",
+            // Old connector version can't handle TRACE messages; disable the
+            // stream status message
+            streamStatus = null,
+        )
+
+        // Second sync
+        val catalog2 =
+            ConfiguredAirbyteCatalog()
+                .withStreams(
+                    listOf(
+                        ConfiguredAirbyteStream()
+                            .withSyncMode(SyncMode.FULL_REFRESH)
+                            .withDestinationSyncMode(DestinationSyncMode.OVERWRITE)
+                            .withSyncId(42L)
+                            .withGenerationId(43L)
+                            .withMinimumGenerationId(43L)
+                            .withStream(
+                                AirbyteStream()
+                                    .withNamespace(streamNamespace)
+                                    .withName(streamName)
+                                    .withJsonSchema(SCHEMA),
+                            ),
+                    ),
+                )
+        val messages2 = readMessages("dat/sync2_messages.jsonl")
+        runSync(catalog2, messages2)
+
+        val expectedRawRecords2 = readRecords("dat/sync2_expectedrecords_overwrite_raw.jsonl")
+        val expectedFinalRecords2 =
+            readRecords("dat/sync2_expectedrecords_fullrefresh_overwrite_final.jsonl")
+        verifySyncResult(expectedRawRecords2, expectedFinalRecords2, disableFinalTableComparison())
+    }
+
     protected open val rawDataset: String
         /**
          * Subclasses using a config with a nonstandard raw table dataset should override this

airbyte-integrations/connectors/destination-bigquery/src/test-integration/kotlin/io/airbyte/integrations/destination/bigquery/typing_deduping/BigQuerySqlGeneratorIntegrationTest.kt

@@ -16,7 +16,6 @@ import io.airbyte.integrations.base.destination.typing_deduping.*
 import io.airbyte.integrations.destination.bigquery.BigQueryConsts
 import io.airbyte.integrations.destination.bigquery.BigQueryDestination.Companion.getBigQuery
 import io.airbyte.integrations.destination.bigquery.migrators.BigQueryDestinationState
-import io.airbyte.protocol.models.v0.DestinationSyncMode
 import java.nio.file.Files
 import java.nio.file.Path
 import java.time.Duration
@@ -509,7 +508,7 @@ class BigQuerySqlGeneratorIntegrationTest :
         val stream =
             StreamConfig(
                 streamId,
-                DestinationSyncMode.APPEND,
+                ImportType.APPEND,
                 emptyList(),
                 Optional.empty(),
                 columns,

airbyte-integrations/connectors/destination-bigquery/src/test-integration/resources/dat/sync2_expectedrecords_overwrite_raw.jsonl

@@ -0,0 +1,4 @@
+// Only sync2 messages present in overwrite mode
+{"_airbyte_extracted_at": "1970-01-01T00:00:02Z", "_airbyte_data": {"id1": 1, "id2": 200, "updated_at": "2000-01-02T00:00:00Z", "_ab_cdc_deleted_at": null, "name": "Alice", "address": {"city": "Seattle", "state": "WA"}}, "_airbyte_meta": {"changes":[],"sync_id":42}, "_airbyte_generation_id":  43}
+{"_airbyte_extracted_at": "1970-01-01T00:00:02Z", "_airbyte_data": {"id1": 1, "id2": 201, "updated_at": "2000-01-02T00:00:00Z", "_ab_cdc_deleted_at": null, "name": "Bob", "address": {"city": "New York", "state": "NY"}}, "_airbyte_meta": {"changes":[],"sync_id":42}, "_airbyte_generation_id":  43}
+{"_airbyte_extracted_at": "1970-01-01T00:00:02Z", "_airbyte_data": {"id1": 1, "id2": 201, "updated_at": "2000-01-02T00:01:00Z", "_ab_cdc_deleted_at": "1970-01-01T00:00:00Z"}, "_airbyte_meta": {"changes":[],"sync_id":42}, "_airbyte_generation_id":  43}

airbyte-integrations/connectors/destination-bigquery/src/test/kotlin/io/airbyte/integrations/destination/bigquery/typing_deduping/BigQuerySqlGeneratorTest.kt

@@ -33,7 +33,7 @@ class BigQuerySqlGeneratorTest {
         Assertions.assertEquals(
             StreamConfig(
                 StreamId("bar", "foo", "airbyte_internal", "bar_raw__stream_foo", "bar", "foo"),
-                DestinationSyncMode.APPEND,
+                ImportType.APPEND,
                 emptyList(),
                 Optional.empty(),
                 columns,