airbyte-cdk 0.72.1__py3-none-any.whl → 6.17.1.dev0__py3-none-any.whl

airbyte_cdk/connector.py

@@ -8,10 +8,15 @@ import logging
 import os
 import pkgutil
 from abc import ABC, abstractmethod
-from typing import Any, Generic, List, Mapping, Optional, Protocol, TypeVar, Union
+from typing import Any, Generic, Mapping, Optional, Protocol, TypeVar
 
 import yaml
-from airbyte_cdk.models import AirbyteConnectionStatus, ConnectorSpecification
+
+from airbyte_cdk.models import (
+    AirbyteConnectionStatus,
+    ConnectorSpecification,
+    ConnectorSpecificationSerializer,
+)
 
 
 def load_optional_package_file(package: str, filename: str) -> Optional[bytes]:
@@ -22,17 +27,6 @@ def load_optional_package_file(package: str, filename: str) -> Optional[bytes]:
         return None
 
 
-class AirbyteSpec(object):
-    @staticmethod
-    def from_file(file_name: str):
-        with open(file_name) as file:
-            spec_text = file.read()
-        return AirbyteSpec(spec_text)
-
-    def __init__(self, spec_string):
-        self.spec_string = spec_string
-
-
 TConfig = TypeVar("TConfig", bound=Mapping[str, Any])
 
 
@@ -57,17 +51,19 @@ class BaseConnector(ABC, Generic[TConfig]):
             )
 
     @staticmethod
-    def _read_json_file(file_path: str) -> Union[None, bool, float, int, str, List[Any], Mapping[str, Any]]:
+    def _read_json_file(file_path: str) -> Any:
         with open(file_path, "r") as file:
             contents = file.read()
 
         try:
             return json.loads(contents)
         except json.JSONDecodeError as error:
-            raise ValueError(f"Could not read json file {file_path}: {error}. Please ensure that it is a valid JSON.")
+            raise ValueError(
+                f"Could not read json file {file_path}: {error}. Please ensure that it is a valid JSON."
+            )
 
     @staticmethod
-    def write_config(config: TConfig, config_path: str):
+    def write_config(config: TConfig, config_path: str) -> None:
         with open(config_path, "w") as fh:
             fh.write(json.dumps(config))
 
@@ -83,7 +79,9 @@ class BaseConnector(ABC, Generic[TConfig]):
         json_spec = load_optional_package_file(package, "spec.json")
 
         if yaml_spec and json_spec:
-            raise RuntimeError("Found multiple spec files in the package. Only one of spec.yaml or spec.json should be provided.")
+            raise RuntimeError(
+                "Found multiple spec files in the package. Only one of spec.yaml or spec.json should be provided."
+            )
 
         if yaml_spec:
             spec_obj = yaml.load(yaml_spec, Loader=yaml.SafeLoader)
@@ -91,11 +89,13 @@ class BaseConnector(ABC, Generic[TConfig]):
             try:
                 spec_obj = json.loads(json_spec)
             except json.JSONDecodeError as error:
-                raise ValueError(f"Could not read json spec file: {error}. Please ensure that it is a valid JSON.")
+                raise ValueError(
+                    f"Could not read json spec file: {error}. Please ensure that it is a valid JSON."
+                )
         else:
             raise FileNotFoundError("Unable to find spec.yaml or spec.json in the package.")
 
-        return ConnectorSpecification.parse_obj(spec_obj)
+        return ConnectorSpecificationSerializer.load(spec_obj)
 
     @abstractmethod
     def check(self, logger: logging.Logger, config: TConfig) -> AirbyteConnectionStatus:
@@ -107,17 +107,17 @@ class BaseConnector(ABC, Generic[TConfig]):
 
 class _WriteConfigProtocol(Protocol):
     @staticmethod
-    def write_config(config: Mapping[str, Any], config_path: str):
-        ...
+    def write_config(config: Mapping[str, Any], config_path: str) -> None: ...
 
 
 class DefaultConnectorMixin:
     # can be overridden to change an input config
-    def configure(self: _WriteConfigProtocol, config: Mapping[str, Any], temp_dir: str) -> Mapping[str, Any]:
+    def configure(
+        self: _WriteConfigProtocol, config: Mapping[str, Any], temp_dir: str
+    ) -> Mapping[str, Any]:
         config_path = os.path.join(temp_dir, "config.json")
         self.write_config(config, config_path)
         return config
 
 
-class Connector(DefaultConnectorMixin, BaseConnector[Mapping[str, Any]], ABC):
-    ...
+class Connector(DefaultConnectorMixin, BaseConnector[Mapping[str, Any]], ABC): ...

airbyte_cdk/connector_builder/README.md

@@ -0,0 +1,53 @@
+# Connector Builder Backend
+
+This is the backend for requests from the
+[Connector Builder](https://docs.airbyte.com/connector-development/connector-builder-ui/overview/).
+
+## Local development
+
+### Locally running the Connector Builder backend
+
+```bash
+python main.py read --config path/to/config --catalog path/to/catalog
+```
+
+Note:
+
+- Requires the keys `__injected_declarative_manifest` and `__command` in its config, where
+  `__injected_declarative_manifest` is a JSON manifest and `__command` is one of the commands
+  handled by the ConnectorBuilderHandler (`stream_read` or `resolve_manifest`), i.e.
+
+```json
+{
+  "config": <normal config>,
+  "__injected_declarative_manifest": {...},
+  "__command": <"resolve_manifest" | "test_read">
+}
+```
+
+\*See
+[ConnectionSpecification](https://docs.airbyte.com/understanding-airbyte/airbyte-protocol/#actor-specification)
+for details on the `"config"` key if needed.
+
+- When the `__command` is `resolve_manifest`, the argument to `catalog` should be an empty string.
+- The config can optionally contain an object under the `__test_read_config` key which can define
+  custom test read limits with `max_records`, `max_slices`, and `max_pages_per_slice` properties.
+  All custom limits are optional; a default will be used for any limit that is not provided.
+
+### Locally running the docker image
+
+#### Build
+
+First, make sure you build the latest Docker image:
+
+```bash
+docker build -t airbyte/source-declarative-manifest:dev .
+```
+
+#### Run
+
+Then run any of the connector commands as follows:
+
+```bash
+docker run --rm -v $(pwd)/secrets:/secrets airbyte/source-declarative-manifest:dev read --config /secrets/config.json
+```

airbyte_cdk/connector_builder/connector_builder_handler.py

@@ -4,15 +4,23 @@
 
 import dataclasses
 from datetime import datetime
-from typing import Any, Mapping
+from typing import Any, List, Mapping
 
 from airbyte_cdk.connector_builder.message_grouper import MessageGrouper
-from airbyte_cdk.models import AirbyteMessage, AirbyteRecordMessage, ConfiguredAirbyteCatalog
-from airbyte_cdk.models import Type
+from airbyte_cdk.models import (
+    AirbyteMessage,
+    AirbyteRecordMessage,
+    AirbyteStateMessage,
+    ConfiguredAirbyteCatalog,
+    Type,
+)
 from airbyte_cdk.models import Type as MessageType
 from airbyte_cdk.sources.declarative.declarative_source import DeclarativeSource
 from airbyte_cdk.sources.declarative.manifest_declarative_source import ManifestDeclarativeSource
-from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import ModelToComponentFactory
+from airbyte_cdk.sources.declarative.parsers.model_to_component_factory import (
+    ModelToComponentFactory,
+)
+from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
 DEFAULT_MAXIMUM_NUMBER_OF_PAGES_PER_SLICE = 5
@@ -33,7 +41,9 @@ class TestReadLimits:
 
 def get_limits(config: Mapping[str, Any]) -> TestReadLimits:
     command_config = config.get("__test_read_config", {})
-    max_pages_per_slice = command_config.get(MAX_PAGES_PER_SLICE_KEY) or DEFAULT_MAXIMUM_NUMBER_OF_PAGES_PER_SLICE
+    max_pages_per_slice = (
+        command_config.get(MAX_PAGES_PER_SLICE_KEY) or DEFAULT_MAXIMUM_NUMBER_OF_PAGES_PER_SLICE
+    )
     max_slices = command_config.get(MAX_SLICES_KEY) or DEFAULT_MAXIMUM_NUMBER_OF_SLICES
     max_records = command_config.get(MAX_RECORDS_KEY) or DEFAULT_MAXIMUM_RECORDS
     return TestReadLimits(max_records, max_pages_per_slice, max_slices)
@@ -49,24 +59,38 @@ def create_source(config: Mapping[str, Any], limits: TestReadLimits) -> Manifest
             limit_pages_fetched_per_slice=limits.max_pages_per_slice,
             limit_slices_fetched=limits.max_slices,
             disable_retries=True,
+            disable_cache=True,
         ),
     )
 
 
 def read_stream(
-    source: DeclarativeSource, config: Mapping[str, Any], configured_catalog: ConfiguredAirbyteCatalog, limits: TestReadLimits
+    source: DeclarativeSource,
+    config: Mapping[str, Any],
+    configured_catalog: ConfiguredAirbyteCatalog,
+    state: List[AirbyteStateMessage],
+    limits: TestReadLimits,
 ) -> AirbyteMessage:
     try:
         handler = MessageGrouper(limits.max_pages_per_slice, limits.max_slices, limits.max_records)
-        stream_name = configured_catalog.streams[0].stream.name  # The connector builder only supports a single stream
-        stream_read = handler.get_message_groups(source, config, configured_catalog, limits.max_records)
+        stream_name = configured_catalog.streams[
+            0
+        ].stream.name  # The connector builder only supports a single stream
+        stream_read = handler.get_message_groups(
+            source, config, configured_catalog, state, limits.max_records
+        )
         return AirbyteMessage(
             type=MessageType.RECORD,
-            record=AirbyteRecordMessage(data=dataclasses.asdict(stream_read), stream=stream_name, emitted_at=_emitted_at()),
+            record=AirbyteRecordMessage(
+                data=dataclasses.asdict(stream_read), stream=stream_name, emitted_at=_emitted_at()
+            ),
         )
     except Exception as exc:
         error = AirbyteTracedException.from_exception(
-            exc, message=f"Error reading stream with config={config} and catalog={configured_catalog}: {str(exc)}"
+            exc,
+            message=filter_secrets(
+                f"Error reading stream with config={config} and catalog={configured_catalog}: {str(exc)}"
+            ),
         )
         return error.as_airbyte_message()
 
@@ -82,7 +106,9 @@ def resolve_manifest(source: ManifestDeclarativeSource) -> AirbyteMessage:
             ),
         )
     except Exception as exc:
-        error = AirbyteTracedException.from_exception(exc, message=f"Error resolving manifest: {str(exc)}")
+        error = AirbyteTracedException.from_exception(
+            exc, message=f"Error resolving manifest: {str(exc)}"
+        )
         return error.as_airbyte_message()
 
 

airbyte_cdk/connector_builder/main.py

@@ -6,19 +6,40 @@
 import sys
 from typing import Any, List, Mapping, Optional, Tuple
 
+import orjson
+
 from airbyte_cdk.connector import BaseConnector
-from airbyte_cdk.connector_builder.connector_builder_handler import TestReadLimits, create_source, get_limits, read_stream, resolve_manifest
+from airbyte_cdk.connector_builder.connector_builder_handler import (
+    TestReadLimits,
+    create_source,
+    get_limits,
+    read_stream,
+    resolve_manifest,
+)
 from airbyte_cdk.entrypoint import AirbyteEntrypoint
-from airbyte_cdk.models import AirbyteMessage, ConfiguredAirbyteCatalog
+from airbyte_cdk.models import (
+    AirbyteMessage,
+    AirbyteMessageSerializer,
+    AirbyteStateMessage,
+    ConfiguredAirbyteCatalog,
+    ConfiguredAirbyteCatalogSerializer,
+)
 from airbyte_cdk.sources.declarative.manifest_declarative_source import ManifestDeclarativeSource
+from airbyte_cdk.sources.source import Source
 from airbyte_cdk.utils.traced_exception import AirbyteTracedException
 
 
-def get_config_and_catalog_from_args(args: List[str]) -> Tuple[str, Mapping[str, Any], Optional[ConfiguredAirbyteCatalog]]:
+def get_config_and_catalog_from_args(
+    args: List[str],
+) -> Tuple[str, Mapping[str, Any], Optional[ConfiguredAirbyteCatalog], Any]:
     # TODO: Add functionality for the `debug` logger.
     #  Currently, no one `debug` level log will be displayed during `read` a stream for a connector created through `connector-builder`.
     parsed_args = AirbyteEntrypoint.parse_args(args)
-    config_path, catalog_path = parsed_args.config, parsed_args.catalog
+    config_path, catalog_path, state_path = (
+        parsed_args.config,
+        parsed_args.catalog,
+        parsed_args.state,
+    )
     if parsed_args.command != "read":
         raise ValueError("Only read commands are allowed for Connector Builder requests.")
 
@@ -31,16 +52,18 @@ def get_config_and_catalog_from_args(args: List[str]) -> Tuple[str, Mapping[str,
 
     command = config["__command"]
     if command == "test_read":
-        catalog = ConfiguredAirbyteCatalog.parse_obj(BaseConnector.read_config(catalog_path))
+        catalog = ConfiguredAirbyteCatalogSerializer.load(BaseConnector.read_config(catalog_path))
+        state = Source.read_state(state_path)
     else:
         catalog = None
+        state = []
 
     if "__injected_declarative_manifest" not in config:
         raise ValueError(
             f"Invalid config: `__injected_declarative_manifest` should be provided at the root of the config but config only has keys {list(config.keys())}"
         )
 
-    return command, config, catalog
+    return command, config, catalog, state
 
 
 def handle_connector_builder_request(
@@ -48,28 +71,37 @@ def handle_connector_builder_request(
     command: str,
     config: Mapping[str, Any],
     catalog: Optional[ConfiguredAirbyteCatalog],
+    state: List[AirbyteStateMessage],
     limits: TestReadLimits,
 ) -> AirbyteMessage:
     if command == "resolve_manifest":
         return resolve_manifest(source)
     elif command == "test_read":
-        assert catalog is not None, "`test_read` requires a valid `ConfiguredAirbyteCatalog`, got None."
-        return read_stream(source, config, catalog, limits)
+        assert (
+            catalog is not None
+        ), "`test_read` requires a valid `ConfiguredAirbyteCatalog`, got None."
+        return read_stream(source, config, catalog, state, limits)
     else:
         raise ValueError(f"Unrecognized command {command}.")
 
 
-def handle_request(args: List[str]) -> AirbyteMessage:
-    command, config, catalog = get_config_and_catalog_from_args(args)
+def handle_request(args: List[str]) -> str:
+    command, config, catalog, state = get_config_and_catalog_from_args(args)
     limits = get_limits(config)
     source = create_source(config, limits)
-    return handle_connector_builder_request(source, command, config, catalog, limits).json(exclude_unset=True)
+    return orjson.dumps(
+        AirbyteMessageSerializer.dump(
+            handle_connector_builder_request(source, command, config, catalog, state, limits)
+        )
+    ).decode()  # type: ignore[no-any-return] # Serializer.dump() always returns AirbyteMessage
 
 
 if __name__ == "__main__":
     try:
         print(handle_request(sys.argv[1:]))
     except Exception as exc:
-        error = AirbyteTracedException.from_exception(exc, message=f"Error handling request: {str(exc)}")
+        error = AirbyteTracedException.from_exception(
+            exc, message=f"Error handling request: {str(exc)}"
+        )
         m = error.as_airbyte_message()
-        print(error.as_airbyte_message().json(exclude_unset=True))
+        print(orjson.dumps(AirbyteMessageSerializer.dump(m)).decode())