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

airbyte_cdk/sql/secrets.py

@@ -0,0 +1,123 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+"""Base classes and methods for working with secrets in Airbyte."""
+
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING, Any
+
+from pydantic_core import CoreSchema, core_schema
+
+from airbyte_cdk.sql import exceptions as exc
+
+if TYPE_CHECKING:
+    from pydantic import GetCoreSchemaHandler, GetJsonSchemaHandler, ValidationInfo
+    from pydantic.json_schema import JsonSchemaValue
+
+
+class SecretString(str):
+    """A string that represents a secret.
+
+    This class is used to mark a string as a secret. When a secret is printed, it
+    will be masked to prevent accidental exposure of sensitive information when debugging
+    or when printing containing objects like dictionaries.
+
+    To create a secret string, simply instantiate the class with any string value:
+
+        ```python
+        secret = SecretString("my_secret_password")
+        ```
+
+    """
+
+    __slots__ = ()
+
+    def __repr__(self) -> str:
+        """Override the representation of the secret string to return a masked value.
+
+        The secret string is always masked with `****` to prevent accidental exposure, unless
+        explicitly converted to a string. For instance, printing a config dictionary that contains
+        a secret will automatically mask the secret value instead of printing it in plain text.
+
+        However, if you explicitly convert the cast the secret as a string, such as when used
+        in an f-string, the secret will be exposed. This is the desired behavior to allow
+        secrets to be used in a controlled manner.
+        """
+        return "<SecretString: ****>"
+
+    def is_empty(self) -> bool:
+        """Check if the secret is an empty string."""
+        return len(self) == 0
+
+    def is_json(self) -> bool:
+        """Check if the secret string is a valid JSON string."""
+        try:
+            json.loads(self)
+        except (json.JSONDecodeError, Exception):
+            return False
+
+        return True
+
+    def __bool__(self) -> bool:
+        """Override the boolean value of the secret string.
+
+        Always returns `True` without inspecting contents.
+        """
+        return True
+
+    def parse_json(self) -> Any:
+        """Parse the secret string as JSON."""
+        try:
+            return json.loads(self)
+        except json.JSONDecodeError as ex:
+            raise exc.AirbyteInputError(
+                message="Failed to parse secret as JSON.",
+                context={
+                    "Message": ex.msg,
+                    "Position": ex.pos,
+                    "SecretString_Length": len(self),  # Debug secret blank or an unexpected format.
+                },
+            ) from None
+
+    # Pydantic compatibility
+
+    @classmethod
+    def validate(
+        cls,
+        v: Any,  # noqa: ANN401  # Must allow `Any` to match Pydantic signature
+        info: ValidationInfo,
+    ) -> SecretString:
+        """Validate the input value is valid as a secret string."""
+        _ = info  # Unused
+        if not isinstance(v, str):
+            raise exc.AirbyteInputError(
+                message="A valid `str` or `SecretString` object is required.",
+            )
+        return cls(v)
+
+    @classmethod
+    def __get_pydantic_core_schema__(  # noqa: PLW3201  # Pydantic dunder
+        cls,
+        source_type: Any,  # noqa: ANN401  # Must allow `Any` to match Pydantic signature
+        handler: GetCoreSchemaHandler,
+    ) -> CoreSchema:
+        """Return a modified core schema for the secret string."""
+        return core_schema.with_info_after_validator_function(
+            function=cls.validate, schema=handler(str), field_name=handler.field_name
+        )
+
+    @classmethod
+    def __get_pydantic_json_schema__(  # noqa: PLW3201  # Pydantic dunder method
+        cls, _core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
+    ) -> JsonSchemaValue:
+        """Return a modified JSON schema for the secret string.
+
+        - `writeOnly=True` is the official way to prevent secrets from being exposed inadvertently.
+        - `Format=password` is a popular and readable convention to indicate the field is sensitive.
+        """
+        _ = _core_schema, handler  # Unused
+        return {
+            "type": "string",
+            "format": "password",
+            "writeOnly": True,
+        }

airbyte_cdk/sql/shared/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+"""Module for future CDK components.
+
+Components here are planned to move to the CDK.
+
+TODO!: Add GitHub link here before merging.
+"""
+
+from __future__ import annotations
+
+from airbyte_cdk.sql.shared.sql_processor import SqlProcessorBase
+
+__all__ = [
+    "SqlProcessorBase",
+]

airbyte_cdk/sql/shared/catalog_providers.py

@@ -0,0 +1,145 @@
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+"""Catalog provider implementation.
+
+A catalog provider wraps a configured catalog and configured streams. This class is responsible for
+providing information about the catalog and streams. A catalog provider can also be updated with new
+streams as they are discovered, providing a thin layer of abstraction over the configured catalog.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, cast, final
+
+from airbyte_cdk.models import ConfiguredAirbyteCatalog
+from airbyte_cdk.sql import exceptions as exc
+from airbyte_cdk.sql._util.name_normalizers import LowerCaseNormalizer
+
+if TYPE_CHECKING:
+    from airbyte_cdk.models import ConfiguredAirbyteStream
+
+
+class CatalogProvider:
+    """A catalog provider wraps a configured catalog and configured streams.
+
+    This class is responsible for providing information about the catalog and streams.
+
+    Note:
+    - The catalog provider is not responsible for managing the catalog or streams but it may
+      be updated with new streams as they are discovered.
+    """
+
+    def __init__(
+        self,
+        configured_catalog: ConfiguredAirbyteCatalog,
+    ) -> None:
+        """Initialize the catalog manager with a catalog object reference.
+
+        Since the catalog is passed by reference, the catalog manager may be updated with new
+        streams as they are discovered.
+        """
+        self._catalog: ConfiguredAirbyteCatalog = self.validate_catalog(configured_catalog)
+
+    @staticmethod
+    def validate_catalog(catalog: ConfiguredAirbyteCatalog) -> Any:
+        """Validate the catalog to ensure it is valid.
+
+        This requires ensuring that `generationId` and `minGenerationId` are both set. If
+        not, both values will be set to `1`.
+        """
+        for stream in catalog.streams:
+            if stream.generation_id is None:
+                stream.generation_id = 1
+            if stream.minimum_generation_id is None:
+                stream.minimum_generation_id = 1
+            if stream.sync_id is None:
+                stream.sync_id = 1  # This should ideally increment monotonically with each sync.
+
+        return catalog
+
+    @property
+    def configured_catalog(self) -> ConfiguredAirbyteCatalog:
+        """Return the configured catalog."""
+        return self._catalog
+
+    @property
+    def stream_names(self) -> list[str]:
+        """Return the names of the streams in the catalog."""
+        return list({stream.stream.name for stream in self.configured_catalog.streams})
+
+    def get_configured_stream_info(
+        self,
+        stream_name: str,
+    ) -> ConfiguredAirbyteStream:
+        """Return the column definitions for the given stream."""
+        if not self.configured_catalog:
+            raise exc.AirbyteInternalError(
+                message="Cannot get stream JSON schema without a catalog.",
+            )
+
+        matching_streams: list[ConfiguredAirbyteStream] = [
+            stream
+            for stream in self.configured_catalog.streams
+            if stream.stream.name == stream_name
+        ]
+        if not matching_streams:
+            raise exc.AirbyteStreamNotFoundError(
+                stream_name=stream_name,
+                context={
+                    "available_streams": [
+                        stream.stream.name for stream in self.configured_catalog.streams
+                    ],
+                },
+            )
+
+        if len(matching_streams) > 1:
+            raise exc.AirbyteInternalError(
+                message="Multiple streams found with same name.",
+                context={
+                    "stream_name": stream_name,
+                },
+            )
+
+        return matching_streams[0]
+
+    @final
+    def get_stream_json_schema(
+        self,
+        stream_name: str,
+    ) -> dict[str, Any]:
+        """Return the column definitions for the given stream."""
+        return cast(dict[str, Any], self.get_configured_stream_info(stream_name).stream.json_schema)
+
+    def get_stream_properties(
+        self,
+        stream_name: str,
+    ) -> dict[str, dict[str, Any]]:
+        """Return the names of the top-level properties for the given stream."""
+        return cast(dict[str, Any], self.get_stream_json_schema(stream_name)["properties"])
+
+    def get_primary_keys(
+        self,
+        stream_name: str,
+    ) -> list[str]:
+        """Return the primary keys for the given stream."""
+        pks = self.get_configured_stream_info(stream_name).primary_key
+        if not pks:
+            return []
+
+        normalized_pks: list[list[str]] = [
+            [LowerCaseNormalizer.normalize(c) for c in pk] for pk in pks
+        ]
+
+        for pk_nodes in normalized_pks:
+            if len(pk_nodes) != 1:
+                raise exc.AirbyteError(
+                    message=(
+                        "Nested primary keys are not supported. "
+                        "Each PK column should have exactly one node. "
+                    ),
+                    context={
+                        "stream_name": stream_name,
+                        "primary_key_nodes": pk_nodes,
+                    },
+                )
+
+        return [pk_nodes[0] for pk_nodes in normalized_pks]