Embedded CDK: run a check before starting to load (#29079)

airbytehq · Aug 21, 2023 · d293e1c · d293e1c
1 parent dd170e2
commit d293e1c
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 2 deletions.
diff --git a/airbyte-cdk/python/airbyte_cdk/sources/embedded/base_integration.py b/airbyte-cdk/python/airbyte_cdk/sources/embedded/base_integration.py
@@ -9,13 +9,16 @@
 from airbyte_cdk.sources.embedded.catalog import create_configured_catalog, get_stream, get_stream_names
 from airbyte_cdk.sources.embedded.runner import SourceRunner
 from airbyte_cdk.sources.embedded.tools import get_defined_id
+from airbyte_cdk.sources.utils.schema_helpers import check_config_against_spec_or_exit
 from airbyte_protocol.models import AirbyteRecordMessage, AirbyteStateMessage, SyncMode, Type
 
 TOutput = TypeVar("TOutput")
 
 
 class BaseEmbeddedIntegration(ABC, Generic[TConfig, TOutput]):
     def __init__(self, runner: SourceRunner[TConfig], config: TConfig):
+        check_config_against_spec_or_exit(config, runner.spec())
+
         self.source = runner
         self.config = config
 

diff --git a/airbyte-cdk/python/airbyte_cdk/sources/embedded/runner.py b/airbyte-cdk/python/airbyte_cdk/sources/embedded/runner.py
@@ -8,11 +8,15 @@
 from typing import Generic, Iterable, Optional
 
 from airbyte_cdk.connector import TConfig
-from airbyte_cdk.models import AirbyteCatalog, AirbyteMessage, AirbyteStateMessage, ConfiguredAirbyteCatalog
+from airbyte_cdk.models import AirbyteCatalog, AirbyteMessage, AirbyteStateMessage, ConfiguredAirbyteCatalog, ConnectorSpecification
 from airbyte_cdk.sources.source import Source
 
 
 class SourceRunner(ABC, Generic[TConfig]):
+    @abstractmethod
+    def spec(self) -> ConnectorSpecification:
+        pass
+
     @abstractmethod
     def discover(self, config: TConfig) -> AirbyteCatalog:
         pass
@@ -27,6 +31,9 @@ def __init__(self, source: Source, name: str):
         self._source = source
         self._logger = logging.getLogger(name)
 
+    def spec(self) -> ConnectorSpecification:
+        return self._source.spec(self._logger)
+
     def discover(self, config: TConfig) -> AirbyteCatalog:
         return self._source.discover(self._logger, config)
 

diff --git a/airbyte-cdk/python/unit_tests/sources/embedded/test_embedded_integration.py b/airbyte-cdk/python/unit_tests/sources/embedded/test_embedded_integration.py
@@ -7,6 +7,7 @@
 from unittest.mock import MagicMock
 
 from airbyte_cdk.sources.embedded.base_integration import BaseEmbeddedIntegration
+from airbyte_cdk.utils import AirbyteTracedException
 from airbyte_protocol.models import (
     AirbyteCatalog,
     AirbyteLogMessage,
@@ -16,6 +17,7 @@
     AirbyteStream,
     ConfiguredAirbyteCatalog,
     ConfiguredAirbyteStream,
+    ConnectorSpecification,
     DestinationSyncMode,
     Level,
     SyncMode,
@@ -33,7 +35,14 @@ def setUp(self):
         self.source_class = MagicMock()
         self.source = MagicMock()
         self.source_class.return_value = self.source
-        self.config = MagicMock()
+        self.source.spec.return_value = ConnectorSpecification(connectionSpecification={
+            "properties": {
+                "test": {
+                    "type": "string",
+                }
+            }
+        })
+        self.config = {"test": "abc"}
         self.integration = TestIntegration(self.source, self.config)
         self.stream1 = AirbyteStream(
             name="test",
@@ -76,6 +85,12 @@ def test_integration(self):
             None,
         )
 
+    def test_failed_check(self):
+        self.config = {"test": 123}
+        with self.assertRaises(AirbyteTracedException) as error:
+            TestIntegration(self.source, self.config)
+        assert str(error.exception) == "123 is not of type 'string'"
+
     def test_state(self):
         state = AirbyteStateMessage(data={})
         self.source.read.return_value = [