Destination snowflake: pull in cdk update for refreshes bugfix (#42505)

Author: edgao

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

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

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

@@ -5,7 +5,7 @@ data:
   connectorSubtype: database
   connectorType: destination
   definitionId: 424892c4-daac-4491-b35d-c6688ba547ba
-  dockerImageTag: 3.11.7
+  dockerImageTag: 3.11.8
   dockerRepository: airbyte/destination-snowflake
   documentationUrl: https://docs.airbyte.com/integrations/destinations/snowflake
   githubIssueLabel: destination-snowflake

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

@@ -26,14 +26,14 @@ import javax.sql.DataSource
 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 {}
 
-abstract class AbstractSnowflakeTypingDedupingTest : BaseTypingDedupingTest() {
-    private var databaseName: String? = null
+abstract class AbstractSnowflakeTypingDedupingTest(
+    private val forceUppercaseIdentifiers: Boolean = false,
+) : BaseTypingDedupingTest() {
+    private lateinit var databaseName: String
     private var database: JdbcDatabase? = null
     // not super happy with this one, but our test classes are not super kotlin-friendly
     private lateinit var dataSource: DataSource
@@ -55,7 +55,7 @@ abstract class AbstractSnowflakeTypingDedupingTest : BaseTypingDedupingTest() {
         dataSource =
             SnowflakeDatabaseUtils.createDataSource(config, OssCloudEnvVarConsts.AIRBYTE_OSS)
         database = SnowflakeDatabaseUtils.getDatabase(dataSource)
-        cleanAirbyteInternalTable(database)
+        cleanAirbyteInternalTable(databaseName, database, forceUppercaseIdentifiers)
         return config
     }
 
@@ -103,8 +103,15 @@ abstract class AbstractSnowflakeTypingDedupingTest : BaseTypingDedupingTest() {
                 DROP SCHEMA IF EXISTS "%s" CASCADE
                 
                 """.trimIndent(),
-                rawSchema, // Raw table is still lowercase.
-                StreamId.concatenateRawTableName(namespaceOrDefault, streamName),
+                // Raw table is still lowercase by default
+                if (forceUppercaseIdentifiers) {
+                    rawSchema.uppercase()
+                } else {
+                    rawSchema
+                },
+                StreamId.concatenateRawTableName(namespaceOrDefault, streamName).let {
+                    if (forceUppercaseIdentifiers) it.uppercase() else it
+                },
                 namespaceOrDefault.uppercase(Locale.getDefault()),
             ),
         )
@@ -383,18 +390,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()
 
@@ -419,17 +414,45 @@ abstract class AbstractSnowflakeTypingDedupingTest : BaseTypingDedupingTest() {
         @Volatile private var cleanedAirbyteInternalTable = false
 
         @Throws(SQLException::class)
-        private fun cleanAirbyteInternalTable(database: JdbcDatabase?) {
+        private fun cleanAirbyteInternalTable(
+            databaseName: String,
+            database: JdbcDatabase?,
+            forceUppercase: Boolean,
+        ) {
             if (!cleanedAirbyteInternalTable) {
                 synchronized(AbstractSnowflakeTypingDedupingTest::class.java) {
                     if (!cleanedAirbyteInternalTable) {
-                        database!!.execute(
-                            "DELETE FROM \"airbyte_internal\".\"_airbyte_destination_state\" WHERE \"updated_at\" < current_date() - 7",
-                        )
+                        val destinationStateTableExists =
+                            database!!.executeMetadataQuery {
+                                it.getTables(
+                                        databaseName,
+                                        if (forceUppercase) {
+                                            "AIRBYTE_INTERNAL"
+                                        } else {
+                                            "airbyte_internal"
+                                        },
+                                        if (forceUppercase) {
+                                            "_AIRBYTE_DESTINATION_STATE"
+                                        } else {
+                                            "_airbyte_destination_state"
+                                        },
+                                        null
+                                    )
+                                    .next()
+                            }
+                        if (destinationStateTableExists) {
+                            database.execute(
+                                """DELETE FROM "airbyte_internal"."_airbyte_destination_state" WHERE "updated_at" < current_date() - 7""",
+                            )
+                        }
                         cleanedAirbyteInternalTable = true
                     }
                 }
             }
         }
     }
 }
+
+open class Batch(val name: String)
+
+class LocalFileBatch(name: String) : Batch(name)

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

@@ -9,7 +9,7 @@ import io.airbyte.commons.json.Jsons.emptyObject
 import java.util.*
 
 class SnowflakeInternalStagingCaseInsensitiveTypingDedupingTest :
-    AbstractSnowflakeTypingDedupingTest() {
+    AbstractSnowflakeTypingDedupingTest(forceUppercaseIdentifiers = true) {
     override val configPath: String
         get() = "secrets/1s1t_case_insensitive.json"
 

docs/integrations/destinations/snowflake.md

@@ -268,6 +268,7 @@ desired namespace.
 
 | Version         | Date       | Pull Request                                               | Subject                                                                                                                                                                          |
 |:----------------|:-----------|:-----------------------------------------------------------|:---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| 3.11.8          | 2024-08-16 | [\#42505](https://github.com/airbytehq/airbyte/pull/42505) | Fix bug in refreshes logic (already mitigated in platform, just fixing protocol compliance)                                                                                      |
 | 3.11.7          | 2024-08-09 | [\#43440](https://github.com/airbytehq/airbyte/pull/43440) | remove contention on state table by deleting rows ony once every 100 updates |
 | 3.11.6          | 2024-08-09 | [\#43332](https://github.com/airbytehq/airbyte/pull/43332) | bump Java CDK |
 | 3.11.5          | 2024-08-07 | [\#43348](https://github.com/airbytehq/airbyte/pull/43348) | SnowflakeSqlGen cleanup to Kotlin string interpolation                                                                                                                           |