bulk-cdk: DiscoverOperation supports isResumable (#43406)

Marius Posta · web-flow · commit a16cc58fbb89 · 2024-08-08T20:47:37.000Z
diff --git a/airbyte-cdk/bulk/core/base/build.gradle b/airbyte-cdk/bulk/core/base/build.gradle
@@ -5,7 +5,7 @@ dependencies {
     api 'com.fasterxml.jackson.core:jackson-databind'
     api 'com.fasterxml.jackson.datatype:jackson-datatype-jsr310'
     api 'com.kjetland:mbknor-jackson-jsonschema_2.13:1.0.39'
-    api('io.airbyte.airbyte-protocol:protocol-models:0.9.0') {
+    api('io.airbyte.airbyte-protocol:protocol-models:0.12.2') {
         exclude group: 'com.google.guava', module: 'guava'
         exclude group: 'com.google.api-client'
         exclude group: 'org.apache.logging.log4j'
diff --git a/airbyte-cdk/bulk/core/extract/src/main/kotlin/io/airbyte/cdk/operation/DiscoverOperation.kt b/airbyte-cdk/bulk/core/extract/src/main/kotlin/io/airbyte/cdk/operation/DiscoverOperation.kt
@@ -39,9 +39,8 @@ class DiscoverOperation(
                         }
                         continue
                     }
-                    val primaryKeys: List<List<String>> =
-                        metadataQuerier.primaryKeys(name, namespace)
-                    val discoveredStream = DiscoveredStream(name, namespace, fields, primaryKeys)
+                    val primaryKey: List<List<String>> = metadataQuerier.primaryKey(name, namespace)
+                    val discoveredStream = DiscoveredStream(name, namespace, fields, primaryKey)
                     airbyteStreams.add(toAirbyteStream(discoveredStream))
                 }
             }
@@ -68,6 +67,7 @@ class DiscoverOperation(
         airbyteStream.withSourceDefinedPrimaryKey(
             if (isValidPK) discoveredStream.primaryKeyColumnIDs else listOf(),
         )
+        airbyteStream.isResumable = airbyteStream.sourceDefinedPrimaryKey.isNotEmpty()
         if (config.global) {
             // There is a global feed of incremental records, like CDC.
             airbyteStreamDecorator.decorateGlobal(airbyteStream)
diff --git a/airbyte-cdk/bulk/core/extract/src/main/kotlin/io/airbyte/cdk/source/AirbyteStreamDecorator.kt b/airbyte-cdk/bulk/core/extract/src/main/kotlin/io/airbyte/cdk/source/AirbyteStreamDecorator.kt
@@ -34,7 +34,7 @@ interface AirbyteStreamDecorator {
      * message.
      *
      * This method does not determine (1), of course, because the source keys are defined in the
-     * source database itself and are retrieved via [MetadataQuerier.primaryKeys]. Instead, this
+     * source database itself and are retrieved via [MetadataQuerier.primaryKey]. Instead, this
      * method determines (2) based on the type information of the field, typically the [FieldType]
      * objects. For instance if the [Field.type] does not map to a [LosslessFieldType] then the
      * field can't reliably round-trip checkpoint values during a resumable initial sync.
diff --git a/airbyte-cdk/bulk/core/extract/src/main/kotlin/io/airbyte/cdk/source/MetadataQuerier.kt b/airbyte-cdk/bulk/core/extract/src/main/kotlin/io/airbyte/cdk/source/MetadataQuerier.kt
@@ -19,8 +19,8 @@ interface MetadataQuerier : AutoCloseable {
         streamNamespace: String?,
     ): List<Field>
 
-    /** Queries the information_schema for all primary keys for the given table. */
-    fun primaryKeys(
+    /** Queries the information_schema for any primary key on the given table. */
+    fun primaryKey(
         streamName: String,
         streamNamespace: String?,
     ): List<List<String>>
diff --git a/airbyte-cdk/bulk/core/extract/src/test/kotlin/io/airbyte/cdk/jdbc/JdbcMetadataQuerierTest.kt b/airbyte-cdk/bulk/core/extract/src/test/kotlin/io/airbyte/cdk/jdbc/JdbcMetadataQuerierTest.kt
@@ -72,7 +72,7 @@ class JdbcMetadataQuerierTest {
             val tableName = (mdq as JdbcMetadataQuerier).findTableName("KV", "PUBLIC")
             Assertions.assertNotNull(tableName)
             Assertions.assertEquals(expectedColumnMetadata, mdq.columnMetadata(tableName!!))
-            Assertions.assertEquals(listOf(listOf("K")), mdq.primaryKeys("KV", "PUBLIC"))
+            Assertions.assertEquals(listOf(listOf("K")), mdq.primaryKey("KV", "PUBLIC"))
         }
     }
 }
diff --git a/airbyte-cdk/bulk/core/extract/src/test/kotlin/io/airbyte/cdk/test/source/FakeSourceDiscoverTest.kt b/airbyte-cdk/bulk/core/extract/src/test/kotlin/io/airbyte/cdk/test/source/FakeSourceDiscoverTest.kt
@@ -34,13 +34,15 @@ class FakeSourceDiscoverTest {
                 .withJsonSchema(Jsons.readTree(EVENTS_SCHEMA))
                 .withSupportedSyncModes(listOf(SyncMode.FULL_REFRESH, SyncMode.INCREMENTAL))
                 .withSourceDefinedPrimaryKey(listOf(listOf("ID")))
+                .withIsResumable(true)
         val kv =
             AirbyteStream()
                 .withName("KV")
                 .withNamespace("PUBLIC")
                 .withJsonSchema(Jsons.readTree(KV_SCHEMA))
                 .withSupportedSyncModes(listOf(SyncMode.FULL_REFRESH, SyncMode.INCREMENTAL))
                 .withSourceDefinedPrimaryKey(listOf(listOf("K")))
+                .withIsResumable(true)
         val expected = AirbyteCatalog().withStreams(listOf(events, kv))
         discoverOperation.execute()
         Assertions.assertEquals(listOf(expected), outputConsumer.catalogs())
@@ -59,13 +61,15 @@ class FakeSourceDiscoverTest {
                 .withJsonSchema(Jsons.readTree(EVENTS_SCHEMA))
                 .withSupportedSyncModes(listOf(SyncMode.FULL_REFRESH, SyncMode.INCREMENTAL))
                 .withSourceDefinedPrimaryKey(listOf(listOf("ID")))
+                .withIsResumable(true)
         val kv =
             AirbyteStream()
                 .withName("KV")
                 .withNamespace("PUBLIC")
                 .withJsonSchema(Jsons.readTree(KV_SCHEMA))
                 .withSupportedSyncModes(listOf(SyncMode.FULL_REFRESH, SyncMode.INCREMENTAL))
                 .withSourceDefinedPrimaryKey(listOf(listOf("K")))
+                .withIsResumable(true)
         val expected = AirbyteCatalog().withStreams(listOf(events, kv))
         discoverOperation.execute()
         Assertions.assertEquals(listOf(expected), outputConsumer.catalogs())
diff --git a/airbyte-cdk/bulk/core/extract/src/test/resources/read/cdc-catalog.json b/airbyte-cdk/bulk/core/extract/src/test/resources/read/cdc-catalog.json
@@ -22,6 +22,7 @@
         "supported_sync_modes": ["full_refresh", "incremental"],
         "default_cursor_field": ["ID", "TS"],
         "source_defined_primary_key": [["ID"]],
+        "is_resumable": true,
         "namespace": "PUBLIC"
       },
       "sync_mode": "full_refresh",
@@ -47,6 +48,7 @@
         "supported_sync_modes": ["full_refresh", "incremental"],
         "default_cursor_field": ["K"],
         "source_defined_primary_key": [["K"]],
+        "is_resumable": true,
         "namespace": "PUBLIC"
       },
       "sync_mode": "incremental",
diff --git a/airbyte-cdk/bulk/core/extract/src/test/resources/read/cursor-catalog.json b/airbyte-cdk/bulk/core/extract/src/test/resources/read/cursor-catalog.json
@@ -22,6 +22,7 @@
         "supported_sync_modes": ["full_refresh", "incremental"],
         "default_cursor_field": [],
         "source_defined_primary_key": [["ID"]],
+        "is_resumable": true,
         "namespace": "PUBLIC"
       },
       "sync_mode": "incremental",
@@ -47,6 +48,7 @@
         "supported_sync_modes": ["full_refresh", "incremental"],
         "default_cursor_field": [],
         "source_defined_primary_key": [["K"]],
+        "is_resumable": true,
         "namespace": "PUBLIC"
       },
       "sync_mode": "full_refresh",
diff --git a/airbyte-cdk/bulk/core/extract/src/test/resources/test/source/cdc-catalog.json b/airbyte-cdk/bulk/core/extract/src/test/resources/test/source/cdc-catalog.json
@@ -23,6 +23,7 @@
         "supported_sync_modes": ["full_refresh", "incremental"],
         "default_cursor_field": ["ID", "TS"],
         "source_defined_primary_key": [["ID"]],
+        "is_resumable": true,
         "namespace": "PUBLIC"
       },
       "sync_mode": "full_refresh",
@@ -48,6 +49,7 @@
         "supported_sync_modes": ["full_refresh", "incremental"],
         "default_cursor_field": ["K"],
         "source_defined_primary_key": [["K"]],
+        "is_resumable": true,
         "namespace": "PUBLIC"
       },
       "sync_mode": "incremental",
diff --git a/airbyte-cdk/bulk/core/extract/src/test/resources/test/source/cursor-catalog.json b/airbyte-cdk/bulk/core/extract/src/test/resources/test/source/cursor-catalog.json
@@ -23,6 +23,7 @@
         "supported_sync_modes": ["full_refresh", "incremental"],
         "default_cursor_field": [],
         "source_defined_primary_key": [["ID"]],
+        "is_resumable": true,
         "namespace": "PUBLIC"
       },
       "sync_mode": "incremental",
@@ -48,6 +49,7 @@
         "supported_sync_modes": ["full_refresh", "incremental"],
         "default_cursor_field": [],
         "source_defined_primary_key": [["K"]],
+        "is_resumable": true,
         "namespace": "PUBLIC"
       },
       "sync_mode": "full_refresh",
diff --git a/airbyte-cdk/bulk/core/extract/src/test/resources/test/source/expected-cdc-catalog.json b/airbyte-cdk/bulk/core/extract/src/test/resources/test/source/expected-cdc-catalog.json
@@ -40,6 +40,7 @@
       "source_defined_cursor": true,
       "default_cursor_field": ["_ab_cdc_lsn"],
       "source_defined_primary_key": [["ID"]],
+      "is_resumable": true,
       "namespace": "PUBLIC"
     },
     {
@@ -77,6 +78,7 @@
       "source_defined_cursor": true,
       "default_cursor_field": ["_ab_cdc_lsn"],
       "source_defined_primary_key": [["K"]],
+      "is_resumable": true,
       "namespace": "PUBLIC"
     }
   ]
diff --git a/airbyte-cdk/bulk/core/extract/src/test/resources/test/source/expected-cursor-catalog.json b/airbyte-cdk/bulk/core/extract/src/test/resources/test/source/expected-cursor-catalog.json
@@ -26,6 +26,7 @@
       "supported_sync_modes": ["full_refresh", "incremental"],
       "default_cursor_field": [],
       "source_defined_primary_key": [["ID"]],
+      "is_resumable": true,
       "namespace": "PUBLIC"
     },
     {
@@ -49,6 +50,7 @@
       "supported_sync_modes": ["full_refresh", "incremental"],
       "default_cursor_field": [],
       "source_defined_primary_key": [["K"]],
+      "is_resumable": true,
       "namespace": "PUBLIC"
     }
   ]
diff --git a/airbyte-cdk/bulk/core/extract/src/testFixtures/kotlin/io/airbyte/cdk/discover/ResourceDrivenMetadataQuerierFactory.kt b/airbyte-cdk/bulk/core/extract/src/testFixtures/kotlin/io/airbyte/cdk/discover/ResourceDrivenMetadataQuerierFactory.kt
@@ -66,7 +66,7 @@ class ResourceDrivenMetadataQuerierFactory(
                     ?: throw SQLException("query failed", "tbl")
             }
 
-            override fun primaryKeys(
+            override fun primaryKey(
                 streamName: String,
                 streamNamespace: String?,
             ): List<List<String>> {
diff --git a/airbyte-cdk/bulk/toolkits/extract-jdbc/src/main/kotlin/io/airbyte/cdk/jdbc/JdbcMetadataQuerier.kt b/airbyte-cdk/bulk/toolkits/extract-jdbc/src/main/kotlin/io/airbyte/cdk/jdbc/JdbcMetadataQuerier.kt
@@ -272,7 +272,7 @@ class JdbcMetadataQuerier(
 
     val memoizedPrimaryKeys = mutableMapOf<TableName, List<List<String>>>()
 
-    override fun primaryKeys(
+    override fun primaryKey(
         streamName: String,
         streamNamespace: String?,
     ): List<List<String>> {

Original file line number	Diff line number	Diff line change
`@@ -39,9 +39,8 @@ class DiscoverOperation(`
`39`	`39`	`}`
`40`	`40`	`continue`
`41`	`41`	`}`
`42`		`- val primaryKeys: List<List<String>> =`
`43`		`- metadataQuerier.primaryKeys(name, namespace)`
`44`		`- val discoveredStream = DiscoveredStream(name, namespace, fields, primaryKeys)`
	`42`	`+ val primaryKey: List<List<String>> = metadataQuerier.primaryKey(name, namespace)`
	`43`	`+ val discoveredStream = DiscoveredStream(name, namespace, fields, primaryKey)`
`45`	`44`	`airbyteStreams.add(toAirbyteStream(discoveredStream))`
`46`	`45`	`}`
`47`	`46`	`}`
`@@ -68,6 +67,7 @@ class DiscoverOperation(`
`68`	`67`	`airbyteStream.withSourceDefinedPrimaryKey(`
`69`	`68`	`if (isValidPK) discoveredStream.primaryKeyColumnIDs else listOf(),`
`70`	`69`	`)`
	`70`	`+ airbyteStream.isResumable = airbyteStream.sourceDefinedPrimaryKey.isNotEmpty()`
`71`	`71`	`if (config.global) {`
`72`	`72`	`// There is a global feed of incremental records, like CDC.`
`73`	`73`	`airbyteStreamDecorator.decorateGlobal(airbyteStream)`
Original file line number	Diff line number	Diff line change
`@@ -34,7 +34,7 @@ interface AirbyteStreamDecorator {`
`34`	`34`	`* message.`
`35`	`35`	`*`
`36`	`36`	`* This method does not determine (1), of course, because the source keys are defined in the`
`37`		`- * source database itself and are retrieved via [MetadataQuerier.primaryKeys]. Instead, this`
	`37`	`+ * source database itself and are retrieved via [MetadataQuerier.primaryKey]. Instead, this`
`38`	`38`	`* method determines (2) based on the type information of the field, typically the [FieldType]`
`39`	`39`	`* objects. For instance if the [Field.type] does not map to a [LosslessFieldType] then the`
`40`	`40`	`* field can't reliably round-trip checkpoint values during a resumable initial sync.`
Original file line number	Diff line number	Diff line change
`@@ -72,7 +72,7 @@ class JdbcMetadataQuerierTest {`
`72`	`72`	`val tableName = (mdq as JdbcMetadataQuerier).findTableName("KV", "PUBLIC")`
`73`	`73`	`Assertions.assertNotNull(tableName)`
`74`	`74`	`Assertions.assertEquals(expectedColumnMetadata, mdq.columnMetadata(tableName!!))`
`75`		`- Assertions.assertEquals(listOf(listOf("K")), mdq.primaryKeys("KV", "PUBLIC"))`
	`75`	`+ Assertions.assertEquals(listOf(listOf("K")), mdq.primaryKey("KV", "PUBLIC"))`
`76`	`76`	`}`
`77`	`77`	`}`
`78`	`78`	`}`
Original file line number	Diff line number	Diff line change
`@@ -40,6 +40,7 @@`
`40`	`40`	`"source_defined_cursor": true,`
`41`	`41`	`"default_cursor_field": ["_ab_cdc_lsn"],`
`42`	`42`	`"source_defined_primary_key": [["ID"]],`
	`43`	`+ "is_resumable": true,`
`43`	`44`	`"namespace": "PUBLIC"`
`44`	`45`	`},`
`45`	`46`	`{`
`@@ -77,6 +78,7 @@`
`77`	`78`	`"source_defined_cursor": true,`
`78`	`79`	`"default_cursor_field": ["_ab_cdc_lsn"],`
`79`	`80`	`"source_defined_primary_key": [["K"]],`
	`81`	`+ "is_resumable": true,`
`80`	`82`	`"namespace": "PUBLIC"`
`81`	`83`	`}`
`82`	`84`	`]`
Original file line number	Diff line number	Diff line change
`@@ -26,6 +26,7 @@`
`26`	`26`	`"supported_sync_modes": ["full_refresh", "incremental"],`
`27`	`27`	`"default_cursor_field": [],`
`28`	`28`	`"source_defined_primary_key": [["ID"]],`
	`29`	`+ "is_resumable": true,`
`29`	`30`	`"namespace": "PUBLIC"`
`30`	`31`	`},`
`31`	`32`	`{`
`@@ -49,6 +50,7 @@`
`49`	`50`	`"supported_sync_modes": ["full_refresh", "incremental"],`
`50`	`51`	`"default_cursor_field": [],`
`51`	`52`	`"source_defined_primary_key": [["K"]],`
	`53`	`+ "is_resumable": true,`
`52`	`54`	`"namespace": "PUBLIC"`
`53`	`55`	`}`
`54`	`56`	`]`
Original file line number	Diff line number	Diff line change
`@@ -66,7 +66,7 @@ class ResourceDrivenMetadataQuerierFactory(`
`66`	`66`	`?: throw SQLException("query failed", "tbl")`
`67`	`67`	`}`
`68`	`68`
`69`		`- override fun primaryKeys(`
	`69`	`+ override fun primaryKey(`
`70`	`70`	`streamName: String,`
`71`	`71`	`streamNamespace: String?,`
`72`	`72`	`): List<List<String>> {`