Feat: Adds config_change_callback to Destinations and Sources (#440)

airbyte/_connector_base.py

@@ -16,6 +16,7 @@
 from airbyte_protocol.models import (
     AirbyteMessage,
     ConnectorSpecification,
+    OrchestratorType,
     Status,
     TraceType,
     Type,
@@ -35,7 +36,7 @@
 
 if TYPE_CHECKING:
     import logging
-    from collections.abc import Generator
+    from collections.abc import Callable, Generator
     from typing import IO
 
     from airbyte._executors.base import Executor
@@ -56,13 +57,15 @@ def __init__(
         executor: Executor,
         name: str,
         config: dict[str, Any] | None = None,
+        config_change_callback: Callable[[dict[str, Any]], None] | None = None,
         *,
         validate: bool = False,
     ) -> None:
         """Initialize the source.
 
         If config is provided, it will be validated against the spec if validate is True.
         """
+        self.config_change_callback = config_change_callback
         self.executor = executor
         self._name = name
         self._config_dict: dict[str, Any] | None = None
@@ -361,7 +364,8 @@ def _peek_airbyte_message(
 
         This method handles reading Airbyte messages and taking action, if needed, based on the
         message type. For instance, log messages are logged, records are tallied, and errors are
-        raised as exceptions if `raise_on_error` is True.
+        raised as exceptions if `raise_on_error` is True. If a config change message is received,
+        the config change callback is called.
 
         Raises:
             AirbyteConnectorFailedError: If a TRACE message of type ERROR is emitted.
@@ -380,6 +384,14 @@ def _peek_airbyte_message(
                 )
             return
 
+        if (
+            message.type == Type.CONTROL
+            and message.control.type == OrchestratorType.CONNECTOR_CONFIG
+            and self.config_change_callback is not None
+        ):
+            self.config_change_callback(message.control.connectorConfig.config)
+            return
+
     def _execute(
         self,
         args: list[str],

airbyte/destinations/base.py

@@ -33,6 +33,8 @@
 
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from airbyte._executors.base import Executor
     from airbyte.caches.base import CacheBase
     from airbyte.shared.state_writers import StateWriterBase
@@ -48,6 +50,7 @@ def __init__(
         executor: Executor,
         name: str,
         config: dict[str, Any] | None = None,
+        config_change_callback: Callable[[dict[str, Any]], None] | None = None,
         *,
         validate: bool = False,
     ) -> None:
@@ -59,6 +62,7 @@ def __init__(
             executor=executor,
             name=name,
             config=config,
+            config_change_callback=config_change_callback,
             validate=validate,
         )
 

airbyte/destinations/util.py

@@ -13,12 +13,14 @@
 
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
     from pathlib import Path
 
 
-def get_destination(
+def get_destination(  # noqa: PLR0913 # Too many arguments
     name: str,
     config: dict[str, Any] | None = None,
+    config_change_callback: Callable[[dict[str, Any]], None] | None = None,
     *,
     version: str | None = None,
     pip_url: str | None = None,
@@ -33,6 +35,7 @@ def get_destination(
         name: connector name
         config: connector config - if not provided, you need to set it later via the set_config
             method.
+        config_change_callback: callback function to be called when the connector config changes.
         streams: list of stream names to select for reading. If set to "*", all streams will be
             selected. If not provided, you can set it later via the `select_streams()` or
             `select_all_streams()` method.
@@ -58,6 +61,7 @@ def get_destination(
     return Destination(
         name=name,
         config=config,
+        config_change_callback=config_change_callback,
         executor=get_connector_executor(
             name=name,
             version=version,

airbyte/sources/base.py

@@ -36,7 +36,7 @@
 
 
 if TYPE_CHECKING:
-    from collections.abc import Generator, Iterable, Iterator
+    from collections.abc import Callable, Generator, Iterable, Iterator
 
     from airbyte_cdk import ConnectorSpecification
     from airbyte_protocol.models import AirbyteStream
@@ -58,6 +58,7 @@ def __init__(
         executor: Executor,
         name: str,
         config: dict[str, Any] | None = None,
+        config_change_callback: Callable[[dict[str, Any]], None] | None = None,
         streams: str | list[str] | None = None,
         *,
         validate: bool = False,
@@ -73,6 +74,7 @@ def __init__(
             executor=executor,
             name=name,
             config=config,
+            config_change_callback=config_change_callback,
             validate=validate,
         )
         self._config_dict: dict[str, Any] | None = None

airbyte/sources/util.py

@@ -13,6 +13,7 @@
 
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
     from pathlib import Path
 
 
@@ -45,6 +46,7 @@ def get_connector(
 def get_source(  # noqa: PLR0913 # Too many arguments
     name: str,
     config: dict[str, Any] | None = None,
+    config_change_callback: Callable[[dict[str, Any]], None] | None = None,
     *,
     streams: str | list[str] | None = None,
     version: str | None = None,
@@ -72,6 +74,7 @@ def get_source(  # noqa: PLR0913 # Too many arguments
         name: connector name
         config: connector config - if not provided, you need to set it later via the set_config
             method.
+        config_change_callback: callback function to be called when the connector config changes.
         streams: list of stream names to select for reading. If set to "*", all streams will be
             selected. If not provided, you can set it later via the `select_streams()` or
             `select_all_streams()` method.
@@ -103,6 +106,7 @@ def get_source(  # noqa: PLR0913 # Too many arguments
     return Source(
         name=name,
         config=config,
+        config_change_callback=config_change_callback,
         streams=streams,
         executor=get_connector_executor(
             name=name,

tests/integration_tests/test_config_change_callback.py

@@ -0,0 +1,103 @@
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+
+"""Integration tests which test destination capabilities using the JSONL destination (docker-based)."""
+
+from __future__ import annotations
+
+import pytest
+from typing import Any
+from unittest.mock import patch
+from airbyte import get_source, get_destination, Source, Destination
+from airbyte_protocol.models import (
+    AirbyteMessage,
+    Type,
+    AirbyteControlMessage,
+    OrchestratorType,
+    AirbyteControlConnectorConfigMessage,
+)
+
+
+def config_change_callback(config: dict[str, Any]) -> None:
+    print(f"Updated config: {config}")
+
+
+@pytest.fixture
+def new_duckdb_destination() -> Destination:
+    """Return a new JSONL destination."""
+    return get_destination(
+        name="destination-duckdb",
+        config={
+            # This path is relative to the container:
+            "destination_path": "/local/temp/db.duckdb",
+        },
+        config_change_callback=config_change_callback,
+    )
+
+
+@pytest.fixture
+def new_source_faker() -> Source:
+    return get_source(
+        "source-faker",
+        config={
+            "count": 100,
+            "seed": 1234,
+            "parallelism": 16,
+        },
+        install_if_missing=True,
+        streams=["products"],
+        config_change_callback=config_change_callback,
+    )
+
+
+def test_source_config_callback(
+    new_duckdb_destination: Destination,
+    new_source_faker: Source,
+) -> None:
+    with patch.object(
+        new_source_faker, "config_change_callback"
+    ) as mock_config_change_callback:
+        updated_config = {
+            "count": 1000,
+            "seed": 1234,
+            "parallelism": 16,
+        }
+        airbyte_source_control_message = AirbyteMessage(
+            type=Type.CONTROL,
+            control=AirbyteControlMessage(
+                type=OrchestratorType.CONNECTOR_CONFIG,
+                emitted_at=0,
+                connectorConfig=AirbyteControlConnectorConfigMessage(
+                    config=updated_config
+                ),
+            ),
+        )
+
+        new_source_faker._peek_airbyte_message(airbyte_source_control_message)
+        mock_config_change_callback.assert_called_once_with(updated_config)
+
+
+def test_destination_config_callback(
+    new_duckdb_destination: Destination,
+    new_source_faker: Source,
+) -> None:
+    with patch.object(
+        new_duckdb_destination, "config_change_callback"
+    ) as mock_config_change_callback:
+        updated_config = {
+            "destination_path": "/local/temp/db.duckdb",
+        }
+        airbyte_destination_control_message = AirbyteMessage(
+            type=Type.CONTROL,
+            control=AirbyteControlMessage(
+                type=OrchestratorType.CONNECTOR_CONFIG,
+                emitted_at=0,
+                connectorConfig=AirbyteControlConnectorConfigMessage(
+                    config=updated_config
+                ),
+            ),
+        )
+
+        new_duckdb_destination._peek_airbyte_message(
+            airbyte_destination_control_message
+        )
+        mock_config_change_callback.assert_called_once_with(updated_config)