snowflake cdk bump

Author: edgao

airbyte-cdk/java/airbyte-cdk/typing-deduping/src/testFixtures/kotlin/io/airbyte/integrations/base/destination/typing_deduping/BaseTypingDedupingTest.kt

@@ -1117,9 +1117,7 @@ abstract class BaseTypingDedupingTest {
         val expectedRawRecords0 = readRecords("dat/sync2_expectedrecords_raw.jsonl")
         val expectedFinalRecords0 =
             readRecords("dat/sync2_expectedrecords_fullrefresh_append_final.jsonl")
-        (expectedFinalRecords0 + expectedRawRecords0).forEach {
-            (it as ObjectNode).put(JavaBaseConstants.COLUMN_NAME_AB_GENERATION_ID, 41)
-        }
+        fixGenerationId(expectedRawRecords0, expectedFinalRecords0, 41)
         verifySyncResult(expectedRawRecords0, expectedFinalRecords0, disableFinalTableComparison())
 
         val catalog =
@@ -1324,7 +1322,10 @@ abstract class BaseTypingDedupingTest {
                 val sync2Records =
                     baseRecords.subList(expectedFinalRecords1.size, baseRecords.size).onEach {
                         (it as ObjectNode).put(
-                            JavaBaseConstants.COLUMN_NAME_AB_GENERATION_ID,
+                            finalMetadataColumnNames.getOrDefault(
+                                JavaBaseConstants.COLUMN_NAME_AB_GENERATION_ID,
+                                JavaBaseConstants.COLUMN_NAME_AB_GENERATION_ID
+                            ),
                             44,
                         )
                     }

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

@@ -5,7 +5,7 @@ plugins {
 airbyteJavaConnector {
     cdkVersionRequired = '0.44.3'
     features = ['db-destinations', 's3-destinations', 'typing-deduping']
-    useLocalCdk = false
+    useLocalCdk = true
 }
 
 java {

airbyte-integrations/connectors/destination-snowflake/src/main/kotlin/io/airbyte/integrations/destination/snowflake/operation/SnowflakeStorageOperation.kt

@@ -100,7 +100,7 @@ class SnowflakeStorageOperation(
             generationIdNode =
                 results.first().get(JavaBaseConstants.COLUMN_NAME_AB_GENERATION_ID.uppercase())
         }
-        return generationIdNode?.asLong()
+        return generationIdNode?.asLong() ?: 0
     }
 
     internal fun createTableQuery(streamId: StreamId, suffix: String): String {

airbyte-integrations/connectors/destination-snowflake/src/main/kotlin/io/airbyte/integrations/destination/snowflake/typing_deduping/SnowflakeDestinationHandler.kt

@@ -27,6 +27,7 @@ import io.airbyte.integrations.base.destination.typing_deduping.Union
 import io.airbyte.integrations.base.destination.typing_deduping.UnsupportedOneOf
 import io.airbyte.integrations.destination.snowflake.SnowflakeDatabaseUtils
 import io.airbyte.integrations.destination.snowflake.migrations.SnowflakeState
+import io.airbyte.integrations.destination.snowflake.typing_deduping.SnowflakeSqlGenerator.Companion.QUOTE
 import java.sql.Connection
 import java.sql.DatabaseMetaData
 import java.sql.ResultSet
@@ -57,7 +58,7 @@ class SnowflakeDestinationHandler(
                     namespace: String,
                     name: String
                 ): Long? {
-                    return null
+                    throw NotImplementedError()
                 }
             }
     ) {
@@ -421,6 +422,12 @@ class SnowflakeDestinationHandler(
                             streamConfig.id.asPair(),
                             toDestinationState(emptyObject())
                         )
+                    val finalTableGenerationId =
+                        if (isFinalTablePresent && !isFinalTableEmpty) {
+                            getFinalTableGenerationId(streamConfig.id)
+                        } else {
+                            null
+                        }
                     return@map DestinationInitialStatus<SnowflakeState>(
                         streamConfig,
                         isFinalTablePresent,
@@ -429,7 +436,11 @@ class SnowflakeDestinationHandler(
                         isSchemaMismatch,
                         isFinalTableEmpty,
                         destinationState,
-                        finalTableGenerationId = null,
+                        finalTableGenerationId = finalTableGenerationId,
+                        // I think the temp final table gen is always null?
+                        // since the only time we T+D into the temp table
+                        // is when we're committing the sync anyway
+                        // (i.e. we'll immediately rename it to the real table)
                         finalTempTableGenerationId = null,
                     )
                 } catch (e: Exception) {
@@ -439,6 +450,44 @@ class SnowflakeDestinationHandler(
             .collect(Collectors.toList())
     }
 
+    /**
+     * Query the final table to find the generation ID of any record. Assumes that the table exists
+     * and is nonempty.
+     */
+    private fun getFinalTableGenerationId(streamId: StreamId): Long? {
+        val tableExistsWithGenerationId =
+            jdbcDatabase.executeMetadataQuery {
+                // Find a column named _airbyte_generation_id
+                // in the relevant table.
+                val resultSet =
+                    it.getColumns(
+                        databaseName,
+                        streamId.finalNamespace,
+                        streamId.finalName,
+                        JavaBaseConstants.COLUMN_NAME_AB_GENERATION_ID.uppercase()
+                    )
+                // Check if there were any such columns.
+                resultSet.next()
+            }
+        // The table doesn't exist, or exists but doesn't have generation id
+        if (!tableExistsWithGenerationId) {
+            return null
+        }
+
+        return jdbcDatabase
+            .queryJsons(
+                """
+                SELECT ${JavaBaseConstants.COLUMN_NAME_AB_GENERATION_ID.uppercase()}
+                FROM ${streamId.finalNamespace(QUOTE)}.${streamId.finalName(QUOTE)}
+                LIMIT 1
+                """.trimIndent(),
+            )
+            .first()
+            .get(JavaBaseConstants.COLUMN_NAME_AB_GENERATION_ID.uppercase())
+            ?.asLong()
+            ?: 0
+    }
+
     override fun toJdbcTypeName(airbyteType: AirbyteType): String {
         if (airbyteType is AirbyteProtocolType) {
             return toJdbcTypeName(airbyteType)

airbyte-integrations/connectors/destination-snowflake/src/test-integration/kotlin/io/airbyte/integrations/destination/snowflake/operation/SnowflakeStagingClientIntegrationTest.kt

@@ -78,7 +78,12 @@ class SnowflakeStagingClientIntegrationTest {
                     throw NotImplementedError("This method should not be called in this test")
                 }
 
-                override fun getDataRow(id: UUID, recordMessage: AirbyteRecordMessage): List<Any> {
+                override fun getDataRow(
+                    id: UUID,
+                    recordMessage: AirbyteRecordMessage,
+                    generationId: Long,
+                    syncId: Long,
+                ): List<Any> {
                     throw NotImplementedError("This method should not be called in this test")
                 }
 
@@ -129,7 +134,12 @@ class SnowflakeStagingClientIntegrationTest {
                     throw NotImplementedError("This method should not be called in this test")
                 }
 
-                override fun getDataRow(id: UUID, recordMessage: AirbyteRecordMessage): List<Any> {
+                override fun getDataRow(
+                    id: UUID,
+                    recordMessage: AirbyteRecordMessage,
+                    generationId: Long,
+                    syncId: Long,
+                ): List<Any> {
                     throw NotImplementedError("This method should not be called in this test")
                 }
 

airbyte-integrations/connectors/destination-snowflake/src/test-integration/kotlin/io/airbyte/integrations/destination/snowflake/typing_deduping/AbstractSnowflakeTypingDedupingTest.kt

@@ -27,7 +27,6 @@ import kotlin.concurrent.Volatile
 import org.junit.jupiter.api.Assertions
 import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.Assertions.assertThrows
-import org.junit.jupiter.api.Disabled
 import org.junit.jupiter.api.Test
 
 private val LOGGER = KotlinLogging.logger {}
@@ -383,18 +382,6 @@ abstract class AbstractSnowflakeTypingDedupingTest : BaseTypingDedupingTest() {
         verifySyncResult(expectedRawRecords2, expectedFinalRecords2, disableFinalTableComparison())
     }
 
-    @Test
-    @Disabled
-    override fun interruptedTruncateWithPriorData() {
-        super.interruptedTruncateWithPriorData()
-    }
-
-    @Test
-    @Disabled
-    override fun interruptedOverwriteWithoutPriorData() {
-        super.interruptedOverwriteWithoutPriorData()
-    }
-
     private val defaultSchema: String
         get() = config!!["schema"].asText()