airbyte-cdk 0.0.0.dev0__py3-none-any.whl

airbyte_cdk/sources/types.py

@@ -0,0 +1,169 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+from __future__ import annotations
+
+from typing import Any, ItemsView, Iterator, KeysView, List, Mapping, Optional, ValuesView
+
+from airbyte_cdk.models import AirbyteRecordMessageFileReference
+from airbyte_cdk.utils.slice_hasher import SliceHasher
+
+# A FieldPointer designates a path to a field inside a mapping. For example, retrieving ["k1", "k1.2"] in the object {"k1" :{"k1.2":
+# "hello"}] returns "hello"
+FieldPointer = List[str]
+Config = Mapping[str, Any]
+ConnectionDefinition = Mapping[str, Any]
+StreamState = Mapping[str, Any]
+EmptyString = str()
+
+
+class Record(Mapping[str, Any]):
+    def __init__(
+        self,
+        data: Mapping[str, Any],
+        stream_name: str,
+        associated_slice: Optional[StreamSlice] = None,
+        file_reference: Optional[AirbyteRecordMessageFileReference] = None,
+    ):
+        self._data = data
+        self._associated_slice = associated_slice
+        self.stream_name = stream_name
+        self._file_reference = file_reference
+
+    @property
+    def data(self) -> Mapping[str, Any]:
+        return self._data
+
+    @property
+    def associated_slice(self) -> Optional[StreamSlice]:
+        return self._associated_slice
+
+    @property
+    def file_reference(self) -> AirbyteRecordMessageFileReference:
+        return self._file_reference
+
+    @file_reference.setter
+    def file_reference(self, value: AirbyteRecordMessageFileReference) -> None:
+        self._file_reference = value
+
+    def __repr__(self) -> str:
+        return repr(self._data)
+
+    def __getitem__(self, key: str) -> Any:
+        return self._data[key]
+
+    def __len__(self) -> int:
+        return len(self._data)
+
+    def __iter__(self) -> Any:
+        return iter(self._data)
+
+    def __contains__(self, item: object) -> bool:
+        return item in self._data
+
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, Record):
+            # noinspection PyProtectedMember
+            return self._data == other._data
+        return False
+
+    def __ne__(self, other: object) -> bool:
+        return not self.__eq__(other)
+
+
+class StreamSlice(Mapping[str, Any]):
+    def __init__(
+        self,
+        *,
+        partition: Mapping[str, Any],
+        cursor_slice: Mapping[str, Any],
+        extra_fields: Optional[Mapping[str, Any]] = None,
+    ) -> None:
+        """
+        :param partition: The partition keys representing a unique partition in the stream.
+        :param cursor_slice: The incremental cursor slice keys, such as dates or pagination tokens.
+        :param extra_fields: Additional fields that should not be part of the partition but passed along, such as metadata from the parent stream.
+        """
+        self._partition = partition
+        self._cursor_slice = cursor_slice
+        self._extra_fields = extra_fields or {}
+
+        # Ensure that partition keys do not overlap with cursor slice keys
+        if partition.keys() & cursor_slice.keys():
+            raise ValueError("Keys for partition and incremental sync cursor should not overlap")
+
+        self._stream_slice = dict(partition) | dict(cursor_slice)
+
+    @property
+    def partition(self) -> Mapping[str, Any]:
+        """Returns the partition portion of the stream slice."""
+        p = self._partition
+        while isinstance(p, StreamSlice):
+            p = p.partition
+        return p
+
+    @property
+    def cursor_slice(self) -> Mapping[str, Any]:
+        """Returns the cursor slice portion of the stream slice."""
+        c = self._cursor_slice
+        while isinstance(c, StreamSlice):
+            c = c.cursor_slice
+        return c
+
+    @property
+    def extra_fields(self) -> Mapping[str, Any]:
+        """Returns the extra fields that are not part of the partition."""
+        return self._extra_fields
+
+    def __repr__(self) -> str:
+        return repr(self._stream_slice)
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        raise ValueError("StreamSlice is immutable")
+
+    def __getitem__(self, key: str) -> Any:
+        return self._stream_slice[key]
+
+    def __len__(self) -> int:
+        return len(self._stream_slice)
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._stream_slice)
+
+    def __contains__(self, item: Any) -> bool:
+        return item in self._stream_slice
+
+    def keys(self) -> KeysView[str]:
+        return self._stream_slice.keys()
+
+    def items(self) -> ItemsView[str, Any]:
+        return self._stream_slice.items()
+
+    def values(self) -> ValuesView[Any]:
+        return self._stream_slice.values()
+
+    def get(self, key: str, default: Any = None) -> Optional[Any]:
+        return self._stream_slice.get(key, default)
+
+    def __eq__(self, other: Any) -> bool:
+        if isinstance(other, dict):
+            return self._stream_slice == other
+        if isinstance(other, StreamSlice):
+            # noinspection PyProtectedMember
+            return self._partition == other._partition and self._cursor_slice == other._cursor_slice
+        return False
+
+    def __ne__(self, other: Any) -> bool:
+        return not self.__eq__(other)
+
+    def __json_serializable__(self) -> Any:
+        return self._stream_slice
+
+    def __hash__(self) -> int:
+        return SliceHasher.hash(
+            stream_slice=self._stream_slice
+        )  # no need to provide stream_name here as this is used for slicing the cursor
+
+    def __bool__(self) -> bool:
+        return bool(self._stream_slice) or bool(self._extra_fields)

airbyte_cdk/sources/utils/__init__.py

@@ -0,0 +1,7 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+# Initialize Utils Package
+
+__all__ = ["record_helper"]

airbyte_cdk/sources/utils/casing.py

@@ -0,0 +1,12 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+
+import re
+
+
+# https://stackoverflow.com/a/1176023
+def camel_to_snake(s: str) -> str:
+    s = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", s)
+    return re.sub("([a-z0-9])([A-Z])", r"\1_\2", s).lower()

airbyte_cdk/sources/utils/files_directory.py

@@ -0,0 +1,15 @@
+#
+# Copyright (c) 2025 Airbyte, Inc., all rights reserved.
+#
+import os
+
+AIRBYTE_STAGING_DIRECTORY = os.getenv("AIRBYTE_STAGING_DIRECTORY", "/staging/files")
+DEFAULT_LOCAL_DIRECTORY = "/tmp/airbyte-file-transfer"
+
+
+def get_files_directory() -> str:
+    return (
+        AIRBYTE_STAGING_DIRECTORY
+        if os.path.exists(AIRBYTE_STAGING_DIRECTORY)
+        else DEFAULT_LOCAL_DIRECTORY
+    )

airbyte_cdk/sources/utils/record_helper.py

@@ -0,0 +1,53 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+import time
+from collections.abc import Mapping as ABCMapping
+from typing import Any, Mapping, Optional
+
+from airbyte_cdk.models import (
+    AirbyteLogMessage,
+    AirbyteMessage,
+    AirbyteRecordMessage,
+    AirbyteRecordMessageFileReference,
+    AirbyteTraceMessage,
+)
+from airbyte_cdk.models import Type as MessageType
+from airbyte_cdk.sources.streams.core import StreamData
+from airbyte_cdk.sources.utils.transform import TransformConfig, TypeTransformer
+
+
+def stream_data_to_airbyte_message(
+    stream_name: str,
+    data_or_message: StreamData,
+    transformer: TypeTransformer = TypeTransformer(TransformConfig.NoTransform),
+    schema: Optional[Mapping[str, Any]] = None,
+    file_reference: Optional[AirbyteRecordMessageFileReference] = None,
+) -> AirbyteMessage:
+    if schema is None:
+        schema = {}
+
+    match data_or_message:
+        case ABCMapping():
+            data = dict(data_or_message)
+            now_millis = time.time_ns() // 1_000_000
+            # Transform object fields according to config. Most likely you will
+            # need it to normalize values against json schema. By default no action
+            # taken unless configured. See
+            # docs/connector-development/cdk-python/schemas.md for details.
+            transformer.transform(data, schema)
+            message = AirbyteRecordMessage(
+                stream=stream_name,
+                data=data,
+                emitted_at=now_millis,
+                file_reference=file_reference,
+            )
+            return AirbyteMessage(type=MessageType.RECORD, record=message)
+        case AirbyteTraceMessage():
+            return AirbyteMessage(type=MessageType.TRACE, trace=data_or_message)
+        case AirbyteLogMessage():
+            return AirbyteMessage(type=MessageType.LOG, log=data_or_message)
+        case _:
+            raise ValueError(
+                f"Unexpected type for data_or_message: {type(data_or_message)}: {data_or_message}"
+            )

airbyte_cdk/sources/utils/schema_helpers.py

@@ -0,0 +1,230 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+
+import importlib
+import json
+import os
+import pkgutil
+from typing import Any, ClassVar, Dict, List, Mapping, MutableMapping, Optional, Tuple
+
+import jsonref
+from jsonschema import RefResolver, validate
+from jsonschema.exceptions import ValidationError
+from pydantic.v1 import BaseModel, Field
+
+from airbyte_cdk.models import ConnectorSpecification, FailureType
+from airbyte_cdk.utils.traced_exception import AirbyteTracedException
+
+
+class JsonFileLoader:
+    """
+    Custom json file loader to resolve references to resources located in "shared" directory.
+    We need this for compatability with existing schemas cause all of them have references
+    pointing to shared_schema.json file instead of shared/shared_schema.json
+    """
+
+    def __init__(self, uri_base: str, shared: str):
+        self.shared = shared
+        self.uri_base = uri_base
+
+    def __call__(self, uri: str) -> Dict[str, Any]:
+        uri = uri.replace(self.uri_base, f"{self.uri_base}/{self.shared}/")
+        with open(uri) as f:
+            data = json.load(f)
+            if isinstance(data, dict):
+                return data
+            else:
+                raise ValueError(f"Expected to read a dictionary from {uri}. Got: {data}")
+
+
+def resolve_ref_links(obj: Any) -> Any:
+    """
+    Scan resolved schema and convert jsonref.JsonRef object to JSON serializable dict.
+
+    :param obj - jsonschema object with ref field resolved.
+    :return JSON serializable object with references without external dependencies.
+    """
+    if isinstance(obj, jsonref.JsonRef):
+        obj = resolve_ref_links(obj.__subject__)
+        # Omit existing definitions for external resource since
+        # we dont need it anymore.
+        if isinstance(obj, dict):
+            obj.pop("definitions", None)
+            return obj
+        else:
+            raise ValueError(f"Expected obj to be a dict. Got {obj}")
+    elif isinstance(obj, dict):
+        return {k: resolve_ref_links(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [resolve_ref_links(item) for item in obj]
+    else:
+        return obj
+
+
+def _expand_refs(schema: Any, ref_resolver: Optional[RefResolver] = None) -> None:
+    """Internal function to iterate over schema and replace all occurrences of $ref with their definitions. Recursive.
+
+    :param schema: schema that will be patched
+    :param ref_resolver: resolver to get definition from $ref, if None pass it will be instantiated
+    """
+    ref_resolver = ref_resolver or RefResolver.from_schema(schema)
+
+    if isinstance(schema, MutableMapping):
+        if "$ref" in schema:
+            ref_url = schema.pop("$ref")
+            _, definition = ref_resolver.resolve(ref_url)
+            _expand_refs(
+                definition, ref_resolver=ref_resolver
+            )  # expand refs in definitions as well
+            schema.update(definition)
+        else:
+            for key, value in schema.items():
+                _expand_refs(value, ref_resolver=ref_resolver)
+    elif isinstance(schema, List):
+        for value in schema:
+            _expand_refs(value, ref_resolver=ref_resolver)
+
+
+def expand_refs(schema: Any) -> None:
+    """Iterate over schema and replace all occurrences of $ref with their definitions.
+
+    :param schema: schema that will be patched
+    """
+    _expand_refs(schema)
+    schema.pop("definitions", None)  # remove definitions created by $ref
+
+
+def rename_key(schema: Any, old_key: str, new_key: str) -> None:
+    """Iterate over nested dictionary and replace one key with another. Used to replace anyOf with oneOf. Recursive."
+
+    :param schema: schema that will be patched
+    :param old_key: name of the key to replace
+    :param new_key: new name of the key
+    """
+    if not isinstance(schema, MutableMapping):
+        return
+
+    for key, value in schema.items():
+        rename_key(value, old_key, new_key)
+        if old_key in schema:
+            schema[new_key] = schema.pop(old_key)
+
+
+class ResourceSchemaLoader:
+    """JSONSchema loader from package resources"""
+
+    def __init__(self, package_name: str):
+        self.package_name = package_name
+
+    def get_schema(self, name: str) -> dict[str, Any]:
+        """
+        This method retrieves a JSON schema from the schemas/ folder.
+
+
+        The expected file structure is to have all top-level schemas (corresponding to streams) in the "schemas/" folder, with any shared $refs
+        living inside the "schemas/shared/" folder. For example:
+
+        schemas/shared/<shared_definition>.json
+        schemas/<name>.json # contains a $ref to shared_definition
+        schemas/<name2>.json # contains a $ref to shared_definition
+        """
+
+        schema_filename = f"schemas/{name}.json"
+        raw_file = pkgutil.get_data(self.package_name, schema_filename)
+        if not raw_file:
+            raise IOError(f"Cannot find file {schema_filename}")
+        try:
+            raw_schema = json.loads(raw_file)
+        except ValueError as err:
+            raise RuntimeError(f"Invalid JSON file format for file {schema_filename}") from err
+
+        return self._resolve_schema_references(raw_schema)
+
+    def _resolve_schema_references(self, raw_schema: dict[str, Any]) -> dict[str, Any]:
+        """
+        Resolve links to external references and move it to local "definitions" map.
+
+        :param raw_schema jsonschema to lookup for external links.
+        :return JSON serializable object with references without external dependencies.
+        """
+
+        package = importlib.import_module(self.package_name)
+        if package.__file__:
+            base = os.path.dirname(package.__file__) + "/"
+        else:
+            raise ValueError(f"Package {package} does not have a valid __file__ field")
+        resolved = jsonref.JsonRef.replace_refs(
+            raw_schema, loader=JsonFileLoader(base, "schemas/shared"), base_uri=base
+        )
+        resolved = resolve_ref_links(resolved)
+        if isinstance(resolved, dict):
+            return resolved
+        else:
+            raise ValueError(f"Expected resolved to be a dict. Got {resolved}")
+
+
+def check_config_against_spec_or_exit(
+    config: Mapping[str, Any], spec: ConnectorSpecification
+) -> None:
+    """
+    Check config object against spec. In case of spec is invalid, throws
+    an exception with validation error description.
+
+    :param config - config loaded from file specified over command line
+    :param spec - spec object generated by connector
+    """
+    spec_schema = spec.connectionSpecification
+    try:
+        validate(instance=config, schema=spec_schema)
+    except ValidationError as validation_error:
+        raise AirbyteTracedException(
+            message="Config validation error: " + validation_error.message,
+            internal_message=validation_error.message,
+            failure_type=FailureType.config_error,
+        ) from None  # required to prevent logging config secrets from the ValidationError's stacktrace
+
+
+class InternalConfig(BaseModel):
+    KEYWORDS: ClassVar[set[str]] = {"_limit", "_page_size"}
+    limit: int = Field(None, alias="_limit")
+    page_size: int = Field(None, alias="_page_size")
+
+    def dict(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
+        kwargs["by_alias"] = True
+        kwargs["exclude_unset"] = True
+        return super().dict(*args, **kwargs)
+
+    def is_limit_reached(self, records_counter: int) -> bool:
+        """
+        Check if record count reached limit set by internal config.
+        :param records_counter - number of records already red
+        :return True if limit reached, False otherwise
+        """
+        if self.limit:
+            if records_counter >= self.limit:
+                return True
+        return False
+
+
+def split_config(config: Mapping[str, Any]) -> Tuple[dict[str, Any], InternalConfig]:
+    """
+    Break config map object into 2 instances: first is a dict with user defined
+    configuration and second is internal config that contains private keys for
+    acceptance test configuration.
+
+    :param
+     config - Dict object that has been loaded from config file.
+
+    :return tuple of user defined config dict with filtered out internal
+    parameters and connector acceptance test internal config object.
+    """
+    main_config = {}
+    internal_config = {}
+    for k, v in config.items():
+        if k in InternalConfig.KEYWORDS:
+            internal_config[k] = v
+        else:
+            main_config[k] = v
+    return main_config, InternalConfig.parse_obj(internal_config)

airbyte_cdk/sources/utils/slice_logger.py

@@ -0,0 +1,57 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import json
+import logging
+from abc import ABC, abstractmethod
+from typing import Any, Mapping, Optional
+
+from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage, Level
+from airbyte_cdk.models import Type as MessageType
+
+
+class SliceLogger(ABC):
+    """
+    SliceLogger is an interface that allows us to log slices of data in a uniform way.
+    It is responsible for determining whether or not a slice should be logged and for creating the log message.
+    """
+
+    SLICE_LOG_PREFIX = "slice:"
+
+    def create_slice_log_message(self, _slice: Optional[Mapping[str, Any]]) -> AirbyteMessage:
+        """
+        Mapping is an interface that can be implemented in various ways. However, json.dumps will just do a `str(<object>)` if
+        the slice is a class implementing Mapping. Therefore, we want to cast this as a dict before passing this to json.dump
+        """
+        printable_slice = dict(_slice) if _slice else _slice
+        return AirbyteMessage(
+            type=MessageType.LOG,
+            log=AirbyteLogMessage(
+                level=Level.INFO,
+                message=f"{SliceLogger.SLICE_LOG_PREFIX}{json.dumps(printable_slice, default=str)}",
+            ),
+        )
+
+    @abstractmethod
+    def should_log_slice_message(self, logger: logging.Logger) -> bool:
+        """
+
+        :param logger:
+        :return:
+        """
+
+
+class DebugSliceLogger(SliceLogger):
+    def should_log_slice_message(self, logger: logging.Logger) -> bool:
+        """
+
+        :param logger:
+        :return:
+        """
+        return logger.isEnabledFor(logging.DEBUG)
+
+
+class AlwaysLogSliceLogger(SliceLogger):
+    def should_log_slice_message(self, logger: logging.Logger) -> bool:
+        return True