Bigquery / Bulk Load CDK: implement T+D, truncate refresh, etc (#57578)

Author: edgao

airbyte-cdk/bulk/core/base/src/main/kotlin/io/airbyte/cdk/util/CollectionUtils.kt

@@ -0,0 +1,62 @@
+/*
+ * Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.cdk.util
+
+import java.util.Optional
+
+object CollectionUtils {
+    /**
+     * Pass in a collection and search term to determine whether any of the values match ignoring
+     * case
+     *
+     * @param collection the collection of values
+     * @param search the value to look for
+     * @return whether the value matches anything in the collection
+     */
+    @JvmStatic
+    fun containsIgnoreCase(collection: Collection<String>, search: String): Boolean {
+        return matchingKey(collection, search).isPresent
+    }
+
+    /**
+     * Convenience method for when you need to check an entire collection for membership in another
+     * collection.
+     *
+     * @param searchCollection the collection you want to check membership in
+     * @param searchTerms the keys you're looking for
+     * @return whether all searchTerms are in the searchCollection
+     */
+    @JvmStatic
+    fun containsAllIgnoreCase(
+        searchCollection: Collection<String>,
+        searchTerms: Collection<String>
+    ): Boolean {
+        require(!searchTerms.isEmpty()) {
+            // There isn't a good behavior for an empty collection. Without this check, an empty
+            // collection
+            // would always return
+            // true, but it feels misleading to say that the searchCollection does "contain all"
+            // when
+            // searchTerms is empty
+            "Search Terms collection may not be empty"
+        }
+        return searchTerms.all { term: String -> containsIgnoreCase(searchCollection, term) }
+    }
+
+    /**
+     * From a collection of strings, return an entry which matches the search term ignoring case
+     *
+     * @param collection the collection to search
+     * @param search the key you're looking for
+     * @return an Optional value which might contain the key that matches the search
+     */
+    @JvmStatic
+    fun matchingKey(collection: Collection<String>, search: String): Optional<String> {
+        if (collection.contains(search)) {
+            return Optional.of(search)
+        }
+        return Optional.ofNullable(collection.firstOrNull { it.equals(search, ignoreCase = true) })
+    }
+}

airbyte-cdk/bulk/core/base/src/main/kotlin/io/airbyte/cdk/util/ConnectorExceptionUtil.kt

@@ -0,0 +1,9 @@
+/*
+ * Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.cdk.util
+
+object ConnectorExceptionUtil {
+    val HTTP_AUTHENTICATION_ERROR_CODES: List<Int> = listOf(401, 403)
+}

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/command/DestinationStream.kt

@@ -7,6 +7,7 @@ package io.airbyte.cdk.load.command
 import io.airbyte.cdk.load.data.AirbyteType
 import io.airbyte.cdk.load.data.json.AirbyteTypeToJsonSchema
 import io.airbyte.cdk.load.data.json.JsonSchemaToAirbyteType
+import io.airbyte.cdk.load.message.DestinationRecord
 import io.airbyte.protocol.models.v0.AirbyteStream
 import io.airbyte.protocol.models.v0.ConfiguredAirbyteStream
 import io.airbyte.protocol.models.v0.DestinationSyncMode
@@ -136,6 +137,9 @@ data class Dedupe(
     /**
      * theoretically, the path to the cursor. In practice, most destinations only support cursors at
      * the root level, i.e. `listOf(cursorField)`.
+     *
+     * If this is set to an empty list, then the destination should use
+     * [DestinationRecord.message.record.emittedAt] as the cursor.
      */
     val cursor: List<String>,
 ) : ImportType

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/data/AirbyteType.kt

@@ -5,41 +5,127 @@
 package io.airbyte.cdk.load.data
 
 import com.fasterxml.jackson.databind.JsonNode
+import io.airbyte.cdk.load.util.Jsons
+import io.github.oshai.kotlinlogging.KotlinLogging
+
+private val logger = KotlinLogging.logger {}
+
+sealed class AirbyteType {
+    /**
+     * Utility method for database/warehouse destinations, which assume that the top-level schema is
+     * an object.
+     */
+    open fun asColumns(): LinkedHashMap<String, FieldType> {
+        return linkedMapOf()
+    }
 
-sealed interface AirbyteType
+    open val isObject: Boolean = false
+    open val isArray: Boolean = false
+}
 
-data object StringType : AirbyteType
+data object StringType : AirbyteType()
 
-data object BooleanType : AirbyteType
+data object BooleanType : AirbyteType()
 
-data object IntegerType : AirbyteType
+data object IntegerType : AirbyteType()
 
-data object NumberType : AirbyteType
+data object NumberType : AirbyteType()
 
-data object DateType : AirbyteType
+data object DateType : AirbyteType()
 
-data object TimestampTypeWithTimezone : AirbyteType
+data object TimestampTypeWithTimezone : AirbyteType()
 
-data object TimestampTypeWithoutTimezone : AirbyteType
+data object TimestampTypeWithoutTimezone : AirbyteType()
 
-data object TimeTypeWithTimezone : AirbyteType
+data object TimeTypeWithTimezone : AirbyteType()
 
-data object TimeTypeWithoutTimezone : AirbyteType
+data object TimeTypeWithoutTimezone : AirbyteType()
 
-data class ArrayType(val items: FieldType) : AirbyteType
+data class ArrayType(val items: FieldType) : AirbyteType() {
+    override val isArray = true
+}
+
+data object ArrayTypeWithoutSchema : AirbyteType() {
+    override val isArray = true
+}
 
-data object ArrayTypeWithoutSchema : AirbyteType
+data class ObjectType(val properties: LinkedHashMap<String, FieldType>) : AirbyteType() {
+    override fun asColumns(): LinkedHashMap<String, FieldType> {
+        return properties
+    }
 
-data class ObjectType(val properties: LinkedHashMap<String, FieldType>) : AirbyteType
+    override val isObject = true
+}
 
-data object ObjectTypeWithEmptySchema : AirbyteType
+data object ObjectTypeWithEmptySchema : AirbyteType() {
+    override val isObject = true
+}
 
-data object ObjectTypeWithoutSchema : AirbyteType
+data object ObjectTypeWithoutSchema : AirbyteType() {
+    override val isObject = true
+}
 
 data class UnionType(
     val options: Set<AirbyteType>,
     val isLegacyUnion: Boolean,
-) : AirbyteType {
+) : AirbyteType() {
+    /**
+     * This is a hack to handle weird schemas like {type: [object, string]}. If a stream's top-level
+     * schema looks like this, we still want to be able to extract the object properties (i.e. treat
+     * it as though the string option didn't exist).
+     *
+     * @throws IllegalArgumentException if we cannot extract columns from this schema
+     */
+    override fun asColumns(): LinkedHashMap<String, FieldType> {
+        logger.warn { "asColumns options=$options" }
+        val numObjectOptions = options.count { it.isObject }
+        if (numObjectOptions > 1) {
+            logger.error { "Can't extract columns from a schema with multiple object options" }
+            return LinkedHashMap()
+        }
+
+        var retVal: LinkedHashMap<String, FieldType>
+        try {
+            retVal = options.first { it.isObject }.asColumns()
+        } catch (_: NoSuchElementException) {
+            logger.error { "Can't extract columns from a schema with no object options" }
+            retVal = LinkedHashMap()
+        }
+        logger.warn { "Union.asColumns retVal=$retVal" }
+        return retVal
+    }
+
+    /**
+     * This matches legacy behavior. Some destinations handle legacy unions by choosing the "best"
+     * type from amongst the options. This is... not great, but it would be painful to change.
+     */
+    fun chooseType(): AirbyteType {
+        check(isLegacyUnion) { "Cannot chooseType for a non-legacy union type" }
+        if (options.isEmpty()) {
+            return UnknownType(Jsons.createObjectNode())
+        }
+        return options.minBy {
+            when (it) {
+                is ArrayType,
+                ArrayTypeWithoutSchema -> -2
+                is ObjectType,
+                ObjectTypeWithEmptySchema,
+                ObjectTypeWithoutSchema -> -1
+                StringType -> 0
+                DateType -> 1
+                TimeTypeWithoutTimezone -> 2
+                TimeTypeWithTimezone -> 3
+                TimestampTypeWithoutTimezone -> 4
+                TimestampTypeWithTimezone -> 5
+                NumberType -> 6
+                IntegerType -> 7
+                BooleanType -> 8
+                is UnknownType -> 9
+                is UnionType -> Int.MAX_VALUE
+            }
+        }
+    }
+
     companion object {
         fun of(options: Set<AirbyteType>, isLegacyUnion: Boolean = false): AirbyteType {
             if (options.size == 1) {
@@ -56,6 +142,6 @@ data class UnionType(
     }
 }
 
-data class UnknownType(val schema: JsonNode) : AirbyteType
+data class UnknownType(val schema: JsonNode) : AirbyteType()
 
 data class FieldType(val type: AirbyteType, val nullable: Boolean)

airbyte-cdk/bulk/core/load/src/main/kotlin/io/airbyte/cdk/load/message/DestinationMessage.kt

@@ -128,6 +128,9 @@ data class Meta(
                 COLUMN_NAME_AB_GENERATION_ID,
             )
 
+        /** A legacy column name. Used in "DV2" destinations' raw tables. */
+        const val COLUMN_NAME_AB_LOADED_AT = "_airbyte_loaded_at"
+
         fun getMetaValue(metaColumnName: String, value: String): AirbyteValue {
             if (!COLUMN_NAMES.contains(metaColumnName)) {
                 throw IllegalArgumentException("Invalid meta column name: $metaColumnName")

airbyte-cdk/bulk/toolkits/load-db/src/main/kotlin/io/airbyte/cdk/load/orchestration/db/DatabaseHandler.kt

@@ -0,0 +1,17 @@
+/*
+ * Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.cdk.load.orchestration.db
+
+interface DatabaseHandler {
+    fun execute(sql: Sql)
+
+    /**
+     * Create the namespaces (typically something like `create schema`).
+     *
+     * This function should assume that all `namespaces` are valid identifiers, i.e. any special
+     * characters have already been escaped, they respect identifier name length, etc.
+     */
+    suspend fun createNamespaces(namespaces: Collection<String>)
+}

airbyte-cdk/bulk/toolkits/load-db/src/main/kotlin/io/airbyte/cdk/load/orchestration/db/DatabaseInitialStatusGatherer.kt

@@ -0,0 +1,28 @@
+/*
+ * Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.cdk.load.orchestration.db
+
+import io.airbyte.cdk.load.command.DestinationStream
+import io.airbyte.cdk.load.orchestration.db.legacy_typing_deduping.TableCatalog
+
+interface DatabaseInitialStatus
+
+/**
+ * Some destinations can efficiently fetch multiple tables' information in a single query, so this
+ * interface accepts multiple streams in a single method call.
+ *
+ * For destinations which do not support that optimization, a simpler implementation would be
+ * something like this:
+ * ```kotlin
+ * streams.forEach { (stream, (tableNames, columnNames)) ->
+ *   launch {
+ *     // ... gather state...
+ *   }
+ * }
+ * ```
+ */
+fun interface DatabaseInitialStatusGatherer<InitialStatus : DatabaseInitialStatus> {
+    suspend fun gatherInitialStatus(streams: TableCatalog): Map<DestinationStream, InitialStatus>
+}

airbyte-cdk/bulk/toolkits/load-db/src/main/kotlin/io/airbyte/cdk/load/orchestration/db/DestinationNames.kt

@@ -0,0 +1,82 @@
+/*
+ * Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.cdk.load.orchestration.db
+
+import io.airbyte.cdk.load.command.DestinationStream
+
+data class TableNames(
+    // this is pretty dumb, but in theory we could have:
+    // * old-style implementation: raw+final tables both exist
+    // * only the raw table exists (i.e. T+D disabled)
+    // * only the final table exists (i.e. new-style direct-load tables)
+    val rawTableName: TableName?,
+    val finalTableName: TableName?,
+) {
+    init {
+        check(rawTableName != null || finalTableName != null) {
+            "At least one table name should be nonnull"
+        }
+    }
+
+    fun hasNamingConflictWith(other: TableNames) =
+        this.rawTableName.hasNamingConflictWith(other.rawTableName) ||
+            this.finalTableName.hasNamingConflictWith(other.finalTableName)
+
+    fun toPrettyString() =
+        "Raw table: ${rawTableName?.toPrettyString()}; Final table: ${finalTableName?.toPrettyString()}"
+
+    companion object {
+        const val NO_SUFFIX = ""
+        // TODO comment explaining this
+        const val TMP_TABLE_SUFFIX = "_airbyte_tmp"
+        const val SOFT_RESET_SUFFIX = "_ab_soft_reset"
+    }
+}
+
+data class TableName(val namespace: String, val name: String) {
+    fun toPrettyString(quote: String = "", suffix: String = "") =
+        "$quote$namespace$quote.$quote$name$suffix$quote"
+}
+
+fun TableName?.hasNamingConflictWith(other: TableName?): Boolean {
+    if (this == null || other == null) {
+        return false
+    }
+    return this.namespace == other.namespace && this.name == other.name
+}
+
+/**
+ * map from the column name as declared in the schema, to the column name that we'll create in the
+ * final (typed) table.
+ */
+@JvmInline
+value class ColumnNameMapping(private val columnNameMapping: Map<String, String>) :
+    Map<String, String> by columnNameMapping
+
+sealed interface TableNameGenerator {
+    fun getTableName(streamDescriptor: DestinationStream.Descriptor): TableName
+}
+
+fun interface RawTableNameGenerator : TableNameGenerator
+
+fun interface FinalTableNameGenerator : TableNameGenerator
+
+fun interface ColumnNameGenerator {
+    /**
+     * In some database/warehouses, there's a difference between how a name is _displayed_, and how
+     * the underlying engine actually treats it. For example, a column might be displayed as
+     * `CamelCaseColumn`, but the engine actually treats it as lowercase `camelcasecolumn`, or
+     * truncate it to `CamelCas`. Bigquery is an example of this: `create table foo (foo int, FOO
+     * int)` is invalid, because `foo` is duplicated.
+     *
+     * This is relevant for handling collisions between column names. We need to know what name will
+     * be displayed to the user, since that's what we'll use in queries - but we also need to know
+     * the "canonical" name to check whether two columns will collide.
+     */
+    data class ColumnName(val displayName: String, val canonicalName: String)
+    fun getColumnName(column: String): ColumnName
+}
+
+const val CDC_DELETED_AT_COLUMN = "_ab_cdc_deleted_at"