Forward hook config to `OpenSearchServerlessCollectionActiveTrigger`
- Dominant language
- Python
- Stars
- 46.9k
- Forks
- 17.8k
- Avg merge
- 2d 9h
- Merged PRs (30d)
- 472
Description
## Description
Sub-issue of #72144, covering one of the three defer sites set aside for the Contributor's Workshop in [this comment](https://github.com/apache/airflow/issues/72144#issuecomment-5443098667).
| | |
| --- | --- |
| Call site | `providers/amazon/src/airflow/providers/amazon/aws/sensors/opensearch_serverless.py`, in `OpenSearchServerlessCollectionActiveSensor.execute` |
| Trigger | `OpenSearchServerlessCollectionActiveTrigger` in `providers/amazon/src/airflow/providers/amazon/aws/triggers/opensearch_serverless.py` |
| Shape | Full migration, the trigger currently discards the configuration |
This is the largest of the three workshop tasks and the one that shows the whole pattern, so it is worth taking after the other two rather than first.
`OpenSearchServerlessCollectionActiveSensor` is an `AwsBaseSensor`, so it always carries `region_name`, `verify` and `botocore_config`. When it defers it passes none of them, and unlike the batch and SageMaker cases the trigger could not use them anyway.
### What needs to change?
Two things are in the way here, which is what makes this the teaching example.
The trigger's `__init__` is closed. It has no `**kwargs`, so the three parameters cannot even be handed to `AwsBaseWaiterTrigger`:
```python
def __init__(
self,
*,
collection_id: str | None = None,
collection_name: str | None = None,
waiter_delay: int = 60,
waiter_max_attempts: int = 20,
aws_conn_id: str | None = None,
) -> None:
```
And its `hook()` rebuilds from the connection alone, so it would discard the configuration even if the signature accepted it:
```python
def hook(self) -> AwsGenericHook:
return OpenSearchServerlessHook(aws_conn_id=self.aws_conn_id)
```
So the work is to widen the signature so the parameters reach the base class, stop the bespoke `hook()` from throwing the configuration away, and then fix the call site to send all three.
For the `hook()` half there are two possible shapes, and which one applies depends on whether #72171 has landed:
- **After #72171 lands**, that PR gives `AwsBaseWaiterTrigger` a default `hook()` driven by a `hook_class` attribute. The bespoke `hook()` can then be deleted outright and replaced by setting the hook class on the trigger, which is the intended end state described in the parent issue.
- **Before it lands**, `AwsBaseWaiterTrigger.hook()` is still an `@abstractmethod` and the default does not exist. The equivalent fix is to keep `hook()` but have it pass all three through, matching the shape already used by the triggers migrated in #72098, for example `SsmRunCommandTrigger`:
```python
def hook(self) -> AwsGenericHook:
return SsmHook(
aws_conn_id=self.aws_conn_id,
region_name=self.region_name,
verify=self.verify,
config=self.botocore_config,
)
```
Either shape fixes the bug. Coordinating with #72171 first is worth doing so the trigger does not get written one way and immediately rewritten the other.
One detail not to lose while widening the signature: the existing `exactly_one(collection_id is None, collection_name is None)` guard, and the fact that `waiter_args` and `return_key` both branch on which of the two was supplied.
### How to verify it
The parent issue's reproduction applies: construct the sensor with a non-default `region_name`, `verify` and `botocore_config`, trigger the defer, and assert those values survive into the serialized payload.
```python
with pytest.raises(TaskDeferred) as deferred:
sensor.execute(None)
assert deferred.value.trigger.serialize()[1]
```
`AwsBaseWaiterTrigger.serialize` prunes empty values, so use values that are actually distinguishable from the defaults, and assert on the serialized dict rather than on attributes of the trigger object. Because this one changes the trigger's constructor, it is also worth asserting that the trigger still round trips through serialize and deserialize with both the `collection_id` and the `collection_name` form.
## Definition of Done
1. Widen `OpenSearchServerlessCollectionActiveTrigger.__init__` so `region_name`, `verify` and `botocore_config` reach `AwsBaseWaiterTrigger`.
2. Stop `hook()` discarding them, using whichever of the two shapes above matches the state of `main` at the time.
3. Pass all three through at the defer site in `sensors/opensearch_serverless.py`.
4. Add or extend unit tests covering the serialized payload and the existing collection id versus collection name branching.
5. Once #72171 has landed, remove the `("sensors/opensearch_serverless.py", "OpenSearchServerlessCollectionActiveTrigger")` entry from the `PENDING_MIGRATION` allowlist that PR introduces. That allowlist does not exist on `main` yet, so this step only applies after it merges. The invariant test asserts each entry is still needed, so a stale line fails the suite.
6. These should pass:
```bash
breeze testing providers-tests providers/amazon/tests/unit/amazon/aws/sensors/test_opensearch_serverless.py
breeze testing providers-tests providers/amazon/tests/unit/amazon/aws/triggers/test_opensearch_serverless.py
```
---
Drafted-by: Claude Code (Opus 5); reviewed and edited by @jroachgolf84 before posting
Contributor guide
Research direction
Read providers/amazon/src/airflow/providers/amazon/aws/triggers/opensearch_serverless.py and the sensor defer site in providers/amazon/src/airflow/providers/amazon/aws/sensors/opensearch_serverless.py, checking #72171 first. Run the two named unit-test files and use the parent issue's reproduction to verify serialized configuration and both collection_id and collection_name forms. Done means configuration survives deferral and the existing branching and round trips remain covered.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- aws, python
- Domain
- backend, cloud
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100