bulk-cdk: add ExceptionClassifier, revamp exception handling (#44608)

Author: Marius Posta

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

@@ -2,11 +2,12 @@
 package io.airbyte.cdk
 
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings
+import io.airbyte.cdk.output.ExceptionClassifier
 import io.airbyte.cdk.output.OutputConsumer
-import io.airbyte.cdk.util.ApmTraceUtils
 import io.github.oshai.kotlinlogging.KotlinLogging
 import io.micronaut.context.annotation.Value
 import jakarta.inject.Inject
+import jakarta.inject.Provider
 
 private val log = KotlinLogging.logger {}
 
@@ -15,18 +16,31 @@ private val log = KotlinLogging.logger {}
 class AirbyteConnectorRunnable : Runnable {
     @Value("\${airbyte.connector.metadata.docker-repository}") lateinit var connectorName: String
 
-    @Inject lateinit var operation: Operation
+    @Inject lateinit var operationProvider: Provider<Operation>
 
     @Inject lateinit var outputConsumer: OutputConsumer
 
+    @Inject lateinit var exceptionClassifier: ExceptionClassifier
+
     override fun run() {
-        log.info { "Executing ${operation::class} operation." }
+        var operation: Operation? = null
         try {
+            try {
+                operation = operationProvider.get()!!
+            } catch (e: Throwable) {
+                throw ConfigErrorException("Failed to initialize connector operation", e)
+            }
+            log.info { "Executing ${operation::class} operation." }
             operation.execute()
         } catch (e: Throwable) {
-            log.error(e) { "Failed ${operation::class} operation execution." }
-            ApmTraceUtils.addExceptionToTrace(e)
-            outputConsumer.acceptTraceOnConfigError(e)
+            log.error(e) {
+                if (operation == null) {
+                    "Failed connector operation initialization."
+                } else {
+                    "Failed ${operation::class} operation execution."
+                }
+            }
+            outputConsumer.accept(exceptionClassifier.handle(e))
             throw e
         } finally {
             log.info { "Flushing output consumer prior to shutdown." }

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

@@ -0,0 +1,32 @@
+/*
+ * Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+ */
+package io.airbyte.cdk
+
+/**
+ * A [ConnectorErrorException] is an exception which readily maps to an
+ * [io.airbyte.cdk.output.ConnectorError] object.
+ */
+sealed class ConnectorErrorException(
+    displayMessage: String?,
+    exception: Throwable?,
+) : RuntimeException(displayMessage, exception)
+
+/** See [io.airbyte.cdk.output.ConfigError]. */
+class ConfigErrorException(
+    displayMessage: String,
+    exception: Throwable? = null,
+) : ConnectorErrorException(displayMessage, exception)
+
+/** See [io.airbyte.cdk.output.TransientError]. */
+class TransientErrorException(
+    displayMessage: String,
+    exception: Throwable? = null,
+) : ConnectorErrorException(displayMessage, exception)
+
+/** See [io.airbyte.cdk.output.SystemError]. */
+class SystemErrorException
+private constructor(
+    displayMessage: String?,
+    exception: Throwable? = null,
+) : ConnectorErrorException(displayMessage, exception)

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

@@ -0,0 +1,103 @@
+/*
+ * Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+ */
+package io.airbyte.cdk.output
+
+import io.airbyte.cdk.ConfigErrorException
+import io.airbyte.cdk.ConnectorErrorException
+import io.airbyte.cdk.SystemErrorException
+import io.airbyte.cdk.TransientErrorException
+import io.airbyte.cdk.util.ApmTraceUtils
+import io.airbyte.protocol.models.v0.AirbyteErrorTraceMessage
+import io.micronaut.context.annotation.DefaultImplementation
+import jakarta.inject.Singleton
+import org.apache.commons.lang3.exception.ExceptionUtils
+
+@Singleton
+@DefaultImplementation(DefaultExceptionClassifier::class)
+fun interface ExceptionClassifier {
+
+    /** Classifies [e] into a [ConnectorError] if possible, null otherwise. */
+    fun classify(e: Throwable): ConnectorError?
+
+    /** [SystemError] display message for [e] in case it can't be classified. */
+    fun fallbackDisplayMessage(e: Throwable): String? = e.message
+
+    /** Maps [e] to a [AirbyteErrorTraceMessage] to be passed to the [OutputConsumer]. */
+    fun handle(e: Throwable): AirbyteErrorTraceMessage {
+        ApmTraceUtils.addExceptionToTrace(e)
+        val connectorError: ConnectorError =
+            DefaultExceptionClassifier().classify(e)
+                ?: classify(e) ?: SystemError(fallbackDisplayMessage(e) ?: e.message)
+        val errorTraceMessage =
+            AirbyteErrorTraceMessage()
+                .withInternalMessage(e.toString())
+                .withStackTrace(ExceptionUtils.getStackTrace(e))
+        return when (connectorError) {
+            is ConfigError ->
+                errorTraceMessage
+                    .withFailureType(AirbyteErrorTraceMessage.FailureType.CONFIG_ERROR)
+                    .withMessage(connectorError.displayMessage)
+            is TransientError ->
+                errorTraceMessage
+                    .withFailureType(AirbyteErrorTraceMessage.FailureType.TRANSIENT_ERROR)
+                    .withMessage(connectorError.displayMessage)
+            is SystemError ->
+                errorTraceMessage
+                    .withFailureType(AirbyteErrorTraceMessage.FailureType.SYSTEM_ERROR)
+                    .withMessage(connectorError.displayMessage ?: e.message)
+        }
+    }
+}
+
+/** Each [ConnectorError] subtype corresponds to a [AirbyteErrorTraceMessage.FailureType]. */
+sealed interface ConnectorError
+
+/**
+ * A [ConfigError] means there is something wrong with the user's connector configuration or any
+ * other error for which the connector informs the platform that the error is not transient in
+ * nature and thus to not bother retrying.
+ */
+data class ConfigError(val displayMessage: String) : ConnectorError
+
+/**
+ * A [TransientError] means there is something wrong with the user's source or any other error for
+ * which the connector informs the platform that the error is transient in nature.
+ */
+data class TransientError(val displayMessage: String) : ConnectorError
+
+/**
+ * A [SystemError] means there is something wrong with the connector.
+ *
+ * In practice these are also ll errors that are neither [ConfigError] or [TransientError]. This is
+ * whatever falls through the cracks of the [ExceptionClassifier], as such there is a standing goal
+ * to minimize occurrences of these instances.
+ */
+data class SystemError(val displayMessage: String?) : ConnectorError
+
+/** Default implementation of [ExceptionClassifier]. */
+@Singleton
+class DefaultExceptionClassifier : ExceptionClassifier {
+
+    override fun classify(e: Throwable): ConnectorError? {
+        return when (val connectorErrorException: ConnectorErrorException? = unwind(e)) {
+            is ConfigErrorException -> ConfigError(connectorErrorException.message!!)
+            is TransientErrorException -> TransientError(connectorErrorException.message!!)
+            is SystemErrorException -> SystemError(connectorErrorException.message)
+            null -> null
+        }
+    }
+
+    /** Recursively walks the causes of [e] and returns the last [ConnectorErrorException]. */
+    fun unwind(e: Throwable): ConnectorErrorException? {
+        var connectorErrorException: ConnectorErrorException? = null
+        var unwound: Throwable? = e
+        while (unwound != null) {
+            if (unwound is ConnectorErrorException) {
+                connectorErrorException = unwound
+            }
+            unwound = unwound.cause
+        }
+        return connectorErrorException
+    }
+}

airbyte-cdk/bulk/core/base/src/main/kotlin/io/airbyte/cdk/output/ExceptionClassifier.kt

@@ -1,7 +1,6 @@
 /* Copyright (c) 2024 Airbyte, Inc., all rights reserved. */
 package io.airbyte.cdk.output
 
-import io.airbyte.cdk.ConfigErrorException
 import io.airbyte.cdk.util.Jsons
 import io.airbyte.protocol.models.v0.AirbyteAnalyticsTraceMessage
 import io.airbyte.protocol.models.v0.AirbyteCatalog
@@ -21,7 +20,6 @@ import jakarta.inject.Singleton
 import java.io.ByteArrayOutputStream
 import java.time.Instant
 import java.util.function.Consumer
-import org.apache.commons.lang3.exception.ExceptionUtils
 
 /** Emits the [AirbyteMessage] instances produced by the connector. */
 @DefaultImplementation(StdoutOutputConsumer::class)
@@ -89,17 +87,6 @@ interface OutputConsumer : Consumer<AirbyteMessage>, AutoCloseable {
                 .withAnalytics(analytics),
         )
     }
-
-    fun acceptTraceOnConfigError(e: Throwable) {
-        val configErrorException: ConfigErrorException = ConfigErrorException.unwind(e) ?: return
-        accept(
-            AirbyteErrorTraceMessage()
-                .withFailureType(AirbyteErrorTraceMessage.FailureType.CONFIG_ERROR)
-                .withMessage(configErrorException.message)
-                .withInternalMessage(e.toString())
-                .withStackTrace(ExceptionUtils.getStackTrace(e)),
-        )
-    }
 }
 
 // Used for integration tests.

airbyte-cdk/bulk/core/base/src/main/kotlin/io/airbyte/cdk/output/OutputConsumer.kt

@@ -1,18 +1,20 @@
 /* Copyright (c) 2024 Airbyte, Inc., all rights reserved. */
 package io.airbyte.cdk.check
 
+import io.airbyte.cdk.ConfigErrorException
 import io.airbyte.cdk.Operation
-import io.airbyte.cdk.command.*
+import io.airbyte.cdk.command.ConfigurationJsonObjectBase
+import io.airbyte.cdk.command.ConfigurationJsonObjectSupplier
+import io.airbyte.cdk.command.SourceConfiguration
+import io.airbyte.cdk.command.SourceConfigurationFactory
 import io.airbyte.cdk.discover.MetadataQuerier
+import io.airbyte.cdk.output.ExceptionClassifier
 import io.airbyte.cdk.output.OutputConsumer
-import io.airbyte.cdk.util.ApmTraceUtils
 import io.airbyte.protocol.models.v0.AirbyteConnectionStatus
 import io.airbyte.protocol.models.v0.AirbyteErrorTraceMessage
 import io.github.oshai.kotlinlogging.KotlinLogging
 import io.micronaut.context.annotation.Requires
 import jakarta.inject.Singleton
-import java.sql.SQLException
-import org.apache.commons.lang3.exception.ExceptionUtils
 
 @Singleton
 @Requires(property = Operation.PROPERTY, value = "check")
@@ -22,6 +24,7 @@ class CheckOperation<T : ConfigurationJsonObjectBase>(
     val configFactory: SourceConfigurationFactory<T, out SourceConfiguration>,
     val metadataQuerierFactory: MetadataQuerier.Factory<SourceConfiguration>,
     val outputConsumer: OutputConsumer,
+    val exceptionClassifier: ExceptionClassifier,
 ) : Operation {
     private val log = KotlinLogging.logger {}
 
@@ -34,37 +37,15 @@ class CheckOperation<T : ConfigurationJsonObjectBase>(
             val config: SourceConfiguration = configFactory.make(pojo)
             log.info { "Connecting for config check." }
             metadataQuerierFactory.session(config).use { connectionCheck(it) }
-        } catch (e: SQLException) {
-            log.debug(e) { "SQLException while checking config." }
-            val message: String =
-                listOfNotNull(
-                        e.sqlState?.let { "State code: $it" },
-                        e.errorCode.takeIf { it != 0 }?.let { "Error code: $it" },
-                        e.message?.let { "Message: $it" },
-                    )
-                    .joinToString(separator = "; ")
-            ApmTraceUtils.addExceptionToTrace(e)
-            outputConsumer.accept(
-                AirbyteErrorTraceMessage()
-                    .withFailureType(AirbyteErrorTraceMessage.FailureType.CONFIG_ERROR)
-                    .withMessage(message)
-                    .withInternalMessage(e.toString())
-                    .withStackTrace(ExceptionUtils.getStackTrace(e)),
-            )
-            outputConsumer.accept(
-                AirbyteConnectionStatus()
-                    .withMessage(message)
-                    .withStatus(AirbyteConnectionStatus.Status.FAILED),
-            )
-            log.info { "Config check failed." }
-            return
         } catch (e: Exception) {
             log.debug(e) { "Exception while checking config." }
-            ApmTraceUtils.addExceptionToTrace(e)
-            outputConsumer.acceptTraceOnConfigError(e)
+            val errorTraceMessage: AirbyteErrorTraceMessage = exceptionClassifier.handle(e)
+            outputConsumer.accept(errorTraceMessage)
+            val connectionStatusMessage: String =
+                String.format(COMMON_EXCEPTION_MESSAGE_TEMPLATE, errorTraceMessage.message)
             outputConsumer.accept(
                 AirbyteConnectionStatus()
-                    .withMessage(String.format(COMMON_EXCEPTION_MESSAGE_TEMPLATE, e.message))
+                    .withMessage(connectionStatusMessage)
                     .withStatus(AirbyteConnectionStatus.Status.FAILED),
             )
             log.info { "Config check failed." }
@@ -101,9 +82,9 @@ class CheckOperation<T : ConfigurationJsonObjectBase>(
             }
         }
         if (n == 0) {
-            throw RuntimeException("Discovered zero tables.")
+            throw ConfigErrorException("Discovered zero tables.")
         } else {
-            throw RuntimeException("Unable to query any of the $n discovered table(s).")
+            throw ConfigErrorException("Unable to query any of the $n discovered table(s).")
         }
     }
 

airbyte-cdk/bulk/core/extract/src/main/kotlin/io/airbyte/cdk/check/CheckOperation.kt

@@ -0,0 +1,38 @@
+/*
+ * Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+ */
+
+package io.airbyte.cdk.output
+
+import io.airbyte.cdk.Operation
+import io.micronaut.context.annotation.Replaces
+import io.micronaut.context.annotation.Value
+import java.sql.SQLException
+
+@Replaces(ExceptionClassifier::class)
+class JdbcExceptionClassifier(@Value("\${${Operation.PROPERTY}}") operationName: String) :
+    ExceptionClassifier {
+
+    val isCheck: Boolean = operationName == "check"
+
+    override fun classify(e: Throwable): ConnectorError? =
+        if (isCheck && e is SQLException) {
+            ConfigError(sqlExceptionDisplayMessage(e))
+        } else {
+            null
+        }
+
+    override fun fallbackDisplayMessage(e: Throwable): String? =
+        when (e) {
+            is SQLException -> sqlExceptionDisplayMessage(e)
+            else -> null
+        }
+
+    fun sqlExceptionDisplayMessage(e: SQLException): String =
+        listOfNotNull(
+                e.sqlState?.let { "State code: $it" },
+                e.errorCode.takeIf { it != 0 }?.let { "Error code: $it" },
+                e.message?.let { "Message: $it" },
+            )
+            .joinToString(separator = "; ")
+}

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

@@ -35,7 +35,7 @@ class H2SourceIntegrationTest {
                     port = h2.port
                     database = h2.database + "_garbage"
                 }
-            SyncsTestFixture.testCheck(configPojo, "Error code: 90149")
+            SyncsTestFixture.testCheck(configPojo, "Database \"mem:.*_garbage\" not found")
         }
     }