airbyte-cdk 0.72.1__py3-none-any.whl → 6.13.1.dev4106__py3-none-any.whl

airbyte_cdk/sources/declarative/declarative_source.py

@@ -2,8 +2,9 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
+import logging
 from abc import abstractmethod
-from typing import Tuple
+from typing import Any, Mapping, Tuple
 
 from airbyte_cdk.sources.abstract_source import AbstractSource
 from airbyte_cdk.sources.declarative.checks.connection_checker import ConnectionChecker
@@ -19,7 +20,9 @@ class DeclarativeSource(AbstractSource):
     def connection_checker(self) -> ConnectionChecker:
         """Returns the ConnectionChecker to use for the `check` operation"""
 
-    def check_connection(self, logger, config) -> Tuple[bool, any]:
+    def check_connection(
+        self, logger: logging.Logger, config: Mapping[str, Any]
+    ) -> Tuple[bool, Any]:
         """
         :param logger: The source logger
         :param config: The user-provided configuration as specified by the source's spec.

airbyte_cdk/sources/declarative/declarative_stream.py

@@ -1,17 +1,30 @@
 #
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
-
+import logging
 from dataclasses import InitVar, dataclass, field
 from typing import Any, Iterable, List, Mapping, MutableMapping, Optional, Union
 
 from airbyte_cdk.models import SyncMode
+from airbyte_cdk.sources.declarative.incremental import (
+    GlobalSubstreamCursor,
+    PerPartitionCursor,
+    PerPartitionWithGlobalCursor,
+)
 from airbyte_cdk.sources.declarative.interpolation import InterpolatedString
+from airbyte_cdk.sources.declarative.migrations.state_migration import StateMigration
+from airbyte_cdk.sources.declarative.retrievers import SimpleRetriever
 from airbyte_cdk.sources.declarative.retrievers.retriever import Retriever
 from airbyte_cdk.sources.declarative.schema import DefaultSchemaLoader
 from airbyte_cdk.sources.declarative.schema.schema_loader import SchemaLoader
-from airbyte_cdk.sources.declarative.types import Config, StreamSlice
+from airbyte_cdk.sources.streams.checkpoint import (
+    CheckpointMode,
+    CheckpointReader,
+    Cursor,
+    CursorBasedCheckpointReader,
+)
 from airbyte_cdk.sources.streams.core import Stream
+from airbyte_cdk.sources.types import Config, StreamSlice
 
 
 @dataclass
@@ -34,6 +47,7 @@ class DeclarativeStream(Stream):
     parameters: InitVar[Mapping[str, Any]]
     name: str
     primary_key: Optional[Union[str, List[str], List[List[str]]]]
+    state_migrations: List[StateMigration] = field(repr=True, default_factory=list)
     schema_loader: Optional[SchemaLoader] = None
     _name: str = field(init=False, repr=False, default="")
     _primary_key: str = field(init=False, repr=False, default="")
@@ -45,7 +59,11 @@ class DeclarativeStream(Stream):
             if isinstance(self.stream_cursor_field, str)
             else self.stream_cursor_field
         )
-        self._schema_loader = self.schema_loader if self.schema_loader else DefaultSchemaLoader(config=self.config, parameters=parameters)
+        self._schema_loader = (
+            self.schema_loader
+            if self.schema_loader
+            else DefaultSchemaLoader(config=self.config, parameters=parameters)
+        )
 
     @property  # type: ignore
     def primary_key(self) -> Optional[Union[str, List[str], List[List[str]]]]:
@@ -56,6 +74,14 @@ class DeclarativeStream(Stream):
         if not isinstance(value, property):
             self._primary_key = value
 
+    @property
+    def exit_on_rate_limit(self) -> bool:
+        return self.retriever.requester.exit_on_rate_limit  # type: ignore # abstract Retriever class has not requester attribute
+
+    @exit_on_rate_limit.setter
+    def exit_on_rate_limit(self, value: bool) -> None:
+        self.retriever.requester.exit_on_rate_limit = value  # type: ignore[attr-defined]
+
     @property  # type: ignore
     def name(self) -> str:
         """
@@ -75,7 +101,12 @@ class DeclarativeStream(Stream):
     @state.setter
     def state(self, value: MutableMapping[str, Any]) -> None:
         """State setter, accept state serialized by state getter."""
-        self.retriever.state = value
+        state: Mapping[str, Any] = value
+        if self.state_migrations:
+            for migration in self.state_migrations:
+                if migration.should_migrate(state):
+                    state = migration.migrate(state)
+        self.retriever.state = state
 
     def get_updated_state(
         self, current_stream_state: MutableMapping[str, Any], latest_record: Mapping[str, Any]
@@ -88,9 +119,15 @@ class DeclarativeStream(Stream):
         Override to return the default cursor field used by this stream e.g: an API entity might always use created_at as the cursor field.
         :return: The name of the field used as a cursor. If the cursor is nested, return an array consisting of the path to the cursor.
         """
-        cursor = self._stream_cursor_field.eval(self.config)
+        cursor = self._stream_cursor_field.eval(self.config)  # type: ignore # _stream_cursor_field is always cast to interpolated string
         return cursor if cursor else []
 
+    @property
+    def is_resumable(self) -> bool:
+        # Declarative sources always implement state getter/setter, but whether it supports checkpointing is based on
+        # if the retriever has a cursor defined.
+        return self.retriever.cursor is not None if hasattr(self.retriever, "cursor") else False
+
     def read_records(
         self,
         sync_mode: SyncMode,
@@ -101,7 +138,7 @@ class DeclarativeStream(Stream):
         """
         :param: stream_state We knowingly avoid using stream_state as we want cursors to manage their own state.
         """
-        if stream_slice is None:
+        if stream_slice is None or stream_slice == {}:
             # As the parameter is Optional, many would just call `read_records(sync_mode)` during testing without specifying the field
             # As part of the declarative model without custom components, this should never happen as the CDK would wire up a
             # SinglePartitionRouter that would create this StreamSlice properly
@@ -109,8 +146,10 @@ class DeclarativeStream(Stream):
             # empty slice which seems to make sense.
             stream_slice = StreamSlice(partition={}, cursor_slice={})
         if not isinstance(stream_slice, StreamSlice):
-            raise ValueError(f"DeclarativeStream does not support stream_slices that are not StreamSlice. Got {stream_slice}")
-        yield from self.retriever.read_records(self.get_json_schema(), stream_slice)
+            raise ValueError(
+                f"DeclarativeStream does not support stream_slices that are not StreamSlice. Got {stream_slice}"
+            )
+        yield from self.retriever.read_records(self.get_json_schema(), stream_slice)  # type: ignore # records are of the correct type
 
     def get_json_schema(self) -> Mapping[str, Any]:  # type: ignore
         """
@@ -122,7 +161,11 @@ class DeclarativeStream(Stream):
         return self._schema_loader.get_json_schema()
 
     def stream_slices(
-        self, *, sync_mode: SyncMode, cursor_field: Optional[List[str]] = None, stream_state: Optional[Mapping[str, Any]] = None
+        self,
+        *,
+        sync_mode: SyncMode,
+        cursor_field: Optional[List[str]] = None,
+        stream_state: Optional[Mapping[str, Any]] = None,
     ) -> Iterable[Optional[StreamSlice]]:
         """
         Override to define the slices for this stream. See the stream slicing section of the docs for more information.
@@ -144,3 +187,46 @@ class DeclarativeStream(Stream):
             important state is the one at the beginning of the slice
         """
         return None
+
+    def get_cursor(self) -> Optional[Cursor]:
+        if self.retriever and isinstance(self.retriever, SimpleRetriever):
+            return self.retriever.cursor
+        return None
+
+    def _get_checkpoint_reader(
+        self,
+        logger: logging.Logger,
+        cursor_field: Optional[List[str]],
+        sync_mode: SyncMode,
+        stream_state: MutableMapping[str, Any],
+    ) -> CheckpointReader:
+        """
+        This method is overridden to prevent issues with stream slice classification for incremental streams that have parent streams.
+
+        The classification logic, when used with `itertools.tee`, creates a copy of the stream slices. When `stream_slices` is called
+        the second time, the parent records generated during the classification phase are lost. This occurs because `itertools.tee`
+        only buffers the results, meaning the logic in `simple_retriever` that observes and updates the cursor isn't executed again.
+
+        By overriding this method, we ensure that the stream slices are processed correctly and parent records are not lost,
+        allowing the cursor to function as expected.
+        """
+        mappings_or_slices = self.stream_slices(
+            cursor_field=cursor_field,
+            sync_mode=sync_mode,  # todo: change this interface to no longer rely on sync_mode for behavior
+            stream_state=stream_state,
+        )
+
+        cursor = self.get_cursor()
+        checkpoint_mode = self._checkpoint_mode
+
+        if isinstance(
+            cursor, (GlobalSubstreamCursor, PerPartitionCursor, PerPartitionWithGlobalCursor)
+        ):
+            self.has_multiple_slices = True
+            return CursorBasedCheckpointReader(
+                stream_slices=mappings_or_slices,
+                cursor=cursor,
+                read_state_from_cursor=checkpoint_mode == CheckpointMode.RESUMABLE_FULL_REFRESH,
+            )
+
+        return super()._get_checkpoint_reader(logger, cursor_field, sync_mode, stream_state)

airbyte_cdk/sources/declarative/decoders/__init__.py

@@ -2,7 +2,28 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
 
+from airbyte_cdk.sources.declarative.decoders.composite_raw_decoder import CompositeRawDecoder
 from airbyte_cdk.sources.declarative.decoders.decoder import Decoder
-from airbyte_cdk.sources.declarative.decoders.json_decoder import JsonDecoder
+from airbyte_cdk.sources.declarative.decoders.json_decoder import (
+    GzipJsonDecoder,
+    IterableDecoder,
+    JsonDecoder,
+    JsonlDecoder,
+)
+from airbyte_cdk.sources.declarative.decoders.noop_decoder import NoopDecoder
+from airbyte_cdk.sources.declarative.decoders.pagination_decoder_decorator import (
+    PaginationDecoderDecorator,
+)
+from airbyte_cdk.sources.declarative.decoders.xml_decoder import XmlDecoder
 
-__all__ = ["Decoder", "JsonDecoder"]
+__all__ = [
+    "Decoder",
+    "CompositeRawDecoder",
+    "JsonDecoder",
+    "JsonlDecoder",
+    "IterableDecoder",
+    "GzipJsonDecoder",
+    "NoopDecoder",
+    "PaginationDecoderDecorator",
+    "XmlDecoder",
+]

airbyte_cdk/sources/declarative/decoders/composite_raw_decoder.py

@@ -0,0 +1,97 @@
+import csv
+import gzip
+import json
+import logging
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from io import BufferedIOBase, TextIOWrapper
+from typing import Any, Generator, MutableMapping, Optional
+
+import requests
+
+from airbyte_cdk.sources.declarative.decoders.decoder import Decoder
+
+logger = logging.getLogger("airbyte")
+
+
+@dataclass
+class Parser(ABC):
+    @abstractmethod
+    def parse(
+        self,
+        data: BufferedIOBase,
+    ) -> Generator[MutableMapping[str, Any], None, None]:
+        """
+        Parse data and yield dictionaries.
+        """
+        pass
+
+
+@dataclass
+class GzipParser(Parser):
+    inner_parser: Parser
+
+    def parse(
+        self,
+        data: BufferedIOBase,
+    ) -> Generator[MutableMapping[str, Any], None, None]:
+        """
+        Decompress gzipped bytes and pass decompressed data to the inner parser.
+        """
+        with gzip.GzipFile(fileobj=data, mode="rb") as gzipobj:
+            yield from self.inner_parser.parse(gzipobj)
+
+
+@dataclass
+class JsonLineParser(Parser):
+    encoding: Optional[str] = "utf-8"
+
+    def parse(
+        self,
+        data: BufferedIOBase,
+    ) -> Generator[MutableMapping[str, Any], None, None]:
+        for line in data:
+            try:
+                yield json.loads(line.decode(encoding=self.encoding or "utf-8"))
+            except json.JSONDecodeError as e:
+                logger.warning(f"Cannot decode/parse line {line!r} as JSON, error: {e}")
+
+
+@dataclass
+class CsvParser(Parser):
+    # TODO: migrate implementation to re-use file-base classes
+    encoding: Optional[str] = "utf-8"
+    delimiter: Optional[str] = ","
+
+    def parse(
+        self,
+        data: BufferedIOBase,
+    ) -> Generator[MutableMapping[str, Any], None, None]:
+        """
+        Parse CSV data from decompressed bytes.
+        """
+        text_data = TextIOWrapper(data, encoding=self.encoding)  # type: ignore
+        reader = csv.DictReader(text_data, delimiter=self.delimiter or ",")
+        yield from reader
+
+
+@dataclass
+class CompositeRawDecoder(Decoder):
+    """
+    Decoder strategy to transform a requests.Response into a Generator[MutableMapping[str, Any], None, None]
+    passed response.raw to parser(s).
+    Note: response.raw is not decoded/decompressed by default.
+    parsers should be instantiated recursively.
+    Example:
+    composite_raw_decoder = CompositeRawDecoder(parser=GzipParser(inner_parser=JsonLineParser(encoding="iso-8859-1")))
+    """
+
+    parser: Parser
+
+    def is_stream_response(self) -> bool:
+        return True
+
+    def decode(
+        self, response: requests.Response
+    ) -> Generator[MutableMapping[str, Any], None, None]:
+        yield from self.parser.parse(data=response.raw)  # type: ignore[arg-type]

airbyte_cdk/sources/declarative/decoders/decoder.py

@@ -4,7 +4,7 @@
 
 from abc import abstractmethod
 from dataclasses import dataclass
-from typing import Any, List, Mapping, Union
+from typing import Any, Generator, MutableMapping
 
 import requests
 
@@ -16,10 +16,17 @@ class Decoder:
     """
 
     @abstractmethod
-    def decode(self, response: requests.Response) -> Union[Mapping[str, Any], List]:
+    def is_stream_response(self) -> bool:
+        """
+        Set to True if you'd like to use stream=True option in http requester
+        """
+
+    @abstractmethod
+    def decode(
+        self, response: requests.Response
+    ) -> Generator[MutableMapping[str, Any], None, None]:
         """
         Decodes a requests.Response into a Mapping[str, Any] or an array
         :param response: the response to decode
-        :return: Mapping or array describing the response
+        :return: Generator of Mapping describing the response
         """
-        pass

airbyte_cdk/sources/declarative/decoders/json_decoder.py

@@ -1,13 +1,19 @@
 #
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
-
+import codecs
+import logging
 from dataclasses import InitVar, dataclass
-from typing import Any, List, Mapping, Union
+from gzip import decompress
+from typing import Any, Generator, List, Mapping, MutableMapping, Optional
 
+import orjson
 import requests
+
 from airbyte_cdk.sources.declarative.decoders.decoder import Decoder
 
+logger = logging.getLogger("airbyte")
+
 
 @dataclass
 class JsonDecoder(Decoder):
@@ -17,8 +23,89 @@ class JsonDecoder(Decoder):
 
     parameters: InitVar[Mapping[str, Any]]
 
-    def decode(self, response: requests.Response) -> Union[Mapping[str, Any], List]:
+    def is_stream_response(self) -> bool:
+        return False
+
+    def decode(
+        self, response: requests.Response
+    ) -> Generator[MutableMapping[str, Any], None, None]:
+        """
+        Given the response is an empty string or an emtpy list, the function will return a generator with an empty mapping.
+        """
         try:
-            return response.json()
+            body_json = response.json()
+            yield from self.parse_body_json(body_json)
         except requests.exceptions.JSONDecodeError:
-            return {}
+            logger.warning(
+                f"Response cannot be parsed into json: {response.status_code=}, {response.text=}"
+            )
+            yield {}
+
+    @staticmethod
+    def parse_body_json(
+        body_json: MutableMapping[str, Any] | List[MutableMapping[str, Any]],
+    ) -> Generator[MutableMapping[str, Any], None, None]:
+        if not isinstance(body_json, list):
+            body_json = [body_json]
+        if len(body_json) == 0:
+            yield {}
+        else:
+            yield from body_json
+
+
+@dataclass
+class IterableDecoder(Decoder):
+    """
+    Decoder strategy that returns the string content of the response, if any.
+    """
+
+    parameters: InitVar[Mapping[str, Any]]
+
+    def is_stream_response(self) -> bool:
+        return True
+
+    def decode(
+        self, response: requests.Response
+    ) -> Generator[MutableMapping[str, Any], None, None]:
+        for line in response.iter_lines():
+            yield {"record": line.decode()}
+
+
+@dataclass
+class JsonlDecoder(Decoder):
+    """
+    Decoder strategy that returns the json-encoded content of the response, if any.
+    """
+
+    parameters: InitVar[Mapping[str, Any]]
+
+    def is_stream_response(self) -> bool:
+        return True
+
+    def decode(
+        self, response: requests.Response
+    ) -> Generator[MutableMapping[str, Any], None, None]:
+        # TODO???: set delimiter? usually it is `\n` but maybe it would be useful to set optional?
+        #  https://github.com/airbytehq/airbyte-internal-issues/issues/8436
+        for record in response.iter_lines():
+            yield orjson.loads(record)
+
+
+@dataclass
+class GzipJsonDecoder(JsonDecoder):
+    encoding: Optional[str]
+
+    def __post_init__(self, parameters: Mapping[str, Any]) -> None:
+        if self.encoding:
+            try:
+                codecs.lookup(self.encoding)
+            except LookupError:
+                raise ValueError(
+                    f"Invalid encoding '{self.encoding}'. Please check provided encoding"
+                )
+
+    def decode(
+        self, response: requests.Response
+    ) -> Generator[MutableMapping[str, Any], None, None]:
+        raw_string = decompress(response.content).decode(encoding=self.encoding or "utf-8")
+        yield from self.parse_body_json(orjson.loads(raw_string))

airbyte_cdk/sources/declarative/decoders/noop_decoder.py

@@ -0,0 +1,21 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+import logging
+from typing import Any, Generator, Mapping
+
+import requests
+
+from airbyte_cdk.sources.declarative.decoders.decoder import Decoder
+
+logger = logging.getLogger("airbyte")
+
+
+class NoopDecoder(Decoder):
+    def is_stream_response(self) -> bool:
+        return False
+
+    def decode(  # type: ignore[override]  # Signature doesn't match base class
+        self,
+        response: requests.Response,
+    ) -> Generator[Mapping[str, Any], None, None]:
+        yield from [{}]

airbyte_cdk/sources/declarative/decoders/pagination_decoder_decorator.py

@@ -0,0 +1,39 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+import logging
+from dataclasses import dataclass
+from typing import Any, Generator, MutableMapping
+
+import requests
+
+from airbyte_cdk.sources.declarative.decoders import Decoder
+
+logger = logging.getLogger("airbyte")
+
+
+@dataclass
+class PaginationDecoderDecorator(Decoder):
+    """
+    Decoder to wrap other decoders when instantiating a DefaultPaginator in order to bypass decoding if the response is streamed.
+    """
+
+    def __init__(self, decoder: Decoder):
+        self._decoder = decoder
+
+    @property
+    def decoder(self) -> Decoder:
+        return self._decoder
+
+    def is_stream_response(self) -> bool:
+        return self._decoder.is_stream_response()
+
+    def decode(
+        self, response: requests.Response
+    ) -> Generator[MutableMapping[str, Any], None, None]:
+        if self._decoder.is_stream_response():
+            logger.warning("Response is streamed and therefore will not be decoded for pagination.")
+            yield {}
+        else:
+            yield from self._decoder.decode(response)

airbyte_cdk/sources/declarative/decoders/xml_decoder.py

@@ -0,0 +1,98 @@
+#
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+#
+
+import logging
+from dataclasses import InitVar, dataclass
+from typing import Any, Generator, Mapping, MutableMapping
+from xml.parsers.expat import ExpatError
+
+import requests
+import xmltodict
+
+from airbyte_cdk.sources.declarative.decoders.decoder import Decoder
+
+logger = logging.getLogger("airbyte")
+
+
+@dataclass
+class XmlDecoder(Decoder):
+    """
+    XmlDecoder is a decoder strategy that parses the XML content of the resopnse, and converts it to a dict.
+
+    This class handles XML attributes by prefixing them with an '@' symbol and represents XML text content by using the '#text' key if the element has attributes or the element name/tag. It does not currently support XML namespace declarations.
+
+    Example XML Input:
+    <root>
+        <location id="123">
+            San Francisco
+        </location>
+        <item id="1" category="books">
+            <name>Book Title 1</name>
+            <price>10.99</price>
+        </item>
+        <item id="2" category="electronics">
+            <name>Gadget</name>
+            <price>299.99</price>
+            <description>A useful gadget</description>
+        </item>
+    </root>
+
+    Converted Output:
+    {
+        "root": {
+            "location: {
+                "@id": "123,
+                "#text": "San Francisco"
+            },
+            "item": [
+              {
+                "@id": "1",
+                "@category": "books",
+                "name": "Book Title 1",
+                "price": "10.99"
+              },
+              {
+                "@id": "2",
+                "@category": "electronics",
+                "name": "Gadget",
+                "price": "299.99",
+                "description": "A useful gadget"
+              }
+            ]
+        }
+    }
+
+    Notes:
+        - Attributes of an XML element are prefixed with an '@' symbol in the dictionary output.
+        - Text content of an XML element is handled in two different ways, depending on whether
+          the element has attributes.
+                - If the element has attributes, the text content will be
+                  represented by the "#text" key.
+                - If the element does not have any attributes, the text content will be
+                  represented by element name.
+        - Namespace declarations are not supported in the current implementation.
+    """
+
+    parameters: InitVar[Mapping[str, Any]]
+
+    def is_stream_response(self) -> bool:
+        return False
+
+    def decode(
+        self, response: requests.Response
+    ) -> Generator[MutableMapping[str, Any], None, None]:
+        body_xml = response.text
+        try:
+            body_json = xmltodict.parse(body_xml)
+            if not isinstance(body_json, list):
+                body_json = [body_json]
+            if len(body_json) == 0:
+                yield {}
+            else:
+                yield from body_json
+        except ExpatError as exc:
+            logger.warning(
+                f"Response cannot be parsed from XML: {response.status_code=}, {response.text=}, {exc=}"
+            )
+            yield {}

airbyte_cdk/sources/declarative/extractors/__init__.py

@@ -6,5 +6,16 @@ from airbyte_cdk.sources.declarative.extractors.dpath_extractor import DpathExtr
 from airbyte_cdk.sources.declarative.extractors.http_selector import HttpSelector
 from airbyte_cdk.sources.declarative.extractors.record_filter import RecordFilter
 from airbyte_cdk.sources.declarative.extractors.record_selector import RecordSelector
+from airbyte_cdk.sources.declarative.extractors.response_to_file_extractor import (
+    ResponseToFileExtractor,
+)
+from airbyte_cdk.sources.declarative.extractors.type_transformer import TypeTransformer
 
-__all__ = ["HttpSelector", "DpathExtractor", "RecordFilter", "RecordSelector"]
+__all__ = [
+    "TypeTransformer",
+    "HttpSelector",
+    "DpathExtractor",
+    "RecordFilter",
+    "RecordSelector",
+    "ResponseToFileExtractor",
+]