bulk-cdk-core-extract: fix catalog schema bug (#43369)

Author: Marius Posta

airbyte-cdk/bulk/core/extract/src/main/kotlin/io/airbyte/cdk/operation/DiscoverOperation.kt

@@ -59,12 +59,15 @@ class DiscoverOperation(
                     AirbyteField.of(it.id, it.type.airbyteType.asJsonSchemaType())
                 },
             )
-        val pkColumnIDs: List<List<String>> =
-            discoveredStream.primaryKeyColumnIDs.filter { pk: List<String> ->
-                // Only keep PKs whose values can be round-tripped.
-                pk.all { airbyteStreamDecorator.isPossiblePrimaryKeyElement(allColumnsByID[it]!!) }
+        val isValidPK: Boolean =
+            discoveredStream.primaryKeyColumnIDs.all { idComponents: List<String> ->
+                val id: String = idComponents.joinToString(separator = ".")
+                val field: Field? = allColumnsByID[id]
+                field != null && airbyteStreamDecorator.isPossiblePrimaryKeyElement(field)
             }
-        airbyteStream.withSourceDefinedPrimaryKey(pkColumnIDs)
+        airbyteStream.withSourceDefinedPrimaryKey(
+            if (isValidPK) discoveredStream.primaryKeyColumnIDs else listOf(),
+        )
         if (config.global) {
             // There is a global feed of incremental records, like CDC.
             airbyteStreamDecorator.decorateGlobal(airbyteStream)

airbyte-cdk/bulk/core/extract/src/main/kotlin/io/airbyte/cdk/read/Feed.kt

@@ -34,7 +34,6 @@ data class Stream(
     val name: String,
     val namespace: String?,
     val fields: List<Field>,
-    val primaryKeyCandidates: List<List<Field>>,
     val configuredSyncMode: SyncMode,
     val configuredPrimaryKey: List<Field>?,
     val configuredCursor: FieldOrMetaField?,

airbyte-cdk/bulk/core/extract/src/main/kotlin/io/airbyte/cdk/read/StateManagerFactory.kt

@@ -156,27 +156,34 @@ class StateManagerFactory(
             return null
         }
 
-        fun pkOrNull(pkColumnIDs: List<String>): List<Field>? {
+        fun pkOrNull(pkColumnIDComponents: List<List<String>>): List<Field>? {
+            if (pkColumnIDComponents.isEmpty()) {
+                return null
+            }
+            val pkColumnIDs: List<String> =
+                pkColumnIDComponents.map { it.joinToString(separator = ".") }
             val pk: List<Field> = pkColumnIDs.mapNotNull(::dataColumnOrNull)
-            if (pk.isEmpty() || pk.size < pkColumnIDs.size) {
+            if (pk.size < pkColumnIDComponents.size) {
                 handler.accept(InvalidPrimaryKey(name, namespace, pkColumnIDs))
                 return null
             }
             return pk
         }
 
-        fun cursorOrNull(cursorColumnID: String): FieldOrMetaField? {
+        fun cursorOrNull(cursorColumnIDComponents: List<String>): FieldOrMetaField? {
+            if (cursorColumnIDComponents.isEmpty()) {
+                return null
+            }
+            val cursorColumnID: String = cursorColumnIDComponents.joinToString(separator = ".")
             if (cursorColumnID == CommonMetaField.CDC_LSN.id) {
                 return CommonMetaField.CDC_LSN
             }
             return dataColumnOrNull(cursorColumnID)
         }
-        val primaryKeyCandidates: List<List<Field>> =
-            stream.sourceDefinedPrimaryKey.mapNotNull(::pkOrNull)
         val configuredPrimaryKey: List<Field>? =
-            configuredStream.primaryKey?.asSequence()?.mapNotNull(::pkOrNull)?.firstOrNull()
+            configuredStream.primaryKey?.asSequence()?.let { pkOrNull(it.toList()) }
         val configuredCursor: FieldOrMetaField? =
-            configuredStream.cursorField?.asSequence()?.mapNotNull(::cursorOrNull)?.firstOrNull()
+            configuredStream.cursorField?.asSequence()?.let { cursorOrNull(it.toList()) }
         val configuredSyncMode: SyncMode =
             when (configuredStream.syncMode) {
                 SyncMode.INCREMENTAL ->
@@ -192,7 +199,6 @@ class StateManagerFactory(
             name,
             namespace,
             streamFields,
-            primaryKeyCandidates,
             configuredSyncMode,
             configuredPrimaryKey,
             configuredCursor,

airbyte-cdk/bulk/core/extract/src/test/kotlin/io/airbyte/cdk/read/RootReaderIntegrationTest.kt

@@ -214,7 +214,6 @@ data class TestCase(
             name = name,
             namespace = "test",
             fields = listOf(),
-            primaryKeyCandidates = listOf(),
             configuredSyncMode = SyncMode.FULL_REFRESH,
             configuredPrimaryKey = null,
             configuredCursor = null,

airbyte-cdk/bulk/core/extract/src/test/kotlin/io/airbyte/cdk/read/StateManagerGlobalStatesTest.kt

@@ -204,18 +204,12 @@ class StateManagerGlobalStatesTest {
         val kv: Stream = global.streams.first()
         Assertions.assertEquals("KV", kv.name)
         Assertions.assertEquals(listOf("V", "K"), kv.fields.map { it.id })
-        Assertions.assertEquals(
-            listOf(listOf("K")),
-            kv.primaryKeyCandidates.map { col -> col.map { it.id } },
-        )
+        Assertions.assertEquals(listOf("K"), kv.configuredPrimaryKey?.map { it.id })
         Assertions.assertEquals(SyncMode.INCREMENTAL, kv.configuredSyncMode)
         val events: Stream = streams.filter { it.namePair != kv.namePair }.first()
         Assertions.assertEquals("EVENTS", events.name)
         Assertions.assertEquals(listOf("MSG", "ID", "TS"), events.fields.map { it.id })
-        Assertions.assertEquals(
-            listOf(listOf("ID")),
-            events.primaryKeyCandidates.map { col -> col.map { it.id } },
-        )
+        Assertions.assertEquals(listOf("ID"), events.configuredPrimaryKey?.map { it.id })
         Assertions.assertEquals(SyncMode.FULL_REFRESH, events.configuredSyncMode)
         return Streams(global, kv, events)
     }

airbyte-cdk/bulk/core/extract/src/test/kotlin/io/airbyte/cdk/read/StateManagerStreamStatesTest.kt

@@ -185,10 +185,6 @@ class StateManagerStreamStatesTest {
         val eventsStream: Stream = stateManager.feeds.mapNotNull { it as? Stream }.first()
         Assertions.assertEquals("EVENTS", eventsStream.name)
         Assertions.assertEquals(listOf("MSG", "ID", "TS"), eventsStream.fields.map { it.id })
-        Assertions.assertEquals(
-            listOf(listOf("ID")),
-            eventsStream.primaryKeyCandidates.map { col -> col.map { it.id } },
-        )
         Assertions.assertEquals(expectedSyncMode, eventsStream.configuredSyncMode)
         Assertions.assertEquals(
             expectedPrimaryKey,

airbyte-cdk/bulk/core/extract/src/test/kotlin/io/airbyte/cdk/read/streams/StreamPartitionsCreatorUtilsTest.kt

@@ -60,7 +60,6 @@ class StreamPartitionsCreatorUtilsTest {
             name = "kv",
             namespace = "public",
             fields = listOf(k, v),
-            primaryKeyCandidates = listOf(listOf(k)),
             configuredSyncMode = SyncMode.FULL_REFRESH,
             configuredPrimaryKey = listOf(k),
             configuredCursor = null,

airbyte-cdk/bulk/toolkits/extract-jdbc/src/main/kotlin/io/airbyte/cdk/jdbc/JdbcMetadataQuerier.kt

@@ -299,13 +299,10 @@ class JdbcMetadataQuerier(
         } catch (e: Exception) {
             throw RuntimeException("Primary key discovery query failed: ${e.message}", e)
         }
-        return results
-            .groupBy { it.name }
-            .values
-            .map { rowsByPK: List<PrimaryKeyRow> ->
-                rowsByPK.sortedBy { it.ordinal }.map { it.columnName }
-            }
-            .also { memoizedPrimaryKeys[table] = it }
+        val rows: List<PrimaryKeyRow> = results.groupBy { it.name }.values.firstOrNull() ?: listOf()
+        val pk: List<List<String>> = rows.sortedBy { it.ordinal }.map { listOf(it.columnName) }
+        memoizedPrimaryKeys[table] = pk
+        return pk
     }
 
     private data class PrimaryKeyRow(

airbyte-cdk/bulk/toolkits/extract-jdbc/src/main/kotlin/io/airbyte/cdk/read/stream/CheckpointStreamState.kt

@@ -82,11 +82,8 @@ private fun StreamStateJsonValue.checkpoint(ctx: StreamReadContext): CheckpointS
         if (primaryKey.isEmpty()) {
             return@run mapOf()
         }
-        val keys: List<Field>? =
-            ctx.stream.primaryKeyCandidates.find { pk: List<Field> ->
-                pk.map { it.id }.toSet() == primaryKey.keys
-            }
-        if (keys == null) {
+        val pk: List<Field> = ctx.stream.configuredPrimaryKey ?: listOf()
+        if (primaryKey.keys != pk.map { it.id }.toSet()) {
             ctx.handler.accept(
                 InvalidPrimaryKey(
                     ctx.stream.name,
@@ -96,7 +93,7 @@ private fun StreamStateJsonValue.checkpoint(ctx: StreamReadContext): CheckpointS
             )
             return null
         }
-        keys.associateWith { primaryKey[it.id]!! }
+        pk.associateWith { primaryKey[it.id]!! }
     }
     val cursorPair: Pair<Field, JsonNode>? = run {
         if (cursors.isEmpty()) {

airbyte-cdk/bulk/toolkits/extract-jdbc/src/main/kotlin/io/airbyte/cdk/read/stream/StreamPartitionsCreator.kt

@@ -162,9 +162,7 @@ fun CheckpointStreamState?.streamPartitionsCreatorInput(
     ctx: StreamReadContext,
 ): StreamPartitionsCreator.Input {
     if (this == null) {
-        val pkChosenFromCatalog: List<Field> =
-            ctx.stream.configuredPrimaryKey
-                ?: ctx.stream.primaryKeyCandidates.firstOrNull() ?: listOf()
+        val pkChosenFromCatalog: List<Field> = ctx.stream.configuredPrimaryKey ?: listOf()
         if (ctx.stream.configuredSyncMode == SyncMode.FULL_REFRESH || ctx.configuration.global) {
             return StreamPartitionsCreator.SnapshotColdStart(pkChosenFromCatalog)
         }