airbyte-cdk 6.21.1.dev0__py3-none-any.whl → 6.26.0.dev4103__py3-none-any.whl

airbyte_cdk/sources/file_based/file_based_source.py

@@ -33,6 +33,9 @@ from airbyte_cdk.sources.file_based.config.file_based_stream_config import (
     FileBasedStreamConfig,
     ValidationPolicy,
 )
+from airbyte_cdk.sources.file_based.config.identities_based_stream_config import (
+    IdentitiesStreamConfig,
+)
 from airbyte_cdk.sources.file_based.discovery_policy import (
     AbstractDiscoveryPolicy,
     DefaultDiscoveryPolicy,
@@ -49,7 +52,11 @@ from airbyte_cdk.sources.file_based.schema_validation_policies import (
     DEFAULT_SCHEMA_VALIDATION_POLICIES,
     AbstractSchemaValidationPolicy,
 )
-from airbyte_cdk.sources.file_based.stream import AbstractFileBasedStream, DefaultFileBasedStream
+from airbyte_cdk.sources.file_based.stream import (
+    AbstractFileBasedStream,
+    DefaultFileBasedStream,
+    IdentitiesStream,
+)
 from airbyte_cdk.sources.file_based.stream.concurrent.adapters import FileBasedStreamFacade
 from airbyte_cdk.sources.file_based.stream.concurrent.cursor import (
     AbstractConcurrentFileBasedCursor,
@@ -157,6 +164,9 @@ class FileBasedSource(ConcurrentSourceAdapter, ABC):
         errors = []
         tracebacks = []
         for stream in streams:
+            if isinstance(stream, IdentitiesStream):
+                # Probably need to check identities endpoint/api access but will skip for now.
+                continue
             if not isinstance(stream, AbstractFileBasedStream):
                 raise ValueError(f"Stream {stream} is not a file-based stream.")
             try:
@@ -164,6 +174,7 @@ class FileBasedSource(ConcurrentSourceAdapter, ABC):
                 availability_method = (
                     stream.availability_strategy.check_availability
                     if self._use_file_transfer(parsed_config)
+                    or self._sync_acl_permissions(parsed_config)
                     else stream.availability_strategy.check_availability_and_parsability
                 )
                 (
@@ -289,6 +300,13 @@ class FileBasedSource(ConcurrentSourceAdapter, ABC):
                     )
 
                 streams.append(stream)
+
+            identities_stream_config = self._get_identities_stream_config(parsed_config)
+            if identities_stream_config:
+                identities_stream = self._make_identities_stream(
+                    stream_config=identities_stream_config
+                )
+                streams.append(identities_stream)
             return streams
 
         except ValidationError as exc:
@@ -312,6 +330,19 @@ class FileBasedSource(ConcurrentSourceAdapter, ABC):
             cursor=cursor,
             use_file_transfer=self._use_file_transfer(parsed_config),
             preserve_directory_structure=self._preserve_directory_structure(parsed_config),
+            sync_acl_permissions=self._sync_acl_permissions(parsed_config),
+        )
+
+    def _make_identities_stream(
+        self,
+        stream_config: IdentitiesStreamConfig,
+    ) -> Stream:
+        return IdentitiesStream(
+            config=stream_config,
+            catalog_schema=self.stream_schemas.get(stream_config.name),
+            stream_reader=self.stream_reader,
+            discovery_policy=self.discovery_policy,
+            errors_collector=self.errors_collector,
         )
 
     def _get_stream_from_catalog(
@@ -387,6 +418,14 @@ class FileBasedSource(ConcurrentSourceAdapter, ABC):
         )
         return use_file_transfer
 
+    @staticmethod
+    def _use_records_transfer(parsed_config: AbstractFileBasedSpec) -> bool:
+        use_records_transfer = (
+            hasattr(parsed_config.delivery_method, "delivery_type")
+            and parsed_config.delivery_method.delivery_type == "use_records_transfer"
+        )
+        return use_records_transfer
+
     @staticmethod
     def _preserve_directory_structure(parsed_config: AbstractFileBasedSpec) -> bool:
         """
@@ -408,3 +447,28 @@ class FileBasedSource(ConcurrentSourceAdapter, ABC):
         ):
             return parsed_config.delivery_method.preserve_directory_structure
         return True
+
+    @staticmethod
+    def _sync_acl_permissions(parsed_config: AbstractFileBasedSpec) -> bool:
+        if (
+            FileBasedSource._use_records_transfer(parsed_config)
+            and hasattr(parsed_config.delivery_method, "sync_acl_permissions")
+            and parsed_config.delivery_method.sync_acl_permissions is not None
+        ):
+            return parsed_config.delivery_method.sync_acl_permissions
+        return False
+
+    @staticmethod
+    def _get_identities_stream_config(
+        parsed_config: AbstractFileBasedSpec,
+    ) -> Optional[IdentitiesStreamConfig]:
+        identities_stream_config = None
+        if (
+            FileBasedSource._sync_acl_permissions(parsed_config)
+            and hasattr(parsed_config.delivery_method, "identities")
+            and parsed_config.delivery_method.identities is not None
+            and isinstance(parsed_config.delivery_method.identities, IdentitiesStreamConfig)
+            and parsed_config.delivery_method.identities.domain
+        ):
+            identities_stream_config = parsed_config.delivery_method.identities
+        return identities_stream_config

airbyte_cdk/sources/file_based/file_based_stream_reader.py

@@ -135,6 +135,15 @@ class AbstractFileBasedStreamReader(ABC):
             return use_file_transfer
         return False
 
+    def use_records_transfer(self) -> bool:
+        if self.config:
+            use_records_transfer = (
+                hasattr(self.config.delivery_method, "delivery_type")
+                and self.config.delivery_method.delivery_type == "use_records_transfer"
+            )
+            return use_records_transfer
+        return False
+
     def preserve_directory_structure(self) -> bool:
         # fall back to preserve subdirectories if config is not present or incomplete
         if (
@@ -146,6 +155,16 @@ class AbstractFileBasedStreamReader(ABC):
             return self.config.delivery_method.preserve_directory_structure
         return True
 
+    def sync_acl_permissions(self) -> bool:
+        if (
+            self.config
+            and self.use_records_transfer()
+            and hasattr(self.config.delivery_method, "sync_acl_permissions")
+            and self.config.delivery_method.sync_acl_permissions is not None
+        ):
+            return self.config.delivery_method.sync_acl_permissions
+        return False
+
     @abstractmethod
     def get_file(
         self, file: RemoteFile, local_directory: str, logger: logging.Logger
@@ -183,3 +202,17 @@ class AbstractFileBasedStreamReader(ABC):
         makedirs(path.dirname(local_file_path), exist_ok=True)
         absolute_file_path = path.abspath(local_file_path)
         return [file_relative_path, local_file_path, absolute_file_path]
+
+    def get_file_acl_permissions(self, file: RemoteFile, logger: logging.Logger) -> Dict[str, Any]:
+        """
+        This is required for connectors that will support syncing
+        ACL Permissions from files.
+        """
+        return {}
+
+    def load_identity_groups(self, logger: logging.Logger) -> Iterable[Dict[str, Any]]:
+        """
+        This is required for connectors that will support syncing
+        identities.
+        """
+        yield {}

airbyte_cdk/sources/file_based/schema_helpers.py

@@ -23,6 +23,31 @@ file_transfer_schema = {
     "properties": {"data": {"type": "object"}, "file": {"type": "object"}},
 }
 
+remote_file_permissions_schema = {
+    "type": "object",
+    "properties": {
+        "id": {"type": "string"},
+        "file_path": {"type": "string"},
+        "allowed_identity_remote_ids": {"type": "array", "items": {"type": "string"}},
+        "publicly_accessible": {"type": "boolean"},
+    },
+}
+
+remote_file_identity_schema = {
+    "type": "object",
+    "properties": {
+        "id": {"type": "string"},
+        "remote_id": {"type": "string"},
+        "parent_id": {"type": ["null", "string"]},
+        "name": {"type": ["null", "string"]},
+        "description": {"type": ["null", "string"]},
+        "email_address": {"type": ["null", "string"]},
+        "member_email_addresses": {"type": ["null", "array"]},
+        "type": {"type": "string"},
+        "modified_at": {"type": "string"},
+    },
+}
+
 
 @total_ordering
 class ComparableType(Enum):

airbyte_cdk/sources/file_based/stream/__init__.py

@@ -1,4 +1,5 @@
 from airbyte_cdk.sources.file_based.stream.abstract_file_based_stream import AbstractFileBasedStream
 from airbyte_cdk.sources.file_based.stream.default_file_based_stream import DefaultFileBasedStream
+from airbyte_cdk.sources.file_based.stream.identities_stream import IdentitiesStream
 
-__all__ = ["AbstractFileBasedStream", "DefaultFileBasedStream"]
+__all__ = ["AbstractFileBasedStream", "DefaultFileBasedStream", "IdentitiesStream"]

airbyte_cdk/sources/file_based/stream/default_file_based_stream.py

@@ -29,6 +29,7 @@ from airbyte_cdk.sources.file_based.schema_helpers import (
     SchemaType,
     file_transfer_schema,
     merge_schemas,
+    remote_file_permissions_schema,
     schemaless_schema,
 )
 from airbyte_cdk.sources.file_based.stream import AbstractFileBasedStream
@@ -47,6 +48,7 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
 
     FILE_TRANSFER_KW = "use_file_transfer"
     PRESERVE_DIRECTORY_STRUCTURE_KW = "preserve_directory_structure"
+    SYNC_ACL_PERMISSIONS_KW = "sync_acl_permissions"
     FILES_KEY = "files"
     DATE_TIME_FORMAT = "%Y-%m-%dT%H:%M:%S.%fZ"
     ab_last_mod_col = "_ab_source_file_last_modified"
@@ -56,6 +58,7 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
     airbyte_columns = [ab_last_mod_col, ab_file_name_col]
     use_file_transfer = False
     preserve_directory_structure = True
+    sync_acl_permissions = False
 
     def __init__(self, **kwargs: Any):
         if self.FILE_TRANSFER_KW in kwargs:
@@ -64,6 +67,8 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
             self.preserve_directory_structure = kwargs.pop(
                 self.PRESERVE_DIRECTORY_STRUCTURE_KW, True
             )
+        if self.SYNC_ACL_PERMISSIONS_KW in kwargs:
+            self.sync_acl_permissions = kwargs.pop(self.SYNC_ACL_PERMISSIONS_KW, False)
         super().__init__(**kwargs)
 
     @property
@@ -105,6 +110,8 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
                     self.ab_file_name_col: {"type": "string"},
                 },
             }
+        elif self.sync_acl_permissions:
+            return remote_file_permissions_schema
         else:
             return super()._filter_schema_invalid_properties(configured_catalog_json_schema)
 
@@ -187,6 +194,26 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
                         yield stream_data_to_airbyte_message(
                             self.name, record, is_file_transfer_message=True
                         )
+                elif self.sync_acl_permissions:
+                    try:
+                        metadata_record = self.stream_reader.get_file_acl_permissions(
+                            file, logger=self.logger
+                        )
+                        yield stream_data_to_airbyte_message(
+                            self.name, metadata_record, is_file_transfer_message=False
+                        )
+                    except Exception as e:
+                        self.logger.error(
+                            f"Failed to retrieve metadata for file {file.uri}: {str(e)}"
+                        )
+                        yield AirbyteMessage(
+                            type=MessageType.LOG,
+                            log=AirbyteLogMessage(
+                                level=Level.ERROR,
+                                message=f"Error retrieving metadata: stream={self.name} file={file.uri}",
+                                stack_trace=traceback.format_exc(),
+                            ),
+                        )
                 else:
                     for record in parser.parse_records(
                         self.config, file, self.stream_reader, self.logger, schema
@@ -284,6 +311,8 @@ class DefaultFileBasedStream(AbstractFileBasedStream, IncrementalMixin):
     def _get_raw_json_schema(self) -> JsonSchema:
         if self.use_file_transfer:
             return file_transfer_schema
+        elif self.sync_acl_permissions:
+            return remote_file_permissions_schema
         elif self.config.input_schema:
             return self.config.get_input_schema()  # type: ignore
         elif self.config.schemaless:

airbyte_cdk/sources/file_based/stream/identities_stream.py

@@ -0,0 +1,99 @@
+#
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+#
+
+import traceback
+from functools import cache
+from typing import Any, Iterable, List, Mapping, MutableMapping, Optional
+
+from airbyte_protocol_dataclasses.models import SyncMode
+
+from airbyte_cdk.models import AirbyteLogMessage, AirbyteMessage, Level
+from airbyte_cdk.models import Type as MessageType
+from airbyte_cdk.sources.file_based.config.file_based_stream_config import PrimaryKeyType
+from airbyte_cdk.sources.file_based.config.identities_based_stream_config import (
+    IdentitiesStreamConfig,
+)
+from airbyte_cdk.sources.file_based.discovery_policy import AbstractDiscoveryPolicy
+from airbyte_cdk.sources.file_based.exceptions import FileBasedErrorsCollector, FileBasedSourceError
+from airbyte_cdk.sources.file_based.file_based_stream_reader import AbstractFileBasedStreamReader
+from airbyte_cdk.sources.file_based.schema_helpers import remote_file_identity_schema
+from airbyte_cdk.sources.file_based.types import StreamSlice
+from airbyte_cdk.sources.streams import Stream
+from airbyte_cdk.sources.streams.checkpoint import Cursor
+from airbyte_cdk.sources.streams.core import JsonSchema
+from airbyte_cdk.sources.utils.record_helper import stream_data_to_airbyte_message
+from airbyte_cdk.utils.traced_exception import AirbyteTracedException
+
+
+class IdentitiesStream(Stream):
+    """
+    The identities stream. A full refresh stream to sync identities from a certain domain.
+    The stream reader manage the logic to get such data, which is implemented on connector side.
+    """
+
+    is_resumable = False
+
+    def __init__(
+        self,
+        config: IdentitiesStreamConfig,
+        catalog_schema: Optional[Mapping[str, Any]],
+        stream_reader: AbstractFileBasedStreamReader,
+        discovery_policy: AbstractDiscoveryPolicy,
+        errors_collector: FileBasedErrorsCollector,
+    ):
+        super().__init__()
+        self.config = config
+        self.catalog_schema = catalog_schema
+        self.stream_reader = stream_reader
+        self._discovery_policy = discovery_policy
+        self.errors_collector = errors_collector
+        self._cursor: MutableMapping[str, Any] = {}
+
+    @property
+    def state(self) -> MutableMapping[str, Any]:
+        return self._cursor
+
+    @state.setter
+    def state(self, value: MutableMapping[str, Any]) -> None:
+        """State setter, accept state serialized by state getter."""
+        self._cursor = value
+
+    @property
+    def primary_key(self) -> PrimaryKeyType:
+        return None
+
+    def read_records(
+        self,
+        sync_mode: SyncMode,
+        cursor_field: Optional[List[str]] = None,
+        stream_slice: Optional[StreamSlice] = None,
+        stream_state: Optional[Mapping[str, Any]] = None,
+    ) -> Iterable[Mapping[str, Any] | AirbyteMessage]:
+        try:
+            identity_groups = self.stream_reader.load_identity_groups(logger=self.logger)
+            for record in identity_groups:
+                yield stream_data_to_airbyte_message(self.name, record)
+        except AirbyteTracedException as exc:
+            # Re-raise the exception to stop the whole sync immediately as this is a fatal error
+            raise exc
+        except Exception:
+            yield AirbyteMessage(
+                type=MessageType.LOG,
+                log=AirbyteLogMessage(
+                    level=Level.ERROR,
+                    message=f"{FileBasedSourceError.ERROR_PARSING_RECORD.value} stream={self.name}",
+                    stack_trace=traceback.format_exc(),
+                ),
+            )
+
+    @cache
+    def get_json_schema(self) -> JsonSchema:
+        return remote_file_identity_schema
+
+    @property
+    def name(self) -> str:
+        return self.config.name
+
+    def get_cursor(self) -> Optional[Cursor]:
+        return None

airbyte_cdk/sources/http_logger.py

@@ -45,7 +45,7 @@ def format_http_message(
         log_message["http"]["is_auxiliary"] = is_auxiliary  # type: ignore [index]
     if stream_name:
         log_message["airbyte_cdk"] = {"stream": {"name": stream_name}}
-    return log_message  # type: ignore [return-value]  # got "dict[str, object]", expected "dict[str, JsonType]"
+    return log_message  # type: ignore[return-value]  # got "dict[str, object]", expected "dict[str, JsonType]"
 
 
 def _normalize_body_string(body_str: Optional[Union[str, bytes]]) -> Optional[str]:

airbyte_cdk/sources/streams/concurrent/clamping.py

@@ -0,0 +1,99 @@
+from abc import ABC
+from datetime import datetime, timedelta
+from enum import Enum
+from typing import Callable
+
+from airbyte_cdk.sources.streams.concurrent.cursor_types import CursorValueType
+
+
+class ClampingStrategy(ABC):
+    def clamp(self, value: CursorValueType) -> CursorValueType:
+        raise NotImplementedError()
+
+
+class NoClamping(ClampingStrategy):
+    def clamp(self, value: CursorValueType) -> CursorValueType:
+        return value
+
+
+class ClampingEndProvider:
+    def __init__(
+        self,
+        clamping_strategy: ClampingStrategy,
+        end_provider: Callable[[], CursorValueType],
+        granularity: timedelta,
+    ) -> None:
+        self._clamping_strategy = clamping_strategy
+        self._end_provider = end_provider
+        self._granularity = granularity
+
+    def __call__(self) -> CursorValueType:
+        return self._clamping_strategy.clamp(self._end_provider()) - self._granularity
+
+
+class DayClampingStrategy(ClampingStrategy):
+    def __init__(self, is_ceiling: bool = True) -> None:
+        self._is_ceiling = is_ceiling
+
+    def clamp(self, value: datetime) -> datetime:  # type: ignore  # datetime implements method from CursorValueType
+        return_value = value.replace(hour=0, minute=0, second=0, microsecond=0)
+        if self._is_ceiling:
+            return return_value + timedelta(days=1)
+        return return_value
+
+
+class MonthClampingStrategy(ClampingStrategy):
+    def __init__(self, is_ceiling: bool = True) -> None:
+        self._is_ceiling = is_ceiling
+
+    def clamp(self, value: datetime) -> datetime:  # type: ignore  # datetime implements method from CursorValueType
+        return_value = value.replace(hour=0, minute=0, second=0, microsecond=0)
+        needs_to_round = value.day != 1
+        if not needs_to_round:
+            return return_value
+
+        return self._ceil(return_value) if self._is_ceiling else return_value.replace(day=1)
+
+    def _ceil(self, value: datetime) -> datetime:
+        return value.replace(
+            year=value.year + 1 if value.month == 12 else value.year,
+            month=(value.month % 12) + 1,
+            day=1,
+            hour=0,
+            minute=0,
+            second=0,
+            microsecond=0,
+        )
+
+
+class Weekday(Enum):
+    """
+    These integer values map to the same ones used by the Datetime.date.weekday() implementation
+    """
+
+    MONDAY = 0
+    TUESDAY = 1
+    WEDNESDAY = 2
+    THURSDAY = 3
+    FRIDAY = 4
+    SATURDAY = 5
+    SUNDAY = 6
+
+
+class WeekClampingStrategy(ClampingStrategy):
+    def __init__(self, day_of_week: Weekday, is_ceiling: bool = True) -> None:
+        self._day_of_week = day_of_week.value
+        self._is_ceiling = is_ceiling
+
+    def clamp(self, value: datetime) -> datetime:  # type: ignore  # datetime implements method from CursorValueType
+        days_diff_to_ceiling = (
+            7 - (value.weekday() - self._day_of_week)
+            if value.weekday() > self._day_of_week
+            else abs(value.weekday() - self._day_of_week)
+        )
+        delta = (
+            timedelta(days_diff_to_ceiling)
+            if self._is_ceiling
+            else timedelta(days_diff_to_ceiling - 7)
+        )
+        return value.replace(hour=0, minute=0, second=0, microsecond=0) + delta