You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "vandonr-amz (via GitHub)" <gi...@apache.org> on 2023/06/29 23:51:42 UTC
[GitHub] [airflow] vandonr-amz opened a new pull request, #32274: WIP: introduce a base class for aws triggers
vandonr-amz opened a new pull request, #32274:
URL: https://github.com/apache/airflow/pull/32274
There has been a lot of triggers added recently, and they all reuse more or less the same code, with some variations.
This leads to many (small) inconsistencies in trigger code, and also in a lot of duplicated code.
We can package all this in a single base class and reuse it as much as possible, which is what I'm doing here.
Adapting all triggers to this base class is also a good way to resolve some of the inconsistencies.
This is still a work in progress, but I'm putting it out there as a marker of "I am working on this" 😄
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ferruzzi commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253588089
##########
airflow/providers/amazon/aws/triggers/ecs.py:
##########
@@ -22,68 +22,89 @@
from botocore.exceptions import ClientError, WaiterError
+from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
from airflow.providers.amazon.aws.hooks.ecs import EcsHook
from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
+from airflow.providers.amazon.aws.triggers.base_trigger import AwsBaseWaiterTrigger
from airflow.providers.amazon.aws.utils.task_log_fetcher import AwsTaskLogFetcher
-from airflow.providers.amazon.aws.utils.waiter_with_logging import async_wait
from airflow.triggers.base import BaseTrigger, TriggerEvent
-class ClusterWaiterTrigger(BaseTrigger):
+class ClusterActiveTrigger(AwsBaseWaiterTrigger):
"""
- Polls the status of a cluster using a given waiter. Can be used to poll for an active or inactive cluster.
+ Polls the status of a cluster until it's ready.
- :param waiter_name: Name of the waiter to use, for instance 'cluster_active' or 'cluster_inactive'
:param cluster_arn: ARN of the cluster to watch.
:param waiter_delay: The amount of time in seconds to wait between attempts.
:param waiter_max_attempts: The number of times to ping for status.
Will fail after that many unsuccessful attempts.
:param aws_conn_id: The Airflow connection used for AWS credentials.
- :param region: The AWS region where the cluster is located.
+ :param region_name: The AWS region where the cluster is located.
"""
def __init__(
self,
- waiter_name: str,
cluster_arn: str,
- waiter_delay: int | None,
- waiter_max_attempts: int | None,
+ waiter_delay: int,
+ waiter_max_attempts: int,
aws_conn_id: str | None,
- region: str | None,
+ region_name: str | None,
):
- self.cluster_arn = cluster_arn
- self.waiter_name = waiter_name
- self.waiter_delay = waiter_delay if waiter_delay is not None else 15 # written like this to allow 0
- self.attempts = waiter_max_attempts or 999999999
- self.aws_conn_id = aws_conn_id
- self.region = region
+ super().__init__(
+ serialized_fields={"cluster_arn": cluster_arn},
+ waiter_name="cluster_active",
+ waiter_args={"clusters": [cluster_arn]},
+ failure_message="Failure while waiting for cluster to be available",
+ status_message="Cluster is not ready yet",
+ status_queries=["clusters[].status", "failures"],
+ return_key="arn",
+ return_value=cluster_arn,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ aws_conn_id=aws_conn_id,
+ region_name=region_name,
+ )
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "waiter_name": self.waiter_name,
- "cluster_arn": self.cluster_arn,
- "waiter_delay": self.waiter_delay,
- "waiter_max_attempts": self.attempts,
- "aws_conn_id": self.aws_conn_id,
- "region": self.region,
- },
+ def hook(self) -> AwsGenericHook:
+ return EcsHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
+
+
+class ClusterInactiveTrigger(AwsBaseWaiterTrigger):
+ """
+ Polls the status of a cluster until it's ready.
Review Comment:
Possible copypasta, I think this should be "until it is inactive"??
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] syedahsn commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "syedahsn (via GitHub)" <gi...@apache.org>.
syedahsn commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1254852770
##########
airflow/providers/amazon/aws/triggers/emr.py:
##########
@@ -102,274 +102,184 @@ async def run(self):
yield TriggerEvent({"status": "success", "message": "Steps completed", "step_ids": self.step_ids})
-class EmrCreateJobFlowTrigger(BaseTrigger):
+class EmrCreateJobFlowTrigger(AwsBaseWaiterTrigger):
"""
Asynchronously poll the boto3 API and wait for the JobFlow to finish executing.
:param job_flow_id: The id of the job flow to wait for.
- :param poll_interval: The amount of time in seconds to wait between attempts.
- :param max_attempts: The maximum number of attempts to be made.
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
:param aws_conn_id: The Airflow connection used for AWS credentials.
"""
def __init__(
self,
job_flow_id: str,
- poll_interval: int,
- max_attempts: int,
- aws_conn_id: str,
+ poll_interval: int | None = None, # deprecated
+ max_attempts: int | None = None, # deprecated
+ aws_conn_id: str | None = None,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
):
- self.job_flow_id = job_flow_id
- self.poll_interval = poll_interval
- self.max_attempts = max_attempts
- self.aws_conn_id = aws_conn_id
-
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "job_flow_id": self.job_flow_id,
- "poll_interval": str(self.poll_interval),
- "max_attempts": str(self.max_attempts),
- "aws_conn_id": self.aws_conn_id,
- },
+ if poll_interval is not None or max_attempts is not None:
+ warnings.warn(
+ "please use waiter_delay instead of poll_interval "
+ "and waiter_max_attempts instead of max_attempts",
+ AirflowProviderDeprecationWarning,
+ stacklevel=2,
+ )
+ waiter_delay = poll_interval or waiter_delay
+ waiter_max_attempts = max_attempts or waiter_max_attempts
+ super().__init__(
+ serialized_fields={"job_flow_id": job_flow_id},
+ waiter_name="job_flow_waiting",
+ waiter_args={"ClusterId": job_flow_id},
+ failure_message="JobFlow creation failed",
+ status_message="JobFlow creation in progress",
+ status_queries=[
+ "Cluster.Status.State",
+ "Cluster.Status.StateChangeReason",
+ "Cluster.Status.ErrorDetails",
+ ],
+ return_key="job_flow_id",
+ return_value=job_flow_id,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ aws_conn_id=aws_conn_id,
)
- async def run(self):
- self.hook = EmrHook(aws_conn_id=self.aws_conn_id)
- async with self.hook.async_conn as client:
- attempt = 0
- waiter = self.hook.get_waiter("job_flow_waiting", deferrable=True, client=client)
- while attempt < int(self.max_attempts):
- attempt = attempt + 1
- try:
- await waiter.wait(
- ClusterId=self.job_flow_id,
- WaiterConfig=prune_dict(
- {
- "Delay": self.poll_interval,
- "MaxAttempts": 1,
- }
- ),
- )
- break
- except WaiterError as error:
- if "terminal failure" in str(error):
- raise AirflowException(f"JobFlow creation failed: {error}")
- self.log.info(
- "Status of jobflow is %s - %s",
- error.last_response["Cluster"]["Status"]["State"],
- error.last_response["Cluster"]["Status"]["StateChangeReason"],
- )
- await asyncio.sleep(int(self.poll_interval))
- if attempt >= int(self.max_attempts):
- raise AirflowException(f"JobFlow creation failed - max attempts reached: {self.max_attempts}")
- else:
- yield TriggerEvent(
- {
- "status": "success",
- "message": "JobFlow completed successfully",
- "job_flow_id": self.job_flow_id,
- }
- )
+ def hook(self) -> AwsGenericHook:
+ return EmrHook(aws_conn_id=self.aws_conn_id)
-class EmrTerminateJobFlowTrigger(BaseTrigger):
+class EmrTerminateJobFlowTrigger(AwsBaseWaiterTrigger):
"""
Asynchronously poll the boto3 API and wait for the JobFlow to finish terminating.
:param job_flow_id: ID of the EMR Job Flow to terminate
- :param poll_interval: The amount of time in seconds to wait between attempts.
- :param max_attempts: The maximum number of attempts to be made.
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
:param aws_conn_id: The Airflow connection used for AWS credentials.
"""
def __init__(
self,
job_flow_id: str,
- poll_interval: int,
- max_attempts: int,
- aws_conn_id: str,
+ poll_interval: int | None = None, # deprecated
+ max_attempts: int | None = None, # deprecated
+ aws_conn_id: str | None = None,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
):
- self.job_flow_id = job_flow_id
- self.poll_interval = poll_interval
- self.max_attempts = max_attempts
- self.aws_conn_id = aws_conn_id
-
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "job_flow_id": self.job_flow_id,
- "poll_interval": str(self.poll_interval),
- "max_attempts": str(self.max_attempts),
- "aws_conn_id": self.aws_conn_id,
- },
+ if poll_interval is not None or max_attempts is not None:
+ warnings.warn(
+ "please use waiter_delay instead of poll_interval "
+ "and waiter_max_attempts instead of max_attempts",
+ AirflowProviderDeprecationWarning,
+ stacklevel=2,
+ )
+ waiter_delay = poll_interval or waiter_delay
+ waiter_max_attempts = max_attempts or waiter_max_attempts
+ super().__init__(
+ serialized_fields={"job_flow_id": job_flow_id},
+ waiter_name="job_flow_terminated",
+ waiter_args={"ClusterId": job_flow_id},
+ failure_message="JobFlow termination failed",
+ status_message="JobFlow termination in progress",
+ status_queries=[
+ "Cluster.Status.State",
+ "Cluster.Status.StateChangeReason",
+ "Cluster.Status.ErrorDetails",
+ ],
+ return_value=None,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ aws_conn_id=aws_conn_id,
)
- async def run(self):
- self.hook = EmrHook(aws_conn_id=self.aws_conn_id)
- async with self.hook.async_conn as client:
- attempt = 0
- waiter = self.hook.get_waiter("job_flow_terminated", deferrable=True, client=client)
- while attempt < int(self.max_attempts):
- attempt = attempt + 1
- try:
- await waiter.wait(
- ClusterId=self.job_flow_id,
- WaiterConfig=prune_dict(
- {
- "Delay": self.poll_interval,
- "MaxAttempts": 1,
- }
- ),
- )
- break
- except WaiterError as error:
- if "terminal failure" in str(error):
- raise AirflowException(f"JobFlow termination failed: {error}")
- self.log.info(
- "Status of jobflow is %s - %s",
- error.last_response["Cluster"]["Status"]["State"],
- error.last_response["Cluster"]["Status"]["StateChangeReason"],
- )
- await asyncio.sleep(int(self.poll_interval))
- if attempt >= int(self.max_attempts):
- raise AirflowException(f"JobFlow termination failed - max attempts reached: {self.max_attempts}")
- else:
- yield TriggerEvent(
- {
- "status": "success",
- "message": "JobFlow terminated successfully",
- }
- )
+ def hook(self) -> AwsGenericHook:
+ return EmrHook(aws_conn_id=self.aws_conn_id)
-class EmrContainerTrigger(BaseTrigger):
+class EmrContainerTrigger(AwsBaseWaiterTrigger):
"""
Poll for the status of EMR container until reaches terminal state.
:param virtual_cluster_id: Reference Emr cluster id
:param job_id: job_id to check the state
:param aws_conn_id: Reference to AWS connection id
- :param poll_interval: polling period in seconds to check for the status
+ :param waiter_delay: polling period in seconds to check for the status
"""
def __init__(
self,
virtual_cluster_id: str,
job_id: str,
aws_conn_id: str = "aws_default",
- poll_interval: int = 30,
- **kwargs: Any,
+ poll_interval: int | None = None, # deprecated
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 600,
Review Comment:
I think a good compromise is to keep it the same as the default one used by boto. That way, we don't inadvertently mess things up for users who might be relying on a Task to fail after a certain amount of time, because they didn't change the default value
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vandonr-amz commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "vandonr-amz (via GitHub)" <gi...@apache.org>.
vandonr-amz commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253446469
##########
airflow/providers/amazon/aws/triggers/athena.py:
##########
@@ -16,61 +16,43 @@
# under the License.
from __future__ import annotations
-from typing import Any
-
from airflow.providers.amazon.aws.hooks.athena import AthenaHook
-from airflow.providers.amazon.aws.utils.waiter_with_logging import async_wait
-from airflow.triggers.base import BaseTrigger, TriggerEvent
+from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+from airflow.providers.amazon.aws.triggers.base_trigger import AwsBaseWaiterTrigger
-class AthenaTrigger(BaseTrigger):
+class AthenaTrigger(AwsBaseWaiterTrigger):
"""
Trigger for RedshiftCreateClusterOperator.
The trigger will asynchronously poll the boto3 API and wait for the
Redshift cluster to be in the `available` state.
:param query_execution_id: ID of the Athena query execution to watch
- :param poll_interval: The amount of time in seconds to wait between attempts.
- :param max_attempt: The maximum number of attempts to be made.
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
:param aws_conn_id: The Airflow connection used for AWS credentials.
"""
def __init__(
self,
query_execution_id: str,
- poll_interval: int,
- max_attempt: int,
+ waiter_delay: int,
+ waiter_max_attempts: int,
Review Comment:
this trigger was added in #32186 merged on June 27th, last provider release was on June 20th, so this breaking change is OK.
##########
airflow/providers/amazon/aws/triggers/ecs.py:
##########
@@ -22,68 +22,89 @@
from botocore.exceptions import ClientError, WaiterError
+from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
from airflow.providers.amazon.aws.hooks.ecs import EcsHook
from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
+from airflow.providers.amazon.aws.triggers.base_trigger import AwsBaseWaiterTrigger
from airflow.providers.amazon.aws.utils.task_log_fetcher import AwsTaskLogFetcher
-from airflow.providers.amazon.aws.utils.waiter_with_logging import async_wait
from airflow.triggers.base import BaseTrigger, TriggerEvent
-class ClusterWaiterTrigger(BaseTrigger):
+class ClusterActiveTrigger(AwsBaseWaiterTrigger):
Review Comment:
this trigger was added in #31881 merged on June 23rd, last provider release was on June 20th, so this breaking change is OK.
##########
airflow/providers/amazon/aws/triggers/batch.py:
##########
@@ -189,56 +193,78 @@ async def run(self):
)
-class BatchCreateComputeEnvironmentTrigger(BaseTrigger):
+class BatchJobTrigger(AwsBaseWaiterTrigger):
+ """
+ Checks for the status of a submitted job_id to AWS Batch until it reaches a failure or a success state.
+
+ :param job_id: the job ID, to poll for job completion or not
+ :param region_name: AWS region name to use
+ Override the region_name in connection (if provided)
+ :param aws_conn_id: connection id of AWS credentials / region name. If None,
+ credential boto3 strategy will be used
+ :param waiter_delay: polling period in seconds to check for the status of the job
+ :param waiter_max_attempts: The maximum number of attempts to be made.
+ """
+
+ def __init__(
+ self,
+ job_id: str | None,
+ region_name: str | None,
+ aws_conn_id: str | None = "aws_default",
+ waiter_delay: int = 5,
+ waiter_max_attempts: int = 720,
+ ):
+ super().__init__(
+ serialized_fields={"job_id": job_id},
+ waiter_name="batch_job_complete",
+ waiter_args={"jobs": [job_id]},
+ failure_message=f"Failure while running batch job {job_id}",
+ status_message=f"Batch job {job_id} not ready yet",
+ status_queries=["jobs[].status", "computeEnvironments[].statusReason"],
+ return_key="job_id",
+ return_value=job_id,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ aws_conn_id=aws_conn_id,
+ region_name=region_name,
+ )
+
+ def hook(self) -> AwsGenericHook:
+ return BatchClientHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
+
+
+class BatchCreateComputeEnvironmentTrigger(AwsBaseWaiterTrigger):
"""
Asynchronously poll the boto3 API and wait for the compute environment to be ready.
- :param job_id: A unique identifier for the cluster.
- :param max_retries: The maximum number of attempts to be made.
+ :param compute_env_arn: The ARN of the compute env.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
:param aws_conn_id: The Airflow connection used for AWS credentials.
:param region_name: region name to use in AWS Hook
- :param poll_interval: The amount of time in seconds to wait between attempts.
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
"""
def __init__(
self,
- compute_env_arn: str | None = None,
- poll_interval: int = 30,
- max_retries: int = 10,
+ compute_env_arn: str,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 10,
Review Comment:
this trigger was added in #32036 merged on June 27th, last provider release was on June 20th, so this breaking change is OK.
##########
tests/providers/amazon/aws/triggers/test_athena.py:
##########
@@ -16,38 +16,20 @@
# under the License.
from __future__ import annotations
-from unittest import mock
-from unittest.mock import AsyncMock
-
-import pytest
-from botocore.exceptions import WaiterError
-
-from airflow.providers.amazon.aws.hooks.athena import AthenaHook
from airflow.providers.amazon.aws.triggers.athena import AthenaTrigger
class TestAthenaTrigger:
- @pytest.mark.asyncio
- @mock.patch.object(AthenaHook, "get_waiter")
- @mock.patch.object(AthenaHook, "async_conn") # LatestBoto step of CI fails without this
- async def test_run_with_error(self, conn_mock, waiter_mock):
- waiter_mock.side_effect = WaiterError("name", "reason", {})
-
- trigger = AthenaTrigger("query_id", 0, 5, None)
+ def test_serialize_recreate(self):
Review Comment:
This is probably best viewed in side-by-side diff. I removed existing tests because there is no logic anymore in individual triggers.
Instead, I'm testing the only thing that can be broken, which is the serialization/deserialization.
To do that, I do a cycle of serialize-deserialize-reserialize and I compare the serialized data. It'd probably be better to compare the instances, but at least comparing the serialized output can be done with a simple `==`
I copy-pasted the same test for all triggers inheriting from the base, because I think it's better to have it in their respective files ? It could also be a parametrized test with many cases in test_base_trigger to avoid the code duplication, open to hear opinion about it.
##########
airflow/providers/amazon/aws/triggers/eks.py:
##########
@@ -115,59 +90,37 @@ def __init__(
waiter_max_attempts: int,
aws_conn_id: str,
region: str | None = None,
+ region_name: str | None = None,
):
- self.cluster_name = cluster_name
- self.fargate_profile_name = fargate_profile_name
- self.waiter_delay = waiter_delay
- self.waiter_max_attempts = waiter_max_attempts
- self.aws_conn_id = aws_conn_id
- self.region = region
-
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "cluster_name": self.cluster_name,
- "fargate_profile_name": self.fargate_profile_name,
- "waiter_delay": str(self.waiter_delay),
- "waiter_max_attempts": str(self.waiter_max_attempts),
- "aws_conn_id": self.aws_conn_id,
- "region": self.region,
- },
+ if region is not None:
+ warnings.warn(
+ "please use region_name param instead of region",
+ AirflowProviderDeprecationWarning,
+ stacklevel=2,
+ )
+ region_name = region
+
+ super().__init__(
+ serialized_fields={"cluster_name": cluster_name, "fargate_profile_name": fargate_profile_name},
+ waiter_name="fargate_profile_deleted",
+ waiter_args={"clusterName": cluster_name, "fargateProfileName": fargate_profile_name},
+ failure_message="Failure while deleting Fargate profile",
+ status_message="Fargate profile not deleted yet",
+ status_queries=["fargateProfile.status"],
+ return_value=None,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ aws_conn_id=aws_conn_id,
+ region_name=region_name,
)
- async def run(self):
- self.hook = EksHook(aws_conn_id=self.aws_conn_id, region_name=self.region)
- async with self.hook.async_conn as client:
- attempt = 0
- waiter = client.get_waiter("fargate_profile_deleted")
- while attempt < int(self.waiter_max_attempts):
- attempt += 1
- try:
- await waiter.wait(
- clusterName=self.cluster_name,
- fargateProfileName=self.fargate_profile_name,
- WaiterConfig={"Delay": int(self.waiter_delay), "MaxAttempts": 1},
- )
- break
- except WaiterError as error:
- if "terminal failure" in str(error):
- raise AirflowException(f"Delete Fargate Profile failed: {error}")
- self.log.info(
- "Status of fargate profile is %s", error.last_response["fargateProfile"]["status"]
- )
- await asyncio.sleep(int(self.waiter_delay))
- if attempt >= int(self.waiter_max_attempts):
- raise AirflowException(
- f"Delete Fargate Profile failed - max attempts reached: {self.waiter_max_attempts}"
- )
- else:
- yield TriggerEvent({"status": "success", "message": "Fargate Profile Deleted"})
+ def hook(self) -> AwsGenericHook:
+ return EksHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
-class EksNodegroupTrigger(BaseTrigger):
+class EksCreateNodegroupTrigger(AwsBaseWaiterTrigger):
Review Comment:
this trigger was added in #32165 merged on June 26th, last provider release was on June 20th, so this breaking change is OK.
##########
airflow/providers/amazon/aws/hooks/athena.py:
##########
@@ -253,7 +253,7 @@ def poll_query_status(
try:
wait(
waiter=self.get_waiter("query_complete"),
- waiter_delay=sleep_time or self.sleep_time,
+ waiter_delay=self.sleep_time if sleep_time is None else sleep_time,
Review Comment:
this is a somewhat unrelated fix that allows specifying a sleep time of 0 in unit tests. Without this, athena unit tests were taking 30s each
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vandonr-amz commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "vandonr-amz (via GitHub)" <gi...@apache.org>.
vandonr-amz commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253484051
##########
airflow/providers/amazon/aws/triggers/base_trigger.py:
##########
@@ -0,0 +1,127 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import Any, AsyncIterator
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+from airflow.providers.amazon.aws.utils.waiter_with_logging import async_wait
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+
+
+class AwsBaseWaiterTrigger(BaseTrigger):
+ """
+ Base class for all AWS Triggers that follow the "standard" model of just waiting on a waiter.
+
+ Subclasses need to implement the hook() method.
+
+ :param serialized_fields: Fields that are specific to the subclass trigger and need to be serialized
+ to be passed to the __init__ method on deserialization.
+ format: {<parameter_name>: <parameter_value>}
+
+ :param waiter_name: The name of the (possibly custom) boto waiter to use.
+
+ :param waiter_args: The arguments to pass to the waiter.
+ :param failure_message: The message to log if a failure state is reached.
+ :param status_message: The message logged when printing the status of the service.
+ :param status_queries: A list containing the JMESPath queries to retrieve status information from
+ the waiter response. See https://jmespath.org/tutorial.html
+
+ :param return_key: The key to use for the return_value in the TriggerEvent this emits on success.
+ :param return_value: A value that'll be returned in the return_key field of the TriggerEvent.
+
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ :param region_name: The AWS region where the resources to watch are.
+ """
+
+ def __init__(
+ self,
+ *,
+ serialized_fields: dict[str, Any],
+ waiter_name: str,
+ waiter_args: dict[str, Any],
+ failure_message: str,
+ status_message: str,
+ status_queries: list[str],
+ return_key: str = "value",
+ return_value: Any | None,
Review Comment:
hmm yeah probably. I'll remove it from here and add it to the doc so that it's more obvious that it's an accepted value.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ferruzzi commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253575212
##########
airflow/providers/amazon/aws/operators/ecs.py:
##########
@@ -139,13 +143,12 @@ def execute(self, context: Context):
self.log.info("Cluster %r in state: %r.", self.cluster_name, cluster_state)
elif self.deferrable:
self.defer(
- trigger=ClusterWaiterTrigger(
- waiter_name="cluster_active",
+ trigger=ClusterActiveTrigger(
cluster_arn=cluster_details["clusterArn"],
Review Comment:
I'll make the same suggestion I made in Syed's PR... feel free to ignore it, but I see a lot of repetition in the calls to the Triggers and the Waiters. It may be convenient to fold the waiter parameters into a new self.waiter_config dict in the Operator's init and just unpack it in the Trigger and Waiter calls, something like:
```
self.waiter_config = {
cluster_arn=cluster_details["clusterArn"],
waiter_delay=self.waiter_delay,
waiter_max_attempts=self.waiter_max_attempts,
aws_conn_id=self.aws_conn_id,
region=self.region,
}
```
then all of the waiters and triggers can reuse that cleanly like:
```trigger=ClusterActiveTrigger(**self.waiter_config)```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ferruzzi commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253549826
##########
airflow/providers/amazon/aws/hooks/athena.py:
##########
@@ -253,7 +253,7 @@ def poll_query_status(
try:
wait(
waiter=self.get_waiter("query_complete"),
- waiter_delay=sleep_time or self.sleep_time,
+ waiter_delay=self.sleep_time if sleep_time is None else sleep_time,
Review Comment:
I'll leave this up to you, but it seems easier to read as
```suggestion
waiter_delay=sleep_time if sleep_time else self.sleep_time,
```
which should also boil down further to just
```suggestion
waiter_delay=sleep_time or self.sleep_time,
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vandonr-amz commented on pull request #32274: Introduce a base class for aws triggers
Posted by "vandonr-amz (via GitHub)" <gi...@apache.org>.
vandonr-amz commented on PR #32274:
URL: https://github.com/apache/airflow/pull/32274#issuecomment-1624303874
this can be merged as is only if https://github.com/apache/airflow/issues/32389 does not proceed.
If it does, I'll have to change the parts that are "breaking" to surround them with deprecation warnings.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ferruzzi commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253590001
##########
airflow/providers/amazon/aws/triggers/ecs.py:
##########
@@ -22,68 +22,89 @@
from botocore.exceptions import ClientError, WaiterError
+from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
from airflow.providers.amazon.aws.hooks.ecs import EcsHook
from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
+from airflow.providers.amazon.aws.triggers.base_trigger import AwsBaseWaiterTrigger
from airflow.providers.amazon.aws.utils.task_log_fetcher import AwsTaskLogFetcher
-from airflow.providers.amazon.aws.utils.waiter_with_logging import async_wait
from airflow.triggers.base import BaseTrigger, TriggerEvent
-class ClusterWaiterTrigger(BaseTrigger):
+class ClusterActiveTrigger(AwsBaseWaiterTrigger):
"""
- Polls the status of a cluster using a given waiter. Can be used to poll for an active or inactive cluster.
+ Polls the status of a cluster until it's ready.
- :param waiter_name: Name of the waiter to use, for instance 'cluster_active' or 'cluster_inactive'
:param cluster_arn: ARN of the cluster to watch.
:param waiter_delay: The amount of time in seconds to wait between attempts.
:param waiter_max_attempts: The number of times to ping for status.
Will fail after that many unsuccessful attempts.
:param aws_conn_id: The Airflow connection used for AWS credentials.
- :param region: The AWS region where the cluster is located.
+ :param region_name: The AWS region where the cluster is located.
"""
def __init__(
self,
- waiter_name: str,
cluster_arn: str,
- waiter_delay: int | None,
- waiter_max_attempts: int | None,
+ waiter_delay: int,
+ waiter_max_attempts: int,
aws_conn_id: str | None,
- region: str | None,
+ region_name: str | None,
):
- self.cluster_arn = cluster_arn
- self.waiter_name = waiter_name
- self.waiter_delay = waiter_delay if waiter_delay is not None else 15 # written like this to allow 0
- self.attempts = waiter_max_attempts or 999999999
- self.aws_conn_id = aws_conn_id
- self.region = region
+ super().__init__(
+ serialized_fields={"cluster_arn": cluster_arn},
+ waiter_name="cluster_active",
+ waiter_args={"clusters": [cluster_arn]},
+ failure_message="Failure while waiting for cluster to be available",
+ status_message="Cluster is not ready yet",
+ status_queries=["clusters[].status", "failures"],
+ return_key="arn",
+ return_value=cluster_arn,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ aws_conn_id=aws_conn_id,
+ region_name=region_name,
+ )
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "waiter_name": self.waiter_name,
- "cluster_arn": self.cluster_arn,
- "waiter_delay": self.waiter_delay,
- "waiter_max_attempts": self.attempts,
- "aws_conn_id": self.aws_conn_id,
- "region": self.region,
- },
+ def hook(self) -> AwsGenericHook:
+ return EcsHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
+
+
+class ClusterInactiveTrigger(AwsBaseWaiterTrigger):
Review Comment:
I may be over-complicating this, but it looooks like the only difference between ClusterActiveTrigger and ClusterInactiveTrigger is the waiter_name value, in which case, why not have them both inherit a ClusterStatusTrigger which accepts the waiter_name as a param?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ferruzzi commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253613019
##########
airflow/providers/amazon/aws/triggers/base_trigger.py:
##########
@@ -0,0 +1,129 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import Any, AsyncIterator
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+from airflow.providers.amazon.aws.utils.waiter_with_logging import async_wait
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+
+
+class AwsBaseWaiterTrigger(BaseTrigger):
+ """
+ Base class for all AWS Triggers that follow the "standard" model of just waiting on a waiter.
+
+ Subclasses need to implement the hook() method.
+
+ :param serialized_fields: Fields that are specific to the subclass trigger and need to be serialized
+ to be passed to the __init__ method on deserialization.
+ format: {<parameter_name>: <parameter_value>}
+
+ :param waiter_name: The name of the (possibly custom) boto waiter to use.
+
+ :param waiter_args: The arguments to pass to the waiter.
+ :param failure_message: The message to log if a failure state is reached.
+ :param status_message: The message logged when printing the status of the service.
+ :param status_queries: A list containing the JMESPath queries to retrieve status information from
+ the waiter response. See https://jmespath.org/tutorial.html
+
+ :param return_key: The key to use for the return_value in the TriggerEvent this emits on success.
+ Defaults to "value".
+ :param return_value: A value that'll be returned in the return_key field of the TriggerEvent.
+ Set to None if there is nothing to return.
+
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
+ :param aws_conn_id: The Airflow connection used for AWS credentials. To be used to build the hook.
+ :param region_name: The AWS region where the resources to watch are. To be used to build the hook.
+ """
+
+ def __init__(
+ self,
+ *,
+ serialized_fields: dict[str, Any],
+ waiter_name: str,
+ waiter_args: dict[str, Any],
+ failure_message: str,
+ status_message: str,
+ status_queries: list[str],
+ return_key: str = "value",
+ return_value: Any,
+ waiter_delay: int,
+ waiter_max_attempts: int,
+ aws_conn_id: str | None,
+ region_name: str | None = None,
+ ):
+ # parameters that should be hardcoded in the child's implem
+ self.serialized_fields = serialized_fields
+
+ self.waiter_name = waiter_name
+ self.waiter_args = waiter_args
+ self.failure_message = failure_message
+ self.status_message = status_message
+ self.status_queries = status_queries
+
+ self.return_key = return_key
+ self.return_value = return_value
+
+ # parameters that should be passed directly from the child's parameters
+ self.waiter_delay = waiter_delay
+ self.attempts = waiter_max_attempts
+ self.aws_conn_id = aws_conn_id
+ self.region_name = region_name
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ # here we put together the "common" params,
+ # and whatever extras we got from the subclass in serialized_fields
+ params = dict(
+ {
+ "waiter_delay": self.waiter_delay,
+ "waiter_max_attempts": self.attempts,
+ "aws_conn_id": self.aws_conn_id,
+ },
+ **self.serialized_fields,
+ )
+ if self.region_name:
+ # if we serialize the None value from this, it breaks subclasses that don't have it in their ctor.
+ params["region_name"] = self.region_name
+ return (
+ # remember that self is an instance of the subclass here, not of this class.
+ self.__class__.__module__ + "." + self.__class__.__qualname__,
+ params,
+ )
+
+ @abstractmethod
+ def hook(self) -> AwsGenericHook:
+ """Override in subclasses to return the right hook."""
+ ...
Review Comment:
oooohhhh. I'll need to look into that more. Sorry about that.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vandonr-amz commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "vandonr-amz (via GitHub)" <gi...@apache.org>.
vandonr-amz commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1254913020
##########
airflow/providers/amazon/aws/triggers/emr.py:
##########
@@ -102,274 +102,184 @@ async def run(self):
yield TriggerEvent({"status": "success", "message": "Steps completed", "step_ids": self.step_ids})
-class EmrCreateJobFlowTrigger(BaseTrigger):
+class EmrCreateJobFlowTrigger(AwsBaseWaiterTrigger):
"""
Asynchronously poll the boto3 API and wait for the JobFlow to finish executing.
:param job_flow_id: The id of the job flow to wait for.
- :param poll_interval: The amount of time in seconds to wait between attempts.
- :param max_attempts: The maximum number of attempts to be made.
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
:param aws_conn_id: The Airflow connection used for AWS credentials.
"""
def __init__(
self,
job_flow_id: str,
- poll_interval: int,
- max_attempts: int,
- aws_conn_id: str,
+ poll_interval: int | None = None, # deprecated
+ max_attempts: int | None = None, # deprecated
+ aws_conn_id: str | None = None,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
):
- self.job_flow_id = job_flow_id
- self.poll_interval = poll_interval
- self.max_attempts = max_attempts
- self.aws_conn_id = aws_conn_id
-
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "job_flow_id": self.job_flow_id,
- "poll_interval": str(self.poll_interval),
- "max_attempts": str(self.max_attempts),
- "aws_conn_id": self.aws_conn_id,
- },
+ if poll_interval is not None or max_attempts is not None:
+ warnings.warn(
+ "please use waiter_delay instead of poll_interval "
+ "and waiter_max_attempts instead of max_attempts",
+ AirflowProviderDeprecationWarning,
+ stacklevel=2,
+ )
+ waiter_delay = poll_interval or waiter_delay
+ waiter_max_attempts = max_attempts or waiter_max_attempts
+ super().__init__(
+ serialized_fields={"job_flow_id": job_flow_id},
+ waiter_name="job_flow_waiting",
+ waiter_args={"ClusterId": job_flow_id},
+ failure_message="JobFlow creation failed",
+ status_message="JobFlow creation in progress",
+ status_queries=[
+ "Cluster.Status.State",
+ "Cluster.Status.StateChangeReason",
+ "Cluster.Status.ErrorDetails",
+ ],
+ return_key="job_flow_id",
+ return_value=job_flow_id,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ aws_conn_id=aws_conn_id,
)
- async def run(self):
- self.hook = EmrHook(aws_conn_id=self.aws_conn_id)
- async with self.hook.async_conn as client:
- attempt = 0
- waiter = self.hook.get_waiter("job_flow_waiting", deferrable=True, client=client)
- while attempt < int(self.max_attempts):
- attempt = attempt + 1
- try:
- await waiter.wait(
- ClusterId=self.job_flow_id,
- WaiterConfig=prune_dict(
- {
- "Delay": self.poll_interval,
- "MaxAttempts": 1,
- }
- ),
- )
- break
- except WaiterError as error:
- if "terminal failure" in str(error):
- raise AirflowException(f"JobFlow creation failed: {error}")
- self.log.info(
- "Status of jobflow is %s - %s",
- error.last_response["Cluster"]["Status"]["State"],
- error.last_response["Cluster"]["Status"]["StateChangeReason"],
- )
- await asyncio.sleep(int(self.poll_interval))
- if attempt >= int(self.max_attempts):
- raise AirflowException(f"JobFlow creation failed - max attempts reached: {self.max_attempts}")
- else:
- yield TriggerEvent(
- {
- "status": "success",
- "message": "JobFlow completed successfully",
- "job_flow_id": self.job_flow_id,
- }
- )
+ def hook(self) -> AwsGenericHook:
+ return EmrHook(aws_conn_id=self.aws_conn_id)
-class EmrTerminateJobFlowTrigger(BaseTrigger):
+class EmrTerminateJobFlowTrigger(AwsBaseWaiterTrigger):
"""
Asynchronously poll the boto3 API and wait for the JobFlow to finish terminating.
:param job_flow_id: ID of the EMR Job Flow to terminate
- :param poll_interval: The amount of time in seconds to wait between attempts.
- :param max_attempts: The maximum number of attempts to be made.
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
:param aws_conn_id: The Airflow connection used for AWS credentials.
"""
def __init__(
self,
job_flow_id: str,
- poll_interval: int,
- max_attempts: int,
- aws_conn_id: str,
+ poll_interval: int | None = None, # deprecated
+ max_attempts: int | None = None, # deprecated
+ aws_conn_id: str | None = None,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
):
- self.job_flow_id = job_flow_id
- self.poll_interval = poll_interval
- self.max_attempts = max_attempts
- self.aws_conn_id = aws_conn_id
-
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "job_flow_id": self.job_flow_id,
- "poll_interval": str(self.poll_interval),
- "max_attempts": str(self.max_attempts),
- "aws_conn_id": self.aws_conn_id,
- },
+ if poll_interval is not None or max_attempts is not None:
+ warnings.warn(
+ "please use waiter_delay instead of poll_interval "
+ "and waiter_max_attempts instead of max_attempts",
+ AirflowProviderDeprecationWarning,
+ stacklevel=2,
+ )
+ waiter_delay = poll_interval or waiter_delay
+ waiter_max_attempts = max_attempts or waiter_max_attempts
+ super().__init__(
+ serialized_fields={"job_flow_id": job_flow_id},
+ waiter_name="job_flow_terminated",
+ waiter_args={"ClusterId": job_flow_id},
+ failure_message="JobFlow termination failed",
+ status_message="JobFlow termination in progress",
+ status_queries=[
+ "Cluster.Status.State",
+ "Cluster.Status.StateChangeReason",
+ "Cluster.Status.ErrorDetails",
+ ],
+ return_value=None,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ aws_conn_id=aws_conn_id,
)
- async def run(self):
- self.hook = EmrHook(aws_conn_id=self.aws_conn_id)
- async with self.hook.async_conn as client:
- attempt = 0
- waiter = self.hook.get_waiter("job_flow_terminated", deferrable=True, client=client)
- while attempt < int(self.max_attempts):
- attempt = attempt + 1
- try:
- await waiter.wait(
- ClusterId=self.job_flow_id,
- WaiterConfig=prune_dict(
- {
- "Delay": self.poll_interval,
- "MaxAttempts": 1,
- }
- ),
- )
- break
- except WaiterError as error:
- if "terminal failure" in str(error):
- raise AirflowException(f"JobFlow termination failed: {error}")
- self.log.info(
- "Status of jobflow is %s - %s",
- error.last_response["Cluster"]["Status"]["State"],
- error.last_response["Cluster"]["Status"]["StateChangeReason"],
- )
- await asyncio.sleep(int(self.poll_interval))
- if attempt >= int(self.max_attempts):
- raise AirflowException(f"JobFlow termination failed - max attempts reached: {self.max_attempts}")
- else:
- yield TriggerEvent(
- {
- "status": "success",
- "message": "JobFlow terminated successfully",
- }
- )
+ def hook(self) -> AwsGenericHook:
+ return EmrHook(aws_conn_id=self.aws_conn_id)
-class EmrContainerTrigger(BaseTrigger):
+class EmrContainerTrigger(AwsBaseWaiterTrigger):
"""
Poll for the status of EMR container until reaches terminal state.
:param virtual_cluster_id: Reference Emr cluster id
:param job_id: job_id to check the state
:param aws_conn_id: Reference to AWS connection id
- :param poll_interval: polling period in seconds to check for the status
+ :param waiter_delay: polling period in seconds to check for the status
"""
def __init__(
self,
virtual_cluster_id: str,
job_id: str,
aws_conn_id: str = "aws_default",
- poll_interval: int = 30,
- **kwargs: Any,
+ poll_interval: int | None = None, # deprecated
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 600,
Review Comment:
ok, but in this particular case, the existing behavior was to wait forever, so....
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ferruzzi commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253597610
##########
tests/providers/amazon/aws/triggers/test_batch.py:
##########
@@ -39,186 +31,34 @@
pytest.importorskip("aiobotocore")
-class TestBatchOperatorTrigger:
- def test_batch_operator_trigger_serialize(self):
- batch_trigger = BatchOperatorTrigger(
- job_id=BATCH_JOB_ID,
- poll_interval=POLL_INTERVAL,
- max_retries=MAX_ATTEMPT,
- aws_conn_id=AWS_CONN_ID,
- region_name=AWS_REGION,
- )
- class_path, args = batch_trigger.serialize()
- assert class_path == "airflow.providers.amazon.aws.triggers.batch.BatchOperatorTrigger"
- assert args["job_id"] == BATCH_JOB_ID
- assert args["poll_interval"] == POLL_INTERVAL
- assert args["max_retries"] == MAX_ATTEMPT
- assert args["aws_conn_id"] == AWS_CONN_ID
- assert args["region_name"] == AWS_REGION
-
- @pytest.mark.asyncio
- @mock.patch("airflow.providers.amazon.aws.hooks.batch_client.BatchClientHook.get_waiter")
- @mock.patch("airflow.providers.amazon.aws.hooks.batch_client.BatchClientHook.async_conn")
- async def test_batch_job_trigger_run(self, mock_async_conn, mock_get_waiter):
- the_mock = mock.MagicMock()
- mock_async_conn.__aenter__.return_value = the_mock
-
- mock_get_waiter().wait = AsyncMock()
-
- batch_trigger = BatchOperatorTrigger(
- job_id=BATCH_JOB_ID,
- poll_interval=POLL_INTERVAL,
- max_retries=MAX_ATTEMPT,
- aws_conn_id=AWS_CONN_ID,
- region_name=AWS_REGION,
- )
-
- generator = batch_trigger.run()
- response = await generator.asend(None)
-
- assert response == TriggerEvent({"status": "success", "job_id": BATCH_JOB_ID})
-
-
-class TestBatchSensorTrigger:
- TRIGGER = BatchSensorTrigger(
- job_id=BATCH_JOB_ID,
- region_name=AWS_REGION,
- aws_conn_id=AWS_CONN_ID,
- poke_interval=POLL_INTERVAL,
+class TestBatchTrigger:
+ @pytest.mark.parametrize(
Review Comment:
:chefkiss:
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ferruzzi commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253596830
##########
tests/providers/amazon/aws/triggers/test_base_trigger.py:
##########
@@ -0,0 +1,89 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from unittest import mock
+from unittest.mock import MagicMock
+
+import pytest
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+from airflow.providers.amazon.aws.triggers.base_trigger import AwsBaseWaiterTrigger
+from airflow.triggers.base import TriggerEvent
+
+
+class TestImplem(AwsBaseWaiterTrigger):
+ """A dumb implementation that allows instantiation for tests."""
Review Comment:
```suggestion
"""An empty implementation that allows instantiation for tests."""
```
Would fail inclusive language check.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] syedahsn commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "syedahsn (via GitHub)" <gi...@apache.org>.
syedahsn commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253633598
##########
airflow/providers/amazon/aws/triggers/eks.py:
##########
@@ -115,59 +90,37 @@ def __init__(
waiter_max_attempts: int,
aws_conn_id: str,
region: str | None = None,
+ region_name: str | None = None,
):
- self.cluster_name = cluster_name
- self.fargate_profile_name = fargate_profile_name
- self.waiter_delay = waiter_delay
- self.waiter_max_attempts = waiter_max_attempts
- self.aws_conn_id = aws_conn_id
- self.region = region
-
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "cluster_name": self.cluster_name,
- "fargate_profile_name": self.fargate_profile_name,
- "waiter_delay": str(self.waiter_delay),
- "waiter_max_attempts": str(self.waiter_max_attempts),
- "aws_conn_id": self.aws_conn_id,
- "region": self.region,
- },
+ if region is not None:
+ warnings.warn(
+ "please use region_name param instead of region",
+ AirflowProviderDeprecationWarning,
+ stacklevel=2,
+ )
+ region_name = region
+
+ super().__init__(
+ serialized_fields={"cluster_name": cluster_name, "fargate_profile_name": fargate_profile_name},
+ waiter_name="fargate_profile_deleted",
+ waiter_args={"clusterName": cluster_name, "fargateProfileName": fargate_profile_name},
+ failure_message="Failure while deleting Fargate profile",
+ status_message="Fargate profile not deleted yet",
+ status_queries=["fargateProfile.status"],
+ return_value=None,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ aws_conn_id=aws_conn_id,
+ region_name=region_name,
)
- async def run(self):
- self.hook = EksHook(aws_conn_id=self.aws_conn_id, region_name=self.region)
- async with self.hook.async_conn as client:
- attempt = 0
- waiter = client.get_waiter("fargate_profile_deleted")
- while attempt < int(self.waiter_max_attempts):
- attempt += 1
- try:
- await waiter.wait(
- clusterName=self.cluster_name,
- fargateProfileName=self.fargate_profile_name,
- WaiterConfig={"Delay": int(self.waiter_delay), "MaxAttempts": 1},
- )
- break
- except WaiterError as error:
- if "terminal failure" in str(error):
- raise AirflowException(f"Delete Fargate Profile failed: {error}")
- self.log.info(
- "Status of fargate profile is %s", error.last_response["fargateProfile"]["status"]
- )
- await asyncio.sleep(int(self.waiter_delay))
- if attempt >= int(self.waiter_max_attempts):
- raise AirflowException(
- f"Delete Fargate Profile failed - max attempts reached: {self.waiter_max_attempts}"
- )
- else:
- yield TriggerEvent({"status": "success", "message": "Fargate Profile Deleted"})
+ def hook(self) -> AwsGenericHook:
+ return EksHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
-class EksNodegroupTrigger(BaseTrigger):
+class EksCreateNodegroupTrigger(AwsBaseWaiterTrigger):
Review Comment:
Ah, I ended up adopting your method haha. Alright, so just to be clear, I am going to be creating a separate Trigger for each operator again.
>Also, given the lower footprint of triggers created this way, I think it's ok to have specific triggers for each thing.
Definitely. Big +1 for reducing the repeated code.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ferruzzi commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253594423
##########
airflow/providers/amazon/aws/triggers/emr.py:
##########
@@ -102,274 +102,184 @@ async def run(self):
yield TriggerEvent({"status": "success", "message": "Steps completed", "step_ids": self.step_ids})
-class EmrCreateJobFlowTrigger(BaseTrigger):
+class EmrCreateJobFlowTrigger(AwsBaseWaiterTrigger):
"""
Asynchronously poll the boto3 API and wait for the JobFlow to finish executing.
:param job_flow_id: The id of the job flow to wait for.
- :param poll_interval: The amount of time in seconds to wait between attempts.
- :param max_attempts: The maximum number of attempts to be made.
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
:param aws_conn_id: The Airflow connection used for AWS credentials.
"""
def __init__(
self,
job_flow_id: str,
- poll_interval: int,
- max_attempts: int,
- aws_conn_id: str,
+ poll_interval: int | None = None, # deprecated
+ max_attempts: int | None = None, # deprecated
+ aws_conn_id: str | None = None,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
):
- self.job_flow_id = job_flow_id
- self.poll_interval = poll_interval
- self.max_attempts = max_attempts
- self.aws_conn_id = aws_conn_id
-
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "job_flow_id": self.job_flow_id,
- "poll_interval": str(self.poll_interval),
- "max_attempts": str(self.max_attempts),
- "aws_conn_id": self.aws_conn_id,
- },
+ if poll_interval is not None or max_attempts is not None:
Review Comment:
Here and elsewhere:
```suggestion
if poll_interval or max_attempts:
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vincbeck commented on pull request #32274: Introduce a base class for aws triggers
Posted by "vincbeck (via GitHub)" <gi...@apache.org>.
vincbeck commented on PR #32274:
URL: https://github.com/apache/airflow/pull/32274#issuecomment-1625897151
Amazon provider package is excluded from RC2, therefore we can merge it. See #32389
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] eladkal commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253525725
##########
airflow/providers/amazon/provider.yaml:
##########
@@ -516,6 +516,9 @@ hooks:
- airflow.providers.amazon.aws.hooks.appflow
triggers:
+ - integration-name: Amazon Web Services
+ python-modules:
+ - airflow.providers.amazon.aws.triggers.base_trigger
Review Comment:
When it comes to naming things my philosophy is the shorter - the better (as long as not compromising name being self explained)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] syedahsn commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "syedahsn (via GitHub)" <gi...@apache.org>.
syedahsn commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253635881
##########
airflow/providers/amazon/aws/triggers/emr.py:
##########
@@ -102,274 +102,184 @@ async def run(self):
yield TriggerEvent({"status": "success", "message": "Steps completed", "step_ids": self.step_ids})
-class EmrCreateJobFlowTrigger(BaseTrigger):
+class EmrCreateJobFlowTrigger(AwsBaseWaiterTrigger):
"""
Asynchronously poll the boto3 API and wait for the JobFlow to finish executing.
:param job_flow_id: The id of the job flow to wait for.
- :param poll_interval: The amount of time in seconds to wait between attempts.
- :param max_attempts: The maximum number of attempts to be made.
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
:param aws_conn_id: The Airflow connection used for AWS credentials.
"""
def __init__(
self,
job_flow_id: str,
- poll_interval: int,
- max_attempts: int,
- aws_conn_id: str,
+ poll_interval: int | None = None, # deprecated
+ max_attempts: int | None = None, # deprecated
+ aws_conn_id: str | None = None,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
):
- self.job_flow_id = job_flow_id
- self.poll_interval = poll_interval
- self.max_attempts = max_attempts
- self.aws_conn_id = aws_conn_id
-
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "job_flow_id": self.job_flow_id,
- "poll_interval": str(self.poll_interval),
- "max_attempts": str(self.max_attempts),
- "aws_conn_id": self.aws_conn_id,
- },
+ if poll_interval is not None or max_attempts is not None:
+ warnings.warn(
+ "please use waiter_delay instead of poll_interval "
+ "and waiter_max_attempts instead of max_attempts",
+ AirflowProviderDeprecationWarning,
+ stacklevel=2,
+ )
+ waiter_delay = poll_interval or waiter_delay
+ waiter_max_attempts = max_attempts or waiter_max_attempts
+ super().__init__(
+ serialized_fields={"job_flow_id": job_flow_id},
+ waiter_name="job_flow_waiting",
+ waiter_args={"ClusterId": job_flow_id},
+ failure_message="JobFlow creation failed",
+ status_message="JobFlow creation in progress",
+ status_queries=[
+ "Cluster.Status.State",
+ "Cluster.Status.StateChangeReason",
+ "Cluster.Status.ErrorDetails",
+ ],
+ return_key="job_flow_id",
+ return_value=job_flow_id,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ aws_conn_id=aws_conn_id,
)
- async def run(self):
- self.hook = EmrHook(aws_conn_id=self.aws_conn_id)
- async with self.hook.async_conn as client:
- attempt = 0
- waiter = self.hook.get_waiter("job_flow_waiting", deferrable=True, client=client)
- while attempt < int(self.max_attempts):
- attempt = attempt + 1
- try:
- await waiter.wait(
- ClusterId=self.job_flow_id,
- WaiterConfig=prune_dict(
- {
- "Delay": self.poll_interval,
- "MaxAttempts": 1,
- }
- ),
- )
- break
- except WaiterError as error:
- if "terminal failure" in str(error):
- raise AirflowException(f"JobFlow creation failed: {error}")
- self.log.info(
- "Status of jobflow is %s - %s",
- error.last_response["Cluster"]["Status"]["State"],
- error.last_response["Cluster"]["Status"]["StateChangeReason"],
- )
- await asyncio.sleep(int(self.poll_interval))
- if attempt >= int(self.max_attempts):
- raise AirflowException(f"JobFlow creation failed - max attempts reached: {self.max_attempts}")
- else:
- yield TriggerEvent(
- {
- "status": "success",
- "message": "JobFlow completed successfully",
- "job_flow_id": self.job_flow_id,
- }
- )
+ def hook(self) -> AwsGenericHook:
+ return EmrHook(aws_conn_id=self.aws_conn_id)
-class EmrTerminateJobFlowTrigger(BaseTrigger):
+class EmrTerminateJobFlowTrigger(AwsBaseWaiterTrigger):
"""
Asynchronously poll the boto3 API and wait for the JobFlow to finish terminating.
:param job_flow_id: ID of the EMR Job Flow to terminate
- :param poll_interval: The amount of time in seconds to wait between attempts.
- :param max_attempts: The maximum number of attempts to be made.
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
:param aws_conn_id: The Airflow connection used for AWS credentials.
"""
def __init__(
self,
job_flow_id: str,
- poll_interval: int,
- max_attempts: int,
- aws_conn_id: str,
+ poll_interval: int | None = None, # deprecated
+ max_attempts: int | None = None, # deprecated
+ aws_conn_id: str | None = None,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
):
- self.job_flow_id = job_flow_id
- self.poll_interval = poll_interval
- self.max_attempts = max_attempts
- self.aws_conn_id = aws_conn_id
-
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "job_flow_id": self.job_flow_id,
- "poll_interval": str(self.poll_interval),
- "max_attempts": str(self.max_attempts),
- "aws_conn_id": self.aws_conn_id,
- },
+ if poll_interval is not None or max_attempts is not None:
+ warnings.warn(
+ "please use waiter_delay instead of poll_interval "
+ "and waiter_max_attempts instead of max_attempts",
+ AirflowProviderDeprecationWarning,
+ stacklevel=2,
+ )
+ waiter_delay = poll_interval or waiter_delay
+ waiter_max_attempts = max_attempts or waiter_max_attempts
+ super().__init__(
+ serialized_fields={"job_flow_id": job_flow_id},
+ waiter_name="job_flow_terminated",
+ waiter_args={"ClusterId": job_flow_id},
+ failure_message="JobFlow termination failed",
+ status_message="JobFlow termination in progress",
+ status_queries=[
+ "Cluster.Status.State",
+ "Cluster.Status.StateChangeReason",
+ "Cluster.Status.ErrorDetails",
+ ],
+ return_value=None,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ aws_conn_id=aws_conn_id,
)
- async def run(self):
- self.hook = EmrHook(aws_conn_id=self.aws_conn_id)
- async with self.hook.async_conn as client:
- attempt = 0
- waiter = self.hook.get_waiter("job_flow_terminated", deferrable=True, client=client)
- while attempt < int(self.max_attempts):
- attempt = attempt + 1
- try:
- await waiter.wait(
- ClusterId=self.job_flow_id,
- WaiterConfig=prune_dict(
- {
- "Delay": self.poll_interval,
- "MaxAttempts": 1,
- }
- ),
- )
- break
- except WaiterError as error:
- if "terminal failure" in str(error):
- raise AirflowException(f"JobFlow termination failed: {error}")
- self.log.info(
- "Status of jobflow is %s - %s",
- error.last_response["Cluster"]["Status"]["State"],
- error.last_response["Cluster"]["Status"]["StateChangeReason"],
- )
- await asyncio.sleep(int(self.poll_interval))
- if attempt >= int(self.max_attempts):
- raise AirflowException(f"JobFlow termination failed - max attempts reached: {self.max_attempts}")
- else:
- yield TriggerEvent(
- {
- "status": "success",
- "message": "JobFlow terminated successfully",
- }
- )
+ def hook(self) -> AwsGenericHook:
+ return EmrHook(aws_conn_id=self.aws_conn_id)
-class EmrContainerTrigger(BaseTrigger):
+class EmrContainerTrigger(AwsBaseWaiterTrigger):
"""
Poll for the status of EMR container until reaches terminal state.
:param virtual_cluster_id: Reference Emr cluster id
:param job_id: job_id to check the state
:param aws_conn_id: Reference to AWS connection id
- :param poll_interval: polling period in seconds to check for the status
+ :param waiter_delay: polling period in seconds to check for the status
"""
def __init__(
self,
virtual_cluster_id: str,
job_id: str,
aws_conn_id: str = "aws_default",
- poll_interval: int = 30,
- **kwargs: Any,
+ poll_interval: int | None = None, # deprecated
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 600,
Review Comment:
is this supposed to be 60?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vandonr-amz commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "vandonr-amz (via GitHub)" <gi...@apache.org>.
vandonr-amz commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253674300
##########
airflow/providers/amazon/aws/triggers/emr.py:
##########
@@ -102,274 +102,184 @@ async def run(self):
yield TriggerEvent({"status": "success", "message": "Steps completed", "step_ids": self.step_ids})
-class EmrCreateJobFlowTrigger(BaseTrigger):
+class EmrCreateJobFlowTrigger(AwsBaseWaiterTrigger):
"""
Asynchronously poll the boto3 API and wait for the JobFlow to finish executing.
:param job_flow_id: The id of the job flow to wait for.
- :param poll_interval: The amount of time in seconds to wait between attempts.
- :param max_attempts: The maximum number of attempts to be made.
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
:param aws_conn_id: The Airflow connection used for AWS credentials.
"""
def __init__(
self,
job_flow_id: str,
- poll_interval: int,
- max_attempts: int,
- aws_conn_id: str,
+ poll_interval: int | None = None, # deprecated
+ max_attempts: int | None = None, # deprecated
+ aws_conn_id: str | None = None,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
):
- self.job_flow_id = job_flow_id
- self.poll_interval = poll_interval
- self.max_attempts = max_attempts
- self.aws_conn_id = aws_conn_id
-
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "job_flow_id": self.job_flow_id,
- "poll_interval": str(self.poll_interval),
- "max_attempts": str(self.max_attempts),
- "aws_conn_id": self.aws_conn_id,
- },
+ if poll_interval is not None or max_attempts is not None:
+ warnings.warn(
+ "please use waiter_delay instead of poll_interval "
+ "and waiter_max_attempts instead of max_attempts",
+ AirflowProviderDeprecationWarning,
+ stacklevel=2,
+ )
+ waiter_delay = poll_interval or waiter_delay
+ waiter_max_attempts = max_attempts or waiter_max_attempts
+ super().__init__(
+ serialized_fields={"job_flow_id": job_flow_id},
+ waiter_name="job_flow_waiting",
+ waiter_args={"ClusterId": job_flow_id},
+ failure_message="JobFlow creation failed",
+ status_message="JobFlow creation in progress",
+ status_queries=[
+ "Cluster.Status.State",
+ "Cluster.Status.StateChangeReason",
+ "Cluster.Status.ErrorDetails",
+ ],
+ return_key="job_flow_id",
+ return_value=job_flow_id,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ aws_conn_id=aws_conn_id,
)
- async def run(self):
- self.hook = EmrHook(aws_conn_id=self.aws_conn_id)
- async with self.hook.async_conn as client:
- attempt = 0
- waiter = self.hook.get_waiter("job_flow_waiting", deferrable=True, client=client)
- while attempt < int(self.max_attempts):
- attempt = attempt + 1
- try:
- await waiter.wait(
- ClusterId=self.job_flow_id,
- WaiterConfig=prune_dict(
- {
- "Delay": self.poll_interval,
- "MaxAttempts": 1,
- }
- ),
- )
- break
- except WaiterError as error:
- if "terminal failure" in str(error):
- raise AirflowException(f"JobFlow creation failed: {error}")
- self.log.info(
- "Status of jobflow is %s - %s",
- error.last_response["Cluster"]["Status"]["State"],
- error.last_response["Cluster"]["Status"]["StateChangeReason"],
- )
- await asyncio.sleep(int(self.poll_interval))
- if attempt >= int(self.max_attempts):
- raise AirflowException(f"JobFlow creation failed - max attempts reached: {self.max_attempts}")
- else:
- yield TriggerEvent(
- {
- "status": "success",
- "message": "JobFlow completed successfully",
- "job_flow_id": self.job_flow_id,
- }
- )
+ def hook(self) -> AwsGenericHook:
+ return EmrHook(aws_conn_id=self.aws_conn_id)
-class EmrTerminateJobFlowTrigger(BaseTrigger):
+class EmrTerminateJobFlowTrigger(AwsBaseWaiterTrigger):
"""
Asynchronously poll the boto3 API and wait for the JobFlow to finish terminating.
:param job_flow_id: ID of the EMR Job Flow to terminate
- :param poll_interval: The amount of time in seconds to wait between attempts.
- :param max_attempts: The maximum number of attempts to be made.
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
:param aws_conn_id: The Airflow connection used for AWS credentials.
"""
def __init__(
self,
job_flow_id: str,
- poll_interval: int,
- max_attempts: int,
- aws_conn_id: str,
+ poll_interval: int | None = None, # deprecated
+ max_attempts: int | None = None, # deprecated
+ aws_conn_id: str | None = None,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
):
- self.job_flow_id = job_flow_id
- self.poll_interval = poll_interval
- self.max_attempts = max_attempts
- self.aws_conn_id = aws_conn_id
-
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "job_flow_id": self.job_flow_id,
- "poll_interval": str(self.poll_interval),
- "max_attempts": str(self.max_attempts),
- "aws_conn_id": self.aws_conn_id,
- },
+ if poll_interval is not None or max_attempts is not None:
+ warnings.warn(
+ "please use waiter_delay instead of poll_interval "
+ "and waiter_max_attempts instead of max_attempts",
+ AirflowProviderDeprecationWarning,
+ stacklevel=2,
+ )
+ waiter_delay = poll_interval or waiter_delay
+ waiter_max_attempts = max_attempts or waiter_max_attempts
+ super().__init__(
+ serialized_fields={"job_flow_id": job_flow_id},
+ waiter_name="job_flow_terminated",
+ waiter_args={"ClusterId": job_flow_id},
+ failure_message="JobFlow termination failed",
+ status_message="JobFlow termination in progress",
+ status_queries=[
+ "Cluster.Status.State",
+ "Cluster.Status.StateChangeReason",
+ "Cluster.Status.ErrorDetails",
+ ],
+ return_value=None,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ aws_conn_id=aws_conn_id,
)
- async def run(self):
- self.hook = EmrHook(aws_conn_id=self.aws_conn_id)
- async with self.hook.async_conn as client:
- attempt = 0
- waiter = self.hook.get_waiter("job_flow_terminated", deferrable=True, client=client)
- while attempt < int(self.max_attempts):
- attempt = attempt + 1
- try:
- await waiter.wait(
- ClusterId=self.job_flow_id,
- WaiterConfig=prune_dict(
- {
- "Delay": self.poll_interval,
- "MaxAttempts": 1,
- }
- ),
- )
- break
- except WaiterError as error:
- if "terminal failure" in str(error):
- raise AirflowException(f"JobFlow termination failed: {error}")
- self.log.info(
- "Status of jobflow is %s - %s",
- error.last_response["Cluster"]["Status"]["State"],
- error.last_response["Cluster"]["Status"]["StateChangeReason"],
- )
- await asyncio.sleep(int(self.poll_interval))
- if attempt >= int(self.max_attempts):
- raise AirflowException(f"JobFlow termination failed - max attempts reached: {self.max_attempts}")
- else:
- yield TriggerEvent(
- {
- "status": "success",
- "message": "JobFlow terminated successfully",
- }
- )
+ def hook(self) -> AwsGenericHook:
+ return EmrHook(aws_conn_id=self.aws_conn_id)
-class EmrContainerTrigger(BaseTrigger):
+class EmrContainerTrigger(AwsBaseWaiterTrigger):
"""
Poll for the status of EMR container until reaches terminal state.
:param virtual_cluster_id: Reference Emr cluster id
:param job_id: job_id to check the state
:param aws_conn_id: Reference to AWS connection id
- :param poll_interval: polling period in seconds to check for the status
+ :param waiter_delay: polling period in seconds to check for the status
"""
def __init__(
self,
virtual_cluster_id: str,
job_id: str,
aws_conn_id: str = "aws_default",
- poll_interval: int = 30,
- **kwargs: Any,
+ poll_interval: int | None = None, # deprecated
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 600,
Review Comment:
it's supposed to be ∞ lol
we don't really have a clear posture on this, some triggers have sensible max attempts values, some others wait forever...
What I don't like with setting a "low" value like 60 is that if the user sets the poll interval to 1 for instance, it ends in 1 minute, which might be annoying.
Overall, I think it should be a really high value, and users should set it themselves if they care.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vandonr-amz commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "vandonr-amz (via GitHub)" <gi...@apache.org>.
vandonr-amz commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253602343
##########
airflow/providers/amazon/aws/hooks/athena.py:
##########
@@ -253,7 +253,7 @@ def poll_query_status(
try:
wait(
waiter=self.get_waiter("query_complete"),
- waiter_delay=sleep_time or self.sleep_time,
+ waiter_delay=self.sleep_time if sleep_time is None else sleep_time,
Review Comment:
I think just setting sleep_time=0 is sooo much simpler & cleaner & easier to read
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vandonr-amz commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "vandonr-amz (via GitHub)" <gi...@apache.org>.
vandonr-amz commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253609094
##########
airflow/providers/amazon/aws/sensors/batch.py:
##########
@@ -113,9 +114,10 @@ def execute_complete(self, context: Context, event: dict[str, Any]) -> None:
Relies on trigger to throw an exception, otherwise it assumes execution was successful.
"""
- if "status" in event and event["status"] == "failure":
Review Comment:
no I don't think so. The event is something we control and that we send to ourselves. If it doesn't contain a "status" key, it means something bad happened, either the event is not formatted properly, or an event was routed to the wrong place, and we definitely want the error to blow up.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ferruzzi commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253577445
##########
airflow/providers/amazon/aws/operators/ecs.py:
##########
@@ -139,13 +143,12 @@ def execute(self, context: Context):
self.log.info("Cluster %r in state: %r.", self.cluster_name, cluster_state)
elif self.deferrable:
self.defer(
- trigger=ClusterWaiterTrigger(
- waiter_name="cluster_active",
+ trigger=ClusterActiveTrigger(
cluster_arn=cluster_details["clusterArn"],
Review Comment:
oorrrrrrr maybe that's another use for a BaseAwsOperator class... it could have a helper method that generates a `self.trigger_config` and/or a `self.waiter_config` in the constructor...
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vincbeck merged pull request #32274: Introduce a base class for aws triggers
Posted by "vincbeck (via GitHub)" <gi...@apache.org>.
vincbeck merged PR #32274:
URL: https://github.com/apache/airflow/pull/32274
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ferruzzi commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253582036
##########
airflow/providers/amazon/aws/triggers/base_trigger.py:
##########
@@ -0,0 +1,129 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import Any, AsyncIterator
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+from airflow.providers.amazon.aws.utils.waiter_with_logging import async_wait
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+
+
+class AwsBaseWaiterTrigger(BaseTrigger):
+ """
+ Base class for all AWS Triggers that follow the "standard" model of just waiting on a waiter.
+
+ Subclasses need to implement the hook() method.
+
+ :param serialized_fields: Fields that are specific to the subclass trigger and need to be serialized
+ to be passed to the __init__ method on deserialization.
+ format: {<parameter_name>: <parameter_value>}
+
+ :param waiter_name: The name of the (possibly custom) boto waiter to use.
+
+ :param waiter_args: The arguments to pass to the waiter.
+ :param failure_message: The message to log if a failure state is reached.
+ :param status_message: The message logged when printing the status of the service.
+ :param status_queries: A list containing the JMESPath queries to retrieve status information from
+ the waiter response. See https://jmespath.org/tutorial.html
+
+ :param return_key: The key to use for the return_value in the TriggerEvent this emits on success.
+ Defaults to "value".
+ :param return_value: A value that'll be returned in the return_key field of the TriggerEvent.
+ Set to None if there is nothing to return.
+
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
+ :param aws_conn_id: The Airflow connection used for AWS credentials. To be used to build the hook.
+ :param region_name: The AWS region where the resources to watch are. To be used to build the hook.
+ """
+
+ def __init__(
+ self,
+ *,
+ serialized_fields: dict[str, Any],
+ waiter_name: str,
+ waiter_args: dict[str, Any],
+ failure_message: str,
+ status_message: str,
+ status_queries: list[str],
+ return_key: str = "value",
+ return_value: Any,
+ waiter_delay: int,
+ waiter_max_attempts: int,
+ aws_conn_id: str | None,
+ region_name: str | None = None,
+ ):
+ # parameters that should be hardcoded in the child's implem
+ self.serialized_fields = serialized_fields
+
+ self.waiter_name = waiter_name
+ self.waiter_args = waiter_args
+ self.failure_message = failure_message
+ self.status_message = status_message
+ self.status_queries = status_queries
+
+ self.return_key = return_key
+ self.return_value = return_value
+
+ # parameters that should be passed directly from the child's parameters
+ self.waiter_delay = waiter_delay
+ self.attempts = waiter_max_attempts
+ self.aws_conn_id = aws_conn_id
+ self.region_name = region_name
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ # here we put together the "common" params,
+ # and whatever extras we got from the subclass in serialized_fields
+ params = dict(
+ {
+ "waiter_delay": self.waiter_delay,
+ "waiter_max_attempts": self.attempts,
+ "aws_conn_id": self.aws_conn_id,
+ },
+ **self.serialized_fields,
+ )
+ if self.region_name:
+ # if we serialize the None value from this, it breaks subclasses that don't have it in their ctor.
+ params["region_name"] = self.region_name
+ return (
+ # remember that self is an instance of the subclass here, not of this class.
+ self.__class__.__module__ + "." + self.__class__.__qualname__,
+ params,
+ )
+
+ @abstractmethod
+ def hook(self) -> AwsGenericHook:
+ """Override in subclasses to return the right hook."""
+ ...
Review Comment:
Not positive, but should this raise NotImplementedException like they did [here](https://github.com/apache/airflow/blob/main/airflow/providers/amazon/aws/sensors/sagemaker.py#L75) to enforce that?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vandonr-amz commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "vandonr-amz (via GitHub)" <gi...@apache.org>.
vandonr-amz commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253598328
##########
airflow/providers/amazon/aws/triggers/eks.py:
##########
@@ -115,59 +90,37 @@ def __init__(
waiter_max_attempts: int,
aws_conn_id: str,
region: str | None = None,
+ region_name: str | None = None,
):
- self.cluster_name = cluster_name
- self.fargate_profile_name = fargate_profile_name
- self.waiter_delay = waiter_delay
- self.waiter_max_attempts = waiter_max_attempts
- self.aws_conn_id = aws_conn_id
- self.region = region
-
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "cluster_name": self.cluster_name,
- "fargate_profile_name": self.fargate_profile_name,
- "waiter_delay": str(self.waiter_delay),
- "waiter_max_attempts": str(self.waiter_max_attempts),
- "aws_conn_id": self.aws_conn_id,
- "region": self.region,
- },
+ if region is not None:
+ warnings.warn(
+ "please use region_name param instead of region",
+ AirflowProviderDeprecationWarning,
+ stacklevel=2,
+ )
+ region_name = region
+
+ super().__init__(
+ serialized_fields={"cluster_name": cluster_name, "fargate_profile_name": fargate_profile_name},
+ waiter_name="fargate_profile_deleted",
+ waiter_args={"clusterName": cluster_name, "fargateProfileName": fargate_profile_name},
+ failure_message="Failure while deleting Fargate profile",
+ status_message="Fargate profile not deleted yet",
+ status_queries=["fargateProfile.status"],
+ return_value=None,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ aws_conn_id=aws_conn_id,
+ region_name=region_name,
)
- async def run(self):
- self.hook = EksHook(aws_conn_id=self.aws_conn_id, region_name=self.region)
- async with self.hook.async_conn as client:
- attempt = 0
- waiter = client.get_waiter("fargate_profile_deleted")
- while attempt < int(self.waiter_max_attempts):
- attempt += 1
- try:
- await waiter.wait(
- clusterName=self.cluster_name,
- fargateProfileName=self.fargate_profile_name,
- WaiterConfig={"Delay": int(self.waiter_delay), "MaxAttempts": 1},
- )
- break
- except WaiterError as error:
- if "terminal failure" in str(error):
- raise AirflowException(f"Delete Fargate Profile failed: {error}")
- self.log.info(
- "Status of fargate profile is %s", error.last_response["fargateProfile"]["status"]
- )
- await asyncio.sleep(int(self.waiter_delay))
- if attempt >= int(self.waiter_max_attempts):
- raise AirflowException(
- f"Delete Fargate Profile failed - max attempts reached: {self.waiter_max_attempts}"
- )
- else:
- yield TriggerEvent({"status": "success", "message": "Fargate Profile Deleted"})
+ def hook(self) -> AwsGenericHook:
+ return EksHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
-class EksNodegroupTrigger(BaseTrigger):
+class EksCreateNodegroupTrigger(AwsBaseWaiterTrigger):
Review Comment:
I remember you said back in the day that you preferred separate triggers so that the status/failure messages can be more descriptive ;)
Also, given the lower footprint of triggers created this way, I think it's ok to have specific triggers for each thing.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vandonr-amz commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "vandonr-amz (via GitHub)" <gi...@apache.org>.
vandonr-amz commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253601492
##########
airflow/providers/amazon/aws/hooks/athena.py:
##########
@@ -253,7 +253,7 @@ def poll_query_status(
try:
wait(
waiter=self.get_waiter("query_complete"),
- waiter_delay=sleep_time or self.sleep_time,
+ waiter_delay=self.sleep_time if sleep_time is None else sleep_time,
Review Comment:
didn't I put a self-review comment about that ? 🤔
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vandonr-amz commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "vandonr-amz (via GitHub)" <gi...@apache.org>.
vandonr-amz commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253678597
##########
tests/providers/amazon/aws/triggers/test_base_trigger.py:
##########
@@ -0,0 +1,89 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from unittest import mock
+from unittest.mock import MagicMock
+
+import pytest
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
+from airflow.triggers.base import TriggerEvent
+
+
+class TestImplem(AwsBaseWaiterTrigger):
+ """An empty implementation that allows instantiation for tests."""
+
+ __test__ = False
+
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+
+ def hook(self) -> AwsGenericHook:
+ return MagicMock()
+
+
+class TestAwsBaseWaiterTrigger:
+ def setup_method(self):
+ self.trigger = TestImplem(
+ serialized_fields={},
+ waiter_name="",
+ waiter_args={},
+ failure_message="",
+ status_message="",
+ status_queries=[],
+ return_value=None,
+ waiter_delay=0,
+ waiter_max_attempts=0,
+ aws_conn_id="",
+ )
+
+ def test_region_serialized(self):
+ self.trigger.region_name = "my_region"
+ _, args = self.trigger.serialize()
+
+ assert "region_name" in args
+ assert args["region_name"] == "my_region"
+
+ def test_region_not_serialized_if_omitted(self):
+ _, args = self.trigger.serialize()
+
+ assert "region_name" not in args
+
+ def test_serialize_extra_fields(self):
+ self.trigger.serialized_fields = {"foo": "bar", "foz": "baz"}
+
+ _, args = self.trigger.serialize()
+
+ assert "foo" in args
+ assert args["foo"] == "bar"
+ assert "foz" in args
+ assert args["foz"] == "baz"
+
+ @pytest.mark.asyncio
Review Comment:
yes indeed, this code is not in the base trigger, so it doesn't have to be tested here imho
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] eladkal commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253520217
##########
airflow/providers/amazon/provider.yaml:
##########
@@ -516,6 +516,9 @@ hooks:
- airflow.providers.amazon.aws.hooks.appflow
triggers:
+ - integration-name: Amazon Web Services
+ python-modules:
+ - airflow.providers.amazon.aws.triggers.base_trigger
Review Comment:
Do we need the `_trigger` suffix?
For example we have:
`from airflow.hooks.base import BaseHook`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ferruzzi commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253612171
##########
airflow/providers/amazon/aws/operators/ecs.py:
##########
@@ -139,13 +143,12 @@ def execute(self, context: Context):
self.log.info("Cluster %r in state: %r.", self.cluster_name, cluster_state)
elif self.deferrable:
self.defer(
- trigger=ClusterWaiterTrigger(
- waiter_name="cluster_active",
+ trigger=ClusterActiveTrigger(
cluster_arn=cluster_details["clusterArn"],
Review Comment:
I'll buy that. We should keep a list somewhere of all these little "would be nice in a BaseAwsOperator" things... I can never think of one big convincing reason to put in the effort, but there have been a few smaller conveniences here and there that I don't remember.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vandonr-amz commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "vandonr-amz (via GitHub)" <gi...@apache.org>.
vandonr-amz commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253609846
##########
airflow/providers/amazon/aws/triggers/base_trigger.py:
##########
@@ -0,0 +1,129 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import Any, AsyncIterator
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+from airflow.providers.amazon.aws.utils.waiter_with_logging import async_wait
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+
+
+class AwsBaseWaiterTrigger(BaseTrigger):
+ """
+ Base class for all AWS Triggers that follow the "standard" model of just waiting on a waiter.
+
+ Subclasses need to implement the hook() method.
+
+ :param serialized_fields: Fields that are specific to the subclass trigger and need to be serialized
+ to be passed to the __init__ method on deserialization.
+ format: {<parameter_name>: <parameter_value>}
+
+ :param waiter_name: The name of the (possibly custom) boto waiter to use.
+
+ :param waiter_args: The arguments to pass to the waiter.
+ :param failure_message: The message to log if a failure state is reached.
+ :param status_message: The message logged when printing the status of the service.
+ :param status_queries: A list containing the JMESPath queries to retrieve status information from
+ the waiter response. See https://jmespath.org/tutorial.html
+
+ :param return_key: The key to use for the return_value in the TriggerEvent this emits on success.
+ Defaults to "value".
+ :param return_value: A value that'll be returned in the return_key field of the TriggerEvent.
+ Set to None if there is nothing to return.
+
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
+ :param aws_conn_id: The Airflow connection used for AWS credentials. To be used to build the hook.
+ :param region_name: The AWS region where the resources to watch are. To be used to build the hook.
+ """
+
+ def __init__(
+ self,
+ *,
+ serialized_fields: dict[str, Any],
+ waiter_name: str,
+ waiter_args: dict[str, Any],
+ failure_message: str,
+ status_message: str,
+ status_queries: list[str],
+ return_key: str = "value",
+ return_value: Any,
+ waiter_delay: int,
+ waiter_max_attempts: int,
+ aws_conn_id: str | None,
+ region_name: str | None = None,
+ ):
+ # parameters that should be hardcoded in the child's implem
+ self.serialized_fields = serialized_fields
+
+ self.waiter_name = waiter_name
+ self.waiter_args = waiter_args
+ self.failure_message = failure_message
+ self.status_message = status_message
+ self.status_queries = status_queries
+
+ self.return_key = return_key
+ self.return_value = return_value
+
+ # parameters that should be passed directly from the child's parameters
+ self.waiter_delay = waiter_delay
+ self.attempts = waiter_max_attempts
+ self.aws_conn_id = aws_conn_id
+ self.region_name = region_name
+
+ def serialize(self) -> tuple[str, dict[str, Any]]:
+ # here we put together the "common" params,
+ # and whatever extras we got from the subclass in serialized_fields
+ params = dict(
+ {
+ "waiter_delay": self.waiter_delay,
+ "waiter_max_attempts": self.attempts,
+ "aws_conn_id": self.aws_conn_id,
+ },
+ **self.serialized_fields,
+ )
+ if self.region_name:
+ # if we serialize the None value from this, it breaks subclasses that don't have it in their ctor.
+ params["region_name"] = self.region_name
+ return (
+ # remember that self is an instance of the subclass here, not of this class.
+ self.__class__.__module__ + "." + self.__class__.__qualname__,
+ params,
+ )
+
+ @abstractmethod
+ def hook(self) -> AwsGenericHook:
+ """Override in subclasses to return the right hook."""
+ ...
Review Comment:
the `@abstractmethod` annotation is taking care of that (and more)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vandonr-amz commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "vandonr-amz (via GitHub)" <gi...@apache.org>.
vandonr-amz commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253612474
##########
airflow/providers/amazon/aws/triggers/ecs.py:
##########
@@ -22,68 +22,89 @@
from botocore.exceptions import ClientError, WaiterError
+from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
from airflow.providers.amazon.aws.hooks.ecs import EcsHook
from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
+from airflow.providers.amazon.aws.triggers.base_trigger import AwsBaseWaiterTrigger
from airflow.providers.amazon.aws.utils.task_log_fetcher import AwsTaskLogFetcher
-from airflow.providers.amazon.aws.utils.waiter_with_logging import async_wait
from airflow.triggers.base import BaseTrigger, TriggerEvent
-class ClusterWaiterTrigger(BaseTrigger):
+class ClusterActiveTrigger(AwsBaseWaiterTrigger):
"""
- Polls the status of a cluster using a given waiter. Can be used to poll for an active or inactive cluster.
+ Polls the status of a cluster until it's ready.
- :param waiter_name: Name of the waiter to use, for instance 'cluster_active' or 'cluster_inactive'
:param cluster_arn: ARN of the cluster to watch.
:param waiter_delay: The amount of time in seconds to wait between attempts.
:param waiter_max_attempts: The number of times to ping for status.
Will fail after that many unsuccessful attempts.
:param aws_conn_id: The Airflow connection used for AWS credentials.
- :param region: The AWS region where the cluster is located.
+ :param region_name: The AWS region where the cluster is located.
"""
def __init__(
self,
- waiter_name: str,
cluster_arn: str,
- waiter_delay: int | None,
- waiter_max_attempts: int | None,
+ waiter_delay: int,
+ waiter_max_attempts: int,
aws_conn_id: str | None,
- region: str | None,
+ region_name: str | None,
):
- self.cluster_arn = cluster_arn
- self.waiter_name = waiter_name
- self.waiter_delay = waiter_delay if waiter_delay is not None else 15 # written like this to allow 0
- self.attempts = waiter_max_attempts or 999999999
- self.aws_conn_id = aws_conn_id
- self.region = region
+ super().__init__(
+ serialized_fields={"cluster_arn": cluster_arn},
+ waiter_name="cluster_active",
+ waiter_args={"clusters": [cluster_arn]},
+ failure_message="Failure while waiting for cluster to be available",
+ status_message="Cluster is not ready yet",
+ status_queries=["clusters[].status", "failures"],
+ return_key="arn",
+ return_value=cluster_arn,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ aws_conn_id=aws_conn_id,
+ region_name=region_name,
+ )
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "waiter_name": self.waiter_name,
- "cluster_arn": self.cluster_arn,
- "waiter_delay": self.waiter_delay,
- "waiter_max_attempts": self.attempts,
- "aws_conn_id": self.aws_conn_id,
- "region": self.region,
- },
+ def hook(self) -> AwsGenericHook:
+ return EcsHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
+
+
+class ClusterInactiveTrigger(AwsBaseWaiterTrigger):
Review Comment:
the status/failure messages change as well.
But you have the same comment as syed here https://github.com/apache/airflow/pull/32274#discussion_r1253537693
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vincbeck commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "vincbeck (via GitHub)" <gi...@apache.org>.
vincbeck commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253466965
##########
airflow/providers/amazon/aws/triggers/base_trigger.py:
##########
@@ -0,0 +1,127 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import Any, AsyncIterator
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+from airflow.providers.amazon.aws.utils.waiter_with_logging import async_wait
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+
+
+class AwsBaseWaiterTrigger(BaseTrigger):
+ """
+ Base class for all AWS Triggers that follow the "standard" model of just waiting on a waiter.
+
+ Subclasses need to implement the hook() method.
+
+ :param serialized_fields: Fields that are specific to the subclass trigger and need to be serialized
+ to be passed to the __init__ method on deserialization.
+ format: {<parameter_name>: <parameter_value>}
+
+ :param waiter_name: The name of the (possibly custom) boto waiter to use.
+
+ :param waiter_args: The arguments to pass to the waiter.
+ :param failure_message: The message to log if a failure state is reached.
+ :param status_message: The message logged when printing the status of the service.
+ :param status_queries: A list containing the JMESPath queries to retrieve status information from
+ the waiter response. See https://jmespath.org/tutorial.html
+
+ :param return_key: The key to use for the return_value in the TriggerEvent this emits on success.
+ :param return_value: A value that'll be returned in the return_key field of the TriggerEvent.
+
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ :param region_name: The AWS region where the resources to watch are.
+ """
+
+ def __init__(
+ self,
+ *,
+ serialized_fields: dict[str, Any],
+ waiter_name: str,
+ waiter_args: dict[str, Any],
+ failure_message: str,
+ status_message: str,
+ status_queries: list[str],
+ return_key: str = "value",
+ return_value: Any | None,
Review Comment:
Is not `None` included in `Any`?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ferruzzi commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253602175
##########
tests/providers/amazon/aws/triggers/test_eks.py:
##########
@@ -40,416 +34,50 @@
TEST_REGION = "test-region"
-class TestEksCreateFargateProfileTrigger:
- def test_eks_create_fargate_profile_serialize(self):
- eks_create_fargate_profile_trigger = EksCreateFargateProfileTrigger(
- cluster_name=TEST_CLUSTER_IDENTIFIER,
- fargate_profile_name=TEST_FARGATE_PROFILE_NAME,
- aws_conn_id=TEST_AWS_CONN_ID,
- waiter_delay=TEST_WAITER_DELAY,
- waiter_max_attempts=TEST_WAITER_MAX_ATTEMPTS,
- )
-
- class_path, args = eks_create_fargate_profile_trigger.serialize()
- assert class_path == "airflow.providers.amazon.aws.triggers.eks.EksCreateFargateProfileTrigger"
- assert args["cluster_name"] == TEST_CLUSTER_IDENTIFIER
- assert args["fargate_profile_name"] == TEST_FARGATE_PROFILE_NAME
- assert args["aws_conn_id"] == TEST_AWS_CONN_ID
- assert args["waiter_delay"] == str(TEST_WAITER_DELAY)
- assert args["waiter_max_attempts"] == str(TEST_WAITER_MAX_ATTEMPTS)
-
- @pytest.mark.asyncio
- @mock.patch.object(EksHook, "async_conn")
- async def test_eks_create_fargate_profile_trigger_run(self, mock_async_conn):
- a_mock = mock.MagicMock()
- mock_async_conn.__aenter__.return_value = a_mock
-
- a_mock.get_waiter().wait = AsyncMock()
-
- eks_create_fargate_profile_trigger = EksCreateFargateProfileTrigger(
- cluster_name=TEST_CLUSTER_IDENTIFIER,
- fargate_profile_name=TEST_FARGATE_PROFILE_NAME,
- aws_conn_id=TEST_AWS_CONN_ID,
- waiter_delay=TEST_WAITER_DELAY,
- waiter_max_attempts=TEST_WAITER_MAX_ATTEMPTS,
- )
-
- generator = eks_create_fargate_profile_trigger.run()
- response = await generator.asend(None)
-
- assert response == TriggerEvent({"status": "success", "message": "Fargate Profile Created"})
-
- @pytest.mark.asyncio
- @mock.patch("asyncio.sleep")
- @mock.patch.object(EksHook, "async_conn")
- async def test_eks_create_fargate_profile_trigger_run_multiple_attempts(
- self, mock_async_conn, mock_sleep
- ):
- a_mock = mock.MagicMock()
- mock_async_conn.__aenter__.return_value = a_mock
- error = WaiterError(
- name="test_name",
- reason="test_reason",
- last_response={"fargateProfile": {"status": "CREATING"}},
- )
- a_mock.get_waiter().wait = AsyncMock(side_effect=[error, error, True])
- mock_sleep.return_value = True
-
- eks_create_fargate_profile_trigger = EksCreateFargateProfileTrigger(
- cluster_name=TEST_CLUSTER_IDENTIFIER,
- fargate_profile_name=TEST_FARGATE_PROFILE_NAME,
- aws_conn_id=TEST_AWS_CONN_ID,
- waiter_delay=TEST_WAITER_DELAY,
- waiter_max_attempts=TEST_WAITER_MAX_ATTEMPTS,
- )
-
- generator = eks_create_fargate_profile_trigger.run()
- response = await generator.asend(None)
-
- assert a_mock.get_waiter().wait.call_count == 3
- assert response == TriggerEvent({"status": "success", "message": "Fargate Profile Created"})
-
- @pytest.mark.asyncio
- @mock.patch("asyncio.sleep")
- @mock.patch.object(EksHook, "async_conn")
- async def test_eks_create_fargate_profile_trigger_run_attempts_exceeded(
- self, mock_async_conn, mock_sleep
- ):
- a_mock = mock.MagicMock()
- mock_async_conn.__aenter__.return_value = a_mock
- error = WaiterError(
- name="test_name",
- reason="test_reason",
- last_response={"fargateProfile": {"status": "CREATING"}},
- )
- a_mock.get_waiter().wait = AsyncMock(side_effect=[error, error, True])
- mock_sleep.return_value = True
-
- eks_create_fargate_profile_trigger = EksCreateFargateProfileTrigger(
- cluster_name=TEST_CLUSTER_IDENTIFIER,
- fargate_profile_name=TEST_FARGATE_PROFILE_NAME,
- aws_conn_id=TEST_AWS_CONN_ID,
- waiter_delay=TEST_WAITER_DELAY,
- waiter_max_attempts=2,
- )
- with pytest.raises(AirflowException) as exc:
- generator = eks_create_fargate_profile_trigger.run()
- await generator.asend(None)
- assert "Create Fargate Profile failed - max attempts reached:" in str(exc.value)
- assert a_mock.get_waiter().wait.call_count == 2
-
- @pytest.mark.asyncio
- @mock.patch("asyncio.sleep")
- @mock.patch.object(EksHook, "async_conn")
- async def test_eks_create_fargate_profile_trigger_run_attempts_failed(self, mock_async_conn, mock_sleep):
- a_mock = mock.MagicMock()
- mock_async_conn.__aenter__.return_value = a_mock
- error_creating = WaiterError(
- name="test_name",
- reason="test_reason",
- last_response={"fargateProfile": {"status": "CREATING"}},
- )
- error_failed = WaiterError(
- name="test_name",
- reason="Waiter encountered a terminal failure state:",
- last_response={"fargateProfile": {"status": "CREATE_FAILED"}},
- )
- a_mock.get_waiter().wait = AsyncMock(side_effect=[error_creating, error_creating, error_failed])
- mock_sleep.return_value = True
-
- eks_create_fargate_profile_trigger = EksCreateFargateProfileTrigger(
- cluster_name=TEST_CLUSTER_IDENTIFIER,
- fargate_profile_name=TEST_FARGATE_PROFILE_NAME,
- aws_conn_id=TEST_AWS_CONN_ID,
- waiter_delay=TEST_WAITER_DELAY,
- waiter_max_attempts=TEST_WAITER_MAX_ATTEMPTS,
- )
-
- with pytest.raises(AirflowException) as exc:
- generator = eks_create_fargate_profile_trigger.run()
- await generator.asend(None)
- assert f"Create Fargate Profile failed: {error_failed}" in str(exc.value)
- assert a_mock.get_waiter().wait.call_count == 3
-
-
-class TestEksDeleteFargateProfileTrigger:
- def test_eks_delete_fargate_profile_serialize(self):
- eks_delete_fargate_profile_trigger = EksDeleteFargateProfileTrigger(
- cluster_name=TEST_CLUSTER_IDENTIFIER,
- fargate_profile_name=TEST_FARGATE_PROFILE_NAME,
- aws_conn_id=TEST_AWS_CONN_ID,
- waiter_delay=TEST_WAITER_DELAY,
- waiter_max_attempts=TEST_WAITER_MAX_ATTEMPTS,
- )
-
- class_path, args = eks_delete_fargate_profile_trigger.serialize()
- assert class_path == "airflow.providers.amazon.aws.triggers.eks.EksDeleteFargateProfileTrigger"
- assert args["cluster_name"] == TEST_CLUSTER_IDENTIFIER
- assert args["fargate_profile_name"] == TEST_FARGATE_PROFILE_NAME
- assert args["aws_conn_id"] == TEST_AWS_CONN_ID
- assert args["waiter_delay"] == str(TEST_WAITER_DELAY)
- assert args["waiter_max_attempts"] == str(TEST_WAITER_MAX_ATTEMPTS)
-
- @pytest.mark.asyncio
- @mock.patch.object(EksHook, "async_conn")
- async def test_eks_delete_fargate_profile_trigger_run(self, mock_async_conn):
- a_mock = mock.MagicMock()
- mock_async_conn.__aenter__.return_value = a_mock
-
- a_mock.get_waiter().wait = AsyncMock()
-
- eks_delete_fargate_profile_trigger = EksDeleteFargateProfileTrigger(
- cluster_name=TEST_CLUSTER_IDENTIFIER,
- fargate_profile_name=TEST_FARGATE_PROFILE_NAME,
- aws_conn_id=TEST_AWS_CONN_ID,
- waiter_delay=TEST_WAITER_DELAY,
- waiter_max_attempts=TEST_WAITER_MAX_ATTEMPTS,
- )
-
- generator = eks_delete_fargate_profile_trigger.run()
- response = await generator.asend(None)
-
- assert response == TriggerEvent({"status": "success", "message": "Fargate Profile Deleted"})
-
- @pytest.mark.asyncio
- @mock.patch("asyncio.sleep")
- @mock.patch.object(EksHook, "async_conn")
- async def test_eks_delete_fargate_profile_trigger_run_multiple_attempts(
- self, mock_async_conn, mock_sleep
- ):
- a_mock = mock.MagicMock()
- mock_async_conn.__aenter__.return_value = a_mock
- error = WaiterError(
- name="test_name",
- reason="test_reason",
- last_response={"fargateProfile": {"status": "DELETING"}},
- )
- a_mock.get_waiter().wait = AsyncMock(side_effect=[error, error, True])
- mock_sleep.return_value = True
-
- eks_delete_fargate_profile_trigger = EksDeleteFargateProfileTrigger(
- cluster_name=TEST_CLUSTER_IDENTIFIER,
- fargate_profile_name=TEST_FARGATE_PROFILE_NAME,
- aws_conn_id=TEST_AWS_CONN_ID,
- waiter_delay=TEST_WAITER_DELAY,
- waiter_max_attempts=TEST_WAITER_MAX_ATTEMPTS,
- )
-
- generator = eks_delete_fargate_profile_trigger.run()
- response = await generator.asend(None)
- assert a_mock.get_waiter().wait.call_count == 3
- assert response == TriggerEvent({"status": "success", "message": "Fargate Profile Deleted"})
-
- @pytest.mark.asyncio
- @mock.patch("asyncio.sleep")
- @mock.patch.object(EksHook, "async_conn")
- async def test_eks_delete_fargate_profile_trigger_run_attempts_exceeded(
- self, mock_async_conn, mock_sleep
- ):
- a_mock = mock.MagicMock()
- mock_async_conn.__aenter__.return_value = a_mock
- error = WaiterError(
- name="test_name",
- reason="test_reason",
- last_response={"fargateProfile": {"status": "DELETING"}},
- )
- a_mock.get_waiter().wait = AsyncMock(side_effect=[error, error, error, True])
- mock_sleep.return_value = True
-
- eks_delete_fargate_profile_trigger = EksDeleteFargateProfileTrigger(
- cluster_name=TEST_CLUSTER_IDENTIFIER,
- fargate_profile_name=TEST_FARGATE_PROFILE_NAME,
- aws_conn_id=TEST_AWS_CONN_ID,
- waiter_delay=TEST_WAITER_DELAY,
- waiter_max_attempts=2,
- )
- with pytest.raises(AirflowException) as exc:
- generator = eks_delete_fargate_profile_trigger.run()
- await generator.asend(None)
- assert "Delete Fargate Profile failed - max attempts reached: 2" in str(exc.value)
- assert a_mock.get_waiter().wait.call_count == 2
-
- @pytest.mark.asyncio
- @mock.patch("asyncio.sleep")
- @mock.patch.object(EksHook, "async_conn")
- async def test_eks_delete_fargate_profile_trigger_run_attempts_failed(self, mock_async_conn, mock_sleep):
- a_mock = mock.MagicMock()
- mock_async_conn.__aenter__.return_value = a_mock
- error_creating = WaiterError(
- name="test_name",
- reason="test_reason",
- last_response={"fargateProfile": {"status": "DELETING"}},
- )
- error_failed = WaiterError(
- name="test_name",
- reason="Waiter encountered a terminal failure state:",
- last_response={"fargateProfile": {"status": "DELETE_FAILED"}},
- )
- a_mock.get_waiter().wait = AsyncMock(side_effect=[error_creating, error_creating, error_failed])
- mock_sleep.return_value = True
-
- eks_delete_fargate_profile_trigger = EksDeleteFargateProfileTrigger(
- cluster_name=TEST_CLUSTER_IDENTIFIER,
- fargate_profile_name=TEST_FARGATE_PROFILE_NAME,
- aws_conn_id=TEST_AWS_CONN_ID,
- waiter_delay=TEST_WAITER_DELAY,
- waiter_max_attempts=TEST_WAITER_MAX_ATTEMPTS,
- )
- with pytest.raises(AirflowException) as exc:
- generator = eks_delete_fargate_profile_trigger.run()
- await generator.asend(None)
- assert f"Delete Fargate Profile failed: {error_failed}" in str(exc.value)
- assert a_mock.get_waiter().wait.call_count == 3
-
-
-class TestEksNodegroupTrigger:
- def test_eks_nodegroup_trigger_serialize(self):
- eks_nodegroup_trigger = EksNodegroupTrigger(
- waiter_name="test_waiter_name",
- cluster_name=TEST_CLUSTER_IDENTIFIER,
- nodegroup_name=TEST_NODEGROUP_NAME,
- aws_conn_id=TEST_AWS_CONN_ID,
- waiter_delay=TEST_WAITER_DELAY,
- waiter_max_attempts=TEST_WAITER_MAX_ATTEMPTS,
- region=TEST_REGION,
- )
-
- class_path, args = eks_nodegroup_trigger.serialize()
- assert class_path == "airflow.providers.amazon.aws.triggers.eks.EksNodegroupTrigger"
- assert args["waiter_name"] == "test_waiter_name"
- assert args["cluster_name"] == TEST_CLUSTER_IDENTIFIER
- assert args["nodegroup_name"] == TEST_NODEGROUP_NAME
- assert args["aws_conn_id"] == TEST_AWS_CONN_ID
- assert args["waiter_delay"] == str(TEST_WAITER_DELAY)
- assert args["waiter_max_attempts"] == str(TEST_WAITER_MAX_ATTEMPTS)
- assert args["region"] == TEST_REGION
-
- @pytest.mark.asyncio
- @mock.patch.object(EksHook, "async_conn")
- async def test_eks_nodegroup_trigger_run(self, mock_async_conn):
- a_mock = mock.MagicMock()
- mock_async_conn.__aenter__.return_value = a_mock
-
- a_mock.get_waiter().wait = AsyncMock()
-
- eks_nodegroup_trigger = EksNodegroupTrigger(
- waiter_name="test_waiter_name",
- cluster_name=TEST_CLUSTER_IDENTIFIER,
- nodegroup_name=TEST_NODEGROUP_NAME,
- aws_conn_id=TEST_AWS_CONN_ID,
- waiter_delay=TEST_WAITER_DELAY,
- waiter_max_attempts=TEST_WAITER_MAX_ATTEMPTS,
- region=TEST_REGION,
- )
-
- generator = eks_nodegroup_trigger.run()
- response = await generator.asend(None)
-
- assert response == TriggerEvent(
- {
- "status": "success",
- "cluster_name": TEST_CLUSTER_IDENTIFIER,
- "nodegroup_name": TEST_NODEGROUP_NAME,
- }
- )
-
- @pytest.mark.asyncio
- @mock.patch("asyncio.sleep")
- @mock.patch.object(EksHook, "async_conn")
- async def test_eks_nodegroup_trigger_run_multiple_attempts(self, mock_async_conn, mock_sleep):
- mock_sleep.return_value = True
- a_mock = mock.MagicMock()
- mock_async_conn.__aenter__.return_value = a_mock
- error = WaiterError(
- name="test_name",
- reason="test_reason",
- last_response={"nodegroup": {"status": "CREATING"}},
- )
- a_mock.get_waiter().wait = AsyncMock(side_effect=[error, error, error, True])
-
- eks_nodegroup_trigger = EksNodegroupTrigger(
- waiter_name="test_waiter_name",
- cluster_name=TEST_CLUSTER_IDENTIFIER,
- nodegroup_name=TEST_NODEGROUP_NAME,
- aws_conn_id=TEST_AWS_CONN_ID,
- waiter_delay=TEST_WAITER_DELAY,
- waiter_max_attempts=TEST_WAITER_MAX_ATTEMPTS,
- region=TEST_REGION,
- )
-
- generator = eks_nodegroup_trigger.run()
- response = await generator.asend(None)
- assert a_mock.get_waiter().wait.call_count == 4
- assert response == TriggerEvent(
- {
- "status": "success",
- "cluster_name": TEST_CLUSTER_IDENTIFIER,
- "nodegroup_name": TEST_NODEGROUP_NAME,
- }
- )
-
- @pytest.mark.asyncio
- @mock.patch("asyncio.sleep")
- @mock.patch.object(EksHook, "async_conn")
- async def test_eks_nodegroup_trigger_run_attempts_exceeded(self, mock_async_conn, mock_sleep):
- mock_sleep.return_value = True
- a_mock = mock.MagicMock()
- mock_async_conn.__aenter__.return_value = a_mock
- error = WaiterError(
- name="test_name",
- reason="test_reason",
- last_response={"nodegroup": {"status": "CREATING"}},
- )
- a_mock.get_waiter().wait = AsyncMock(side_effect=[error, error, error, True])
-
- eks_nodegroup_trigger = EksNodegroupTrigger(
- waiter_name="test_waiter_name",
- cluster_name=TEST_CLUSTER_IDENTIFIER,
- nodegroup_name=TEST_NODEGROUP_NAME,
- aws_conn_id=TEST_AWS_CONN_ID,
- waiter_delay=TEST_WAITER_DELAY,
- waiter_max_attempts=2,
- region=TEST_REGION,
- )
-
- with pytest.raises(AirflowException) as exc:
- generator = eks_nodegroup_trigger.run()
- await generator.asend(None)
- assert "Waiter error: max attempts reached" in str(exc.value)
- assert a_mock.get_waiter().wait.call_count == 2
-
- @pytest.mark.asyncio
- @mock.patch("asyncio.sleep")
- @mock.patch.object(EksHook, "async_conn")
- async def test_eks_nodegroup_trigger_run_attempts_failed(self, mock_async_conn, mock_sleep):
- mock_sleep.return_value = True
- a_mock = mock.MagicMock()
- mock_async_conn.__aenter__.return_value = a_mock
- error_creating = WaiterError(
- name="test_name",
- reason="test_reason",
- last_response={"nodegroup": {"status": "CREATING"}},
- )
- error_failed = WaiterError(
- name="test_name",
- reason="Waiter encountered a terminal failure state:",
- last_response={"nodegroup": {"status": "DELETE_FAILED"}},
- )
- a_mock.get_waiter().wait = AsyncMock(side_effect=[error_creating, error_creating, error_failed])
- mock_sleep.return_value = True
-
- eks_nodegroup_trigger = EksNodegroupTrigger(
- waiter_name="test_waiter_name",
- cluster_name=TEST_CLUSTER_IDENTIFIER,
- nodegroup_name=TEST_NODEGROUP_NAME,
- aws_conn_id=TEST_AWS_CONN_ID,
- waiter_delay=TEST_WAITER_DELAY,
- waiter_max_attempts=TEST_WAITER_MAX_ATTEMPTS,
- region=TEST_REGION,
- )
- with pytest.raises(AirflowException) as exc:
- generator = eks_nodegroup_trigger.run()
- await generator.asend(None)
-
- assert "Error checking nodegroup" in str(exc.value)
- assert a_mock.get_waiter().wait.call_count == 3
+class TestEksTriggers:
+ @pytest.mark.parametrize(
Review Comment:
I do love a well-parameterized test. It's beautiful.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ferruzzi commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253594423
##########
airflow/providers/amazon/aws/triggers/emr.py:
##########
@@ -102,274 +102,184 @@ async def run(self):
yield TriggerEvent({"status": "success", "message": "Steps completed", "step_ids": self.step_ids})
-class EmrCreateJobFlowTrigger(BaseTrigger):
+class EmrCreateJobFlowTrigger(AwsBaseWaiterTrigger):
"""
Asynchronously poll the boto3 API and wait for the JobFlow to finish executing.
:param job_flow_id: The id of the job flow to wait for.
- :param poll_interval: The amount of time in seconds to wait between attempts.
- :param max_attempts: The maximum number of attempts to be made.
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
:param aws_conn_id: The Airflow connection used for AWS credentials.
"""
def __init__(
self,
job_flow_id: str,
- poll_interval: int,
- max_attempts: int,
- aws_conn_id: str,
+ poll_interval: int | None = None, # deprecated
+ max_attempts: int | None = None, # deprecated
+ aws_conn_id: str | None = None,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
):
- self.job_flow_id = job_flow_id
- self.poll_interval = poll_interval
- self.max_attempts = max_attempts
- self.aws_conn_id = aws_conn_id
-
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "job_flow_id": self.job_flow_id,
- "poll_interval": str(self.poll_interval),
- "max_attempts": str(self.max_attempts),
- "aws_conn_id": self.aws_conn_id,
- },
+ if poll_interval is not None or max_attempts is not None:
Review Comment:
```suggestion
if poll_interval or max_attempts:
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ferruzzi commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253587187
##########
airflow/providers/amazon/aws/triggers/ecs.py:
##########
@@ -22,68 +22,89 @@
from botocore.exceptions import ClientError, WaiterError
+from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
from airflow.providers.amazon.aws.hooks.ecs import EcsHook
from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
+from airflow.providers.amazon.aws.triggers.base_trigger import AwsBaseWaiterTrigger
from airflow.providers.amazon.aws.utils.task_log_fetcher import AwsTaskLogFetcher
-from airflow.providers.amazon.aws.utils.waiter_with_logging import async_wait
from airflow.triggers.base import BaseTrigger, TriggerEvent
-class ClusterWaiterTrigger(BaseTrigger):
+class ClusterActiveTrigger(AwsBaseWaiterTrigger):
"""
- Polls the status of a cluster using a given waiter. Can be used to poll for an active or inactive cluster.
+ Polls the status of a cluster until it's ready.
Review Comment:
I'm all for shorter docs, but maybe "until it is `active`"?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] syedahsn commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "syedahsn (via GitHub)" <gi...@apache.org>.
syedahsn commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253530855
##########
airflow/providers/amazon/aws/triggers/base_trigger.py:
##########
@@ -0,0 +1,129 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import Any, AsyncIterator
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+from airflow.providers.amazon.aws.utils.waiter_with_logging import async_wait
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+
+
+class AwsBaseWaiterTrigger(BaseTrigger):
+ """
+ Base class for all AWS Triggers that follow the "standard" model of just waiting on a waiter.
+
+ Subclasses need to implement the hook() method.
+
+ :param serialized_fields: Fields that are specific to the subclass trigger and need to be serialized
+ to be passed to the __init__ method on deserialization.
+ format: {<parameter_name>: <parameter_value>}
Review Comment:
This should specify which fields are serialized by default. For example, we don't want someone passing `aws_conn_id` as a serialized field, because even though it needs to be serialized, it is done by default. What do you think?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] syedahsn commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "syedahsn (via GitHub)" <gi...@apache.org>.
syedahsn commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253537693
##########
airflow/providers/amazon/aws/triggers/eks.py:
##########
@@ -115,59 +90,37 @@ def __init__(
waiter_max_attempts: int,
aws_conn_id: str,
region: str | None = None,
+ region_name: str | None = None,
):
- self.cluster_name = cluster_name
- self.fargate_profile_name = fargate_profile_name
- self.waiter_delay = waiter_delay
- self.waiter_max_attempts = waiter_max_attempts
- self.aws_conn_id = aws_conn_id
- self.region = region
-
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "cluster_name": self.cluster_name,
- "fargate_profile_name": self.fargate_profile_name,
- "waiter_delay": str(self.waiter_delay),
- "waiter_max_attempts": str(self.waiter_max_attempts),
- "aws_conn_id": self.aws_conn_id,
- "region": self.region,
- },
+ if region is not None:
+ warnings.warn(
+ "please use region_name param instead of region",
+ AirflowProviderDeprecationWarning,
+ stacklevel=2,
+ )
+ region_name = region
+
+ super().__init__(
+ serialized_fields={"cluster_name": cluster_name, "fargate_profile_name": fargate_profile_name},
+ waiter_name="fargate_profile_deleted",
+ waiter_args={"clusterName": cluster_name, "fargateProfileName": fargate_profile_name},
+ failure_message="Failure while deleting Fargate profile",
+ status_message="Fargate profile not deleted yet",
+ status_queries=["fargateProfile.status"],
+ return_value=None,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ aws_conn_id=aws_conn_id,
+ region_name=region_name,
)
- async def run(self):
- self.hook = EksHook(aws_conn_id=self.aws_conn_id, region_name=self.region)
- async with self.hook.async_conn as client:
- attempt = 0
- waiter = client.get_waiter("fargate_profile_deleted")
- while attempt < int(self.waiter_max_attempts):
- attempt += 1
- try:
- await waiter.wait(
- clusterName=self.cluster_name,
- fargateProfileName=self.fargate_profile_name,
- WaiterConfig={"Delay": int(self.waiter_delay), "MaxAttempts": 1},
- )
- break
- except WaiterError as error:
- if "terminal failure" in str(error):
- raise AirflowException(f"Delete Fargate Profile failed: {error}")
- self.log.info(
- "Status of fargate profile is %s", error.last_response["fargateProfile"]["status"]
- )
- await asyncio.sleep(int(self.waiter_delay))
- if attempt >= int(self.waiter_max_attempts):
- raise AirflowException(
- f"Delete Fargate Profile failed - max attempts reached: {self.waiter_max_attempts}"
- )
- else:
- yield TriggerEvent({"status": "success", "message": "Fargate Profile Deleted"})
+ def hook(self) -> AwsGenericHook:
+ return EksHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
-class EksNodegroupTrigger(BaseTrigger):
+class EksCreateNodegroupTrigger(AwsBaseWaiterTrigger):
Review Comment:
Why are you making a separate trigger for create/delete here? Is it no longer possible to use a generic Trigger if the responsibility is to just poll for a particular state (depending on `waiter_name`)?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ferruzzi commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253594695
##########
airflow/providers/amazon/aws/triggers/emr.py:
##########
@@ -102,274 +102,184 @@ async def run(self):
yield TriggerEvent({"status": "success", "message": "Steps completed", "step_ids": self.step_ids})
-class EmrCreateJobFlowTrigger(BaseTrigger):
+class EmrCreateJobFlowTrigger(AwsBaseWaiterTrigger):
"""
Asynchronously poll the boto3 API and wait for the JobFlow to finish executing.
:param job_flow_id: The id of the job flow to wait for.
- :param poll_interval: The amount of time in seconds to wait between attempts.
- :param max_attempts: The maximum number of attempts to be made.
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
:param aws_conn_id: The Airflow connection used for AWS credentials.
"""
def __init__(
self,
job_flow_id: str,
- poll_interval: int,
- max_attempts: int,
- aws_conn_id: str,
+ poll_interval: int | None = None, # deprecated
+ max_attempts: int | None = None, # deprecated
+ aws_conn_id: str | None = None,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
):
- self.job_flow_id = job_flow_id
- self.poll_interval = poll_interval
- self.max_attempts = max_attempts
- self.aws_conn_id = aws_conn_id
-
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "job_flow_id": self.job_flow_id,
- "poll_interval": str(self.poll_interval),
- "max_attempts": str(self.max_attempts),
- "aws_conn_id": self.aws_conn_id,
- },
+ if poll_interval is not None or max_attempts is not None:
+ warnings.warn(
+ "please use waiter_delay instead of poll_interval "
+ "and waiter_max_attempts instead of max_attempts",
+ AirflowProviderDeprecationWarning,
+ stacklevel=2,
+ )
+ waiter_delay = poll_interval or waiter_delay
+ waiter_max_attempts = max_attempts or waiter_max_attempts
+ super().__init__(
+ serialized_fields={"job_flow_id": job_flow_id},
+ waiter_name="job_flow_waiting",
+ waiter_args={"ClusterId": job_flow_id},
+ failure_message="JobFlow creation failed",
+ status_message="JobFlow creation in progress",
+ status_queries=[
+ "Cluster.Status.State",
+ "Cluster.Status.StateChangeReason",
+ "Cluster.Status.ErrorDetails",
+ ],
+ return_key="job_flow_id",
+ return_value=job_flow_id,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ aws_conn_id=aws_conn_id,
)
- async def run(self):
- self.hook = EmrHook(aws_conn_id=self.aws_conn_id)
- async with self.hook.async_conn as client:
- attempt = 0
- waiter = self.hook.get_waiter("job_flow_waiting", deferrable=True, client=client)
- while attempt < int(self.max_attempts):
- attempt = attempt + 1
- try:
- await waiter.wait(
- ClusterId=self.job_flow_id,
- WaiterConfig=prune_dict(
- {
- "Delay": self.poll_interval,
- "MaxAttempts": 1,
- }
- ),
- )
- break
- except WaiterError as error:
- if "terminal failure" in str(error):
- raise AirflowException(f"JobFlow creation failed: {error}")
- self.log.info(
- "Status of jobflow is %s - %s",
- error.last_response["Cluster"]["Status"]["State"],
- error.last_response["Cluster"]["Status"]["StateChangeReason"],
- )
- await asyncio.sleep(int(self.poll_interval))
- if attempt >= int(self.max_attempts):
- raise AirflowException(f"JobFlow creation failed - max attempts reached: {self.max_attempts}")
- else:
- yield TriggerEvent(
- {
- "status": "success",
- "message": "JobFlow completed successfully",
- "job_flow_id": self.job_flow_id,
- }
- )
+ def hook(self) -> AwsGenericHook:
+ return EmrHook(aws_conn_id=self.aws_conn_id)
-class EmrTerminateJobFlowTrigger(BaseTrigger):
+class EmrTerminateJobFlowTrigger(AwsBaseWaiterTrigger):
"""
Asynchronously poll the boto3 API and wait for the JobFlow to finish terminating.
:param job_flow_id: ID of the EMR Job Flow to terminate
- :param poll_interval: The amount of time in seconds to wait between attempts.
- :param max_attempts: The maximum number of attempts to be made.
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
:param aws_conn_id: The Airflow connection used for AWS credentials.
"""
def __init__(
self,
job_flow_id: str,
- poll_interval: int,
- max_attempts: int,
- aws_conn_id: str,
+ poll_interval: int | None = None, # deprecated
+ max_attempts: int | None = None, # deprecated
+ aws_conn_id: str | None = None,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
):
- self.job_flow_id = job_flow_id
- self.poll_interval = poll_interval
- self.max_attempts = max_attempts
- self.aws_conn_id = aws_conn_id
-
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "job_flow_id": self.job_flow_id,
- "poll_interval": str(self.poll_interval),
- "max_attempts": str(self.max_attempts),
- "aws_conn_id": self.aws_conn_id,
- },
+ if poll_interval is not None or max_attempts is not None:
Review Comment:
```suggestion
if poll_interval or max_attempts:
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ferruzzi commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253592879
##########
airflow/providers/amazon/aws/triggers/eks.py:
##########
@@ -46,57 +43,35 @@ def __init__(
waiter_max_attempts: int,
aws_conn_id: str,
region: str | None = None,
+ region_name: str | None = None,
):
- self.cluster_name = cluster_name
- self.fargate_profile_name = fargate_profile_name
- self.waiter_delay = waiter_delay
- self.waiter_max_attempts = waiter_max_attempts
- self.aws_conn_id = aws_conn_id
- self.region = region
-
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "cluster_name": self.cluster_name,
- "fargate_profile_name": self.fargate_profile_name,
- "waiter_delay": str(self.waiter_delay),
- "waiter_max_attempts": str(self.waiter_max_attempts),
- "aws_conn_id": self.aws_conn_id,
- "region": self.region,
- },
+ if region is not None:
Review Comment:
```suggestion
if region:
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ferruzzi commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253579117
##########
airflow/providers/amazon/aws/sensors/batch.py:
##########
@@ -113,9 +114,10 @@ def execute_complete(self, context: Context, event: dict[str, Any]) -> None:
Relies on trigger to throw an exception, otherwise it assumes execution was successful.
"""
- if "status" in event and event["status"] == "failure":
Review Comment:
Are we no longer worried about an IndexError here if 'status' is not in event??
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vandonr-amz commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "vandonr-amz (via GitHub)" <gi...@apache.org>.
vandonr-amz commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253606033
##########
airflow/providers/amazon/aws/operators/ecs.py:
##########
@@ -139,13 +143,12 @@ def execute(self, context: Context):
self.log.info("Cluster %r in state: %r.", self.cluster_name, cluster_state)
elif self.deferrable:
self.defer(
- trigger=ClusterWaiterTrigger(
- waiter_name="cluster_active",
+ trigger=ClusterActiveTrigger(
cluster_arn=cluster_details["clusterArn"],
Review Comment:
yeah, maybe something to do in a base operator, because there is usually only one trigger being built per operator, so you'd have to pack those parameters once per operator, and use them once, and it's not really saving anything ? Just making the code harder to read by moving the parameters away from where they are passed ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vandonr-amz commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "vandonr-amz (via GitHub)" <gi...@apache.org>.
vandonr-amz commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253615702
##########
airflow/providers/amazon/aws/triggers/eks.py:
##########
@@ -46,57 +43,35 @@ def __init__(
waiter_max_attempts: int,
aws_conn_id: str,
region: str | None = None,
+ region_name: str | None = None,
):
- self.cluster_name = cluster_name
- self.fargate_profile_name = fargate_profile_name
- self.waiter_delay = waiter_delay
- self.waiter_max_attempts = waiter_max_attempts
- self.aws_conn_id = aws_conn_id
- self.region = region
-
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "cluster_name": self.cluster_name,
- "fargate_profile_name": self.fargate_profile_name,
- "waiter_delay": str(self.waiter_delay),
- "waiter_max_attempts": str(self.waiter_max_attempts),
- "aws_conn_id": self.aws_conn_id,
- "region": self.region,
- },
+ if region is not None:
Review Comment:
I tend to be cautious with checks like this, because it'll trigger on empty string as well (and empty array, empty dict, etc.), and I've been bitten by this several times.
Here it's not a problem, but I prefer to be more explicit in general.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vandonr-amz commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "vandonr-amz (via GitHub)" <gi...@apache.org>.
vandonr-amz commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253523315
##########
airflow/providers/amazon/provider.yaml:
##########
@@ -516,6 +516,9 @@ hooks:
- airflow.providers.amazon.aws.hooks.appflow
triggers:
+ - integration-name: Amazon Web Services
+ python-modules:
+ - airflow.providers.amazon.aws.triggers.base_trigger
Review Comment:
hmm, yeah, I could name the file `base` instead of `base_trigger`, but what would it change ? How is it better ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] syedahsn commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "syedahsn (via GitHub)" <gi...@apache.org>.
syedahsn commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253523543
##########
airflow/providers/amazon/aws/hooks/athena.py:
##########
@@ -253,7 +253,7 @@ def poll_query_status(
try:
wait(
waiter=self.get_waiter("query_complete"),
- waiter_delay=sleep_time or self.sleep_time,
+ waiter_delay=self.sleep_time if sleep_time is None else sleep_time,
Review Comment:
This is fine, but you can also just mock the sleep function no?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ferruzzi commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253590001
##########
airflow/providers/amazon/aws/triggers/ecs.py:
##########
@@ -22,68 +22,89 @@
from botocore.exceptions import ClientError, WaiterError
+from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
from airflow.providers.amazon.aws.hooks.ecs import EcsHook
from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
+from airflow.providers.amazon.aws.triggers.base_trigger import AwsBaseWaiterTrigger
from airflow.providers.amazon.aws.utils.task_log_fetcher import AwsTaskLogFetcher
-from airflow.providers.amazon.aws.utils.waiter_with_logging import async_wait
from airflow.triggers.base import BaseTrigger, TriggerEvent
-class ClusterWaiterTrigger(BaseTrigger):
+class ClusterActiveTrigger(AwsBaseWaiterTrigger):
"""
- Polls the status of a cluster using a given waiter. Can be used to poll for an active or inactive cluster.
+ Polls the status of a cluster until it's ready.
- :param waiter_name: Name of the waiter to use, for instance 'cluster_active' or 'cluster_inactive'
:param cluster_arn: ARN of the cluster to watch.
:param waiter_delay: The amount of time in seconds to wait between attempts.
:param waiter_max_attempts: The number of times to ping for status.
Will fail after that many unsuccessful attempts.
:param aws_conn_id: The Airflow connection used for AWS credentials.
- :param region: The AWS region where the cluster is located.
+ :param region_name: The AWS region where the cluster is located.
"""
def __init__(
self,
- waiter_name: str,
cluster_arn: str,
- waiter_delay: int | None,
- waiter_max_attempts: int | None,
+ waiter_delay: int,
+ waiter_max_attempts: int,
aws_conn_id: str | None,
- region: str | None,
+ region_name: str | None,
):
- self.cluster_arn = cluster_arn
- self.waiter_name = waiter_name
- self.waiter_delay = waiter_delay if waiter_delay is not None else 15 # written like this to allow 0
- self.attempts = waiter_max_attempts or 999999999
- self.aws_conn_id = aws_conn_id
- self.region = region
+ super().__init__(
+ serialized_fields={"cluster_arn": cluster_arn},
+ waiter_name="cluster_active",
+ waiter_args={"clusters": [cluster_arn]},
+ failure_message="Failure while waiting for cluster to be available",
+ status_message="Cluster is not ready yet",
+ status_queries=["clusters[].status", "failures"],
+ return_key="arn",
+ return_value=cluster_arn,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ aws_conn_id=aws_conn_id,
+ region_name=region_name,
+ )
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "waiter_name": self.waiter_name,
- "cluster_arn": self.cluster_arn,
- "waiter_delay": self.waiter_delay,
- "waiter_max_attempts": self.attempts,
- "aws_conn_id": self.aws_conn_id,
- "region": self.region,
- },
+ def hook(self) -> AwsGenericHook:
+ return EcsHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
+
+
+class ClusterInactiveTrigger(AwsBaseWaiterTrigger):
Review Comment:
I may be over-complicating this, but it looooks like the only difference between `ClusterActiveTrigger` and `ClusterInactiveTrigger` is the `waiter_name` value and the two message. In which case, why not have them both inherit a `ClusterStatusTrigger` which accepts those three values and drop all the repetition?
Same below where you split other Triggers out into status-specific ones.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] ferruzzi commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "ferruzzi (via GitHub)" <gi...@apache.org>.
ferruzzi commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253590001
##########
airflow/providers/amazon/aws/triggers/ecs.py:
##########
@@ -22,68 +22,89 @@
from botocore.exceptions import ClientError, WaiterError
+from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
from airflow.providers.amazon.aws.hooks.ecs import EcsHook
from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
+from airflow.providers.amazon.aws.triggers.base_trigger import AwsBaseWaiterTrigger
from airflow.providers.amazon.aws.utils.task_log_fetcher import AwsTaskLogFetcher
-from airflow.providers.amazon.aws.utils.waiter_with_logging import async_wait
from airflow.triggers.base import BaseTrigger, TriggerEvent
-class ClusterWaiterTrigger(BaseTrigger):
+class ClusterActiveTrigger(AwsBaseWaiterTrigger):
"""
- Polls the status of a cluster using a given waiter. Can be used to poll for an active or inactive cluster.
+ Polls the status of a cluster until it's ready.
- :param waiter_name: Name of the waiter to use, for instance 'cluster_active' or 'cluster_inactive'
:param cluster_arn: ARN of the cluster to watch.
:param waiter_delay: The amount of time in seconds to wait between attempts.
:param waiter_max_attempts: The number of times to ping for status.
Will fail after that many unsuccessful attempts.
:param aws_conn_id: The Airflow connection used for AWS credentials.
- :param region: The AWS region where the cluster is located.
+ :param region_name: The AWS region where the cluster is located.
"""
def __init__(
self,
- waiter_name: str,
cluster_arn: str,
- waiter_delay: int | None,
- waiter_max_attempts: int | None,
+ waiter_delay: int,
+ waiter_max_attempts: int,
aws_conn_id: str | None,
- region: str | None,
+ region_name: str | None,
):
- self.cluster_arn = cluster_arn
- self.waiter_name = waiter_name
- self.waiter_delay = waiter_delay if waiter_delay is not None else 15 # written like this to allow 0
- self.attempts = waiter_max_attempts or 999999999
- self.aws_conn_id = aws_conn_id
- self.region = region
+ super().__init__(
+ serialized_fields={"cluster_arn": cluster_arn},
+ waiter_name="cluster_active",
+ waiter_args={"clusters": [cluster_arn]},
+ failure_message="Failure while waiting for cluster to be available",
+ status_message="Cluster is not ready yet",
+ status_queries=["clusters[].status", "failures"],
+ return_key="arn",
+ return_value=cluster_arn,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ aws_conn_id=aws_conn_id,
+ region_name=region_name,
+ )
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "waiter_name": self.waiter_name,
- "cluster_arn": self.cluster_arn,
- "waiter_delay": self.waiter_delay,
- "waiter_max_attempts": self.attempts,
- "aws_conn_id": self.aws_conn_id,
- "region": self.region,
- },
+ def hook(self) -> AwsGenericHook:
+ return EcsHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
+
+
+class ClusterInactiveTrigger(AwsBaseWaiterTrigger):
Review Comment:
I may be over-complicating this, but it looooks like the only difference between `ClusterActiveTrigger` and `ClusterInactiveTrigger` is the `waiter_name` value and the two message. In which case, why not have them both inherit a `ClusterStatusTrigger` which accepts those three values and drop all the repetition?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vandonr-amz commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "vandonr-amz (via GitHub)" <gi...@apache.org>.
vandonr-amz commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253596043
##########
airflow/providers/amazon/aws/triggers/base_trigger.py:
##########
@@ -0,0 +1,129 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import Any, AsyncIterator
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+from airflow.providers.amazon.aws.utils.waiter_with_logging import async_wait
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+
+
+class AwsBaseWaiterTrigger(BaseTrigger):
+ """
+ Base class for all AWS Triggers that follow the "standard" model of just waiting on a waiter.
+
+ Subclasses need to implement the hook() method.
+
+ :param serialized_fields: Fields that are specific to the subclass trigger and need to be serialized
+ to be passed to the __init__ method on deserialization.
+ format: {<parameter_name>: <parameter_value>}
Review Comment:
hmm, yeah. From the way it's used, it'd be mostly harmless, just unnecessary. I can add it to the comments.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] syedahsn commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "syedahsn (via GitHub)" <gi...@apache.org>.
syedahsn commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253675457
##########
tests/providers/amazon/aws/triggers/test_base_trigger.py:
##########
@@ -0,0 +1,89 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from unittest import mock
+from unittest.mock import MagicMock
+
+import pytest
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
+from airflow.triggers.base import TriggerEvent
+
+
+class TestImplem(AwsBaseWaiterTrigger):
+ """An empty implementation that allows instantiation for tests."""
+
+ __test__ = False
+
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+
+ def hook(self) -> AwsGenericHook:
+ return MagicMock()
+
+
+class TestAwsBaseWaiterTrigger:
+ def setup_method(self):
+ self.trigger = TestImplem(
+ serialized_fields={},
+ waiter_name="",
+ waiter_args={},
+ failure_message="",
+ status_message="",
+ status_queries=[],
+ return_value=None,
+ waiter_delay=0,
+ waiter_max_attempts=0,
+ aws_conn_id="",
+ )
+
+ def test_region_serialized(self):
+ self.trigger.region_name = "my_region"
+ _, args = self.trigger.serialize()
+
+ assert "region_name" in args
+ assert args["region_name"] == "my_region"
+
+ def test_region_not_serialized_if_omitted(self):
+ _, args = self.trigger.serialize()
+
+ assert "region_name" not in args
+
+ def test_serialize_extra_fields(self):
+ self.trigger.serialized_fields = {"foo": "bar", "foz": "baz"}
+
+ _, args = self.trigger.serialize()
+
+ assert "foo" in args
+ assert args["foo"] == "bar"
+ assert "foz" in args
+ assert args["foz"] == "baz"
+
+ @pytest.mark.asyncio
Review Comment:
What about testing things like multiple attempts, attempts exceeded etc? I suppose those are covered in [`test_waiter_with_logging.py`](https://github.com/apache/airflow/blob/main/tests/providers/amazon/aws/utils/test_waiter_with_logging.py) so they don't need to be repeated.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vandonr-amz commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "vandonr-amz (via GitHub)" <gi...@apache.org>.
vandonr-amz commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253616672
##########
airflow/providers/amazon/aws/triggers/emr.py:
##########
@@ -102,274 +102,184 @@ async def run(self):
yield TriggerEvent({"status": "success", "message": "Steps completed", "step_ids": self.step_ids})
-class EmrCreateJobFlowTrigger(BaseTrigger):
+class EmrCreateJobFlowTrigger(AwsBaseWaiterTrigger):
"""
Asynchronously poll the boto3 API and wait for the JobFlow to finish executing.
:param job_flow_id: The id of the job flow to wait for.
- :param poll_interval: The amount of time in seconds to wait between attempts.
- :param max_attempts: The maximum number of attempts to be made.
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
:param aws_conn_id: The Airflow connection used for AWS credentials.
"""
def __init__(
self,
job_flow_id: str,
- poll_interval: int,
- max_attempts: int,
- aws_conn_id: str,
+ poll_interval: int | None = None, # deprecated
+ max_attempts: int | None = None, # deprecated
+ aws_conn_id: str | None = None,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
):
- self.job_flow_id = job_flow_id
- self.poll_interval = poll_interval
- self.max_attempts = max_attempts
- self.aws_conn_id = aws_conn_id
-
- def serialize(self) -> tuple[str, dict[str, Any]]:
- return (
- self.__class__.__module__ + "." + self.__class__.__qualname__,
- {
- "job_flow_id": self.job_flow_id,
- "poll_interval": str(self.poll_interval),
- "max_attempts": str(self.max_attempts),
- "aws_conn_id": self.aws_conn_id,
- },
+ if poll_interval is not None or max_attempts is not None:
Review Comment:
same remark as above, I think I'd like the warning to show even if the value is zero, hence the explicit check
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] vandonr-amz commented on a diff in pull request #32274: Introduce a base class for aws triggers
Posted by "vandonr-amz (via GitHub)" <gi...@apache.org>.
vandonr-amz commented on code in PR #32274:
URL: https://github.com/apache/airflow/pull/32274#discussion_r1253600626
##########
airflow/providers/amazon/aws/hooks/athena.py:
##########
@@ -253,7 +253,7 @@ def poll_query_status(
try:
wait(
waiter=self.get_waiter("query_complete"),
- waiter_delay=sleep_time or self.sleep_time,
+ waiter_delay=self.sleep_time if sleep_time is None else sleep_time,
Review Comment:
I think I wrote it this way to allow `sleep_time = 0`. `is None` will keep the zero value, but `or` will replace it with the alternative, and being able to specify a sleep time of 0 is useful in tests.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org