Merge branch 'main' into tolik0/add-api-budget

Author: tolik0

.github/workflows/pypi_publish.yml

@@ -146,7 +146,7 @@ jobs:
           (github.event_name == 'push' &&
            startsWith(github.ref, 'refs/tags/v')
           ) || github.event.inputs.publish_to_pypi == 'true'
-        uses: pypa/gh-action-pypi-publish@v1.12.3
+        uses: pypa/gh-action-pypi-publish@v1.12.4
 
   publish_sdm:
     name: Publish SDM to DockerHub

airbyte_cdk/sources/declarative/auth/oauth.py

@@ -3,7 +3,7 @@
 #
 
 from dataclasses import InitVar, dataclass, field
-from datetime import timedelta
+from datetime import datetime, timedelta
 from typing import Any, List, Mapping, MutableMapping, Optional, Union
 
 from airbyte_cdk.sources.declarative.auth.declarative_authenticator import DeclarativeAuthenticator
@@ -232,8 +232,13 @@ def get_refresh_request_headers(self) -> Mapping[str, Any]:
         return self._refresh_request_headers.eval(self.config)
 
     def get_token_expiry_date(self) -> AirbyteDateTime:
+        if not self._has_access_token_been_initialized():
+            return AirbyteDateTime.from_datetime(datetime.min)
         return self._token_expiry_date  # type: ignore # _token_expiry_date is an AirbyteDateTime. It is never None despite what mypy thinks
 
+    def _has_access_token_been_initialized(self) -> bool:
+        return self._access_token is not None
+
     def set_token_expiry_date(self, value: Union[str, int]) -> None:
         self._token_expiry_date = self._parse_token_expiration_date(value)
 

airbyte_cdk/sources/declarative/concurrent_declarative_source.py

@@ -3,7 +3,7 @@
 #
 
 import logging
-from typing import Any, Generic, Iterator, List, Mapping, Optional, Tuple
+from typing import Any, Generic, Iterator, List, Mapping, MutableMapping, Optional, Tuple
 
 from airbyte_cdk.models import (
     AirbyteCatalog,
@@ -224,6 +224,7 @@ def _group_streams(
                     stream_state = self._connector_state_manager.get_stream_state(
                         stream_name=declarative_stream.name, namespace=declarative_stream.namespace
                     )
+                    stream_state = self._migrate_state(declarative_stream, stream_state)
 
                     retriever = self._get_retriever(declarative_stream, stream_state)
 
@@ -331,6 +332,8 @@ def _group_streams(
                     stream_state = self._connector_state_manager.get_stream_state(
                         stream_name=declarative_stream.name, namespace=declarative_stream.namespace
                     )
+                    stream_state = self._migrate_state(declarative_stream, stream_state)
+
                     partition_router = declarative_stream.retriever.stream_slicer._partition_router
 
                     perpartition_cursor = (
@@ -521,3 +524,14 @@ def _remove_concurrent_streams_from_catalog(
                 if stream.stream.name not in concurrent_stream_names
             ]
         )
+
+    @staticmethod
+    def _migrate_state(
+        declarative_stream: DeclarativeStream, stream_state: MutableMapping[str, Any]
+    ) -> MutableMapping[str, Any]:
+        for state_migration in declarative_stream.state_migrations:
+            if state_migration.should_migrate(stream_state):
+                # The state variable is expected to be mutable but the migrate method returns an immutable mapping.
+                stream_state = dict(state_migration.migrate(stream_state))
+
+        return stream_state

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -968,6 +968,17 @@ def create_concurrency_level(
             parameters={},
         )
 
+    @staticmethod
+    def apply_stream_state_migrations(
+        stream_state_migrations: List[Any] | None, stream_state: MutableMapping[str, Any]
+    ) -> MutableMapping[str, Any]:
+        if stream_state_migrations:
+            for state_migration in stream_state_migrations:
+                if state_migration.should_migrate(stream_state):
+                    # The state variable is expected to be mutable but the migrate method returns an immutable mapping.
+                    stream_state = dict(state_migration.migrate(stream_state))
+        return stream_state
+
     def create_concurrent_cursor_from_datetime_based_cursor(
         self,
         model_type: Type[BaseModel],
@@ -977,6 +988,7 @@ def create_concurrent_cursor_from_datetime_based_cursor(
         config: Config,
         message_repository: Optional[MessageRepository] = None,
         runtime_lookback_window: Optional[datetime.timedelta] = None,
+        stream_state_migrations: Optional[List[Any]] = None,
         **kwargs: Any,
     ) -> ConcurrentCursor:
         # Per-partition incremental streams can dynamically create child cursors which will pass their current
@@ -987,6 +999,7 @@ def create_concurrent_cursor_from_datetime_based_cursor(
             if "stream_state" not in kwargs
             else kwargs["stream_state"]
         )
+        stream_state = self.apply_stream_state_migrations(stream_state_migrations, stream_state)
 
         component_type = component_definition.get("type")
         if component_definition.get("type") != model_type.__name__:
@@ -1222,6 +1235,7 @@ def create_concurrent_cursor_from_perpartition_cursor(
         config: Config,
         stream_state: MutableMapping[str, Any],
         partition_router: PartitionRouter,
+        stream_state_migrations: Optional[List[Any]] = None,
         **kwargs: Any,
     ) -> ConcurrentPerPartitionCursor:
         component_type = component_definition.get("type")
@@ -1270,8 +1284,10 @@ def create_concurrent_cursor_from_perpartition_cursor(
                 stream_namespace=stream_namespace,
                 config=config,
                 message_repository=NoopMessageRepository(),
+                stream_state_migrations=stream_state_migrations,
             )
         )
+        stream_state = self.apply_stream_state_migrations(stream_state_migrations, stream_state)
 
         # Return the concurrent cursor and state converter
         return ConcurrentPerPartitionCursor(
@@ -1780,6 +1796,7 @@ def _merge_stream_slicers(
                     stream_name=model.name or "",
                     stream_namespace=None,
                     config=config or {},
+                    stream_state_migrations=model.state_migrations,
                 )
             return (
                 self._create_component_from_model(model=model.incremental_sync, config=config)

airbyte_cdk/sources/streams/http/requests_native_auth/abstract_oauth.py

@@ -261,6 +261,9 @@ def _parse_token_expiration_date(self, value: Union[str, int]) -> AirbyteDateTim
 
         :return: expiration datetime
         """
+        if not value and not self.token_has_expired():
+            # No expiry token was provided but the previous one is not expired so it's fine
+            return self.get_token_expiry_date()
 
         if self.token_expiry_is_time_of_expiration:
             if not self.token_expiry_date_format:

unit_tests/sources/declarative/auth/test_oauth.py

@@ -301,6 +301,7 @@ def test_initialize_declarative_oauth_with_token_expiry_date_as_timestamp(
             client_id="{{ config['client_id'] }}",
             client_secret="{{ config['client_secret'] }}",
             token_expiry_date=timestamp,
+            access_token_value="some_access_token",
             refresh_token="some_refresh_token",
             config={
                 "refresh_endpoint": "refresh_end",
@@ -313,6 +314,34 @@ def test_initialize_declarative_oauth_with_token_expiry_date_as_timestamp(
         assert isinstance(oauth._token_expiry_date, AirbyteDateTime)
         assert oauth.get_token_expiry_date() == ab_datetime_parse(expected_date)
 
+    def test_given_no_access_token_but_expiry_in_the_future_when_refresh_token_then_fetch_access_token(
+        self,
+    ) -> None:
+        expiry_date = ab_datetime_now().add(timedelta(days=1))
+        oauth = DeclarativeOauth2Authenticator(
+            token_refresh_endpoint="https://refresh_endpoint.com/",
+            client_id="some_client_id",
+            client_secret="some_client_secret",
+            token_expiry_date=expiry_date.isoformat(),
+            refresh_token="some_refresh_token",
+            config={},
+            parameters={},
+            grant_type="client",
+        )
+
+        with HttpMocker() as http_mocker:
+            http_mocker.post(
+                HttpRequest(
+                    url="https://refresh_endpoint.com/",
+                    body="grant_type=client&client_id=some_client_id&client_secret=some_client_secret&refresh_token=some_refresh_token",
+                ),
+                HttpResponse(body=json.dumps({"access_token": "new_access_token"})),
+            )
+            oauth.get_access_token()
+
+        assert oauth.access_token == "new_access_token"
+        assert oauth._token_expiry_date == expiry_date
+
     @pytest.mark.parametrize(
         "expires_in_response, token_expiry_date_format",
         [

unit_tests/sources/declarative/custom_state_migration.py

@@ -0,0 +1,47 @@
+#
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+#
+
+from typing import Any, Mapping
+
+from airbyte_cdk.sources.declarative.declarative_stream import DeclarativeStream
+from airbyte_cdk.sources.declarative.interpolation.interpolated_string import InterpolatedString
+from airbyte_cdk.sources.declarative.migrations.state_migration import StateMigration
+from airbyte_cdk.sources.types import Config
+
+
+class CustomStateMigration(StateMigration):
+    declarative_stream: DeclarativeStream
+    config: Config
+
+    def __init__(self, declarative_stream: DeclarativeStream, config: Config):
+        self._config = config
+        self.declarative_stream = declarative_stream
+        self._cursor = declarative_stream.incremental_sync
+        self._parameters = declarative_stream.parameters
+        self._cursor_field = InterpolatedString.create(
+            self._cursor.cursor_field, parameters=self._parameters
+        ).eval(self._config)
+
+    def should_migrate(self, stream_state: Mapping[str, Any]) -> bool:
+        return True
+
+    def migrate(self, stream_state: Mapping[str, Any]) -> Mapping[str, Any]:
+        if not self.should_migrate(stream_state):
+            return stream_state
+        updated_at = stream_state[self._cursor.cursor_field]
+
+        migrated_stream_state = {
+            "states": [
+                {
+                    "partition": {"type": "type_1"},
+                    "cursor": {self._cursor.cursor_field: updated_at},
+                },
+                {
+                    "partition": {"type": "type_2"},
+                    "cursor": {self._cursor.cursor_field: updated_at},
+                },
+            ]
+        }
+
+        return migrated_stream_state

unit_tests/sources/declarative/decoders/test_composite_decoder.py

@@ -30,35 +30,35 @@ def compress_with_gzip(data: str, encoding: str = "utf-8"):
     return buf.getvalue()
 
 
-def generate_csv(encoding: str) -> bytes:
-    """
-    Generate CSV data with tab-separated values (\t).
-    """
+def generate_csv(
+    encoding: str = "utf-8", delimiter: str = ",", should_compress: bool = False
+) -> bytes:
     data = [
-        {"id": 1, "name": "John", "age": 28},
-        {"id": 2, "name": "Alice", "age": 34},
-        {"id": 3, "name": "Bob", "age": 25},
+        {"id": "1", "name": "John", "age": "28"},
+        {"id": "2", "name": "Alice", "age": "34"},
+        {"id": "3", "name": "Bob", "age": "25"},
     ]
 
     output = StringIO()
-    writer = csv.DictWriter(output, fieldnames=["id", "name", "age"], delimiter="\t")
+    writer = csv.DictWriter(output, fieldnames=["id", "name", "age"], delimiter=delimiter)
     writer.writeheader()
     for row in data:
         writer.writerow(row)
 
-    # Ensure the pointer is at the beginning of the buffer before compressing
     output.seek(0)
+    csv_data = output.read()
 
-    # Compress the CSV data with Gzip
-    compressed_data = compress_with_gzip(output.read(), encoding=encoding)
-
-    return compressed_data
+    if should_compress:
+        return compress_with_gzip(csv_data, encoding=encoding)
+    return csv_data.encode(encoding)
 
 
 @pytest.mark.parametrize("encoding", ["utf-8", "utf", "iso-8859-1"])
 def test_composite_raw_decoder_gzip_csv_parser(requests_mock, encoding: str):
     requests_mock.register_uri(
-        "GET", "https://airbyte.io/", content=generate_csv(encoding=encoding)
+        "GET",
+        "https://airbyte.io/",
+        content=generate_csv(encoding=encoding, delimiter="\t", should_compress=True),
     )
     response = requests.get("https://airbyte.io/", stream=True)
 
@@ -175,3 +175,26 @@ def test_composite_raw_decoder_raises_traced_exception_when_both_parsers_fail(re
         with patch("json.loads", side_effect=Exception("test")):
             with pytest.raises(AirbyteTracedException):
                 list(composite_raw_decoder.decode(response))
+
+
+@pytest.mark.parametrize("encoding", ["utf-8", "utf", "iso-8859-1"])
+@pytest.mark.parametrize("delimiter", [",", "\t", ";"])
+def test_composite_raw_decoder_csv_parser_values(requests_mock, encoding: str, delimiter: str):
+    requests_mock.register_uri(
+        "GET",
+        "https://airbyte.io/",
+        content=generate_csv(encoding=encoding, delimiter=delimiter, should_compress=False),
+    )
+    response = requests.get("https://airbyte.io/", stream=True)
+
+    parser = CsvParser(encoding=encoding, delimiter=delimiter)
+    composite_raw_decoder = CompositeRawDecoder(parser=parser)
+
+    expected_data = [
+        {"id": "1", "name": "John", "age": "28"},
+        {"id": "2", "name": "Alice", "age": "34"},
+        {"id": "3", "name": "Bob", "age": "25"},
+    ]
+
+    parsed_records = list(composite_raw_decoder.decode(response))
+    assert parsed_records == expected_data