This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 0a41e28ec2a Preserve OpenSearch Serverless hook configuration when
deferring (#72472)
0a41e28ec2a is described below
commit 0a41e28ec2a2fd0bd302eeb1994a8cb02bf65677
Author: Zhibo Lin <[email protected]>
AuthorDate: Sun Sep 20 23:54:36 2026 +0800
Preserve OpenSearch Serverless hook configuration when deferring (#72472)
* Preserve OpenSearch Serverless hook configuration when deferring
* Narrow deferral regression assertions
* Autospec the OpenSearchServerlessHook patch in the trigger test
Review nit: with autospec=True a kwarg-name typo in hook() would fail at
the call instead of being swallowed by the mock.
Generated-by: Claude Opus 5
---------
Co-authored-by: Jarek Potiuk <[email protected]>
---
.../amazon/aws/sensors/opensearch_serverless.py | 3 +
.../amazon/aws/triggers/opensearch_serverless.py | 18 ++++-
.../aws/sensors/test_opensearch_serverless.py | 20 +++++-
.../aws/triggers/test_opensearch_serverless.py | 77 +++++++++++++++-------
4 files changed, 93 insertions(+), 25 deletions(-)
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/sensors/opensearch_serverless.py
b/providers/amazon/src/airflow/providers/amazon/aws/sensors/opensearch_serverless.py
index 051df9f26a8..38641f3c447 100644
---
a/providers/amazon/src/airflow/providers/amazon/aws/sensors/opensearch_serverless.py
+++
b/providers/amazon/src/airflow/providers/amazon/aws/sensors/opensearch_serverless.py
@@ -119,6 +119,9 @@ class
OpenSearchServerlessCollectionActiveSensor(AwsBaseSensor[OpenSearchServerl
waiter_delay=int(self.poke_interval),
waiter_max_attempts=self.max_retries,
aws_conn_id=self.aws_conn_id,
+ region_name=self.region_name,
+ verify=self.verify,
+ botocore_config=self.botocore_config,
),
method_name="poke",
)
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/triggers/opensearch_serverless.py
b/providers/amazon/src/airflow/providers/amazon/aws/triggers/opensearch_serverless.py
index 7543ac0d64c..93c90af1762 100644
---
a/providers/amazon/src/airflow/providers/amazon/aws/triggers/opensearch_serverless.py
+++
b/providers/amazon/src/airflow/providers/amazon/aws/triggers/opensearch_serverless.py
@@ -16,7 +16,7 @@
# under the License.
from __future__ import annotations
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
from airflow.providers.amazon.aws.hooks.opensearch_serverless import
OpenSearchServerlessHook
from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
@@ -36,6 +36,9 @@ class
OpenSearchServerlessCollectionActiveTrigger(AwsBaseWaiterTrigger):
:param waiter_delay: The amount of time in seconds to wait between
attempts. (default: 60)
:param waiter_max_attempts: The maximum number of attempts to be made.
(default: 20)
:param aws_conn_id: The Airflow connection used for AWS credentials.
+ :param region_name: The AWS region where the collection is located.
+ :param verify: Whether to verify SSL certificates.
+ :param botocore_config: Configuration dictionary for the botocore client.
"""
def __init__(
@@ -46,6 +49,9 @@ class
OpenSearchServerlessCollectionActiveTrigger(AwsBaseWaiterTrigger):
waiter_delay: int = 60,
waiter_max_attempts: int = 20,
aws_conn_id: str | None = None,
+ region_name: str | None = None,
+ verify: bool | str | None = None,
+ botocore_config: dict[str, Any] | None = None,
) -> None:
if not exactly_one(collection_id is None, collection_name is None):
raise AttributeError("Either collection_ids or collection_names
must be provided, not both.")
@@ -67,7 +73,15 @@ class
OpenSearchServerlessCollectionActiveTrigger(AwsBaseWaiterTrigger):
waiter_delay=waiter_delay,
waiter_max_attempts=waiter_max_attempts,
aws_conn_id=aws_conn_id,
+ region_name=region_name,
+ verify=verify,
+ botocore_config=botocore_config,
)
def hook(self) -> AwsGenericHook:
- return OpenSearchServerlessHook(aws_conn_id=self.aws_conn_id)
+ return OpenSearchServerlessHook(
+ aws_conn_id=self.aws_conn_id,
+ region_name=self.region_name,
+ verify=self.verify,
+ config=self.botocore_config,
+ )
diff --git
a/providers/amazon/tests/unit/amazon/aws/sensors/test_opensearch_serverless.py
b/providers/amazon/tests/unit/amazon/aws/sensors/test_opensearch_serverless.py
index d2f2f09a890..f341ec57106 100644
---
a/providers/amazon/tests/unit/amazon/aws/sensors/test_opensearch_serverless.py
+++
b/providers/amazon/tests/unit/amazon/aws/sensors/test_opensearch_serverless.py
@@ -24,7 +24,7 @@ from airflow.providers.amazon.aws.hooks.opensearch_serverless
import OpenSearchS
from airflow.providers.amazon.aws.sensors.opensearch_serverless import (
OpenSearchServerlessCollectionActiveSensor,
)
-from airflow.providers.common.compat.sdk import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException, TaskDeferred
class TestOpenSearchServerlessCollectionActiveSensor:
@@ -57,6 +57,24 @@ class TestOpenSearchServerlessCollectionActiveSensor:
assert op.hook._config is not None
assert op.hook._config.read_timeout == 42
+ def test_deferrable_sensor_forwards_aws_configuration(self):
+ sensor = OpenSearchServerlessCollectionActiveSensor(
+ **self.default_op_kwargs,
+ deferrable=True,
+ aws_conn_id="test_conn",
+ region_name="eu-west-1",
+ verify=False,
+ botocore_config={"read_timeout": 42},
+ )
+
+ with pytest.raises(TaskDeferred) as deferred:
+ sensor.execute(None)
+
+ serialized_kwargs = deferred.value.trigger.serialize()[1]
+ assert serialized_kwargs["region_name"] == "eu-west-1"
+ assert serialized_kwargs["verify"] is False
+ assert serialized_kwargs["botocore_config"] == {"read_timeout": 42}
+
@pytest.mark.parametrize(
("collection_name", "collection_id", "expected_pass"),
[
diff --git
a/providers/amazon/tests/unit/amazon/aws/triggers/test_opensearch_serverless.py
b/providers/amazon/tests/unit/amazon/aws/triggers/test_opensearch_serverless.py
index 2dbe91091e4..71959e6c2f6 100644
---
a/providers/amazon/tests/unit/amazon/aws/triggers/test_opensearch_serverless.py
+++
b/providers/amazon/tests/unit/amazon/aws/triggers/test_opensearch_serverless.py
@@ -52,31 +52,64 @@ class TestOpenSearchServerlessCollectionActiveTrigger:
COLLECTION_ID = "test_collection_id"
@pytest.mark.parametrize(
- ("collection_name", "collection_id", "expected_pass"),
+ ("collection_name", "collection_id"),
[
- pytest.param(COLLECTION_NAME, COLLECTION_ID, False,
id="both_provided_fails"),
- pytest.param(COLLECTION_NAME, None, True,
id="only_name_provided_passes"),
- pytest.param(None, COLLECTION_ID, True,
id="only_id_provided_passes"),
+ pytest.param(COLLECTION_NAME, None, id="collection_name"),
+ pytest.param(None, COLLECTION_ID, id="collection_id"),
],
)
- def test_serialization(self, collection_name, collection_id,
expected_pass):
- """Assert that arguments and classpath are correctly serialized."""
- call_args = prune_dict({"collection_id": collection_id,
"collection_name": collection_name})
-
- if expected_pass:
- trigger = OpenSearchServerlessCollectionActiveTrigger(**call_args)
- classpath, kwargs = trigger.serialize()
- assert classpath == BASE_TRIGGER_CLASSPATH +
"OpenSearchServerlessCollectionActiveTrigger"
- if call_args.get("collection_name"):
- assert kwargs.get("collection_name") == self.COLLECTION_NAME
- if call_args.get("collection_id"):
- assert kwargs.get("collection_id") == self.COLLECTION_ID
-
- if not expected_pass:
- with pytest.raises(
- AttributeError, match="Either collection_ids or
collection_names must be provided, not both."
- ):
- OpenSearchServerlessCollectionActiveTrigger(**call_args)
+ def test_serialization_round_trip(self, collection_name, collection_id):
+ """Assert that all arguments survive serialization and
reconstruction."""
+ trigger = OpenSearchServerlessCollectionActiveTrigger(
+ **prune_dict({"collection_id": collection_id, "collection_name":
collection_name}),
+ aws_conn_id="test_conn",
+ region_name="eu-west-1",
+ verify=False,
+ botocore_config={"read_timeout": 42},
+ )
+
+ classpath, kwargs = trigger.serialize()
+
+ assert classpath == BASE_TRIGGER_CLASSPATH +
"OpenSearchServerlessCollectionActiveTrigger"
+ assert kwargs == {
+ "collection_id": collection_id,
+ "collection_name": collection_name,
+ "waiter_delay": 60,
+ "waiter_max_attempts": 20,
+ "aws_conn_id": "test_conn",
+ "region_name": "eu-west-1",
+ "verify": False,
+ "botocore_config": {"read_timeout": 42},
+ }
+ restored_trigger =
OpenSearchServerlessCollectionActiveTrigger(**kwargs)
+ assert restored_trigger.serialize() == (classpath, kwargs)
+
+ def test_both_collection_name_and_id_raise(self):
+ with pytest.raises(
+ AttributeError, match="Either collection_ids or collection_names
must be provided, not both."
+ ):
+ OpenSearchServerlessCollectionActiveTrigger(
+ collection_name=self.COLLECTION_NAME,
collection_id=self.COLLECTION_ID
+ )
+
+ @mock.patch(BASE_TRIGGER_CLASSPATH + "OpenSearchServerlessHook",
autospec=True)
+ def test_hook_forwards_aws_configuration(self, mock_hook_class):
+ trigger = OpenSearchServerlessCollectionActiveTrigger(
+ collection_id=self.COLLECTION_ID,
+ aws_conn_id="test_conn",
+ region_name="eu-west-1",
+ verify="/path/to/ca-bundle.crt",
+ botocore_config={"read_timeout": 42},
+ )
+
+ trigger.hook()
+
+ mock_hook_class.assert_called_once_with(
+ aws_conn_id="test_conn",
+ region_name="eu-west-1",
+ verify="/path/to/ca-bundle.crt",
+ config={"read_timeout": 42},
+ )
@pytest.mark.asyncio
@mock.patch.object(OpenSearchServerlessHook, "get_waiter")
