Merge branch 'main' into aj/feat/cli/accept-dynamic-gcp-project-id-from-env-var

Author: aaronsteers

.github/workflows/auto_merge_notification.yml

@@ -0,0 +1,31 @@
+# When a PR is has the auto-merge feature enabled or disabled, this workflow adds or updates
+# warning text at the bottom of the PR description.
+
+name: "Add Auto-Merge Notification Text"
+on:
+  pull_request:
+    types: [auto_merge_enabled, auto_merge_disabled]
+
+jobs:
+  update-description:
+    runs-on: ubuntu-latest
+    permissions:
+      pull-requests: write
+    steps:
+      - name: Add Auto-Merge Notice
+        if: github.event.action == 'auto_merge_enabled'
+        uses: bcgov/action-pr-description-add@v2.0.0
+        with:
+          add_markdown: |
+            > [!IMPORTANT]
+            > **Auto-merge enabled.**
+            > 
+            > _This PR is set to merge automatically when all requirements are met._
+
+      - name: Remove Auto-Merge Notice
+        if: github.event.action == 'auto_merge_disabled'
+        uses: bcgov/action-pr-description-add@v2.0.0
+        with:
+          add_markdown: |
+            > [!NOTE]
+            > **Auto-merge may have been disabled. Please check the PR status to confirm.**

airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py

@@ -229,10 +229,10 @@ def ensure_at_least_one_state_emitted(self) -> None:
 
     def _throttle_state_message(self) -> Optional[float]:
         """
-        Throttles the state message emission to once every 60 seconds.
+        Throttles the state message emission to once every 600 seconds.
         """
         current_time = time.time()
-        if current_time - self._last_emission_time <= 60:
+        if current_time - self._last_emission_time <= 600:
             return None
         return current_time
 

airbyte_cdk/sources/declarative/migrations/legacy_to_per_partition_state_migration.py

@@ -78,16 +78,20 @@ def should_migrate(self, stream_state: Mapping[str, Any]) -> bool:
           "<cursor_field>" : "<cursor_value>"
         }
         """
-        if stream_state:
-            for key, value in stream_state.items():
-                if isinstance(value, dict):
-                    keys = list(value.keys())
-                    if len(keys) != 1:
-                        # The input partitioned state should only have one key
-                        return False
-                    if keys[0] != self._cursor_field:
-                        # Unexpected key. Found {keys[0]}. Expected {self._cursor.cursor_field}
-                        return False
+        if not stream_state:
+            return False
+        for key, value in stream_state.items():
+            # it is expected the internal value to be a dictionary according to docstring
+            if not isinstance(value, dict):
+                return False
+            keys = list(value.keys())
+            if len(keys) != 1:
+                # The input partitioned state should only have one key
+                return False
+            if keys[0] != self._cursor_field:
+                # Unexpected key. Found {keys[0]}. Expected {self._cursor.cursor_field}
+                return False
+
         return True
 
     def migrate(self, stream_state: Mapping[str, Any]) -> Mapping[str, Any]:

airbyte_cdk/sources/declarative/parsers/manifest_component_transformer.py

@@ -120,14 +120,6 @@ def propagate_types_and_parameters(
             if found_type:
                 propagated_component["type"] = found_type
 
-        # When there is no resolved type, we're not processing a component (likely a regular object) and don't need to propagate parameters
-        # When the type refers to a json schema, we're not processing a component as well. This check is currently imperfect as there could
-        # be json_schema are not objects but we believe this is not likely in our case because:
-        # * records are Mapping so objects hence SchemaLoader root should be an object
-        # * connection_specification is a Mapping
-        if "type" not in propagated_component or self._is_json_schema_object(propagated_component):
-            return propagated_component
-
         # Combines parameters defined at the current level with parameters from parent components. Parameters at the current
         # level take precedence
         current_parameters = dict(copy.deepcopy(parent_parameters))
@@ -138,6 +130,27 @@ def propagate_types_and_parameters(
             else {**current_parameters, **component_parameters}
         )
 
+        # When there is no resolved type, we're not processing a component (likely a regular object) and don't need to propagate parameters
+        # When the type refers to a json schema, we're not processing a component as well. This check is currently imperfect as there could
+        # be json_schema are not objects but we believe this is not likely in our case because:
+        # * records are Mapping so objects hence SchemaLoader root should be an object
+        # * connection_specification is a Mapping
+        if self._is_json_schema_object(propagated_component):
+            return propagated_component
+
+        # For objects that don't have type check if their object fields have nested components which should have `$parameters` in it.
+        # For example, QueryProperties in requester.request_parameters, etc.
+        # Update propagated_component value with nested components with parent `$parameters` if needed and return propagated_component.
+        if "type" not in propagated_component:
+            if self._has_nested_components(propagated_component):
+                propagated_component = self._process_nested_components(
+                    propagated_component,
+                    parent_field_identifier,
+                    current_parameters,
+                    use_parent_parameters,
+                )
+            return propagated_component
+
         # Parameters should be applied to the current component fields with the existing field taking precedence over parameters if
         # both exist
         for parameter_key, parameter_value in current_parameters.items():
@@ -181,4 +194,33 @@ def propagate_types_and_parameters(
 
     @staticmethod
     def _is_json_schema_object(propagated_component: Mapping[str, Any]) -> bool:
-        return propagated_component.get("type") == "object"
+        return propagated_component.get("type") == "object" or propagated_component.get("type") == [
+            "null",
+            "object",
+        ]
+
+    @staticmethod
+    def _has_nested_components(propagated_component: Dict[str, Any]) -> bool:
+        for k, v in propagated_component.items():
+            if isinstance(v, dict) and v.get("type"):
+                return True
+        return False
+
+    def _process_nested_components(
+        self,
+        propagated_component: Dict[str, Any],
+        parent_field_identifier: str,
+        current_parameters: Mapping[str, Any],
+        use_parent_parameters: Optional[bool] = None,
+    ) -> Dict[str, Any]:
+        for field_name, field_value in propagated_component.items():
+            if isinstance(field_value, dict) and field_value.get("type"):
+                nested_component_with_parameters = self.propagate_types_and_parameters(
+                    parent_field_identifier,
+                    field_value,
+                    current_parameters,
+                    use_parent_parameters=use_parent_parameters,
+                )
+                propagated_component[field_name] = nested_component_with_parameters
+
+        return propagated_component

airbyte_cdk/sources/declarative/parsers/model_to_component_factory.py

@@ -928,9 +928,9 @@ def create_legacy_to_per_partition_state_migration(
         declarative_stream: DeclarativeStreamModel,
     ) -> LegacyToPerPartitionStateMigration:
         retriever = declarative_stream.retriever
-        if not isinstance(retriever, SimpleRetrieverModel):
+        if not isinstance(retriever, (SimpleRetrieverModel, AsyncRetrieverModel)):
             raise ValueError(
-                f"LegacyToPerPartitionStateMigrations can only be applied on a DeclarativeStream with a SimpleRetriever. Got {type(retriever)}"
+                f"LegacyToPerPartitionStateMigrations can only be applied on a DeclarativeStream with a SimpleRetriever or AsyncRetriever. Got {type(retriever)}"
             )
         partition_router = retriever.partition_router
         if not isinstance(
@@ -1484,6 +1484,7 @@ def create_concurrent_cursor_from_perpartition_cursor(
                 stream_state_migrations=stream_state_migrations,
             )
         )
+
         stream_state = self.apply_stream_state_migrations(stream_state_migrations, stream_state)
         # Per-partition state doesn't make sense for GroupingPartitionRouter, so force the global state
         use_global_cursor = isinstance(
@@ -1993,14 +1994,31 @@ def _build_incremental_cursor(
     ) -> Optional[StreamSlicer]:
         if model.incremental_sync and stream_slicer:
             if model.retriever.type == "AsyncRetriever":
+                stream_name = model.name or ""
+                stream_namespace = None
+                stream_state = self._connector_state_manager.get_stream_state(
+                    stream_name, stream_namespace
+                )
+                state_transformations = (
+                    [
+                        self._create_component_from_model(
+                            state_migration, config, declarative_stream=model
+                        )
+                        for state_migration in model.state_migrations
+                    ]
+                    if model.state_migrations
+                    else []
+                )
+
                 return self.create_concurrent_cursor_from_perpartition_cursor(  # type: ignore # This is a known issue that we are creating and returning a ConcurrentCursor which does not technically implement the (low-code) StreamSlicer. However, (low-code) StreamSlicer and ConcurrentCursor both implement StreamSlicer.stream_slices() which is the primary method needed for checkpointing
                     state_manager=self._connector_state_manager,
                     model_type=DatetimeBasedCursorModel,
                     component_definition=model.incremental_sync.__dict__,
-                    stream_name=model.name or "",
-                    stream_namespace=None,
+                    stream_name=stream_name,
+                    stream_namespace=stream_namespace,
                     config=config or {},
-                    stream_state={},
+                    stream_state=stream_state,
+                    stream_state_migrations=state_transformations,
                     partition_router=stream_slicer,
                 )
 

unit_tests/sources/declarative/incremental/test_concurrent_perpartitioncursor.py

@@ -3265,8 +3265,8 @@ def test_incremental_substream_request_options_provider(
 
 def test_state_throttling(mocker):
     """
-    Verifies that _emit_state_message does not emit a new state if less than 60s
-    have passed since last emission, and does emit once 60s or more have passed.
+    Verifies that _emit_state_message does not emit a new state if less than 600s
+    have passed since last emission, and does emit once 600s or more have passed.
     """
     cursor = ConcurrentPerPartitionCursor(
         cursor_factory=MagicMock(),
@@ -3288,20 +3288,20 @@ def test_state_throttling(mocker):
 
     mock_time = mocker.patch("time.time")
 
-    # First attempt: only 10 seconds passed => NO emission
-    mock_time.return_value = 10
+    # First attempt: only 100 seconds passed => NO emission
+    mock_time.return_value = 100
     cursor._emit_state_message()
     mock_connector_manager.update_state_for_stream.assert_not_called()
     mock_repo.emit_message.assert_not_called()
 
-    # Second attempt: 30 seconds passed => still NO emission
-    mock_time.return_value = 30
+    # Second attempt: 300 seconds passed => still NO emission
+    mock_time.return_value = 300
     cursor._emit_state_message()
     mock_connector_manager.update_state_for_stream.assert_not_called()
     mock_repo.emit_message.assert_not_called()
 
-    # Advance time: 70 seconds => exceed 60s => MUST emit
-    mock_time.return_value = 70
+    # Advance time: 700 seconds => exceed 600s => MUST emit
+    mock_time.return_value = 700
     cursor._emit_state_message()
     mock_connector_manager.update_state_for_stream.assert_called_once()
     mock_repo.emit_message.assert_called_once()

unit_tests/sources/declarative/migrations/test_legacy_to_per_partition_migration.py

@@ -189,6 +189,14 @@ def test_migrate_a_valid_legacy_state_to_per_partition():
             },
             id="test_should_not_migrate_if_the_partitioned_state_key_is_not_the_cursor_field",
         ),
+        pytest.param(
+            {"last_changed": "2022-12-27T08:34:39+00:00"},
+            id="test_should_not_migrate_if_the_partitioned_state_is_not_in_correct_format",
+        ),
+        pytest.param(
+            {},
+            id="test_should_not_migrate_if_not_state_is_passed",
+        ),
     ],
 )
 def test_should_not_migrate(input_state):
@@ -277,7 +285,7 @@ def _migrator_with_multiple_parent_streams():
             CustomPartitionRouter,
             True,
             ValueError,
-            "LegacyToPerPartitionStateMigrations can only be applied on a DeclarativeStream with a SimpleRetriever. Got <class 'unittest.mock.MagicMock'>",
+            "LegacyToPerPartitionStateMigrations can only be applied on a DeclarativeStream with a SimpleRetriever or AsyncRetriever. Got <class 'unittest.mock.MagicMock'>",
         ),
         (
             SimpleRetriever,