From 31b176730df6b0597b8d749e9cbae696b5c654bd Mon Sep 17 00:00:00 2001 From: Syed Hussain Date: Mon, 6 Mar 2023 00:25:08 -0800 Subject: [PATCH 01/17] Change base_aws.py to support async_conn Add async custom waiter support in get_waiter, and base_waiter.py Add Deferrable mode to RedshiftCreateClusterOperator Add RedshiftCreateClusterTrigger and unit test Add README.md for writing Triggers for AMPP --- .../providers/amazon/aws/hooks/base_aws.py | 68 ++++++-- .../amazon/aws/operators/redshift_cluster.py | 19 ++- .../providers/amazon/aws/triggers/README.md | 153 ++++++++++++++++++ .../amazon/aws/triggers/redshift_cluster.py | 55 ++++++- .../amazon/aws/waiters/base_waiter.py | 8 +- airflow/providers/amazon/provider.yaml | 1 + generated/provider_dependencies.json | 1 + .../providers/amazon/aws/triggers/__init__.py | 16 ++ .../aws/triggers/test_redshift_cluster.py | 74 +++++++++ 9 files changed, 376 insertions(+), 19 deletions(-) create mode 100644 airflow/providers/amazon/aws/triggers/README.md create mode 100644 tests/providers/amazon/aws/triggers/__init__.py create mode 100644 tests/providers/amazon/aws/triggers/test_redshift_cluster.py diff --git a/airflow/providers/amazon/aws/hooks/base_aws.py b/airflow/providers/amazon/aws/hooks/base_aws.py index 3395990fc343e..09e7126017756 100644 --- a/airflow/providers/amazon/aws/hooks/base_aws.py +++ b/airflow/providers/amazon/aws/hooks/base_aws.py @@ -43,6 +43,7 @@ import jinja2 import requests import tenacity +from aiobotocore.session import AioSession, get_session as async_get_session from botocore.client import ClientMeta from botocore.config import Config from botocore.credentials import ReadOnlyCredentials @@ -72,7 +73,7 @@ class BaseSessionFactory(LoggingMixin): """ - Base AWS Session Factory class to handle boto3 session creation. + Base AWS Session Factory class to handle synchronous and async boto session creation. It can handle most of the AWS supported authentication methods. User can also derive from this class to have full control of boto3 session @@ -127,17 +128,18 @@ def role_arn(self) -> str | None: """Assume Role ARN from AWS Connection""" return self.conn.role_arn - def create_session(self) -> boto3.session.Session: - """Create boto3 Session from connection config.""" + def create_session(self, deferrable: bool = False) -> boto3.session.Session: + """Create boto3 or aiobotocore Session from connection config.""" if not self.conn: self.log.info( "No connection ID provided. Fallback on boto3 credential strategy (region_name=%r). " "See: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html", self.region_name, ) - return boto3.session.Session(region_name=self.region_name) + return async_get_session() if deferrable else boto3.session.Session(region_name=self.region_name) + elif not self.role_arn: - return self.basic_session + return async_get_session() if deferrable else self.basic_session # Values stored in ``AwsConnectionWrapper.session_kwargs`` are intended to be used only # to create the initial boto3 session. @@ -150,12 +152,16 @@ def create_session(self) -> boto3.session.Session: assume_session_kwargs = {} if self.conn.region_name: assume_session_kwargs["region_name"] = self.conn.region_name - return self._create_session_with_assume_role(session_kwargs=assume_session_kwargs) + return self._create_session_with_assume_role( + session_kwargs=assume_session_kwargs, deferrable=deferrable + ) def _create_basic_session(self, session_kwargs: dict[str, Any]) -> boto3.session.Session: return boto3.session.Session(**session_kwargs) - def _create_session_with_assume_role(self, session_kwargs: dict[str, Any]) -> boto3.session.Session: + def _create_session_with_assume_role( + self, session_kwargs: dict[str, Any], deferrable: bool = False + ) -> boto3.session.Session: if self.conn.assume_role_method == "assume_role_with_web_identity": # Deferred credentials have no initial credentials credential_fetcher = self._get_web_identity_credential_fetcher() @@ -172,10 +178,12 @@ def _create_session_with_assume_role(self, session_kwargs: dict[str, Any]) -> bo method="sts-assume-role", ) - session = botocore.session.get_session() - session._credentials = credentials - region_name = self.basic_session.region_name - session.set_config_variable("region", region_name) + session = async_get_session() if deferrable else botocore.session.get_session() + + session.set_credentials( + access_key=credentials.access_key, secret_key=credentials.secret_key, token=credentials.token + ) + session.set_config_variable("region", self.basic_session.region_name) return boto3.session.Session(botocore_session=session, **session_kwargs) @@ -564,11 +572,11 @@ def verify(self) -> bool | str | None: """Verify or not SSL certificates boto3 client/resource read-only property.""" return self.conn_config.verify - def get_session(self, region_name: str | None = None) -> boto3.session.Session: + def get_session(self, region_name: str | None = None, deferrable: bool = False) -> boto3.session.Session: """Get the underlying boto3.session.Session(region_name=region_name).""" return SessionFactory( conn=self.conn_config, region_name=region_name, config=self.config - ).create_session() + ).create_session(deferrable=deferrable) def _get_config(self, config: Config | None = None) -> Config: """ @@ -591,10 +599,19 @@ def get_client_type( self, region_name: str | None = None, config: Config | None = None, + deferrable: bool = False, ) -> boto3.client: """Get the underlying boto3 client using boto3 session""" client_type = self.client_type - session = self.get_session(region_name=region_name) + session = self.get_session(region_name=region_name, deferrable=deferrable) + if isinstance(session, AioSession): + return session.create_client( + client_type, + endpoint_url=self.conn_config.endpoint_url, + config=self._get_config(config), + verify=self.verify, + ) + return session.client( client_type, endpoint_url=self.conn_config.endpoint_url, @@ -634,6 +651,14 @@ def conn(self) -> BaseAwsConnection: else: return self.get_resource_type(region_name=self.region_name) + @cached_property + def async_conn(self): + """Get an Aiobotocore client to use for async operations (cached).""" + if not self.client_type: + raise ValueError("client_type must be specified.") + + return self.get_client_type(region_name=self.region_name, deferrable=True) + @cached_property def conn_client_meta(self) -> ClientMeta: """Get botocore client metadata from Hook connection (cached).""" @@ -797,18 +822,27 @@ def waiter_path(self) -> PathLike[str] | None: path = Path(__file__).parents[1].joinpath(f"waiters/{self.client_type}.json").resolve() return path if path.exists() else None - def get_waiter(self, waiter_name: str, parameters: dict[str, str] | None = None) -> Waiter: + def get_waiter(self, waiter_name: str, parameters: dict[str, str] | None = None, deferrable: bool = False, client=None) -> Waiter: """ First checks if there is a custom waiter with the provided waiter_name and uses that if it exists, otherwise it will check the service client for a waiter that matches the name and pass that through. + If `deferrable` is True, the waiter will be an AIOWaiter, generated from the + client that is passed as a parameter. If `deferrable` is True, `client` must be + provided. + :param waiter_name: The name of the waiter. The name should exactly match the name of the key in the waiter model file (typically this is CamelCase). :param parameters: will scan the waiter config for the keys of that dict, and replace them with the corresponding value. If a custom waiter has such keys to be expanded, they need to be provided here. + :param deferrable: If True, the waiter is going to be an async custom waiter. + """ + if deferrable and not client: + raise ValueError("client must be provided for a deferrable waiter.") + client = client or self.conn if self.waiter_path and (waiter_name in self._list_custom_waiters()): # Technically if waiter_name is in custom_waiters then self.waiter_path must # exist but MyPy doesn't like the fact that self.waiter_path could be None. @@ -816,7 +850,9 @@ def get_waiter(self, waiter_name: str, parameters: dict[str, str] | None = None) config = json.loads(config_file.read()) config = self._apply_parameters_value(config, waiter_name, parameters) - return BaseBotoWaiter(client=self.conn, model_config=config).waiter(waiter_name) + return BaseBotoWaiter(client=client, model_config=config, deferrable=deferrable).waiter( + waiter_name + ) # If there is no custom waiter found for the provided name, # then try checking the service's official waiters. return self.conn.get_waiter(waiter_name) diff --git a/airflow/providers/amazon/aws/operators/redshift_cluster.py b/airflow/providers/amazon/aws/operators/redshift_cluster.py index 5183dc2e2c120..d005fab09f394 100644 --- a/airflow/providers/amazon/aws/operators/redshift_cluster.py +++ b/airflow/providers/amazon/aws/operators/redshift_cluster.py @@ -22,7 +22,7 @@ from airflow.exceptions import AirflowException from airflow.models import BaseOperator from airflow.providers.amazon.aws.hooks.redshift_cluster import RedshiftHook -from airflow.providers.amazon.aws.triggers.redshift_cluster import RedshiftClusterTrigger +from airflow.providers.amazon.aws.triggers.redshift_cluster import RedshiftClusterTrigger, RedshiftCreateClusterTrigger if TYPE_CHECKING: from airflow.utils.context import Context @@ -140,6 +140,7 @@ def __init__( wait_for_completion: bool = False, max_attempt: int = 5, poll_interval: int = 60, + deferrable: bool = False, **kwargs, ): super().__init__(**kwargs) @@ -180,6 +181,7 @@ def __init__( self.wait_for_completion = wait_for_completion self.max_attempt = max_attempt self.poll_interval = poll_interval + self.deferrable = deferrable self.kwargs = kwargs def execute(self, context: Context): @@ -252,6 +254,16 @@ def execute(self, context: Context): self.master_user_password, params, ) + if self.deferrable: + self.defer( + trigger=RedshiftCreateClusterTrigger( + cluster_identifier=self.cluster_identifier, + poll_interval=self.poll_interval, + max_attempt=self.max_attempt, + aws_conn_id=self.aws_conn_id, + ), + method_name="execute_complete", + ) if self.wait_for_completion: redshift_hook.get_conn().get_waiter("cluster_available").wait( ClusterIdentifier=self.cluster_identifier, @@ -264,6 +276,11 @@ def execute(self, context: Context): self.log.info("Created Redshift cluster %s", self.cluster_identifier) self.log.info(cluster) + def execute_complete(self, context, event=None): + if event["status"] != "success": + raise AirflowException(f"Error creating cluster: {event}") + return + class RedshiftCreateClusterSnapshotOperator(BaseOperator): """ diff --git a/airflow/providers/amazon/aws/triggers/README.md b/airflow/providers/amazon/aws/triggers/README.md new file mode 100644 index 0000000000000..7ebc5fc2a3c57 --- /dev/null +++ b/airflow/providers/amazon/aws/triggers/README.md @@ -0,0 +1,153 @@ + + +# Writing Deferrable Operators for Amazon Provider Package + + +Before writing deferrable operators, it is strongly recommended to read and familiarize yourself with the official [documentation](https://airflow.apache.org/docs/apache-airflow/stable/authoring-and-scheduling/deferring.html) of Deferrable Operators. +The purpose of this guide is to provide a standardized way to convert existing Amazon Provider Package (AMPP) operators to deferrable operators. Due to the varied complexities of available operators, it is impossible to define one method that will work for every operator. +The method described in this guide should work for many of the AMPP operators, but it is important to study each operator before determining whether the steps outlined below are applicable. + +Although it varies from operator to operator, a typical AMPP operator has 3 stages: + +1. A pre-processing stage, where information is looked up via boto3 API calls, parameters are formatted etc. The complexity of this stage depends on the complexity of the task the operator is attempting to do. Some operators (e.g. Sagemaker) have a lot of pre-processing, whereas others require little to no pre-processing. +2. The "main" call to the boto3 API to start an operation. This is the task that the operator is attempting to complete. This could be a request to provision a resource, request to change the state of a resource, start a job on a resource etc. Regardless of the operation, the boto3 API returns a response instantly (ignoring network delays) with a response detailing the results of the query. For example, in the case of a resource provisioning request, although the resource can take significant time to be allocated, the boto3 API returns a response to the caller without waiting for the operation to be completed. +3. The last, often optional, stage is to wait for the operation initiated in stage 2 to be completed. This usually involves polling the boto3 API at set intervals, and waiting for a certain criteria to be met. + +In general, it is the last polling stage where we can defer the operator to a trigger which can handle the polling operation. The botocore library defines waiters for certain services, which are built-in functions that poll a service and wait for a given criteria to be met. +As part of our work for writing deferrable operators, we have extended the built-in waiters to allow custom waiters, which follow the same logic, but for services not included in the botocore library. +We can use these custom waiters, along with the built-in waiters to implement the polling logic of the deferrable operators. + +The first step to making an existing operator deferrable is to add `deferrable` as a parameter to the operator, and initialize it in the constructor of the operator. +The next step is to determine where the operator should be deferred. This will be dependent on what the operator does, and how it is written. Although every operator is different, there are a few guidelines to determine the best place to defer an operator. + +1. If the operator has a `wait_for_completion` parameter, the `self.defer` method should be called right before the check for wait_for_completion . +2. If there is no `wait_for_completion` , look for the "main" task that the operator does. Often, operators will make various describe calls to to the boto3 API to verify certain conditions, or look up some information before performing its "main" task. Often, right after the "main" call to the boto3 API is made is a good place to call `self.defer`. + + +Once the location to defer is decided in the operator, call the `self.defer` method if the `deferrable` flag is `True`. The `self.defer` method takes in several parameters, listed below: + +1. `trigger`: This is the trigger which you want to pass the execution to. We will write this trigger in just a moment. +2. `method_name`: This specifies the name of the method you want to execute once the trigger completes its execution. The trigger cannot pass the execution back to the execute method of the operator. By convention, the name for this method is `execute_complete`. +3. `timeout`: An optional parameter that controls the length of time the Trigger can execute for before timing out. This defaults to `None`, meaning no timeout. +4. `kwargs`: Additional keyword arguments to pass to `method_name`. Default is `{}`. + +The Trigger is the main component of deferrable operators. They must be placed in the `airflow/providers/amazon/aws/triggers/` folder. All Triggers must implement the following 3 methods: + +1. `__init__`: the constructor which receives parameters from the operator. These must be JSON serializable. +2. `serialize`: a function that returns the classpath, as well as keyword arguments to the `__init__` method as a tuple +3. `run` : the asynchronous function that is responsible for awaiting the asynchronous operations. + +Ideally, when the operator has deferred itself, it has already initiated the "main" task of the operator, and is now waiting for a certain resource to reach a certain state. +As mentioned earlier, the botocore library defines a `Waiter` class for many services, which implements a `wait` method that can be configured via a config file to poll the boto3 API at set intervals, and return if the success criteria is met. +The aiobotocore library, which is used to make asynchronous botocore calls, defines an `AIOWaiter` class, which also implements a wait method that behaves identical to the botocore method, except that it works asynchronously. +Therefore, any botocore waiter is available as an aiobotocore waiter, and can be used to asynchronously poll a service until the desired criteria is met. + +To call the asynchronous `wait` function, first create a hook for the particular service. For example, for a Redshift hook, it would look like this: + +```python +self.redshift_hook = RedshiftHook(aws_conn_id=self.aws_conn_id) +``` + +With this hook, we can use the async_conn property to get access to the aiobotocore client: + +```python +async with self.redshift_hook.async_conn as client: + await client.get_waiter("cluster_available").wait( + ClusterIdentifier=self.cluster_identifier, + WaiterConfig={ + "Delay": int(self.poll_interval), + "MaxAttempts": int(self.max_attempt), + }, + ) +``` + +In this case, we are using the built-in cluster_available waiter. If we wanted to use a custom waiter, we would change the code slightly to use the get_waiter function from the hook, rather than the aiobotocore client: + +```python +async with self.redshift_hook.async_conn as client: + waiter = self.redshift_hook.get_waiter("cluster_paused", deferrable=True, client=client) + await waiter.wait( + ClusterIdentifier=self.cluster_identifier, + WaiterConfig={ + "Delay": int(self.poll_interval), + "MaxAttempts": int(self.max_attempt), + }, + ) +``` + +Here, we are calling the get_waiter function defined in base_aws.py which takes an optional argument of deferrable (set to True), and the aiobotocore client. cluster_paused is a custom boto waiter defined in redshift.json in the airflow/providers/amazon/aws/waiters folder. In general, the config file for a custom waiter should be named as .json. The config for cluster_paused is shown below: + +```json +{ + "version": 2, + "waiters": { + "cluster_paused": { + "operation": "DescribeClusters", + "delay": 30, + "maxAttempts": 60, + "acceptors": [ + { + "matcher": "pathAll", + "argument": "Clusters[].ClusterStatus", + "expected": "paused", + "state": "success" + }, + { + "expected": "ClusterNotFound", + "matcher": "error", + "state": "retry" + }, + { + "expected": "deleting", + "matcher": "pathAny", + "state": "failure", + "argument": "Clusters[].ClusterStatus" + } + ] + }, + } +} +``` + +For more information about writing custom waiter, see the [README.md](https://github.com/apache/airflow/blob/main/airflow/providers/amazon/aws/waiters/README.md) for custom waiters. + +In some cases, a built-in or custom waiter may not be able to solve the problem. In such cases, the asynchronous method used to poll the boto3 API would need to be defined in the hook of the service being used. This method is essentially the same as the synchronous version of the method, except that it will use the aiobotocore client, and will be awaited. For the Redshift example, the async describe_clusters method would look as follows: + +```python +async with self.async_conn as client: + response = client.describe_clusters(ClusterIdentifier=self.cluster_identifier) +``` + +This async method can be used in the Trigger to poll the boto3 API. The polling logic will need to be implemented manually, taking care to use asyncio.sleep() rather than time.sleep(). + +The last step in the Trigger is to yield a TriggerEvent that will be used to alert the Triggerer that the Trigger has finished execution. The TriggerEvent can pass information from the trigger to the method_name method named in the self.defer call in the operator. In the Redshift example, the TriggerEvent would look as follows: + +``` +yield TriggerEvent({"status": "success", "message": "Cluster Created"}) +``` + +The object passed through the TrigggerEvent can be captured in the method_name method through an event parameter. This can be used to determine what needs to be done based on the outcome of the Trigger execution. In the Redshift case, we can simply check the status of the event, and raise an Exception if something went wrong. + +```python +def execute_complete(self, context, event=None): + if event["status"] != "success": + raise AirflowException(f"Error creating cluster: {event}") + return +``` diff --git a/airflow/providers/amazon/aws/triggers/redshift_cluster.py b/airflow/providers/amazon/aws/triggers/redshift_cluster.py index a32a6efa19924..7230c14e19c6b 100644 --- a/airflow/providers/amazon/aws/triggers/redshift_cluster.py +++ b/airflow/providers/amazon/aws/triggers/redshift_cluster.py @@ -18,8 +18,11 @@ from typing import Any, AsyncIterator -from airflow.providers.amazon.aws.hooks.redshift_cluster import RedshiftAsyncHook +from airflow.providers.amazon.aws.hooks.redshift_cluster import RedshiftAsyncHook, RedshiftHook from airflow.triggers.base import BaseTrigger, TriggerEvent +from typing import Any + +from airflow.compat.functools import cached_property class RedshiftClusterTrigger(BaseTrigger): @@ -85,3 +88,53 @@ async def run(self) -> AsyncIterator["TriggerEvent"]: except Exception as e: if self.attempts < 1: yield TriggerEvent({"status": "error", "message": str(e)}) + +class RedshiftCreateClusterTrigger(BaseTrigger): + """ + Trigger for RedshiftCreateClusterOperator. + The trigger will asynchronously poll the boto3 API and wait for the + Redshift cluster to be in the `available` state. + + :param cluster_identifier: A unique identifier for the cluster. + :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 aws_conn_id: The Airflow connection used for AWS credentials. + """ + + def __init__( + self, + cluster_identifier: str, + poll_interval: int, + max_attempt: int, + aws_conn_id: str, + ): + self.cluster_identifier = cluster_identifier + self.poll_interval = poll_interval + self.max_attempt = max_attempt + self.aws_conn_id = aws_conn_id + + def serialize(self) -> tuple[str, dict[str, Any]]: + return ( + "airflow.providers.amazon.aws.triggers.redshift_cluster.RedshiftCreateClusterTrigger", + { + "cluster_identifier": str(self.cluster_identifier), + "poll_interval": str(self.poll_interval), + "max_attempt": str(self.max_attempt), + "aws_conn_id": str(self.aws_conn_id), + }, + ) + + @cached_property + def hook(self) -> RedshiftHook: + return RedshiftHook(aws_conn_id=self.aws_conn_id) + + async def run(self): + async with self.hook.async_conn as client: + await client.get_waiter("cluster_available").wait( + ClusterIdentifier=self.cluster_identifier, + WaiterConfig={ + "Delay": int(self.poll_interval), + "MaxAttempts": int(self.max_attempt), + }, + ) + yield TriggerEvent({"status": "success", "message": "Cluster Created"}) diff --git a/airflow/providers/amazon/aws/waiters/base_waiter.py b/airflow/providers/amazon/aws/waiters/base_waiter.py index 0d9f8a1d4e407..b4b8668ea8127 100644 --- a/airflow/providers/amazon/aws/waiters/base_waiter.py +++ b/airflow/providers/amazon/aws/waiters/base_waiter.py @@ -18,6 +18,7 @@ from __future__ import annotations import boto3 +from aiobotocore.waiter import create_waiter_with_client as create_async_waiter_with_client from botocore.waiter import Waiter, WaiterModel, create_waiter_with_client @@ -28,9 +29,14 @@ class BaseBotoWaiter: For more details, see airflow/providers/amazon/aws/waiters/README.md """ - def __init__(self, client: boto3.client, model_config: dict) -> None: + def __init__(self, client: boto3.client, model_config: dict, deferrable: bool = False) -> None: self.model = WaiterModel(model_config) self.client = client + self.deferrable = deferrable def waiter(self, waiter_name: str) -> Waiter: + if self.deferrable: + return create_async_waiter_with_client( + waiter_name=waiter_name, waiter_model=self.model, client=self.client + ) return create_waiter_with_client(waiter_name=waiter_name, waiter_model=self.model, client=self.client) diff --git a/airflow/providers/amazon/provider.yaml b/airflow/providers/amazon/provider.yaml index 05ee54482082f..f2aac81f6e7be 100644 --- a/airflow/providers/amazon/provider.yaml +++ b/airflow/providers/amazon/provider.yaml @@ -71,6 +71,7 @@ dependencies: - mypy-boto3-rds>=1.24.0 - mypy-boto3-redshift-data>=1.24.0 - mypy-boto3-appflow>=1.24.0 + - aiobotocore[boto3] integrations: - integration-name: Amazon Athena diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json index fafcbbb606332..186c189db2899 100644 --- a/generated/provider_dependencies.json +++ b/generated/provider_dependencies.json @@ -17,6 +17,7 @@ }, "amazon": { "deps": [ + "aiobotocore[boto3]", "apache-airflow-providers-common-sql>=1.3.1", "apache-airflow>=2.3.0", "asgiref", diff --git a/tests/providers/amazon/aws/triggers/__init__.py b/tests/providers/amazon/aws/triggers/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/tests/providers/amazon/aws/triggers/__init__.py @@ -0,0 +1,16 @@ +# 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. diff --git a/tests/providers/amazon/aws/triggers/test_redshift_cluster.py b/tests/providers/amazon/aws/triggers/test_redshift_cluster.py new file mode 100644 index 0000000000000..941258659e9ae --- /dev/null +++ b/tests/providers/amazon/aws/triggers/test_redshift_cluster.py @@ -0,0 +1,74 @@ +# 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 + +import sys + +import pytest + +from airflow.providers.amazon.aws.triggers.redshift_cluster import RedshiftCreateClusterTrigger +from airflow.triggers.base import TriggerEvent + +if sys.version_info < (3, 8): + from asynctest import CoroutineMock as AsyncMock, mock as async_mock +else: + from unittest import mock as async_mock + from unittest.mock import AsyncMock + + +TEST_CLUSTER_IDENTIFIER = "test-cluster" +TEST_POLL_INTERVAL = 10 +TEST_MAX_ATTEMPT = 10 +TEST_AWS_CONN_ID = "test-aws-id" + + +class TestRedshiftCreateClusterTrigger: + def test_redshift_create_cluster_trigger_serialize(self): + redshift_create_cluster_trigger = RedshiftCreateClusterTrigger( + cluster_identifier=TEST_CLUSTER_IDENTIFIER, + poll_interval=TEST_POLL_INTERVAL, + max_attempt=TEST_MAX_ATTEMPT, + aws_conn_id=TEST_AWS_CONN_ID, + ) + class_path, args = redshift_create_cluster_trigger.serialize() + assert ( + class_path + == "airflow.providers.amazon.aws.triggers.redshift_cluster.RedshiftCreateClusterTrigger" + ) + assert args["cluster_identifier"] == TEST_CLUSTER_IDENTIFIER + assert args["poll_interval"] == str(TEST_POLL_INTERVAL) + assert args["max_attempt"] == str(TEST_MAX_ATTEMPT) + assert args["aws_conn_id"] == TEST_AWS_CONN_ID + + @pytest.mark.asyncio + @async_mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.async_conn") + async def test_redshift_create_cluster_trigger_run(self, mock_async_conn): + mock = async_mock.MagicMock() + mock_async_conn.__aenter__.return_value = mock + mock.get_waiter().wait = AsyncMock() + + redshift_create_cluster_trigger = RedshiftCreateClusterTrigger( + cluster_identifier=TEST_CLUSTER_IDENTIFIER, + poll_interval=TEST_POLL_INTERVAL, + max_attempt=TEST_MAX_ATTEMPT, + aws_conn_id=TEST_AWS_CONN_ID, + ) + + generator = redshift_create_cluster_trigger.run() + response = await generator.asend(None) + + assert response == TriggerEvent({"status": "success", "message": "Cluster Created"}) From 197d7d1be3fd3d391077a9999352e150474cbad5 Mon Sep 17 00:00:00 2001 From: Syed Hussain Date: Wed, 22 Mar 2023 14:51:41 -0700 Subject: [PATCH 02/17] Fix failing test Rebase from main --- .../providers/amazon/aws/hooks/base_aws.py | 20 +++++++++++-------- airflow/providers/amazon/provider.yaml | 2 +- generated/provider_dependencies.json | 2 +- 3 files changed, 14 insertions(+), 10 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/base_aws.py b/airflow/providers/amazon/aws/hooks/base_aws.py index 09e7126017756..47b431d68adc2 100644 --- a/airflow/providers/amazon/aws/hooks/base_aws.py +++ b/airflow/providers/amazon/aws/hooks/base_aws.py @@ -43,7 +43,7 @@ import jinja2 import requests import tenacity -from aiobotocore.session import AioSession, get_session as async_get_session +from aiobotocore.session import get_session as async_get_session from botocore.client import ClientMeta from botocore.config import Config from botocore.credentials import ReadOnlyCredentials @@ -180,9 +180,7 @@ def _create_session_with_assume_role( session = async_get_session() if deferrable else botocore.session.get_session() - session.set_credentials( - access_key=credentials.access_key, secret_key=credentials.secret_key, token=credentials.token - ) + session._credentials = credentials session.set_config_variable("region", self.basic_session.region_name) return boto3.session.Session(botocore_session=session, **session_kwargs) @@ -653,7 +651,7 @@ def conn(self) -> BaseAwsConnection: @cached_property def async_conn(self): - """Get an Aiobotocore client to use for async operations (cached).""" + """Get an aiobotocore client to use for async operations (cached).""" if not self.client_type: raise ValueError("client_type must be specified.") @@ -822,7 +820,13 @@ def waiter_path(self) -> PathLike[str] | None: path = Path(__file__).parents[1].joinpath(f"waiters/{self.client_type}.json").resolve() return path if path.exists() else None - def get_waiter(self, waiter_name: str, parameters: dict[str, str] | None = None, deferrable: bool = False, client=None) -> Waiter: + def get_waiter( + self, + waiter_name: str, + parameters: dict[str, str] | None = None, + deferrable: bool = False, + client=None, + ) -> Waiter: """ First checks if there is a custom waiter with the provided waiter_name and uses that if it exists, otherwise it will check the service client for a @@ -851,8 +855,8 @@ def get_waiter(self, waiter_name: str, parameters: dict[str, str] | None = None, config = self._apply_parameters_value(config, waiter_name, parameters) return BaseBotoWaiter(client=client, model_config=config, deferrable=deferrable).waiter( - waiter_name - ) + waiter_name + ) # If there is no custom waiter found for the provided name, # then try checking the service's official waiters. return self.conn.get_waiter(waiter_name) diff --git a/airflow/providers/amazon/provider.yaml b/airflow/providers/amazon/provider.yaml index f2aac81f6e7be..bc78bf6702c73 100644 --- a/airflow/providers/amazon/provider.yaml +++ b/airflow/providers/amazon/provider.yaml @@ -71,7 +71,7 @@ dependencies: - mypy-boto3-rds>=1.24.0 - mypy-boto3-redshift-data>=1.24.0 - mypy-boto3-appflow>=1.24.0 - - aiobotocore[boto3] + - aiobotocore[boto3]>=2.2.0 integrations: - integration-name: Amazon Athena diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json index 186c189db2899..655f3d5d66aa7 100644 --- a/generated/provider_dependencies.json +++ b/generated/provider_dependencies.json @@ -17,7 +17,7 @@ }, "amazon": { "deps": [ - "aiobotocore[boto3]", + "aiobotocore[boto3]>=2.2.0", "apache-airflow-providers-common-sql>=1.3.1", "apache-airflow>=2.3.0", "asgiref", From 647b414a1bb1b55d9d58522ce55b08b1d1018196 Mon Sep 17 00:00:00 2001 From: Syed Hussain Date: Thu, 23 Mar 2023 15:48:52 -0700 Subject: [PATCH 03/17] Fix static check failures --- airflow/providers/amazon/aws/hooks/base_aws.py | 2 +- airflow/providers/amazon/aws/hooks/batch_waiters.py | 4 +++- airflow/providers/amazon/aws/operators/redshift_cluster.py | 5 ++++- airflow/providers/amazon/aws/triggers/redshift_cluster.py | 5 ++--- 4 files changed, 10 insertions(+), 6 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/base_aws.py b/airflow/providers/amazon/aws/hooks/base_aws.py index 47b431d68adc2..6ffb7a4ffa1f8 100644 --- a/airflow/providers/amazon/aws/hooks/base_aws.py +++ b/airflow/providers/amazon/aws/hooks/base_aws.py @@ -1025,7 +1025,7 @@ def _basic_session(self) -> AioSession: aio_session.set_config_variable("region", region_name) return aio_session - def create_session(self) -> AioSession: + def create_session(self, deferrable: bool = False) -> AioSession: """Create aiobotocore Session from connection and config.""" if not self._conn: self.log.info("No connection ID provided. Fallback on boto3 credential strategy") diff --git a/airflow/providers/amazon/aws/hooks/batch_waiters.py b/airflow/providers/amazon/aws/hooks/batch_waiters.py index dcf111591c9bf..cb852acf9d8b8 100644 --- a/airflow/providers/amazon/aws/hooks/batch_waiters.py +++ b/airflow/providers/amazon/aws/hooks/batch_waiters.py @@ -138,7 +138,9 @@ def waiter_model(self) -> botocore.waiter.WaiterModel: """ return self._waiter_model - def get_waiter(self, waiter_name: str, _: dict[str, str] | None = None) -> botocore.waiter.Waiter: + def get_waiter( + self, waiter_name: str, _: dict[str, str] | None = None, deferrable: bool = False, client=None + ) -> botocore.waiter.Waiter: """ Get an AWS Batch service waiter, using the configured ``.waiter_model``. diff --git a/airflow/providers/amazon/aws/operators/redshift_cluster.py b/airflow/providers/amazon/aws/operators/redshift_cluster.py index d005fab09f394..dc6835434ff23 100644 --- a/airflow/providers/amazon/aws/operators/redshift_cluster.py +++ b/airflow/providers/amazon/aws/operators/redshift_cluster.py @@ -22,7 +22,10 @@ from airflow.exceptions import AirflowException from airflow.models import BaseOperator from airflow.providers.amazon.aws.hooks.redshift_cluster import RedshiftHook -from airflow.providers.amazon.aws.triggers.redshift_cluster import RedshiftClusterTrigger, RedshiftCreateClusterTrigger +from airflow.providers.amazon.aws.triggers.redshift_cluster import ( + RedshiftClusterTrigger, + RedshiftCreateClusterTrigger, +) if TYPE_CHECKING: from airflow.utils.context import Context diff --git a/airflow/providers/amazon/aws/triggers/redshift_cluster.py b/airflow/providers/amazon/aws/triggers/redshift_cluster.py index 7230c14e19c6b..2f831fa14c2f1 100644 --- a/airflow/providers/amazon/aws/triggers/redshift_cluster.py +++ b/airflow/providers/amazon/aws/triggers/redshift_cluster.py @@ -18,11 +18,9 @@ from typing import Any, AsyncIterator +from airflow.compat.functools import cached_property from airflow.providers.amazon.aws.hooks.redshift_cluster import RedshiftAsyncHook, RedshiftHook from airflow.triggers.base import BaseTrigger, TriggerEvent -from typing import Any - -from airflow.compat.functools import cached_property class RedshiftClusterTrigger(BaseTrigger): @@ -89,6 +87,7 @@ async def run(self) -> AsyncIterator["TriggerEvent"]: if self.attempts < 1: yield TriggerEvent({"status": "error", "message": str(e)}) + class RedshiftCreateClusterTrigger(BaseTrigger): """ Trigger for RedshiftCreateClusterOperator. From 50035cff5860706b98cc572ffb71eea79d85705a Mon Sep 17 00:00:00 2001 From: Syed Hussain Date: Mon, 27 Mar 2023 15:12:44 -0700 Subject: [PATCH 04/17] Skip system tests if aiobotocore is not added. Move import of aiobotocore closer to where they are used to minimize usage --- .../providers/amazon/aws/hooks/base_aws.py | 6 +- .../providers/amazon/aws/example_appflow.py | 3 +- .../providers/amazon/aws/example_athena.py | 3 +- .../providers/amazon/aws/example_batch.py | 3 +- .../amazon/aws/example_cloudformation.py | 3 +- .../providers/amazon/aws/example_datasync.py | 3 +- .../providers/amazon/aws/example_dms.py | 3 +- .../amazon/aws/example_dynamodb_to_s3.py | 3 +- .../providers/amazon/aws/example_ec2.py | 3 +- .../providers/amazon/aws/example_ecs.py | 3 +- .../amazon/aws/example_ecs_fargate.py | 3 +- .../amazon/aws/example_eks_templated.py | 3 +- .../example_eks_with_fargate_in_one_step.py | 3 +- .../aws/example_eks_with_fargate_profile.py | 3 +- .../example_eks_with_nodegroup_in_one_step.py | 3 +- .../amazon/aws/example_eks_with_nodegroups.py | 3 +- .../providers/amazon/aws/example_emr.py | 3 +- .../providers/amazon/aws/example_emr_eks.py | 3 +- .../aws/example_emr_notebook_execution.py | 3 +- .../amazon/aws/example_emr_serverless.py | 3 +- .../providers/amazon/aws/example_ftp_to_s3.py | 3 +- .../providers/amazon/aws/example_gcs_to_s3.py | 3 +- .../amazon/aws/example_glacier_to_gcs.py | 3 +- .../providers/amazon/aws/example_glue.py | 3 +- .../aws/example_google_api_sheets_to_s3.py | 3 +- .../aws/example_google_api_youtube_to_s3.py | 3 +- .../amazon/aws/example_hive_to_dynamodb.py | 3 +- .../aws/example_imap_attachment_to_s3.py | 3 +- .../providers/amazon/aws/example_lambda.py | 3 +- .../amazon/aws/example_local_to_s3.py | 3 +- .../amazon/aws/example_mongo_to_s3.py | 3 +- .../amazon/aws/example_quicksight.py | 3 +- .../providers/amazon/aws/example_rds_event.py | 3 +- .../amazon/aws/example_rds_export.py | 3 +- .../amazon/aws/example_rds_instance.py | 3 +- .../amazon/aws/example_rds_snapshot.py | 3 +- .../providers/amazon/aws/example_redshift.py | 3 +- .../aws/example_redshift_s3_transfers.py | 3 +- .../system/providers/amazon/aws/example_s3.py | 3 +- .../providers/amazon/aws/example_s3_to_ftp.py | 3 +- .../amazon/aws/example_s3_to_sftp.py | 3 +- .../providers/amazon/aws/example_s3_to_sql.py | 3 +- .../providers/amazon/aws/example_sagemaker.py | 3 +- .../amazon/aws/example_sagemaker_endpoint.py | 3 +- .../amazon/aws/example_salesforce_to_s3.py | 3 +- .../amazon/aws/example_sftp_to_s3.py | 3 +- .../providers/amazon/aws/example_sns.py | 3 +- .../providers/amazon/aws/example_sql_to_s3.py | 3 +- .../providers/amazon/aws/example_sqs.py | 3 +- .../amazon/aws/example_step_functions.py | 3 +- .../system/providers/amazon/aws/import_fix.py | 79 +++++++++++++++++++ 51 files changed, 181 insertions(+), 51 deletions(-) create mode 100644 tests/system/providers/amazon/aws/import_fix.py diff --git a/airflow/providers/amazon/aws/hooks/base_aws.py b/airflow/providers/amazon/aws/hooks/base_aws.py index 6ffb7a4ffa1f8..a03dde2ef08b6 100644 --- a/airflow/providers/amazon/aws/hooks/base_aws.py +++ b/airflow/providers/amazon/aws/hooks/base_aws.py @@ -43,7 +43,6 @@ import jinja2 import requests import tenacity -from aiobotocore.session import get_session as async_get_session from botocore.client import ClientMeta from botocore.config import Config from botocore.credentials import ReadOnlyCredentials @@ -59,7 +58,6 @@ ) from airflow.hooks.base import BaseHook from airflow.providers.amazon.aws.utils.connection_wrapper import AwsConnectionWrapper -from airflow.providers.amazon.aws.waiters.base_waiter import BaseBotoWaiter from airflow.providers_manager import ProvidersManager from airflow.utils.helpers import exactly_one from airflow.utils.log.logging_mixin import LoggingMixin @@ -130,6 +128,9 @@ def role_arn(self) -> str | None: def create_session(self, deferrable: bool = False) -> boto3.session.Session: """Create boto3 or aiobotocore Session from connection config.""" + from aiobotocore.session import get_session as async_get_session + from pdb import set_trace + #set_trace() if not self.conn: self.log.info( "No connection ID provided. Fallback on boto3 credential strategy (region_name=%r). " @@ -844,6 +845,7 @@ def get_waiter( :param deferrable: If True, the waiter is going to be an async custom waiter. """ + from airflow.providers.amazon.aws.waiters.base_waiter import BaseBotoWaiter if deferrable and not client: raise ValueError("client must be provided for a deferrable waiter.") client = client or self.conn diff --git a/tests/system/providers/amazon/aws/example_appflow.py b/tests/system/providers/amazon/aws/example_appflow.py index 4469c0290b539..7983c469b85a7 100644 --- a/tests/system/providers/amazon/aws/example_appflow.py +++ b/tests/system/providers/amazon/aws/example_appflow.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime from airflow import DAG diff --git a/tests/system/providers/amazon/aws/example_athena.py b/tests/system/providers/amazon/aws/example_athena.py index 3d2487d7b56a4..07feb4ff01a9b 100644 --- a/tests/system/providers/amazon/aws/example_athena.py +++ b/tests/system/providers/amazon/aws/example_athena.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime import boto3 diff --git a/tests/system/providers/amazon/aws/example_batch.py b/tests/system/providers/amazon/aws/example_batch.py index a035b12f87b6e..9188b13e02ab2 100644 --- a/tests/system/providers/amazon/aws/example_batch.py +++ b/tests/system/providers/amazon/aws/example_batch.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime import boto3 diff --git a/tests/system/providers/amazon/aws/example_cloudformation.py b/tests/system/providers/amazon/aws/example_cloudformation.py index fc6e04d422081..f00cabe1e3b7d 100644 --- a/tests/system/providers/amazon/aws/example_cloudformation.py +++ b/tests/system/providers/amazon/aws/example_cloudformation.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") import json from datetime import datetime diff --git a/tests/system/providers/amazon/aws/example_datasync.py b/tests/system/providers/amazon/aws/example_datasync.py index bead9f8f2f80a..98c8de697fa35 100644 --- a/tests/system/providers/amazon/aws/example_datasync.py +++ b/tests/system/providers/amazon/aws/example_datasync.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime import boto3 diff --git a/tests/system/providers/amazon/aws/example_dms.py b/tests/system/providers/amazon/aws/example_dms.py index 8e13e17c32e79..5e44871c3a62f 100644 --- a/tests/system/providers/amazon/aws/example_dms.py +++ b/tests/system/providers/amazon/aws/example_dms.py @@ -21,7 +21,8 @@ """ from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") import json from datetime import datetime from typing import cast diff --git a/tests/system/providers/amazon/aws/example_dynamodb_to_s3.py b/tests/system/providers/amazon/aws/example_dynamodb_to_s3.py index b56efaf2ce1ba..33b632b8ff00a 100644 --- a/tests/system/providers/amazon/aws/example_dynamodb_to_s3.py +++ b/tests/system/providers/amazon/aws/example_dynamodb_to_s3.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime import boto3 diff --git a/tests/system/providers/amazon/aws/example_ec2.py b/tests/system/providers/amazon/aws/example_ec2.py index 1dd98488c2995..8bda33e8236f7 100644 --- a/tests/system/providers/amazon/aws/example_ec2.py +++ b/tests/system/providers/amazon/aws/example_ec2.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime from operator import itemgetter diff --git a/tests/system/providers/amazon/aws/example_ecs.py b/tests/system/providers/amazon/aws/example_ecs.py index 194b070b51686..d018beff01f86 100644 --- a/tests/system/providers/amazon/aws/example_ecs.py +++ b/tests/system/providers/amazon/aws/example_ecs.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime import boto3 diff --git a/tests/system/providers/amazon/aws/example_ecs_fargate.py b/tests/system/providers/amazon/aws/example_ecs_fargate.py index 40132358ab430..b7da5642ab27e 100644 --- a/tests/system/providers/amazon/aws/example_ecs_fargate.py +++ b/tests/system/providers/amazon/aws/example_ecs_fargate.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime import boto3 diff --git a/tests/system/providers/amazon/aws/example_eks_templated.py b/tests/system/providers/amazon/aws/example_eks_templated.py index d09eabf959827..d7c99775009aa 100644 --- a/tests/system/providers/amazon/aws/example_eks_templated.py +++ b/tests/system/providers/amazon/aws/example_eks_templated.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime from airflow.models.baseoperator import chain diff --git a/tests/system/providers/amazon/aws/example_eks_with_fargate_in_one_step.py b/tests/system/providers/amazon/aws/example_eks_with_fargate_in_one_step.py index 37cba110d9bfd..9084cf20bbf05 100644 --- a/tests/system/providers/amazon/aws/example_eks_with_fargate_in_one_step.py +++ b/tests/system/providers/amazon/aws/example_eks_with_fargate_in_one_step.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime from airflow.models.baseoperator import chain diff --git a/tests/system/providers/amazon/aws/example_eks_with_fargate_profile.py b/tests/system/providers/amazon/aws/example_eks_with_fargate_profile.py index 5792332136deb..925cf2abeff27 100644 --- a/tests/system/providers/amazon/aws/example_eks_with_fargate_profile.py +++ b/tests/system/providers/amazon/aws/example_eks_with_fargate_profile.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime from airflow.models.baseoperator import chain diff --git a/tests/system/providers/amazon/aws/example_eks_with_nodegroup_in_one_step.py b/tests/system/providers/amazon/aws/example_eks_with_nodegroup_in_one_step.py index 3dd2b90649389..3a2fde590fe14 100644 --- a/tests/system/providers/amazon/aws/example_eks_with_nodegroup_in_one_step.py +++ b/tests/system/providers/amazon/aws/example_eks_with_nodegroup_in_one_step.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime import boto3 diff --git a/tests/system/providers/amazon/aws/example_eks_with_nodegroups.py b/tests/system/providers/amazon/aws/example_eks_with_nodegroups.py index 4828d58b4ff66..28b6e6c9ce593 100644 --- a/tests/system/providers/amazon/aws/example_eks_with_nodegroups.py +++ b/tests/system/providers/amazon/aws/example_eks_with_nodegroups.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime import boto3 diff --git a/tests/system/providers/amazon/aws/example_emr.py b/tests/system/providers/amazon/aws/example_emr.py index 792b2b9742637..b2a4b843a9f1e 100644 --- a/tests/system/providers/amazon/aws/example_emr.py +++ b/tests/system/providers/amazon/aws/example_emr.py @@ -17,7 +17,8 @@ # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") import json from datetime import datetime diff --git a/tests/system/providers/amazon/aws/example_emr_eks.py b/tests/system/providers/amazon/aws/example_emr_eks.py index 16d1f3bc2475a..2974ee56b96cb 100644 --- a/tests/system/providers/amazon/aws/example_emr_eks.py +++ b/tests/system/providers/amazon/aws/example_emr_eks.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") import json import subprocess from datetime import datetime diff --git a/tests/system/providers/amazon/aws/example_emr_notebook_execution.py b/tests/system/providers/amazon/aws/example_emr_notebook_execution.py index e24d465832ee1..0cd9a5689a80f 100644 --- a/tests/system/providers/amazon/aws/example_emr_notebook_execution.py +++ b/tests/system/providers/amazon/aws/example_emr_notebook_execution.py @@ -17,7 +17,8 @@ # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime from airflow import DAG diff --git a/tests/system/providers/amazon/aws/example_emr_serverless.py b/tests/system/providers/amazon/aws/example_emr_serverless.py index 6d8a669c3e97e..998a2620d758f 100644 --- a/tests/system/providers/amazon/aws/example_emr_serverless.py +++ b/tests/system/providers/amazon/aws/example_emr_serverless.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime import boto3 diff --git a/tests/system/providers/amazon/aws/example_ftp_to_s3.py b/tests/system/providers/amazon/aws/example_ftp_to_s3.py index ca2e6eb8e5b3b..195b20c6b5a78 100644 --- a/tests/system/providers/amazon/aws/example_ftp_to_s3.py +++ b/tests/system/providers/amazon/aws/example_ftp_to_s3.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime from airflow import models diff --git a/tests/system/providers/amazon/aws/example_gcs_to_s3.py b/tests/system/providers/amazon/aws/example_gcs_to_s3.py index c0182f2d099ff..2bf7080f7c349 100644 --- a/tests/system/providers/amazon/aws/example_gcs_to_s3.py +++ b/tests/system/providers/amazon/aws/example_gcs_to_s3.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime from airflow import DAG diff --git a/tests/system/providers/amazon/aws/example_glacier_to_gcs.py b/tests/system/providers/amazon/aws/example_glacier_to_gcs.py index 43e14907ba799..88a29ae878bcc 100644 --- a/tests/system/providers/amazon/aws/example_glacier_to_gcs.py +++ b/tests/system/providers/amazon/aws/example_glacier_to_gcs.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime import boto3 diff --git a/tests/system/providers/amazon/aws/example_glue.py b/tests/system/providers/amazon/aws/example_glue.py index f010b2dfd9fe0..d23929f655f05 100644 --- a/tests/system/providers/amazon/aws/example_glue.py +++ b/tests/system/providers/amazon/aws/example_glue.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime import boto3 diff --git a/tests/system/providers/amazon/aws/example_google_api_sheets_to_s3.py b/tests/system/providers/amazon/aws/example_google_api_sheets_to_s3.py index 926ab55f8f0c3..e94d82873bd24 100644 --- a/tests/system/providers/amazon/aws/example_google_api_sheets_to_s3.py +++ b/tests/system/providers/amazon/aws/example_google_api_sheets_to_s3.py @@ -19,7 +19,8 @@ You need to set all env variables to request the data. """ from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime from os import getenv diff --git a/tests/system/providers/amazon/aws/example_google_api_youtube_to_s3.py b/tests/system/providers/amazon/aws/example_google_api_youtube_to_s3.py index 64646b910468a..4712ad07da4e4 100644 --- a/tests/system/providers/amazon/aws/example_google_api_youtube_to_s3.py +++ b/tests/system/providers/amazon/aws/example_google_api_youtube_to_s3.py @@ -46,7 +46,8 @@ or by creating a custom connection. """ from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") import json from datetime import datetime diff --git a/tests/system/providers/amazon/aws/example_hive_to_dynamodb.py b/tests/system/providers/amazon/aws/example_hive_to_dynamodb.py index 5a9b62c8663f9..372cb571d3c7c 100644 --- a/tests/system/providers/amazon/aws/example_hive_to_dynamodb.py +++ b/tests/system/providers/amazon/aws/example_hive_to_dynamodb.py @@ -20,7 +20,8 @@ https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/EMRforDynamoDB.Tutorial.html """ from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime from airflow import DAG diff --git a/tests/system/providers/amazon/aws/example_imap_attachment_to_s3.py b/tests/system/providers/amazon/aws/example_imap_attachment_to_s3.py index 7e70fd05e0fda..a92a8b0960d57 100644 --- a/tests/system/providers/amazon/aws/example_imap_attachment_to_s3.py +++ b/tests/system/providers/amazon/aws/example_imap_attachment_to_s3.py @@ -19,7 +19,8 @@ protocol from a mail server to S3 Bucket. """ from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime from airflow import DAG diff --git a/tests/system/providers/amazon/aws/example_lambda.py b/tests/system/providers/amazon/aws/example_lambda.py index b4951799c8b08..7852cc6238f1b 100644 --- a/tests/system/providers/amazon/aws/example_lambda.py +++ b/tests/system/providers/amazon/aws/example_lambda.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") import io import json import zipfile diff --git a/tests/system/providers/amazon/aws/example_local_to_s3.py b/tests/system/providers/amazon/aws/example_local_to_s3.py index 8082b1cf6fc91..43fa77b7a2d82 100644 --- a/tests/system/providers/amazon/aws/example_local_to_s3.py +++ b/tests/system/providers/amazon/aws/example_local_to_s3.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") import os from datetime import datetime diff --git a/tests/system/providers/amazon/aws/example_mongo_to_s3.py b/tests/system/providers/amazon/aws/example_mongo_to_s3.py index 3a3e5103fb0b5..6e91c2a4f4282 100644 --- a/tests/system/providers/amazon/aws/example_mongo_to_s3.py +++ b/tests/system/providers/amazon/aws/example_mongo_to_s3.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from airflow import models from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator diff --git a/tests/system/providers/amazon/aws/example_quicksight.py b/tests/system/providers/amazon/aws/example_quicksight.py index bc3cbd2abd31e..5d4d561ae0411 100644 --- a/tests/system/providers/amazon/aws/example_quicksight.py +++ b/tests/system/providers/amazon/aws/example_quicksight.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") import json from datetime import datetime diff --git a/tests/system/providers/amazon/aws/example_rds_event.py b/tests/system/providers/amazon/aws/example_rds_event.py index 94c39137ba416..70d3c56b1a4c7 100644 --- a/tests/system/providers/amazon/aws/example_rds_event.py +++ b/tests/system/providers/amazon/aws/example_rds_event.py @@ -16,7 +16,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime import boto3 diff --git a/tests/system/providers/amazon/aws/example_rds_export.py b/tests/system/providers/amazon/aws/example_rds_export.py index d103763952417..343d4233390d4 100644 --- a/tests/system/providers/amazon/aws/example_rds_export.py +++ b/tests/system/providers/amazon/aws/example_rds_export.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime from airflow import DAG diff --git a/tests/system/providers/amazon/aws/example_rds_instance.py b/tests/system/providers/amazon/aws/example_rds_instance.py index 4e165aaec63e0..b47a12ca3e08b 100644 --- a/tests/system/providers/amazon/aws/example_rds_instance.py +++ b/tests/system/providers/amazon/aws/example_rds_instance.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime from airflow import DAG diff --git a/tests/system/providers/amazon/aws/example_rds_snapshot.py b/tests/system/providers/amazon/aws/example_rds_snapshot.py index 904deea541430..b139985ad3fdf 100644 --- a/tests/system/providers/amazon/aws/example_rds_snapshot.py +++ b/tests/system/providers/amazon/aws/example_rds_snapshot.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime from airflow import DAG diff --git a/tests/system/providers/amazon/aws/example_redshift.py b/tests/system/providers/amazon/aws/example_redshift.py index 7e3a5a92809c5..a23c09fa1077b 100644 --- a/tests/system/providers/amazon/aws/example_redshift.py +++ b/tests/system/providers/amazon/aws/example_redshift.py @@ -17,7 +17,8 @@ # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime import boto3 diff --git a/tests/system/providers/amazon/aws/example_redshift_s3_transfers.py b/tests/system/providers/amazon/aws/example_redshift_s3_transfers.py index 2608e327b6a5d..3e045b432ff22 100644 --- a/tests/system/providers/amazon/aws/example_redshift_s3_transfers.py +++ b/tests/system/providers/amazon/aws/example_redshift_s3_transfers.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime import boto3 diff --git a/tests/system/providers/amazon/aws/example_s3.py b/tests/system/providers/amazon/aws/example_s3.py index a18e16b79ad61..08ddcd7d9206b 100644 --- a/tests/system/providers/amazon/aws/example_s3.py +++ b/tests/system/providers/amazon/aws/example_s3.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime from airflow.models.baseoperator import chain diff --git a/tests/system/providers/amazon/aws/example_s3_to_ftp.py b/tests/system/providers/amazon/aws/example_s3_to_ftp.py index 984b5e41306f1..6164fcd6ede12 100644 --- a/tests/system/providers/amazon/aws/example_s3_to_ftp.py +++ b/tests/system/providers/amazon/aws/example_s3_to_ftp.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime from airflow import models diff --git a/tests/system/providers/amazon/aws/example_s3_to_sftp.py b/tests/system/providers/amazon/aws/example_s3_to_sftp.py index 1212b32e1607e..d7e027695039a 100644 --- a/tests/system/providers/amazon/aws/example_s3_to_sftp.py +++ b/tests/system/providers/amazon/aws/example_s3_to_sftp.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime from airflow import models diff --git a/tests/system/providers/amazon/aws/example_s3_to_sql.py b/tests/system/providers/amazon/aws/example_s3_to_sql.py index f1b1b97dc36e2..3cc1de8b22ea6 100644 --- a/tests/system/providers/amazon/aws/example_s3_to_sql.py +++ b/tests/system/providers/amazon/aws/example_s3_to_sql.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime import boto3 diff --git a/tests/system/providers/amazon/aws/example_sagemaker.py b/tests/system/providers/amazon/aws/example_sagemaker.py index 9506970446320..745709639d096 100644 --- a/tests/system/providers/amazon/aws/example_sagemaker.py +++ b/tests/system/providers/amazon/aws/example_sagemaker.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") import json import logging import subprocess diff --git a/tests/system/providers/amazon/aws/example_sagemaker_endpoint.py b/tests/system/providers/amazon/aws/example_sagemaker_endpoint.py index b4f9ce2b9e847..9b71a579f58da 100644 --- a/tests/system/providers/amazon/aws/example_sagemaker_endpoint.py +++ b/tests/system/providers/amazon/aws/example_sagemaker_endpoint.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") import json from datetime import datetime diff --git a/tests/system/providers/amazon/aws/example_salesforce_to_s3.py b/tests/system/providers/amazon/aws/example_salesforce_to_s3.py index 30c34461b2c4b..598dd9fd3d792 100644 --- a/tests/system/providers/amazon/aws/example_salesforce_to_s3.py +++ b/tests/system/providers/amazon/aws/example_salesforce_to_s3.py @@ -19,7 +19,8 @@ data and upload it to an Amazon S3 bucket. """ from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime from airflow import DAG diff --git a/tests/system/providers/amazon/aws/example_sftp_to_s3.py b/tests/system/providers/amazon/aws/example_sftp_to_s3.py index de07811af0f66..d5758094d286f 100644 --- a/tests/system/providers/amazon/aws/example_sftp_to_s3.py +++ b/tests/system/providers/amazon/aws/example_sftp_to_s3.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime from airflow import models diff --git a/tests/system/providers/amazon/aws/example_sns.py b/tests/system/providers/amazon/aws/example_sns.py index 41431915b43e6..1df77f867526d 100644 --- a/tests/system/providers/amazon/aws/example_sns.py +++ b/tests/system/providers/amazon/aws/example_sns.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime import boto3 diff --git a/tests/system/providers/amazon/aws/example_sql_to_s3.py b/tests/system/providers/amazon/aws/example_sql_to_s3.py index f983881a3f707..882edc0416dbc 100644 --- a/tests/system/providers/amazon/aws/example_sql_to_s3.py +++ b/tests/system/providers/amazon/aws/example_sql_to_s3.py @@ -17,7 +17,8 @@ # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime import boto3 diff --git a/tests/system/providers/amazon/aws/example_sqs.py b/tests/system/providers/amazon/aws/example_sqs.py index 55e242d71559e..49f967fa9bf8a 100644 --- a/tests/system/providers/amazon/aws/example_sqs.py +++ b/tests/system/providers/amazon/aws/example_sqs.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") from datetime import datetime from airflow import DAG diff --git a/tests/system/providers/amazon/aws/example_step_functions.py b/tests/system/providers/amazon/aws/example_step_functions.py index b33a25e48b1aa..0e2ebac237712 100644 --- a/tests/system/providers/amazon/aws/example_step_functions.py +++ b/tests/system/providers/amazon/aws/example_step_functions.py @@ -15,7 +15,8 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations - +from pytest import importorskip +importorskip("aiobotocore") import json from datetime import datetime diff --git a/tests/system/providers/amazon/aws/import_fix.py b/tests/system/providers/amazon/aws/import_fix.py new file mode 100644 index 0000000000000..d7e1128fd1578 --- /dev/null +++ b/tests/system/providers/amazon/aws/import_fix.py @@ -0,0 +1,79 @@ +files = ['example_rds_event.py', + 'example_local_to_s3.py', + 'example_glue.py', + 'example_sagemaker_endpoint.py', + 'example_athena.py', + 'example_eks_with_fargate_profile.py', + 'example_ecs.py', + 'example_s3_to_sql.py', + 'example_rds_instance.py', + 'example_ecs_fargate.py', + 'example_s3_to_sftp.py', + 'example_rds_snapshot.py', + 'example_mongo_to_s3.py', + 'example_step_functions.py', + 'example_google_api_youtube_to_s3.py', + 'example_rds_export.py', + 'example_sns.py', + 'example_eks_with_nodegroups.py', + 'example_redshift_s3_transfers.py', + 'example_google_api_sheets_to_s3.py', + 'example_quicksight.py', + 'example_datasync.py', + 'example_emr_serverless.py', + 'example_ftp_to_s3.py', + 'example_cloudformation.py', + 'example_hive_to_dynamodb.py', + 'example_emr_notebook_execution.py', + 'example_gcs_to_s3.py', + 'example_ec2.py', + 'import_fix.py', + 'example_emr.py', + 'example_glacier_to_gcs.py', + 'example_s3.py', + 'example_emr_eks.py', + 'example_appflow.py', + 'example_redshift.py', + 'example_lambda.py', + 'example_sql_to_s3.py', + 'example_eks_with_nodegroup_in_one_step.py', + 'example_eks_with_fargate_in_one_step.py', + 'example_imap_attachment_to_s3.py', + 'example_eks_templated.py', + 'example_dynamodb_to_s3.py', + 'example_s3_to_ftp.py', + 'example_dms.py', + 'example_salesforce_to_s3.py', + 'example_sagemaker.py', + 'example_sqs.py', + 'example_sftp_to_s3.py', + 'example_batch.py'] + + +for file in files: + base_path = '/opt/airflow/tests/system/providers/amazon/aws/' + path = base_path + file + ans = input(f"working on {path}: Y or N") + if ans == 'N': + break + with open(path, "r") as f: + contents = f.readlines() + + index = -1 + + for i in range(len(contents)): + if contents[i].startswith("from __future__"): + print(f"index is {i}") + index = i + 1 + break + + + value = 'from pytest import importorskip\nimportorskip("aiobotocore")' + + contents.insert(index, value) + + with open(path, "w") as f: + contents = "".join(contents) + f.write(contents) +from pytest import importorskip +importorskip("aiobotocore") From ec21cf1cf0ac8be55e41c468f946725274745f2a Mon Sep 17 00:00:00 2001 From: Syed Hussain Date: Mon, 27 Mar 2023 15:22:30 -0700 Subject: [PATCH 05/17] Remove debug statement, fix static checks --- .../providers/amazon/aws/hooks/base_aws.py | 9 ++- .../providers/amazon/aws/example_appflow.py | 7 +- .../providers/amazon/aws/example_athena.py | 6 +- .../providers/amazon/aws/example_batch.py | 5 +- .../amazon/aws/example_cloudformation.py | 7 +- .../providers/amazon/aws/example_datasync.py | 5 +- .../providers/amazon/aws/example_dms.py | 5 +- .../amazon/aws/example_dynamodb_to_s3.py | 5 +- .../providers/amazon/aws/example_ec2.py | 5 +- .../providers/amazon/aws/example_ecs.py | 5 +- .../amazon/aws/example_ecs_fargate.py | 5 +- .../amazon/aws/example_eks_templated.py | 7 +- .../example_eks_with_fargate_in_one_step.py | 6 +- .../aws/example_eks_with_fargate_profile.py | 6 +- .../example_eks_with_nodegroup_in_one_step.py | 6 +- .../amazon/aws/example_eks_with_nodegroups.py | 5 +- .../providers/amazon/aws/example_emr.py | 5 +- .../providers/amazon/aws/example_emr_eks.py | 5 +- .../aws/example_emr_notebook_execution.py | 6 +- .../amazon/aws/example_emr_serverless.py | 5 +- .../providers/amazon/aws/example_ftp_to_s3.py | 7 +- .../providers/amazon/aws/example_gcs_to_s3.py | 7 +- .../amazon/aws/example_glacier_to_gcs.py | 6 +- .../providers/amazon/aws/example_glue.py | 5 +- .../aws/example_google_api_sheets_to_s3.py | 7 +- .../aws/example_google_api_youtube_to_s3.py | 5 +- .../amazon/aws/example_hive_to_dynamodb.py | 6 +- .../aws/example_imap_attachment_to_s3.py | 6 +- .../providers/amazon/aws/example_lambda.py | 5 +- .../amazon/aws/example_local_to_s3.py | 7 +- .../amazon/aws/example_mongo_to_s3.py | 4 +- .../amazon/aws/example_quicksight.py | 6 +- .../providers/amazon/aws/example_rds_event.py | 5 +- .../amazon/aws/example_rds_export.py | 6 +- .../amazon/aws/example_rds_instance.py | 7 +- .../amazon/aws/example_rds_snapshot.py | 6 +- .../providers/amazon/aws/example_redshift.py | 5 +- .../aws/example_redshift_s3_transfers.py | 5 +- .../system/providers/amazon/aws/example_s3.py | 6 +- .../providers/amazon/aws/example_s3_to_ftp.py | 7 +- .../amazon/aws/example_s3_to_sftp.py | 7 +- .../providers/amazon/aws/example_s3_to_sql.py | 6 +- .../providers/amazon/aws/example_sagemaker.py | 5 +- .../amazon/aws/example_sagemaker_endpoint.py | 5 +- .../amazon/aws/example_salesforce_to_s3.py | 7 +- .../amazon/aws/example_sftp_to_s3.py | 7 +- .../providers/amazon/aws/example_sns.py | 6 +- .../providers/amazon/aws/example_sql_to_s3.py | 5 +- .../providers/amazon/aws/example_sqs.py | 7 +- .../amazon/aws/example_step_functions.py | 6 +- .../system/providers/amazon/aws/import_fix.py | 79 ------------------- 51 files changed, 175 insertions(+), 198 deletions(-) delete mode 100644 tests/system/providers/amazon/aws/import_fix.py diff --git a/airflow/providers/amazon/aws/hooks/base_aws.py b/airflow/providers/amazon/aws/hooks/base_aws.py index a03dde2ef08b6..a8b88ada117f7 100644 --- a/airflow/providers/amazon/aws/hooks/base_aws.py +++ b/airflow/providers/amazon/aws/hooks/base_aws.py @@ -48,6 +48,7 @@ from botocore.credentials import ReadOnlyCredentials from botocore.waiter import Waiter, WaiterModel from dateutil.tz import tzlocal +from pytest import importorskip from slugify import slugify from airflow.compat.functools import cached_property @@ -63,6 +64,8 @@ from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.log.secrets_masker import mask_secret +importorskip("aiobotocore") + BaseAwsConnection = TypeVar("BaseAwsConnection", bound=Union[boto3.client, boto3.resource]) if TYPE_CHECKING: @@ -129,8 +132,7 @@ def role_arn(self) -> str | None: def create_session(self, deferrable: bool = False) -> boto3.session.Session: """Create boto3 or aiobotocore Session from connection config.""" from aiobotocore.session import get_session as async_get_session - from pdb import set_trace - #set_trace() + if not self.conn: self.log.info( "No connection ID provided. Fallback on boto3 credential strategy (region_name=%r). " @@ -163,6 +165,8 @@ def _create_basic_session(self, session_kwargs: dict[str, Any]) -> boto3.session def _create_session_with_assume_role( self, session_kwargs: dict[str, Any], deferrable: bool = False ) -> boto3.session.Session: + from aiobotocore.session import get_session as async_get_session + if self.conn.assume_role_method == "assume_role_with_web_identity": # Deferred credentials have no initial credentials credential_fetcher = self._get_web_identity_credential_fetcher() @@ -846,6 +850,7 @@ def get_waiter( """ from airflow.providers.amazon.aws.waiters.base_waiter import BaseBotoWaiter + if deferrable and not client: raise ValueError("client must be provided for a deferrable waiter.") client = client or self.conn diff --git a/tests/system/providers/amazon/aws/example_appflow.py b/tests/system/providers/amazon/aws/example_appflow.py index 7983c469b85a7..faf32e4399fee 100644 --- a/tests/system/providers/amazon/aws/example_appflow.py +++ b/tests/system/providers/amazon/aws/example_appflow.py @@ -15,10 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime +from pytest import importorskip + from airflow import DAG from airflow.models.baseoperator import chain from airflow.operators.bash import BashOperator @@ -33,7 +34,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() - +importorskip("aiobotocore") DAG_ID = "example_appflow" with DAG( diff --git a/tests/system/providers/amazon/aws/example_athena.py b/tests/system/providers/amazon/aws/example_athena.py index 07feb4ff01a9b..e4603e9515dab 100644 --- a/tests/system/providers/amazon/aws/example_athena.py +++ b/tests/system/providers/amazon/aws/example_athena.py @@ -15,11 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime import boto3 +from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -36,7 +36,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() - +importorskip("aiobotocore") DAG_ID = "example_athena" SAMPLE_DATA = """"Alice",20 diff --git a/tests/system/providers/amazon/aws/example_batch.py b/tests/system/providers/amazon/aws/example_batch.py index 9188b13e02ab2..9b4af52740910 100644 --- a/tests/system/providers/amazon/aws/example_batch.py +++ b/tests/system/providers/amazon/aws/example_batch.py @@ -15,11 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime import boto3 +from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -38,6 +38,7 @@ split_string, ) +importorskip("aiobotocore") DAG_ID = "example_batch" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_cloudformation.py b/tests/system/providers/amazon/aws/example_cloudformation.py index f00cabe1e3b7d..de0f58b96e885 100644 --- a/tests/system/providers/amazon/aws/example_cloudformation.py +++ b/tests/system/providers/amazon/aws/example_cloudformation.py @@ -15,11 +15,12 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + import json from datetime import datetime +from pytest import importorskip + from airflow import DAG from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.cloud_formation import ( @@ -34,7 +35,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() - +importorskip("aiobotocore") DAG_ID = "example_cloudformation" # The CloudFormation template must have at least one resource to diff --git a/tests/system/providers/amazon/aws/example_datasync.py b/tests/system/providers/amazon/aws/example_datasync.py index 98c8de697fa35..4f4cf6965f4b0 100644 --- a/tests/system/providers/amazon/aws/example_datasync.py +++ b/tests/system/providers/amazon/aws/example_datasync.py @@ -15,11 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime import boto3 +from pytest import importorskip from airflow import models from airflow.decorators import task @@ -29,6 +29,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_datasync" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_dms.py b/tests/system/providers/amazon/aws/example_dms.py index 5e44871c3a62f..168752d103ad8 100644 --- a/tests/system/providers/amazon/aws/example_dms.py +++ b/tests/system/providers/amazon/aws/example_dms.py @@ -21,13 +21,13 @@ """ from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + import json from datetime import datetime from typing import cast import boto3 +from pytest import importorskip from sqlalchemy import Column, MetaData, String, Table, create_engine from airflow import DAG @@ -50,6 +50,7 @@ from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder from tests.system.providers.amazon.aws.utils.ec2 import get_default_vpc_id +importorskip("aiobotocore") DAG_ID = "example_dms" ROLE_ARN_KEY = "ROLE_ARN" diff --git a/tests/system/providers/amazon/aws/example_dynamodb_to_s3.py b/tests/system/providers/amazon/aws/example_dynamodb_to_s3.py index 33b632b8ff00a..951a8a134fe6f 100644 --- a/tests/system/providers/amazon/aws/example_dynamodb_to_s3.py +++ b/tests/system/providers/amazon/aws/example_dynamodb_to_s3.py @@ -15,11 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime import boto3 +from pytest import importorskip from airflow.decorators import task from airflow.models.baseoperator import chain @@ -29,6 +29,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_dynamodb_to_s3" sys_test_context_task = SystemTestContextBuilder().build() diff --git a/tests/system/providers/amazon/aws/example_ec2.py b/tests/system/providers/amazon/aws/example_ec2.py index 8bda33e8236f7..cfdf0ac3d418d 100644 --- a/tests/system/providers/amazon/aws/example_ec2.py +++ b/tests/system/providers/amazon/aws/example_ec2.py @@ -15,12 +15,12 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime from operator import itemgetter import boto3 +from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -35,6 +35,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_ec2" sys_test_context_task = SystemTestContextBuilder().build() diff --git a/tests/system/providers/amazon/aws/example_ecs.py b/tests/system/providers/amazon/aws/example_ecs.py index d018beff01f86..80de03cc95652 100644 --- a/tests/system/providers/amazon/aws/example_ecs.py +++ b/tests/system/providers/amazon/aws/example_ecs.py @@ -15,11 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime import boto3 +from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -40,6 +40,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_ecs" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_ecs_fargate.py b/tests/system/providers/amazon/aws/example_ecs_fargate.py index b7da5642ab27e..7955c9f3eddce 100644 --- a/tests/system/providers/amazon/aws/example_ecs_fargate.py +++ b/tests/system/providers/amazon/aws/example_ecs_fargate.py @@ -15,11 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime import boto3 +from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -28,6 +28,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_ecs_fargate" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_eks_templated.py b/tests/system/providers/amazon/aws/example_eks_templated.py index d7c99775009aa..98331c6f28862 100644 --- a/tests/system/providers/amazon/aws/example_eks_templated.py +++ b/tests/system/providers/amazon/aws/example_eks_templated.py @@ -15,10 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime +from pytest import importorskip + from airflow.models.baseoperator import chain from airflow.models.dag import DAG from airflow.providers.amazon.aws.hooks.eks import ClusterStates, NodegroupStates @@ -36,7 +37,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() - +importorskip("aiobotocore") DAG_ID = "example_eks_templated" # Example Jinja Template format, substitute your values: diff --git a/tests/system/providers/amazon/aws/example_eks_with_fargate_in_one_step.py b/tests/system/providers/amazon/aws/example_eks_with_fargate_in_one_step.py index 9084cf20bbf05..275121fcbc4af 100644 --- a/tests/system/providers/amazon/aws/example_eks_with_fargate_in_one_step.py +++ b/tests/system/providers/amazon/aws/example_eks_with_fargate_in_one_step.py @@ -15,10 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime +from pytest import importorskip + from airflow.models.baseoperator import chain from airflow.models.dag import DAG from airflow.operators.bash import BashOperator @@ -32,6 +33,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_eks_with_fargate_in_one_step" # Externally fetched variables diff --git a/tests/system/providers/amazon/aws/example_eks_with_fargate_profile.py b/tests/system/providers/amazon/aws/example_eks_with_fargate_profile.py index 925cf2abeff27..3d242ba7495a2 100644 --- a/tests/system/providers/amazon/aws/example_eks_with_fargate_profile.py +++ b/tests/system/providers/amazon/aws/example_eks_with_fargate_profile.py @@ -15,10 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime +from pytest import importorskip + from airflow.models.baseoperator import chain from airflow.models.dag import DAG from airflow.operators.bash import BashOperator @@ -37,6 +38,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_eks_with_fargate_profile" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_eks_with_nodegroup_in_one_step.py b/tests/system/providers/amazon/aws/example_eks_with_nodegroup_in_one_step.py index 3a2fde590fe14..53307879d3015 100644 --- a/tests/system/providers/amazon/aws/example_eks_with_nodegroup_in_one_step.py +++ b/tests/system/providers/amazon/aws/example_eks_with_nodegroup_in_one_step.py @@ -15,13 +15,12 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime import boto3 -from airflow.decorators import task +from airflow.adecorators import task from airflow.models.baseoperator import chain from airflow.models.dag import DAG from airflow.operators.bash import BashOperator @@ -35,6 +34,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_eks_with_nodegroup_in_one_step" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_eks_with_nodegroups.py b/tests/system/providers/amazon/aws/example_eks_with_nodegroups.py index 28b6e6c9ce593..e6ba0658807e8 100644 --- a/tests/system/providers/amazon/aws/example_eks_with_nodegroups.py +++ b/tests/system/providers/amazon/aws/example_eks_with_nodegroups.py @@ -15,8 +15,7 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime import boto3 @@ -39,7 +38,7 @@ # Ignore missing args provided by default_args # type: ignore[call-arg] - +importorskip("aiobotocore") DAG_ID = "example_eks_with_nodegroups" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_emr.py b/tests/system/providers/amazon/aws/example_emr.py index b2a4b843a9f1e..4f39488b8b1de 100644 --- a/tests/system/providers/amazon/aws/example_emr.py +++ b/tests/system/providers/amazon/aws/example_emr.py @@ -17,12 +17,12 @@ # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + import json from datetime import datetime import boto3 +from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -38,6 +38,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_emr" CONFIG_NAME = "EMR Runtime Role Security Configuration" EXECUTION_ROLE_ARN_KEY = "EXECUTION_ROLE_ARN" diff --git a/tests/system/providers/amazon/aws/example_emr_eks.py b/tests/system/providers/amazon/aws/example_emr_eks.py index 2974ee56b96cb..e0979f5750794 100644 --- a/tests/system/providers/amazon/aws/example_emr_eks.py +++ b/tests/system/providers/amazon/aws/example_emr_eks.py @@ -15,13 +15,13 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + import json import subprocess from datetime import datetime import boto3 +from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -39,6 +39,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_emr_eks" # Externally fetched variables diff --git a/tests/system/providers/amazon/aws/example_emr_notebook_execution.py b/tests/system/providers/amazon/aws/example_emr_notebook_execution.py index 0cd9a5689a80f..2e4f865c674f7 100644 --- a/tests/system/providers/amazon/aws/example_emr_notebook_execution.py +++ b/tests/system/providers/amazon/aws/example_emr_notebook_execution.py @@ -17,10 +17,11 @@ # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime +from pytest import importorskip + from airflow import DAG from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.emr import ( @@ -30,6 +31,7 @@ from airflow.providers.amazon.aws.sensors.emr import EmrNotebookExecutionSensor from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_emr_notebook" # Externally fetched variables: EDITOR_ID_KEY = "EDITOR_ID" diff --git a/tests/system/providers/amazon/aws/example_emr_serverless.py b/tests/system/providers/amazon/aws/example_emr_serverless.py index 998a2620d758f..de6db77eef8ed 100644 --- a/tests/system/providers/amazon/aws/example_emr_serverless.py +++ b/tests/system/providers/amazon/aws/example_emr_serverless.py @@ -15,11 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime import boto3 +from pytest import importorskip from airflow.models.baseoperator import chain from airflow.models.dag import DAG @@ -33,6 +33,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_emr_serverless" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_ftp_to_s3.py b/tests/system/providers/amazon/aws/example_ftp_to_s3.py index 195b20c6b5a78..07c2e63e7a9d2 100644 --- a/tests/system/providers/amazon/aws/example_ftp_to_s3.py +++ b/tests/system/providers/amazon/aws/example_ftp_to_s3.py @@ -15,10 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime +from pytest import importorskip + from airflow import models from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator @@ -27,7 +28,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() - +importorskip("aiobotocore") DAG_ID = "example_ftp_to_s3" with models.DAG( diff --git a/tests/system/providers/amazon/aws/example_gcs_to_s3.py b/tests/system/providers/amazon/aws/example_gcs_to_s3.py index 2bf7080f7c349..7cb5ce9715c9e 100644 --- a/tests/system/providers/amazon/aws/example_gcs_to_s3.py +++ b/tests/system/providers/amazon/aws/example_gcs_to_s3.py @@ -15,10 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime +from pytest import importorskip + from airflow import DAG from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator @@ -27,7 +28,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() - +importorskip("aiobotocore") DAG_ID = "example_gcs_to_s3" with DAG( diff --git a/tests/system/providers/amazon/aws/example_glacier_to_gcs.py b/tests/system/providers/amazon/aws/example_glacier_to_gcs.py index 88a29ae878bcc..c87f7d30d4a25 100644 --- a/tests/system/providers/amazon/aws/example_glacier_to_gcs.py +++ b/tests/system/providers/amazon/aws/example_glacier_to_gcs.py @@ -15,11 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime import boto3 +from pytest import importorskip from airflow import DAG from airflow.models.baseoperator import chain @@ -34,7 +34,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() - +importorskip("aiobotocore") DAG_ID = "example_glacier_to_gcs" diff --git a/tests/system/providers/amazon/aws/example_glue.py b/tests/system/providers/amazon/aws/example_glue.py index d23929f655f05..cf2b9527c1d8a 100644 --- a/tests/system/providers/amazon/aws/example_glue.py +++ b/tests/system/providers/amazon/aws/example_glue.py @@ -15,12 +15,12 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime import boto3 from botocore.client import BaseClient +from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -37,6 +37,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder, prune_logs +importorskip("aiobotocore") DAG_ID = "example_glue" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_google_api_sheets_to_s3.py b/tests/system/providers/amazon/aws/example_google_api_sheets_to_s3.py index e94d82873bd24..47a10dcffda55 100644 --- a/tests/system/providers/amazon/aws/example_google_api_sheets_to_s3.py +++ b/tests/system/providers/amazon/aws/example_google_api_sheets_to_s3.py @@ -19,11 +19,12 @@ You need to set all env variables to request the data. """ from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime from os import getenv +from pytest import importorskip + from airflow import DAG from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator @@ -32,7 +33,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() - +importorskip("aiobotocore") DAG_ID = "example_google_api_sheets_to_s3" GOOGLE_SHEET_ID = getenv("GOOGLE_SHEET_ID", "test-google-sheet-id") diff --git a/tests/system/providers/amazon/aws/example_google_api_youtube_to_s3.py b/tests/system/providers/amazon/aws/example_google_api_youtube_to_s3.py index 4712ad07da4e4..ebce3fa6bce45 100644 --- a/tests/system/providers/amazon/aws/example_google_api_youtube_to_s3.py +++ b/tests/system/providers/amazon/aws/example_google_api_youtube_to_s3.py @@ -46,12 +46,12 @@ or by creating a custom connection. """ from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + import json from datetime import datetime import boto3 +from pytest import importorskip from airflow import DAG, settings from airflow.decorators import task @@ -62,6 +62,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_google_api_youtube_to_s3" YOUTUBE_CHANNEL_ID = "UCSXwxpWZQ7XZ1WL3wqevChA" diff --git a/tests/system/providers/amazon/aws/example_hive_to_dynamodb.py b/tests/system/providers/amazon/aws/example_hive_to_dynamodb.py index 372cb571d3c7c..c6abf2d771ae5 100644 --- a/tests/system/providers/amazon/aws/example_hive_to_dynamodb.py +++ b/tests/system/providers/amazon/aws/example_hive_to_dynamodb.py @@ -20,10 +20,11 @@ https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/EMRforDynamoDB.Tutorial.html """ from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime +from pytest import importorskip + from airflow import DAG from airflow.decorators import task from airflow.models import Connection @@ -34,6 +35,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_hive_to_dynamodb" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_imap_attachment_to_s3.py b/tests/system/providers/amazon/aws/example_imap_attachment_to_s3.py index a92a8b0960d57..ad4b20709c2ee 100644 --- a/tests/system/providers/amazon/aws/example_imap_attachment_to_s3.py +++ b/tests/system/providers/amazon/aws/example_imap_attachment_to_s3.py @@ -19,10 +19,11 @@ protocol from a mail server to S3 Bucket. """ from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime +from pytest import importorskip + from airflow import DAG from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator @@ -30,6 +31,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_imap_attachment_to_s3" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_lambda.py b/tests/system/providers/amazon/aws/example_lambda.py index 7852cc6238f1b..e587a01e268e0 100644 --- a/tests/system/providers/amazon/aws/example_lambda.py +++ b/tests/system/providers/amazon/aws/example_lambda.py @@ -15,14 +15,14 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + import io import json import zipfile from datetime import datetime import boto3 +from pytest import importorskip from airflow import models from airflow.decorators import task @@ -35,6 +35,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder, prune_logs +importorskip("aiobotocore") DAG_ID = "example_lambda" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_local_to_s3.py b/tests/system/providers/amazon/aws/example_local_to_s3.py index 43fa77b7a2d82..e61496bbb6157 100644 --- a/tests/system/providers/amazon/aws/example_local_to_s3.py +++ b/tests/system/providers/amazon/aws/example_local_to_s3.py @@ -15,11 +15,12 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + import os from datetime import datetime +from pytest import importorskip + from airflow import DAG from airflow.decorators import task from airflow.models.baseoperator import chain @@ -30,7 +31,7 @@ sys_test_context_task = SystemTestContextBuilder().build() - +importorskip("aiobotocore") DAG_ID = "example_local_to_s3" TEMP_FILE_PATH = "/tmp/sample-txt.txt" SAMPLE_TEXT = "This is some sample text." diff --git a/tests/system/providers/amazon/aws/example_mongo_to_s3.py b/tests/system/providers/amazon/aws/example_mongo_to_s3.py index 6e91c2a4f4282..db484ca6a75ca 100644 --- a/tests/system/providers/amazon/aws/example_mongo_to_s3.py +++ b/tests/system/providers/amazon/aws/example_mongo_to_s3.py @@ -15,8 +15,9 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations + from pytest import importorskip -importorskip("aiobotocore") + from airflow import models from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator @@ -25,6 +26,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_mongo_to_s3" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_quicksight.py b/tests/system/providers/amazon/aws/example_quicksight.py index 5d4d561ae0411..1bf13cff710eb 100644 --- a/tests/system/providers/amazon/aws/example_quicksight.py +++ b/tests/system/providers/amazon/aws/example_quicksight.py @@ -15,12 +15,12 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + import json from datetime import datetime import boto3 +from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -48,7 +48,7 @@ into this behavior, changing the template for the ingestion name or the ENV_ID and re-running the test should resolve the issue. """ - +importorskip("aiobotocore") DAG_ID = "example_quicksight" sys_test_context_task = SystemTestContextBuilder().build() diff --git a/tests/system/providers/amazon/aws/example_rds_event.py b/tests/system/providers/amazon/aws/example_rds_event.py index 70d3c56b1a4c7..a17d9f39edccd 100644 --- a/tests/system/providers/amazon/aws/example_rds_event.py +++ b/tests/system/providers/amazon/aws/example_rds_event.py @@ -16,11 +16,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime import boto3 +from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -34,6 +34,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_rds_event" sys_test_context_task = SystemTestContextBuilder().build() diff --git a/tests/system/providers/amazon/aws/example_rds_export.py b/tests/system/providers/amazon/aws/example_rds_export.py index 343d4233390d4..069e45de945e9 100644 --- a/tests/system/providers/amazon/aws/example_rds_export.py +++ b/tests/system/providers/amazon/aws/example_rds_export.py @@ -15,10 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime +from pytest import importorskip + from airflow import DAG from airflow.decorators import task from airflow.models.baseoperator import chain @@ -36,6 +37,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_rds_export" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_rds_instance.py b/tests/system/providers/amazon/aws/example_rds_instance.py index b47a12ca3e08b..016e2d6770d49 100644 --- a/tests/system/providers/amazon/aws/example_rds_instance.py +++ b/tests/system/providers/amazon/aws/example_rds_instance.py @@ -15,10 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime +from pytest import importorskip + from airflow import DAG from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.rds import ( @@ -32,7 +33,7 @@ from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() - +importorskip("aiobotocore") DAG_ID = "example_rds_instance" RDS_USERNAME = "database_username" diff --git a/tests/system/providers/amazon/aws/example_rds_snapshot.py b/tests/system/providers/amazon/aws/example_rds_snapshot.py index b139985ad3fdf..3c00c4f761348 100644 --- a/tests/system/providers/amazon/aws/example_rds_snapshot.py +++ b/tests/system/providers/amazon/aws/example_rds_snapshot.py @@ -15,10 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime +from pytest import importorskip + from airflow import DAG from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.rds import ( @@ -32,6 +33,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_rds_snapshot" sys_test_context_task = SystemTestContextBuilder().build() diff --git a/tests/system/providers/amazon/aws/example_redshift.py b/tests/system/providers/amazon/aws/example_redshift.py index a23c09fa1077b..2d7af63ad1854 100644 --- a/tests/system/providers/amazon/aws/example_redshift.py +++ b/tests/system/providers/amazon/aws/example_redshift.py @@ -17,11 +17,11 @@ # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime import boto3 +from pytest import importorskip from airflow import DAG, settings from airflow.decorators import task @@ -42,6 +42,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_redshift" DB_LOGIN = "adminuser" DB_PASS = "MyAmazonPassword1" diff --git a/tests/system/providers/amazon/aws/example_redshift_s3_transfers.py b/tests/system/providers/amazon/aws/example_redshift_s3_transfers.py index 3e045b432ff22..53fe2c9a74b8a 100644 --- a/tests/system/providers/amazon/aws/example_redshift_s3_transfers.py +++ b/tests/system/providers/amazon/aws/example_redshift_s3_transfers.py @@ -15,11 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime import boto3 +from pytest import importorskip from airflow import DAG, settings from airflow.decorators import task @@ -43,6 +43,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_redshift_to_s3" DB_LOGIN = "adminuser" diff --git a/tests/system/providers/amazon/aws/example_s3.py b/tests/system/providers/amazon/aws/example_s3.py index 08ddcd7d9206b..f1abf624ea57a 100644 --- a/tests/system/providers/amazon/aws/example_s3.py +++ b/tests/system/providers/amazon/aws/example_s3.py @@ -15,10 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime +from pytest import importorskip + from airflow.models.baseoperator import chain from airflow.models.dag import DAG from airflow.operators.python import BranchPythonOperator @@ -39,6 +40,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_s3" sys_test_context_task = SystemTestContextBuilder().build() diff --git a/tests/system/providers/amazon/aws/example_s3_to_ftp.py b/tests/system/providers/amazon/aws/example_s3_to_ftp.py index 6164fcd6ede12..a81db7d2bcf91 100644 --- a/tests/system/providers/amazon/aws/example_s3_to_ftp.py +++ b/tests/system/providers/amazon/aws/example_s3_to_ftp.py @@ -15,10 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime +from pytest import importorskip + from airflow import models from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator @@ -27,7 +28,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() - +importorskip("aiobotocore") DAG_ID = "example_s3_to_ftp" with models.DAG( diff --git a/tests/system/providers/amazon/aws/example_s3_to_sftp.py b/tests/system/providers/amazon/aws/example_s3_to_sftp.py index d7e027695039a..54b8239398571 100644 --- a/tests/system/providers/amazon/aws/example_s3_to_sftp.py +++ b/tests/system/providers/amazon/aws/example_s3_to_sftp.py @@ -15,10 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime +from pytest import importorskip + from airflow import models from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator @@ -27,7 +28,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() - +importorskip("aiobotocore") DAG_ID = "example_s3_to_sftp" with models.DAG( diff --git a/tests/system/providers/amazon/aws/example_s3_to_sql.py b/tests/system/providers/amazon/aws/example_s3_to_sql.py index 3cc1de8b22ea6..58fad3bbc0aea 100644 --- a/tests/system/providers/amazon/aws/example_s3_to_sql.py +++ b/tests/system/providers/amazon/aws/example_s3_to_sql.py @@ -15,11 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime import boto3 +from pytest import importorskip from airflow import DAG, settings from airflow.decorators import task @@ -44,7 +44,7 @@ from tests.system.utils.watcher import watcher sys_test_context_task = SystemTestContextBuilder().build() - +importorskip("aiobotocore") DAG_ID = "example_s3_to_sql" DB_LOGIN = "adminuser" diff --git a/tests/system/providers/amazon/aws/example_sagemaker.py b/tests/system/providers/amazon/aws/example_sagemaker.py index 745709639d096..900da98eaa1f8 100644 --- a/tests/system/providers/amazon/aws/example_sagemaker.py +++ b/tests/system/providers/amazon/aws/example_sagemaker.py @@ -15,8 +15,7 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + import json import logging import subprocess @@ -24,6 +23,7 @@ from tempfile import NamedTemporaryFile import boto3 +from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -58,6 +58,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder, prune_logs +importorskip("aiobotocore") DAG_ID = "example_sagemaker" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_sagemaker_endpoint.py b/tests/system/providers/amazon/aws/example_sagemaker_endpoint.py index 9b71a579f58da..8a85b3f803c5e 100644 --- a/tests/system/providers/amazon/aws/example_sagemaker_endpoint.py +++ b/tests/system/providers/amazon/aws/example_sagemaker_endpoint.py @@ -15,12 +15,12 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + import json from datetime import datetime import boto3 +from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -41,6 +41,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder, prune_logs +importorskip("aiobotocore") DAG_ID = "example_sagemaker_endpoint" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_salesforce_to_s3.py b/tests/system/providers/amazon/aws/example_salesforce_to_s3.py index 598dd9fd3d792..63ce2c160af53 100644 --- a/tests/system/providers/amazon/aws/example_salesforce_to_s3.py +++ b/tests/system/providers/amazon/aws/example_salesforce_to_s3.py @@ -19,10 +19,11 @@ data and upload it to an Amazon S3 bucket. """ from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime +from pytest import importorskip + from airflow import DAG from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator @@ -31,7 +32,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() - +importorskip("aiobotocore") DAG_ID = "example_salesforce_to_s3" with DAG( diff --git a/tests/system/providers/amazon/aws/example_sftp_to_s3.py b/tests/system/providers/amazon/aws/example_sftp_to_s3.py index d5758094d286f..1b2e9d4be50f2 100644 --- a/tests/system/providers/amazon/aws/example_sftp_to_s3.py +++ b/tests/system/providers/amazon/aws/example_sftp_to_s3.py @@ -15,10 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime +from pytest import importorskip + from airflow import models from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator @@ -27,7 +28,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() - +importorskip("aiobotocore") DAG_ID = "example_sftp_to_s3" with models.DAG( diff --git a/tests/system/providers/amazon/aws/example_sns.py b/tests/system/providers/amazon/aws/example_sns.py index 1df77f867526d..d414f53a8a96e 100644 --- a/tests/system/providers/amazon/aws/example_sns.py +++ b/tests/system/providers/amazon/aws/example_sns.py @@ -15,11 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime import boto3 +from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -29,7 +29,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() - +importorskip("aiobotocore") DAG_ID = "example_sns" diff --git a/tests/system/providers/amazon/aws/example_sql_to_s3.py b/tests/system/providers/amazon/aws/example_sql_to_s3.py index 882edc0416dbc..149f12eddedb2 100644 --- a/tests/system/providers/amazon/aws/example_sql_to_s3.py +++ b/tests/system/providers/amazon/aws/example_sql_to_s3.py @@ -17,11 +17,11 @@ # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime import boto3 +from pytest import importorskip from airflow import DAG, settings from airflow.decorators import task @@ -39,6 +39,7 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_sql_to_s3" DB_LOGIN = "adminuser" DB_PASS = "MyAmazonPassword1" diff --git a/tests/system/providers/amazon/aws/example_sqs.py b/tests/system/providers/amazon/aws/example_sqs.py index 49f967fa9bf8a..b5945b5c821cc 100644 --- a/tests/system/providers/amazon/aws/example_sqs.py +++ b/tests/system/providers/amazon/aws/example_sqs.py @@ -15,10 +15,11 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + from datetime import datetime +from pytest import importorskip + from airflow import DAG from airflow.decorators import task from airflow.models.baseoperator import chain @@ -29,7 +30,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() - +importorskip("aiobotocore") DAG_ID = "example_sqs" diff --git a/tests/system/providers/amazon/aws/example_step_functions.py b/tests/system/providers/amazon/aws/example_step_functions.py index 0e2ebac237712..ecb050374fde8 100644 --- a/tests/system/providers/amazon/aws/example_step_functions.py +++ b/tests/system/providers/amazon/aws/example_step_functions.py @@ -15,11 +15,12 @@ # specific language governing permissions and limitations # under the License. from __future__ import annotations -from pytest import importorskip -importorskip("aiobotocore") + import json from datetime import datetime +from pytest import importorskip + from airflow import DAG from airflow.decorators import task from airflow.models.baseoperator import chain @@ -31,6 +32,7 @@ from airflow.providers.amazon.aws.sensors.step_function import StepFunctionExecutionSensor from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder +importorskip("aiobotocore") DAG_ID = "example_step_functions" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/import_fix.py b/tests/system/providers/amazon/aws/import_fix.py deleted file mode 100644 index d7e1128fd1578..0000000000000 --- a/tests/system/providers/amazon/aws/import_fix.py +++ /dev/null @@ -1,79 +0,0 @@ -files = ['example_rds_event.py', - 'example_local_to_s3.py', - 'example_glue.py', - 'example_sagemaker_endpoint.py', - 'example_athena.py', - 'example_eks_with_fargate_profile.py', - 'example_ecs.py', - 'example_s3_to_sql.py', - 'example_rds_instance.py', - 'example_ecs_fargate.py', - 'example_s3_to_sftp.py', - 'example_rds_snapshot.py', - 'example_mongo_to_s3.py', - 'example_step_functions.py', - 'example_google_api_youtube_to_s3.py', - 'example_rds_export.py', - 'example_sns.py', - 'example_eks_with_nodegroups.py', - 'example_redshift_s3_transfers.py', - 'example_google_api_sheets_to_s3.py', - 'example_quicksight.py', - 'example_datasync.py', - 'example_emr_serverless.py', - 'example_ftp_to_s3.py', - 'example_cloudformation.py', - 'example_hive_to_dynamodb.py', - 'example_emr_notebook_execution.py', - 'example_gcs_to_s3.py', - 'example_ec2.py', - 'import_fix.py', - 'example_emr.py', - 'example_glacier_to_gcs.py', - 'example_s3.py', - 'example_emr_eks.py', - 'example_appflow.py', - 'example_redshift.py', - 'example_lambda.py', - 'example_sql_to_s3.py', - 'example_eks_with_nodegroup_in_one_step.py', - 'example_eks_with_fargate_in_one_step.py', - 'example_imap_attachment_to_s3.py', - 'example_eks_templated.py', - 'example_dynamodb_to_s3.py', - 'example_s3_to_ftp.py', - 'example_dms.py', - 'example_salesforce_to_s3.py', - 'example_sagemaker.py', - 'example_sqs.py', - 'example_sftp_to_s3.py', - 'example_batch.py'] - - -for file in files: - base_path = '/opt/airflow/tests/system/providers/amazon/aws/' - path = base_path + file - ans = input(f"working on {path}: Y or N") - if ans == 'N': - break - with open(path, "r") as f: - contents = f.readlines() - - index = -1 - - for i in range(len(contents)): - if contents[i].startswith("from __future__"): - print(f"index is {i}") - index = i + 1 - break - - - value = 'from pytest import importorskip\nimportorskip("aiobotocore")' - - contents.insert(index, value) - - with open(path, "w") as f: - contents = "".join(contents) - f.write(contents) -from pytest import importorskip -importorskip("aiobotocore") From 05d226129c0ddcd2cc75ab91fa8ec740aacf9c65 Mon Sep 17 00:00:00 2001 From: Syed Hussain Date: Tue, 28 Mar 2023 13:54:32 -0700 Subject: [PATCH 06/17] Put import of aiobotocore into function so it only imports when it is used --- airflow/providers/amazon/aws/hooks/base_aws.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/base_aws.py b/airflow/providers/amazon/aws/hooks/base_aws.py index a8b88ada117f7..2425e4881ad93 100644 --- a/airflow/providers/amazon/aws/hooks/base_aws.py +++ b/airflow/providers/amazon/aws/hooks/base_aws.py @@ -129,20 +129,27 @@ def role_arn(self) -> str | None: """Assume Role ARN from AWS Connection""" return self.conn.role_arn - def create_session(self, deferrable: bool = False) -> boto3.session.Session: - """Create boto3 or aiobotocore Session from connection config.""" + def get_async_session(self): from aiobotocore.session import get_session as async_get_session + return async_get_session() + + def create_session(self, deferrable: bool = False) -> boto3.session.Session: + """Create boto3 or aiobotocore Session from connection config.""" if not self.conn: self.log.info( "No connection ID provided. Fallback on boto3 credential strategy (region_name=%r). " "See: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html", self.region_name, ) - return async_get_session() if deferrable else boto3.session.Session(region_name=self.region_name) + return ( + self.get_async_session() + if deferrable + else boto3.session.Session(region_name=self.region_name) + ) elif not self.role_arn: - return async_get_session() if deferrable else self.basic_session + return self.get_async_session() if deferrable else self.basic_session # Values stored in ``AwsConnectionWrapper.session_kwargs`` are intended to be used only # to create the initial boto3 session. From c73f7b6230d0596e1f4b8ec30f18cc524a3b22f7 Mon Sep 17 00:00:00 2001 From: Syed Hussain Date: Thu, 30 Mar 2023 14:22:23 -0700 Subject: [PATCH 07/17] Remove importorskip from system tests and base_aws --- airflow/providers/amazon/aws/hooks/base_aws.py | 5 +---- airflow/providers/amazon/aws/waiters/base_waiter.py | 13 +++++++++---- .../system/providers/amazon/aws/example_appflow.py | 4 +--- tests/system/providers/amazon/aws/example_athena.py | 3 +-- tests/system/providers/amazon/aws/example_batch.py | 2 -- .../providers/amazon/aws/example_cloudformation.py | 4 +--- .../system/providers/amazon/aws/example_datasync.py | 2 -- tests/system/providers/amazon/aws/example_dms.py | 2 -- .../providers/amazon/aws/example_dynamodb_to_s3.py | 2 -- tests/system/providers/amazon/aws/example_ec2.py | 2 -- tests/system/providers/amazon/aws/example_ecs.py | 2 -- .../providers/amazon/aws/example_ecs_fargate.py | 2 -- .../providers/amazon/aws/example_eks_templated.py | 4 +--- .../aws/example_eks_with_fargate_in_one_step.py | 3 --- .../amazon/aws/example_eks_with_fargate_profile.py | 3 --- .../aws/example_eks_with_nodegroup_in_one_step.py | 1 - .../amazon/aws/example_eks_with_nodegroups.py | 2 +- tests/system/providers/amazon/aws/example_emr.py | 2 -- .../system/providers/amazon/aws/example_emr_eks.py | 2 -- .../amazon/aws/example_emr_notebook_execution.py | 3 --- .../providers/amazon/aws/example_emr_serverless.py | 2 -- .../providers/amazon/aws/example_ftp_to_s3.py | 4 +--- .../providers/amazon/aws/example_gcs_to_s3.py | 4 +--- .../providers/amazon/aws/example_glacier_to_gcs.py | 3 +-- tests/system/providers/amazon/aws/example_glue.py | 2 -- .../amazon/aws/example_google_api_sheets_to_s3.py | 4 +--- .../amazon/aws/example_google_api_youtube_to_s3.py | 2 -- .../amazon/aws/example_hive_to_dynamodb.py | 3 --- .../amazon/aws/example_imap_attachment_to_s3.py | 3 --- tests/system/providers/amazon/aws/example_lambda.py | 2 -- .../providers/amazon/aws/example_local_to_s3.py | 4 +--- .../providers/amazon/aws/example_mongo_to_s3.py | 3 --- .../providers/amazon/aws/example_quicksight.py | 3 +-- .../providers/amazon/aws/example_rds_event.py | 2 -- .../providers/amazon/aws/example_rds_export.py | 3 --- .../providers/amazon/aws/example_rds_instance.py | 4 +--- .../providers/amazon/aws/example_rds_snapshot.py | 3 --- .../system/providers/amazon/aws/example_redshift.py | 2 -- .../amazon/aws/example_redshift_s3_transfers.py | 2 -- tests/system/providers/amazon/aws/example_s3.py | 3 --- .../providers/amazon/aws/example_s3_to_ftp.py | 4 +--- .../providers/amazon/aws/example_s3_to_sftp.py | 4 +--- .../providers/amazon/aws/example_s3_to_sql.py | 3 +-- .../providers/amazon/aws/example_sagemaker.py | 2 -- .../amazon/aws/example_sagemaker_endpoint.py | 2 -- .../amazon/aws/example_salesforce_to_s3.py | 4 +--- .../providers/amazon/aws/example_sftp_to_s3.py | 4 +--- tests/system/providers/amazon/aws/example_sns.py | 3 +-- .../providers/amazon/aws/example_sql_to_s3.py | 2 -- tests/system/providers/amazon/aws/example_sqs.py | 4 +--- .../providers/amazon/aws/example_step_functions.py | 3 --- 51 files changed, 29 insertions(+), 127 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/base_aws.py b/airflow/providers/amazon/aws/hooks/base_aws.py index 2425e4881ad93..96e0528407133 100644 --- a/airflow/providers/amazon/aws/hooks/base_aws.py +++ b/airflow/providers/amazon/aws/hooks/base_aws.py @@ -48,7 +48,6 @@ from botocore.credentials import ReadOnlyCredentials from botocore.waiter import Waiter, WaiterModel from dateutil.tz import tzlocal -from pytest import importorskip from slugify import slugify from airflow.compat.functools import cached_property @@ -64,8 +63,6 @@ from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.log.secrets_masker import mask_secret -importorskip("aiobotocore") - BaseAwsConnection = TypeVar("BaseAwsConnection", bound=Union[boto3.client, boto3.resource]) if TYPE_CHECKING: @@ -614,7 +611,7 @@ def get_client_type( """Get the underlying boto3 client using boto3 session""" client_type = self.client_type session = self.get_session(region_name=region_name, deferrable=deferrable) - if isinstance(session, AioSession): + if not isinstance(session, boto3.session.Session): return session.create_client( client_type, endpoint_url=self.conn_config.endpoint_url, diff --git a/airflow/providers/amazon/aws/waiters/base_waiter.py b/airflow/providers/amazon/aws/waiters/base_waiter.py index b4b8668ea8127..488767a084a21 100644 --- a/airflow/providers/amazon/aws/waiters/base_waiter.py +++ b/airflow/providers/amazon/aws/waiters/base_waiter.py @@ -18,7 +18,6 @@ from __future__ import annotations import boto3 -from aiobotocore.waiter import create_waiter_with_client as create_async_waiter_with_client from botocore.waiter import Waiter, WaiterModel, create_waiter_with_client @@ -34,9 +33,15 @@ def __init__(self, client: boto3.client, model_config: dict, deferrable: bool = self.client = client self.deferrable = deferrable + def _get_async_waiter_with_client(self, waiter_name: str): + from aiobotocore.waiter import create_waiter_with_client as create_async_waiter_with_client + + return create_async_waiter_with_client( + waiter_name=waiter_name, waiter_model=self.model, client=self.client + ) + def waiter(self, waiter_name: str) -> Waiter: if self.deferrable: - return create_async_waiter_with_client( - waiter_name=waiter_name, waiter_model=self.model, client=self.client - ) + return self._get_async_waiter_with_client(waiter_name=waiter_name) + return create_waiter_with_client(waiter_name=waiter_name, waiter_model=self.model, client=self.client) diff --git a/tests/system/providers/amazon/aws/example_appflow.py b/tests/system/providers/amazon/aws/example_appflow.py index faf32e4399fee..4469c0290b539 100644 --- a/tests/system/providers/amazon/aws/example_appflow.py +++ b/tests/system/providers/amazon/aws/example_appflow.py @@ -18,8 +18,6 @@ from datetime import datetime -from pytest import importorskip - from airflow import DAG from airflow.models.baseoperator import chain from airflow.operators.bash import BashOperator @@ -34,7 +32,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() -importorskip("aiobotocore") + DAG_ID = "example_appflow" with DAG( diff --git a/tests/system/providers/amazon/aws/example_athena.py b/tests/system/providers/amazon/aws/example_athena.py index e4603e9515dab..3d2487d7b56a4 100644 --- a/tests/system/providers/amazon/aws/example_athena.py +++ b/tests/system/providers/amazon/aws/example_athena.py @@ -19,7 +19,6 @@ from datetime import datetime import boto3 -from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -36,7 +35,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() -importorskip("aiobotocore") + DAG_ID = "example_athena" SAMPLE_DATA = """"Alice",20 diff --git a/tests/system/providers/amazon/aws/example_batch.py b/tests/system/providers/amazon/aws/example_batch.py index 9b4af52740910..a035b12f87b6e 100644 --- a/tests/system/providers/amazon/aws/example_batch.py +++ b/tests/system/providers/amazon/aws/example_batch.py @@ -19,7 +19,6 @@ from datetime import datetime import boto3 -from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -38,7 +37,6 @@ split_string, ) -importorskip("aiobotocore") DAG_ID = "example_batch" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_cloudformation.py b/tests/system/providers/amazon/aws/example_cloudformation.py index de0f58b96e885..fc6e04d422081 100644 --- a/tests/system/providers/amazon/aws/example_cloudformation.py +++ b/tests/system/providers/amazon/aws/example_cloudformation.py @@ -19,8 +19,6 @@ import json from datetime import datetime -from pytest import importorskip - from airflow import DAG from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.cloud_formation import ( @@ -35,7 +33,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() -importorskip("aiobotocore") + DAG_ID = "example_cloudformation" # The CloudFormation template must have at least one resource to diff --git a/tests/system/providers/amazon/aws/example_datasync.py b/tests/system/providers/amazon/aws/example_datasync.py index 4f4cf6965f4b0..bead9f8f2f80a 100644 --- a/tests/system/providers/amazon/aws/example_datasync.py +++ b/tests/system/providers/amazon/aws/example_datasync.py @@ -19,7 +19,6 @@ from datetime import datetime import boto3 -from pytest import importorskip from airflow import models from airflow.decorators import task @@ -29,7 +28,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_datasync" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_dms.py b/tests/system/providers/amazon/aws/example_dms.py index 168752d103ad8..8e13e17c32e79 100644 --- a/tests/system/providers/amazon/aws/example_dms.py +++ b/tests/system/providers/amazon/aws/example_dms.py @@ -27,7 +27,6 @@ from typing import cast import boto3 -from pytest import importorskip from sqlalchemy import Column, MetaData, String, Table, create_engine from airflow import DAG @@ -50,7 +49,6 @@ from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder from tests.system.providers.amazon.aws.utils.ec2 import get_default_vpc_id -importorskip("aiobotocore") DAG_ID = "example_dms" ROLE_ARN_KEY = "ROLE_ARN" diff --git a/tests/system/providers/amazon/aws/example_dynamodb_to_s3.py b/tests/system/providers/amazon/aws/example_dynamodb_to_s3.py index 951a8a134fe6f..b56efaf2ce1ba 100644 --- a/tests/system/providers/amazon/aws/example_dynamodb_to_s3.py +++ b/tests/system/providers/amazon/aws/example_dynamodb_to_s3.py @@ -19,7 +19,6 @@ from datetime import datetime import boto3 -from pytest import importorskip from airflow.decorators import task from airflow.models.baseoperator import chain @@ -29,7 +28,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_dynamodb_to_s3" sys_test_context_task = SystemTestContextBuilder().build() diff --git a/tests/system/providers/amazon/aws/example_ec2.py b/tests/system/providers/amazon/aws/example_ec2.py index cfdf0ac3d418d..1dd98488c2995 100644 --- a/tests/system/providers/amazon/aws/example_ec2.py +++ b/tests/system/providers/amazon/aws/example_ec2.py @@ -20,7 +20,6 @@ from operator import itemgetter import boto3 -from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -35,7 +34,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_ec2" sys_test_context_task = SystemTestContextBuilder().build() diff --git a/tests/system/providers/amazon/aws/example_ecs.py b/tests/system/providers/amazon/aws/example_ecs.py index 80de03cc95652..194b070b51686 100644 --- a/tests/system/providers/amazon/aws/example_ecs.py +++ b/tests/system/providers/amazon/aws/example_ecs.py @@ -19,7 +19,6 @@ from datetime import datetime import boto3 -from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -40,7 +39,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_ecs" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_ecs_fargate.py b/tests/system/providers/amazon/aws/example_ecs_fargate.py index 7955c9f3eddce..40132358ab430 100644 --- a/tests/system/providers/amazon/aws/example_ecs_fargate.py +++ b/tests/system/providers/amazon/aws/example_ecs_fargate.py @@ -19,7 +19,6 @@ from datetime import datetime import boto3 -from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -28,7 +27,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_ecs_fargate" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_eks_templated.py b/tests/system/providers/amazon/aws/example_eks_templated.py index 98331c6f28862..d09eabf959827 100644 --- a/tests/system/providers/amazon/aws/example_eks_templated.py +++ b/tests/system/providers/amazon/aws/example_eks_templated.py @@ -18,8 +18,6 @@ from datetime import datetime -from pytest import importorskip - from airflow.models.baseoperator import chain from airflow.models.dag import DAG from airflow.providers.amazon.aws.hooks.eks import ClusterStates, NodegroupStates @@ -37,7 +35,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() -importorskip("aiobotocore") + DAG_ID = "example_eks_templated" # Example Jinja Template format, substitute your values: diff --git a/tests/system/providers/amazon/aws/example_eks_with_fargate_in_one_step.py b/tests/system/providers/amazon/aws/example_eks_with_fargate_in_one_step.py index 275121fcbc4af..37cba110d9bfd 100644 --- a/tests/system/providers/amazon/aws/example_eks_with_fargate_in_one_step.py +++ b/tests/system/providers/amazon/aws/example_eks_with_fargate_in_one_step.py @@ -18,8 +18,6 @@ from datetime import datetime -from pytest import importorskip - from airflow.models.baseoperator import chain from airflow.models.dag import DAG from airflow.operators.bash import BashOperator @@ -33,7 +31,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_eks_with_fargate_in_one_step" # Externally fetched variables diff --git a/tests/system/providers/amazon/aws/example_eks_with_fargate_profile.py b/tests/system/providers/amazon/aws/example_eks_with_fargate_profile.py index 3d242ba7495a2..5792332136deb 100644 --- a/tests/system/providers/amazon/aws/example_eks_with_fargate_profile.py +++ b/tests/system/providers/amazon/aws/example_eks_with_fargate_profile.py @@ -18,8 +18,6 @@ from datetime import datetime -from pytest import importorskip - from airflow.models.baseoperator import chain from airflow.models.dag import DAG from airflow.operators.bash import BashOperator @@ -38,7 +36,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_eks_with_fargate_profile" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_eks_with_nodegroup_in_one_step.py b/tests/system/providers/amazon/aws/example_eks_with_nodegroup_in_one_step.py index 53307879d3015..00c6dc1286476 100644 --- a/tests/system/providers/amazon/aws/example_eks_with_nodegroup_in_one_step.py +++ b/tests/system/providers/amazon/aws/example_eks_with_nodegroup_in_one_step.py @@ -34,7 +34,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_eks_with_nodegroup_in_one_step" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_eks_with_nodegroups.py b/tests/system/providers/amazon/aws/example_eks_with_nodegroups.py index e6ba0658807e8..4828d58b4ff66 100644 --- a/tests/system/providers/amazon/aws/example_eks_with_nodegroups.py +++ b/tests/system/providers/amazon/aws/example_eks_with_nodegroups.py @@ -38,7 +38,7 @@ # Ignore missing args provided by default_args # type: ignore[call-arg] -importorskip("aiobotocore") + DAG_ID = "example_eks_with_nodegroups" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_emr.py b/tests/system/providers/amazon/aws/example_emr.py index 4f39488b8b1de..792b2b9742637 100644 --- a/tests/system/providers/amazon/aws/example_emr.py +++ b/tests/system/providers/amazon/aws/example_emr.py @@ -22,7 +22,6 @@ from datetime import datetime import boto3 -from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -38,7 +37,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_emr" CONFIG_NAME = "EMR Runtime Role Security Configuration" EXECUTION_ROLE_ARN_KEY = "EXECUTION_ROLE_ARN" diff --git a/tests/system/providers/amazon/aws/example_emr_eks.py b/tests/system/providers/amazon/aws/example_emr_eks.py index e0979f5750794..16d1f3bc2475a 100644 --- a/tests/system/providers/amazon/aws/example_emr_eks.py +++ b/tests/system/providers/amazon/aws/example_emr_eks.py @@ -21,7 +21,6 @@ from datetime import datetime import boto3 -from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -39,7 +38,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_emr_eks" # Externally fetched variables diff --git a/tests/system/providers/amazon/aws/example_emr_notebook_execution.py b/tests/system/providers/amazon/aws/example_emr_notebook_execution.py index 2e4f865c674f7..e24d465832ee1 100644 --- a/tests/system/providers/amazon/aws/example_emr_notebook_execution.py +++ b/tests/system/providers/amazon/aws/example_emr_notebook_execution.py @@ -20,8 +20,6 @@ from datetime import datetime -from pytest import importorskip - from airflow import DAG from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.emr import ( @@ -31,7 +29,6 @@ from airflow.providers.amazon.aws.sensors.emr import EmrNotebookExecutionSensor from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_emr_notebook" # Externally fetched variables: EDITOR_ID_KEY = "EDITOR_ID" diff --git a/tests/system/providers/amazon/aws/example_emr_serverless.py b/tests/system/providers/amazon/aws/example_emr_serverless.py index de6db77eef8ed..6d8a669c3e97e 100644 --- a/tests/system/providers/amazon/aws/example_emr_serverless.py +++ b/tests/system/providers/amazon/aws/example_emr_serverless.py @@ -19,7 +19,6 @@ from datetime import datetime import boto3 -from pytest import importorskip from airflow.models.baseoperator import chain from airflow.models.dag import DAG @@ -33,7 +32,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_emr_serverless" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_ftp_to_s3.py b/tests/system/providers/amazon/aws/example_ftp_to_s3.py index 07c2e63e7a9d2..ca2e6eb8e5b3b 100644 --- a/tests/system/providers/amazon/aws/example_ftp_to_s3.py +++ b/tests/system/providers/amazon/aws/example_ftp_to_s3.py @@ -18,8 +18,6 @@ from datetime import datetime -from pytest import importorskip - from airflow import models from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator @@ -28,7 +26,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() -importorskip("aiobotocore") + DAG_ID = "example_ftp_to_s3" with models.DAG( diff --git a/tests/system/providers/amazon/aws/example_gcs_to_s3.py b/tests/system/providers/amazon/aws/example_gcs_to_s3.py index 7cb5ce9715c9e..c0182f2d099ff 100644 --- a/tests/system/providers/amazon/aws/example_gcs_to_s3.py +++ b/tests/system/providers/amazon/aws/example_gcs_to_s3.py @@ -18,8 +18,6 @@ from datetime import datetime -from pytest import importorskip - from airflow import DAG from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator @@ -28,7 +26,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() -importorskip("aiobotocore") + DAG_ID = "example_gcs_to_s3" with DAG( diff --git a/tests/system/providers/amazon/aws/example_glacier_to_gcs.py b/tests/system/providers/amazon/aws/example_glacier_to_gcs.py index c87f7d30d4a25..43e14907ba799 100644 --- a/tests/system/providers/amazon/aws/example_glacier_to_gcs.py +++ b/tests/system/providers/amazon/aws/example_glacier_to_gcs.py @@ -19,7 +19,6 @@ from datetime import datetime import boto3 -from pytest import importorskip from airflow import DAG from airflow.models.baseoperator import chain @@ -34,7 +33,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() -importorskip("aiobotocore") + DAG_ID = "example_glacier_to_gcs" diff --git a/tests/system/providers/amazon/aws/example_glue.py b/tests/system/providers/amazon/aws/example_glue.py index cf2b9527c1d8a..f010b2dfd9fe0 100644 --- a/tests/system/providers/amazon/aws/example_glue.py +++ b/tests/system/providers/amazon/aws/example_glue.py @@ -20,7 +20,6 @@ import boto3 from botocore.client import BaseClient -from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -37,7 +36,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder, prune_logs -importorskip("aiobotocore") DAG_ID = "example_glue" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_google_api_sheets_to_s3.py b/tests/system/providers/amazon/aws/example_google_api_sheets_to_s3.py index 47a10dcffda55..926ab55f8f0c3 100644 --- a/tests/system/providers/amazon/aws/example_google_api_sheets_to_s3.py +++ b/tests/system/providers/amazon/aws/example_google_api_sheets_to_s3.py @@ -23,8 +23,6 @@ from datetime import datetime from os import getenv -from pytest import importorskip - from airflow import DAG from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator @@ -33,7 +31,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() -importorskip("aiobotocore") + DAG_ID = "example_google_api_sheets_to_s3" GOOGLE_SHEET_ID = getenv("GOOGLE_SHEET_ID", "test-google-sheet-id") diff --git a/tests/system/providers/amazon/aws/example_google_api_youtube_to_s3.py b/tests/system/providers/amazon/aws/example_google_api_youtube_to_s3.py index ebce3fa6bce45..64646b910468a 100644 --- a/tests/system/providers/amazon/aws/example_google_api_youtube_to_s3.py +++ b/tests/system/providers/amazon/aws/example_google_api_youtube_to_s3.py @@ -51,7 +51,6 @@ from datetime import datetime import boto3 -from pytest import importorskip from airflow import DAG, settings from airflow.decorators import task @@ -62,7 +61,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_google_api_youtube_to_s3" YOUTUBE_CHANNEL_ID = "UCSXwxpWZQ7XZ1WL3wqevChA" diff --git a/tests/system/providers/amazon/aws/example_hive_to_dynamodb.py b/tests/system/providers/amazon/aws/example_hive_to_dynamodb.py index c6abf2d771ae5..5a9b62c8663f9 100644 --- a/tests/system/providers/amazon/aws/example_hive_to_dynamodb.py +++ b/tests/system/providers/amazon/aws/example_hive_to_dynamodb.py @@ -23,8 +23,6 @@ from datetime import datetime -from pytest import importorskip - from airflow import DAG from airflow.decorators import task from airflow.models import Connection @@ -35,7 +33,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_hive_to_dynamodb" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_imap_attachment_to_s3.py b/tests/system/providers/amazon/aws/example_imap_attachment_to_s3.py index ad4b20709c2ee..7e70fd05e0fda 100644 --- a/tests/system/providers/amazon/aws/example_imap_attachment_to_s3.py +++ b/tests/system/providers/amazon/aws/example_imap_attachment_to_s3.py @@ -22,8 +22,6 @@ from datetime import datetime -from pytest import importorskip - from airflow import DAG from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator @@ -31,7 +29,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_imap_attachment_to_s3" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_lambda.py b/tests/system/providers/amazon/aws/example_lambda.py index e587a01e268e0..b4951799c8b08 100644 --- a/tests/system/providers/amazon/aws/example_lambda.py +++ b/tests/system/providers/amazon/aws/example_lambda.py @@ -22,7 +22,6 @@ from datetime import datetime import boto3 -from pytest import importorskip from airflow import models from airflow.decorators import task @@ -35,7 +34,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder, prune_logs -importorskip("aiobotocore") DAG_ID = "example_lambda" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_local_to_s3.py b/tests/system/providers/amazon/aws/example_local_to_s3.py index e61496bbb6157..8082b1cf6fc91 100644 --- a/tests/system/providers/amazon/aws/example_local_to_s3.py +++ b/tests/system/providers/amazon/aws/example_local_to_s3.py @@ -19,8 +19,6 @@ import os from datetime import datetime -from pytest import importorskip - from airflow import DAG from airflow.decorators import task from airflow.models.baseoperator import chain @@ -31,7 +29,7 @@ sys_test_context_task = SystemTestContextBuilder().build() -importorskip("aiobotocore") + DAG_ID = "example_local_to_s3" TEMP_FILE_PATH = "/tmp/sample-txt.txt" SAMPLE_TEXT = "This is some sample text." diff --git a/tests/system/providers/amazon/aws/example_mongo_to_s3.py b/tests/system/providers/amazon/aws/example_mongo_to_s3.py index db484ca6a75ca..3a3e5103fb0b5 100644 --- a/tests/system/providers/amazon/aws/example_mongo_to_s3.py +++ b/tests/system/providers/amazon/aws/example_mongo_to_s3.py @@ -16,8 +16,6 @@ # under the License. from __future__ import annotations -from pytest import importorskip - from airflow import models from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator @@ -26,7 +24,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_mongo_to_s3" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_quicksight.py b/tests/system/providers/amazon/aws/example_quicksight.py index 1bf13cff710eb..bc3cbd2abd31e 100644 --- a/tests/system/providers/amazon/aws/example_quicksight.py +++ b/tests/system/providers/amazon/aws/example_quicksight.py @@ -20,7 +20,6 @@ from datetime import datetime import boto3 -from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -48,7 +47,7 @@ into this behavior, changing the template for the ingestion name or the ENV_ID and re-running the test should resolve the issue. """ -importorskip("aiobotocore") + DAG_ID = "example_quicksight" sys_test_context_task = SystemTestContextBuilder().build() diff --git a/tests/system/providers/amazon/aws/example_rds_event.py b/tests/system/providers/amazon/aws/example_rds_event.py index a17d9f39edccd..94c39137ba416 100644 --- a/tests/system/providers/amazon/aws/example_rds_event.py +++ b/tests/system/providers/amazon/aws/example_rds_event.py @@ -20,7 +20,6 @@ from datetime import datetime import boto3 -from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -34,7 +33,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_rds_event" sys_test_context_task = SystemTestContextBuilder().build() diff --git a/tests/system/providers/amazon/aws/example_rds_export.py b/tests/system/providers/amazon/aws/example_rds_export.py index 069e45de945e9..d103763952417 100644 --- a/tests/system/providers/amazon/aws/example_rds_export.py +++ b/tests/system/providers/amazon/aws/example_rds_export.py @@ -18,8 +18,6 @@ from datetime import datetime -from pytest import importorskip - from airflow import DAG from airflow.decorators import task from airflow.models.baseoperator import chain @@ -37,7 +35,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_rds_export" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_rds_instance.py b/tests/system/providers/amazon/aws/example_rds_instance.py index 016e2d6770d49..4e165aaec63e0 100644 --- a/tests/system/providers/amazon/aws/example_rds_instance.py +++ b/tests/system/providers/amazon/aws/example_rds_instance.py @@ -18,8 +18,6 @@ from datetime import datetime -from pytest import importorskip - from airflow import DAG from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.rds import ( @@ -33,7 +31,7 @@ from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() -importorskip("aiobotocore") + DAG_ID = "example_rds_instance" RDS_USERNAME = "database_username" diff --git a/tests/system/providers/amazon/aws/example_rds_snapshot.py b/tests/system/providers/amazon/aws/example_rds_snapshot.py index 3c00c4f761348..904deea541430 100644 --- a/tests/system/providers/amazon/aws/example_rds_snapshot.py +++ b/tests/system/providers/amazon/aws/example_rds_snapshot.py @@ -18,8 +18,6 @@ from datetime import datetime -from pytest import importorskip - from airflow import DAG from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.rds import ( @@ -33,7 +31,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_rds_snapshot" sys_test_context_task = SystemTestContextBuilder().build() diff --git a/tests/system/providers/amazon/aws/example_redshift.py b/tests/system/providers/amazon/aws/example_redshift.py index 2d7af63ad1854..7e3a5a92809c5 100644 --- a/tests/system/providers/amazon/aws/example_redshift.py +++ b/tests/system/providers/amazon/aws/example_redshift.py @@ -21,7 +21,6 @@ from datetime import datetime import boto3 -from pytest import importorskip from airflow import DAG, settings from airflow.decorators import task @@ -42,7 +41,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_redshift" DB_LOGIN = "adminuser" DB_PASS = "MyAmazonPassword1" diff --git a/tests/system/providers/amazon/aws/example_redshift_s3_transfers.py b/tests/system/providers/amazon/aws/example_redshift_s3_transfers.py index 53fe2c9a74b8a..2608e327b6a5d 100644 --- a/tests/system/providers/amazon/aws/example_redshift_s3_transfers.py +++ b/tests/system/providers/amazon/aws/example_redshift_s3_transfers.py @@ -19,7 +19,6 @@ from datetime import datetime import boto3 -from pytest import importorskip from airflow import DAG, settings from airflow.decorators import task @@ -43,7 +42,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_redshift_to_s3" DB_LOGIN = "adminuser" diff --git a/tests/system/providers/amazon/aws/example_s3.py b/tests/system/providers/amazon/aws/example_s3.py index f1abf624ea57a..a18e16b79ad61 100644 --- a/tests/system/providers/amazon/aws/example_s3.py +++ b/tests/system/providers/amazon/aws/example_s3.py @@ -18,8 +18,6 @@ from datetime import datetime -from pytest import importorskip - from airflow.models.baseoperator import chain from airflow.models.dag import DAG from airflow.operators.python import BranchPythonOperator @@ -40,7 +38,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_s3" sys_test_context_task = SystemTestContextBuilder().build() diff --git a/tests/system/providers/amazon/aws/example_s3_to_ftp.py b/tests/system/providers/amazon/aws/example_s3_to_ftp.py index a81db7d2bcf91..984b5e41306f1 100644 --- a/tests/system/providers/amazon/aws/example_s3_to_ftp.py +++ b/tests/system/providers/amazon/aws/example_s3_to_ftp.py @@ -18,8 +18,6 @@ from datetime import datetime -from pytest import importorskip - from airflow import models from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator @@ -28,7 +26,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() -importorskip("aiobotocore") + DAG_ID = "example_s3_to_ftp" with models.DAG( diff --git a/tests/system/providers/amazon/aws/example_s3_to_sftp.py b/tests/system/providers/amazon/aws/example_s3_to_sftp.py index 54b8239398571..1212b32e1607e 100644 --- a/tests/system/providers/amazon/aws/example_s3_to_sftp.py +++ b/tests/system/providers/amazon/aws/example_s3_to_sftp.py @@ -18,8 +18,6 @@ from datetime import datetime -from pytest import importorskip - from airflow import models from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator @@ -28,7 +26,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() -importorskip("aiobotocore") + DAG_ID = "example_s3_to_sftp" with models.DAG( diff --git a/tests/system/providers/amazon/aws/example_s3_to_sql.py b/tests/system/providers/amazon/aws/example_s3_to_sql.py index 58fad3bbc0aea..f1b1b97dc36e2 100644 --- a/tests/system/providers/amazon/aws/example_s3_to_sql.py +++ b/tests/system/providers/amazon/aws/example_s3_to_sql.py @@ -19,7 +19,6 @@ from datetime import datetime import boto3 -from pytest import importorskip from airflow import DAG, settings from airflow.decorators import task @@ -44,7 +43,7 @@ from tests.system.utils.watcher import watcher sys_test_context_task = SystemTestContextBuilder().build() -importorskip("aiobotocore") + DAG_ID = "example_s3_to_sql" DB_LOGIN = "adminuser" diff --git a/tests/system/providers/amazon/aws/example_sagemaker.py b/tests/system/providers/amazon/aws/example_sagemaker.py index 900da98eaa1f8..9506970446320 100644 --- a/tests/system/providers/amazon/aws/example_sagemaker.py +++ b/tests/system/providers/amazon/aws/example_sagemaker.py @@ -23,7 +23,6 @@ from tempfile import NamedTemporaryFile import boto3 -from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -58,7 +57,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder, prune_logs -importorskip("aiobotocore") DAG_ID = "example_sagemaker" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_sagemaker_endpoint.py b/tests/system/providers/amazon/aws/example_sagemaker_endpoint.py index 8a85b3f803c5e..b4f9ce2b9e847 100644 --- a/tests/system/providers/amazon/aws/example_sagemaker_endpoint.py +++ b/tests/system/providers/amazon/aws/example_sagemaker_endpoint.py @@ -20,7 +20,6 @@ from datetime import datetime import boto3 -from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -41,7 +40,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder, prune_logs -importorskip("aiobotocore") DAG_ID = "example_sagemaker_endpoint" # Externally fetched variables: diff --git a/tests/system/providers/amazon/aws/example_salesforce_to_s3.py b/tests/system/providers/amazon/aws/example_salesforce_to_s3.py index 63ce2c160af53..30c34461b2c4b 100644 --- a/tests/system/providers/amazon/aws/example_salesforce_to_s3.py +++ b/tests/system/providers/amazon/aws/example_salesforce_to_s3.py @@ -22,8 +22,6 @@ from datetime import datetime -from pytest import importorskip - from airflow import DAG from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator @@ -32,7 +30,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() -importorskip("aiobotocore") + DAG_ID = "example_salesforce_to_s3" with DAG( diff --git a/tests/system/providers/amazon/aws/example_sftp_to_s3.py b/tests/system/providers/amazon/aws/example_sftp_to_s3.py index 1b2e9d4be50f2..de07811af0f66 100644 --- a/tests/system/providers/amazon/aws/example_sftp_to_s3.py +++ b/tests/system/providers/amazon/aws/example_sftp_to_s3.py @@ -18,8 +18,6 @@ from datetime import datetime -from pytest import importorskip - from airflow import models from airflow.models.baseoperator import chain from airflow.providers.amazon.aws.operators.s3 import S3CreateBucketOperator, S3DeleteBucketOperator @@ -28,7 +26,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() -importorskip("aiobotocore") + DAG_ID = "example_sftp_to_s3" with models.DAG( diff --git a/tests/system/providers/amazon/aws/example_sns.py b/tests/system/providers/amazon/aws/example_sns.py index d414f53a8a96e..41431915b43e6 100644 --- a/tests/system/providers/amazon/aws/example_sns.py +++ b/tests/system/providers/amazon/aws/example_sns.py @@ -19,7 +19,6 @@ from datetime import datetime import boto3 -from pytest import importorskip from airflow import DAG from airflow.decorators import task @@ -29,7 +28,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() -importorskip("aiobotocore") + DAG_ID = "example_sns" diff --git a/tests/system/providers/amazon/aws/example_sql_to_s3.py b/tests/system/providers/amazon/aws/example_sql_to_s3.py index 149f12eddedb2..f983881a3f707 100644 --- a/tests/system/providers/amazon/aws/example_sql_to_s3.py +++ b/tests/system/providers/amazon/aws/example_sql_to_s3.py @@ -21,7 +21,6 @@ from datetime import datetime import boto3 -from pytest import importorskip from airflow import DAG, settings from airflow.decorators import task @@ -39,7 +38,6 @@ from airflow.utils.trigger_rule import TriggerRule from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_sql_to_s3" DB_LOGIN = "adminuser" DB_PASS = "MyAmazonPassword1" diff --git a/tests/system/providers/amazon/aws/example_sqs.py b/tests/system/providers/amazon/aws/example_sqs.py index b5945b5c821cc..55e242d71559e 100644 --- a/tests/system/providers/amazon/aws/example_sqs.py +++ b/tests/system/providers/amazon/aws/example_sqs.py @@ -18,8 +18,6 @@ from datetime import datetime -from pytest import importorskip - from airflow import DAG from airflow.decorators import task from airflow.models.baseoperator import chain @@ -30,7 +28,7 @@ from tests.system.providers.amazon.aws.utils import SystemTestContextBuilder sys_test_context_task = SystemTestContextBuilder().build() -importorskip("aiobotocore") + DAG_ID = "example_sqs" diff --git a/tests/system/providers/amazon/aws/example_step_functions.py b/tests/system/providers/amazon/aws/example_step_functions.py index ecb050374fde8..b33a25e48b1aa 100644 --- a/tests/system/providers/amazon/aws/example_step_functions.py +++ b/tests/system/providers/amazon/aws/example_step_functions.py @@ -19,8 +19,6 @@ import json from datetime import datetime -from pytest import importorskip - from airflow import DAG from airflow.decorators import task from airflow.models.baseoperator import chain @@ -32,7 +30,6 @@ from airflow.providers.amazon.aws.sensors.step_function import StepFunctionExecutionSensor from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder -importorskip("aiobotocore") DAG_ID = "example_step_functions" # Externally fetched variables: From 03d6cbc9e49758b14a19d575149e980079c63c50 Mon Sep 17 00:00:00 2001 From: Syed Hussain Date: Tue, 4 Apr 2023 18:03:25 -0700 Subject: [PATCH 08/17] mock isinstance call to allow tests to pass Add unit tests for async client --- .../providers/amazon/aws/hooks/base_aws.py | 37 ++++++-- .../amazon/aws/hooks/test_base_aws.py | 88 ++++++++++++++++--- 2 files changed, 105 insertions(+), 20 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/base_aws.py b/airflow/providers/amazon/aws/hooks/base_aws.py index 96e0528407133..3814ea668a8e4 100644 --- a/airflow/providers/amazon/aws/hooks/base_aws.py +++ b/airflow/providers/amazon/aws/hooks/base_aws.py @@ -126,6 +126,24 @@ def role_arn(self) -> str | None: """Assume Role ARN from AWS Connection""" return self.conn.role_arn + def _apply_session_kwargs(self, session): + if self.conn.session_kwargs.get("profile_name", None) is not None: + session.set_config_variable("profile", self.conn.session_kwargs["profile_name"]) + + if ( + self.conn.session_kwargs.get("aws_access_key_id", None) + or self.conn.session_kwargs.get("aws_secret_access_key", None) + or self.conn.session_kwargs.get("aws_session_token", None) + ): + session.set_credentials( + self.conn.session_kwargs["aws_access_key_id"], + self.conn.session_kwargs["aws_secret_access_key"], + self.conn.session_kwargs["aws_session_token"], + ) + + if self.conn.session_kwargs.get("region_name", None) is not None: + session.set_config_variable("region", self.conn.session_kwargs["region_name"]) + def get_async_session(self): from aiobotocore.session import get_session as async_get_session @@ -139,14 +157,19 @@ def create_session(self, deferrable: bool = False) -> boto3.session.Session: "See: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html", self.region_name, ) - return ( - self.get_async_session() - if deferrable - else boto3.session.Session(region_name=self.region_name) - ) - + if deferrable: + session = self.get_async_session() + self._apply_session_kwargs(session) + return session + else: + return boto3.session.Session(region_name=self.region_name) elif not self.role_arn: - return self.get_async_session() if deferrable else self.basic_session + if deferrable: + session = self.get_async_session() + self._apply_session_kwargs(session) + return session + else: + return self.basic_session # Values stored in ``AwsConnectionWrapper.session_kwargs`` are intended to be used only # to create the initial boto3 session. diff --git a/tests/providers/amazon/aws/hooks/test_base_aws.py b/tests/providers/amazon/aws/hooks/test_base_aws.py index 6540e7b69b506..bfd7f2856620b 100644 --- a/tests/providers/amazon/aws/hooks/test_base_aws.py +++ b/tests/providers/amazon/aws/hooks/test_base_aws.py @@ -227,20 +227,39 @@ def test_create_session_from_credentials(self, mock_boto3_session, region_name, mock_boto3_session.assert_called_once_with(**expected_arguments) assert session == MOCK_BOTO3_SESSION + @pytest.mark.parametrize("region_name", ["eu-central-1", None]) + @pytest.mark.parametrize("profile_name", ["default", None]) + def test_async_create_session_from_credentials(self, region_name, profile_name): + mock_conn = Connection( + conn_type=MOCK_CONN_TYPE, conn_id=MOCK_AWS_CONN_ID, extra={"profile_name": profile_name} + ) + mock_conn_config = AwsConnectionWrapper(conn=mock_conn) + sf = BaseSessionFactory(conn=mock_conn_config, region_name=region_name, config=None) + async_session = sf.create_session(deferrable=True) + if region_name: + session_region = async_session.get_config_variable("region") + assert session_region == region_name + + session_profile = async_session.get_config_variable("profile") + + assert session_profile == profile_name + + config_for_credentials_test = [ + ( + "assume-with-initial-creds", + { + "aws_access_key_id": "mock_aws_access_key_id", + "aws_secret_access_key": "mock_aws_access_key_id", + "aws_session_token": "mock_aws_session_token", + }, + ), + ("assume-without-initial-creds", {}), + ] + @mock_sts @pytest.mark.parametrize( "conn_id, conn_extra", - [ - ( - "assume-with-initial-creds", - { - "aws_access_key_id": "mock_aws_access_key_id", - "aws_secret_access_key": "mock_aws_access_key_id", - "aws_session_token": "mock_aws_session_token", - }, - ), - ("assume-without-initial-creds", {}), - ], + config_for_credentials_test, ) @pytest.mark.parametrize("region_name", ["ap-southeast-2", "sa-east-1"]) def test_get_credentials_from_role_arn(self, conn_id, conn_extra, region_name): @@ -258,6 +277,39 @@ def test_get_credentials_from_role_arn(self, conn_id, conn_extra, region_name): # It shouldn't be 'explicit' which refers in this case to initial credentials. assert session.get_credentials().method == "sts-assume-role" + @pytest.mark.asyncio + @pytest.mark.parametrize( + "conn_id, conn_extra", + config_for_credentials_test, + ) + @pytest.mark.parametrize("region_name", ["ap-southeast-2", "sa-east-1"]) + @mock.patch("airflow.providers.amazon.aws.hooks.base_aws.BaseSessionFactory._refresh_credentials") + async def test_async_get_credentials_from_role_arn(self, mock_refresh, conn_id, conn_extra, region_name): + """Test RefreshableCredentials with assume_role for async_conn""" + + def side_effect(): + return { + "access_key": "mock-AccessKeyId", + "secret_key": "mock-SecretAccessKey", + "token": "mock-SessionToken", + "expiry_time": datetime.now(timezone.utc).isoformat(), + } + + mock_refresh.side_effect = side_effect + extra = { + **conn_extra, + "role_arn": "arn:aws:iam::123456:role/role_arn", + "region_name": region_name, + } + conn = AwsConnectionWrapper.from_connection_metadata(conn_id=conn_id, extra=extra) + sf = BaseSessionFactory(conn=conn) + session = sf.create_session(deferrable=True) + assert session.region_name == region_name + # Validate method of botocore credentials provider. + # It shouldn't be 'explicit' which refers in this case to initial credentials. + credentials = await session.get_credentials() + assert credentials.method == "sts-assume-role" + class TestAwsBaseHook: @mock_emr @@ -394,7 +446,12 @@ def mock_assume_role(**kwargs): with mock.patch( "airflow.providers.amazon.aws.hooks.base_aws.requests.Session.get" - ) as mock_get, mock.patch("airflow.providers.amazon.aws.hooks.base_aws.boto3") as mock_boto3: + ) as mock_get, mock.patch( + "airflow.providers.amazon.aws.hooks.base_aws.boto3" + ) as mock_boto3, mock.patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = False mock_get.return_value.ok = True mock_client = mock_boto3.session.Session.return_value.client @@ -589,7 +646,12 @@ def mock_assume_role_with_saml(**kwargs): with mock.patch("builtins.__import__", side_effect=import_mock), mock.patch( "airflow.providers.amazon.aws.hooks.base_aws.requests.Session.get" - ) as mock_get, mock.patch("airflow.providers.amazon.aws.hooks.base_aws.boto3") as mock_boto3: + ) as mock_get, mock.patch( + "airflow.providers.amazon.aws.hooks.base_aws.boto3" + ) as mock_boto3, mock.patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = False mock_get.return_value.ok = True mock_client = mock_boto3.session.Session.return_value.client From 579e59115e5fb0f87beadb0e00fb129337a5fa51 Mon Sep 17 00:00:00 2001 From: Syed Hussain Date: Wed, 5 Apr 2023 05:52:05 -0700 Subject: [PATCH 09/17] Add mocking to isinstance where it is needed --- .../amazon/aws/hooks/test_base_aws.py | 4 ++-- .../amazon/aws/hooks/test_emr_containers.py | 12 +++++++---- .../aws/operators/test_cloud_formation.py | 10 ++++++++-- .../aws/operators/test_emr_add_steps.py | 15 +++++++++++--- .../aws/operators/test_emr_containers.py | 10 ++++++++-- .../aws/operators/test_emr_create_job_flow.py | 15 +++++++++++--- .../aws/operators/test_emr_modify_cluster.py | 10 ++++++++-- .../operators/test_emr_terminate_job_flow.py | 5 ++++- .../aws/sensors/test_cloud_formation.py | 20 +++++++++++++++---- .../amazon/aws/sensors/test_emr_job_flow.py | 15 +++++++++++--- .../amazon/aws/sensors/test_emr_step.py | 20 +++++++++++++++---- 11 files changed, 106 insertions(+), 30 deletions(-) diff --git a/tests/providers/amazon/aws/hooks/test_base_aws.py b/tests/providers/amazon/aws/hooks/test_base_aws.py index bfd7f2856620b..9cb67a4cc013a 100644 --- a/tests/providers/amazon/aws/hooks/test_base_aws.py +++ b/tests/providers/amazon/aws/hooks/test_base_aws.py @@ -451,7 +451,7 @@ def mock_assume_role(**kwargs): ) as mock_boto3, mock.patch( "airflow.providers.amazon.aws.hooks.base_aws.isinstance" ) as mock_isinstance: - mock_isinstance.return_value = False + mock_isinstance.return_value = True mock_get.return_value.ok = True mock_client = mock_boto3.session.Session.return_value.client @@ -651,7 +651,7 @@ def mock_assume_role_with_saml(**kwargs): ) as mock_boto3, mock.patch( "airflow.providers.amazon.aws.hooks.base_aws.isinstance" ) as mock_isinstance: - mock_isinstance.return_value = False + mock_isinstance.return_value = True mock_get.return_value.ok = True mock_client = mock_boto3.session.Session.return_value.client diff --git a/tests/providers/amazon/aws/hooks/test_emr_containers.py b/tests/providers/amazon/aws/hooks/test_emr_containers.py index 8a5f1303a6921..9be48a08ba04a 100644 --- a/tests/providers/amazon/aws/hooks/test_emr_containers.py +++ b/tests/providers/amazon/aws/hooks/test_emr_containers.py @@ -54,8 +54,9 @@ def test_init(self): assert self.emr_containers.aws_conn_id == "aws_default" assert self.emr_containers.virtual_cluster_id == "vc1234" + @mock.patch("airflow.providers.amazon.aws.hooks.base_aws.isinstance", return_value=True) @mock.patch("boto3.session.Session") - def test_create_emr_on_eks_cluster(self, mock_session): + def test_create_emr_on_eks_cluster(self, mock_session, mock_isinstance): emr_client_mock = mock.MagicMock() emr_client_mock.create_virtual_cluster.return_value = CREATE_EMR_ON_EKS_CLUSTER_RETURN emr_session_mock = mock.MagicMock() @@ -69,8 +70,9 @@ def test_create_emr_on_eks_cluster(self, mock_session): ) assert emr_on_eks_create_cluster_response == "vc1234" + @mock.patch("airflow.providers.amazon.aws.hooks.base_aws.isinstance", return_value=True) @mock.patch("boto3.session.Session") - def test_submit_job(self, mock_session): + def test_submit_job(self, mock_session, mock_isinstance): # Mock out the emr_client creator emr_client_mock = mock.MagicMock() emr_client_mock.start_job_run.return_value = SUBMIT_JOB_SUCCESS_RETURN @@ -88,8 +90,9 @@ def test_submit_job(self, mock_session): ) assert emr_containers_job == "job123456" + @mock.patch("airflow.providers.amazon.aws.hooks.base_aws.isinstance", return_value=True) @mock.patch("boto3.session.Session") - def test_query_status_polling_when_terminal(self, mock_session): + def test_query_status_polling_when_terminal(self, mock_session, mock_isinstance): emr_client_mock = mock.MagicMock() emr_session_mock = mock.MagicMock() emr_session_mock.client.return_value = emr_client_mock @@ -101,8 +104,9 @@ def test_query_status_polling_when_terminal(self, mock_session): emr_client_mock.describe_job_run.assert_called_once() assert query_status == "COMPLETED" + @mock.patch("airflow.providers.amazon.aws.hooks.base_aws.isinstance", return_value=True) @mock.patch("boto3.session.Session") - def test_query_status_polling_with_timeout(self, mock_session): + def test_query_status_polling_with_timeout(self, mock_session, mock_isinstance): emr_client_mock = mock.MagicMock() emr_session_mock = mock.MagicMock() emr_session_mock.client.return_value = emr_client_mock diff --git a/tests/providers/amazon/aws/operators/test_cloud_formation.py b/tests/providers/amazon/aws/operators/test_cloud_formation.py index 2600096df2852..1a1088adab057 100644 --- a/tests/providers/amazon/aws/operators/test_cloud_formation.py +++ b/tests/providers/amazon/aws/operators/test_cloud_formation.py @@ -55,7 +55,10 @@ def test_create_stack(self): dag=DAG("test_dag_id", default_args=DEFAULT_ARGS), ) - with mock.patch("boto3.session.Session", self.boto3_session_mock): + with mock.patch("boto3.session.Session", self.boto3_session_mock), mock.patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True operator.execute(self.mock_context) self.cloudformation_client_mock.create_stack.assert_any_call( @@ -84,7 +87,10 @@ def test_delete_stack(self): dag=DAG("test_dag_id", default_args=DEFAULT_ARGS), ) - with mock.patch("boto3.session.Session", self.boto3_session_mock): + with mock.patch("boto3.session.Session", self.boto3_session_mock), mock.patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True operator.execute(self.mock_context) self.cloudformation_client_mock.delete_stack.assert_any_call(StackName=stack_name) diff --git a/tests/providers/amazon/aws/operators/test_emr_add_steps.py b/tests/providers/amazon/aws/operators/test_emr_add_steps.py index 6f9c1c1b45922..67a4090563067 100644 --- a/tests/providers/amazon/aws/operators/test_emr_add_steps.py +++ b/tests/providers/amazon/aws/operators/test_emr_add_steps.py @@ -151,7 +151,10 @@ def test_render_template_from_file(self): assert json.loads(test_task.steps) == file_steps # String in job_flow_overrides (i.e. from loaded as a file) is not "parsed" until inside execute() - with patch("boto3.session.Session", self.boto3_session_mock): + with patch("boto3.session.Session", self.boto3_session_mock), patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True test_task.execute(None) self.emr_client_mock.add_job_flow_steps.assert_called_once_with( @@ -161,7 +164,10 @@ def test_render_template_from_file(self): def test_execute_returns_step_id(self): self.emr_client_mock.add_job_flow_steps.return_value = ADD_STEPS_SUCCESS_RETURN - with patch("boto3.session.Session", self.boto3_session_mock): + with patch("boto3.session.Session", self.boto3_session_mock), patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True assert self.operator.execute(self.mock_context) == ["s-2LH3R5GW3A53T"] def test_init_with_cluster_name(self): @@ -169,7 +175,10 @@ def test_init_with_cluster_name(self): self.emr_client_mock.add_job_flow_steps.return_value = ADD_STEPS_SUCCESS_RETURN - with patch("boto3.session.Session", self.boto3_session_mock): + with patch("boto3.session.Session", self.boto3_session_mock), patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True with patch( "airflow.providers.amazon.aws.hooks.emr.EmrHook.get_cluster_id_by_name" ) as mock_get_cluster_id_by_name: diff --git a/tests/providers/amazon/aws/operators/test_emr_containers.py b/tests/providers/amazon/aws/operators/test_emr_containers.py index 7f0e2942f6c98..ddc11b15c56ce 100644 --- a/tests/providers/amazon/aws/operators/test_emr_containers.py +++ b/tests/providers/amazon/aws/operators/test_emr_containers.py @@ -85,7 +85,10 @@ def test_execute_with_polling(self, mock_check_query_status): emr_session_mock.client.return_value = emr_client_mock boto3_session_mock = MagicMock(return_value=emr_session_mock) - with patch("boto3.session.Session", boto3_session_mock): + with patch("boto3.session.Session", boto3_session_mock), patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True assert self.emr_container.execute(None) == "job123456" assert mock_check_query_status.call_count == 5 @@ -130,7 +133,10 @@ def test_execute_with_polling_timeout(self, mock_check_query_status): max_polling_attempts=3, ) - with patch("boto3.session.Session", boto3_session_mock): + with patch("boto3.session.Session", boto3_session_mock), patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True with pytest.raises(AirflowException) as ctx: timeout_container.execute(None) diff --git a/tests/providers/amazon/aws/operators/test_emr_create_job_flow.py b/tests/providers/amazon/aws/operators/test_emr_create_job_flow.py index 1920f915b88db..cac14ddf59cd1 100644 --- a/tests/providers/amazon/aws/operators/test_emr_create_job_flow.py +++ b/tests/providers/amazon/aws/operators/test_emr_create_job_flow.py @@ -129,7 +129,10 @@ def test_render_template_from_file(self): boto3_session_mock = MagicMock(return_value=emr_session_mock) # String in job_flow_overrides (i.e. from loaded as a file) is not "parsed" until inside execute() - with patch("boto3.session.Session", boto3_session_mock): + with patch("boto3.session.Session", boto3_session_mock), patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True self.operator.execute(self.mock_context) expected_args = { @@ -161,7 +164,10 @@ def test_execute_returns_job_id(self): emr_session_mock.client.return_value = self.emr_client_mock boto3_session_mock = MagicMock(return_value=emr_session_mock) - with patch("boto3.session.Session", boto3_session_mock): + with patch("boto3.session.Session", boto3_session_mock), patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True assert self.operator.execute(self.mock_context) == JOB_FLOW_ID @mock.patch("botocore.waiter.get_service_module_name", return_value="emr") @@ -175,7 +181,10 @@ def test_execute_with_wait(self, mock_waiter, _): boto3_session_mock = MagicMock(return_value=emr_session_mock) self.operator.wait_for_completion = True - with patch("boto3.session.Session", boto3_session_mock): + with patch("boto3.session.Session", boto3_session_mock), patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True assert self.operator.execute(self.mock_context) == JOB_FLOW_ID mock_waiter.assert_called_once_with(mock.ANY, ClusterId=JOB_FLOW_ID, WaiterConfig=mock.ANY) assert_expected_waiter_type(mock_waiter, "job_flow_waiting") diff --git a/tests/providers/amazon/aws/operators/test_emr_modify_cluster.py b/tests/providers/amazon/aws/operators/test_emr_modify_cluster.py index 12efbcf198326..6fc20ed430b1d 100644 --- a/tests/providers/amazon/aws/operators/test_emr_modify_cluster.py +++ b/tests/providers/amazon/aws/operators/test_emr_modify_cluster.py @@ -63,12 +63,18 @@ def test_init(self): def test_execute_returns_step_concurrency(self): self.emr_client_mock.modify_cluster.return_value = MODIFY_CLUSTER_SUCCESS_RETURN - with patch("boto3.session.Session", self.boto3_session_mock): + with patch("boto3.session.Session", self.boto3_session_mock), patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True assert self.operator.execute(self.mock_context) == 1 def test_execute_returns_error(self): self.emr_client_mock.modify_cluster.return_value = MODIFY_CLUSTER_ERROR_RETURN - with patch("boto3.session.Session", self.boto3_session_mock): + with patch("boto3.session.Session", self.boto3_session_mock), patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True with pytest.raises(AirflowException): self.operator.execute(self.mock_context) diff --git a/tests/providers/amazon/aws/operators/test_emr_terminate_job_flow.py b/tests/providers/amazon/aws/operators/test_emr_terminate_job_flow.py index 594c6d775b9f9..4a639cc154719 100644 --- a/tests/providers/amazon/aws/operators/test_emr_terminate_job_flow.py +++ b/tests/providers/amazon/aws/operators/test_emr_terminate_job_flow.py @@ -37,7 +37,10 @@ def setup_method(self): self.boto3_session_mock = MagicMock(return_value=mock_emr_session) def test_execute_terminates_the_job_flow_and_does_not_error(self): - with patch("boto3.session.Session", self.boto3_session_mock): + with patch("boto3.session.Session", self.boto3_session_mock), patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True operator = EmrTerminateJobFlowOperator( task_id="test_task", job_flow_id="j-8989898989", aws_conn_id="aws_default" ) diff --git a/tests/providers/amazon/aws/sensors/test_cloud_formation.py b/tests/providers/amazon/aws/sensors/test_cloud_formation.py index 14610df2671b5..9aef7fae6dd87 100644 --- a/tests/providers/amazon/aws/sensors/test_cloud_formation.py +++ b/tests/providers/amazon/aws/sensors/test_cloud_formation.py @@ -51,7 +51,10 @@ def test_poke(self): assert op.poke({}) def test_poke_false(self): - with patch("boto3.session.Session", self.boto3_session_mock): + with patch("boto3.session.Session", self.boto3_session_mock), patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True self.cloudformation_client_mock.describe_stacks.return_value = { "Stacks": [{"StackStatus": "CREATE_IN_PROGRESS"}] } @@ -59,7 +62,10 @@ def test_poke_false(self): assert not op.poke({}) def test_poke_stack_in_unsuccessful_state(self): - with patch("boto3.session.Session", self.boto3_session_mock): + with patch("boto3.session.Session", self.boto3_session_mock), patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True self.cloudformation_client_mock.describe_stacks.return_value = { "Stacks": [{"StackStatus": "bar"}] } @@ -91,7 +97,10 @@ def test_poke(self): assert op.poke({}) def test_poke_false(self): - with patch("boto3.session.Session", self.boto3_session_mock): + with patch("boto3.session.Session", self.boto3_session_mock), patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True self.cloudformation_client_mock.describe_stacks.return_value = { "Stacks": [{"StackStatus": "DELETE_IN_PROGRESS"}] } @@ -99,7 +108,10 @@ def test_poke_false(self): assert not op.poke({}) def test_poke_stack_in_unsuccessful_state(self): - with patch("boto3.session.Session", self.boto3_session_mock): + with patch("boto3.session.Session", self.boto3_session_mock), patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True self.cloudformation_client_mock.describe_stacks.return_value = { "Stacks": [{"StackStatus": "bar"}] } diff --git a/tests/providers/amazon/aws/sensors/test_emr_job_flow.py b/tests/providers/amazon/aws/sensors/test_emr_job_flow.py index a10e68abb8483..c81d9d4855747 100644 --- a/tests/providers/amazon/aws/sensors/test_emr_job_flow.py +++ b/tests/providers/amazon/aws/sensors/test_emr_job_flow.py @@ -208,7 +208,10 @@ def test_execute_calls_with_the_job_flow_id_until_it_reaches_a_target_state(self DESCRIBE_CLUSTER_RUNNING_RETURN, DESCRIBE_CLUSTER_TERMINATED_RETURN, ] - with patch("boto3.session.Session", self.boto3_session_mock): + with patch("boto3.session.Session", self.boto3_session_mock), patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True operator = EmrJobFlowSensor( task_id="test_task", poke_interval=0, job_flow_id="j-8989898989", aws_conn_id="aws_default" ) @@ -227,7 +230,10 @@ def test_execute_calls_with_the_job_flow_id_until_it_reaches_failed_state_with_e DESCRIBE_CLUSTER_RUNNING_RETURN, DESCRIBE_CLUSTER_TERMINATED_WITH_ERRORS_RETURN, ] - with patch("boto3.session.Session", self.boto3_session_mock): + with patch("boto3.session.Session", self.boto3_session_mock), patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True operator = EmrJobFlowSensor( task_id="test_task", poke_interval=0, job_flow_id="j-8989898989", aws_conn_id="aws_default" ) @@ -250,7 +256,10 @@ def test_different_target_states(self): DESCRIBE_CLUSTER_TERMINATED_RETURN, # will not be used DESCRIBE_CLUSTER_TERMINATED_WITH_ERRORS_RETURN, # will not be used ] - with patch("boto3.session.Session", self.boto3_session_mock): + with patch("boto3.session.Session", self.boto3_session_mock), patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True operator = EmrJobFlowSensor( task_id="test_task", poke_interval=0, diff --git a/tests/providers/amazon/aws/sensors/test_emr_step.py b/tests/providers/amazon/aws/sensors/test_emr_step.py index d053bda97c00f..cca5f417f83d0 100644 --- a/tests/providers/amazon/aws/sensors/test_emr_step.py +++ b/tests/providers/amazon/aws/sensors/test_emr_step.py @@ -165,7 +165,10 @@ def test_step_completed(self): DESCRIBE_JOB_STEP_COMPLETED_RETURN, ] - with patch("boto3.session.Session", self.boto3_session_mock): + with patch("boto3.session.Session", self.boto3_session_mock), patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True self.sensor.execute(None) assert self.emr_client_mock.describe_step.call_count == 2 @@ -181,7 +184,10 @@ def test_step_cancelled(self): DESCRIBE_JOB_STEP_CANCELLED_RETURN, ] - with patch("boto3.session.Session", self.boto3_session_mock): + with patch("boto3.session.Session", self.boto3_session_mock), patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True with pytest.raises(AirflowException): self.sensor.execute(None) @@ -191,7 +197,10 @@ def test_step_failed(self): DESCRIBE_JOB_STEP_FAILED_RETURN, ] - with patch("boto3.session.Session", self.boto3_session_mock): + with patch("boto3.session.Session", self.boto3_session_mock), patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True with pytest.raises(AirflowException): self.sensor.execute(None) @@ -201,6 +210,9 @@ def test_step_interrupted(self): DESCRIBE_JOB_STEP_INTERRUPTED_RETURN, ] - with patch("boto3.session.Session", self.boto3_session_mock): + with patch("boto3.session.Session", self.boto3_session_mock), patch( + "airflow.providers.amazon.aws.hooks.base_aws.isinstance" + ) as mock_isinstance: + mock_isinstance.return_value = True with pytest.raises(AirflowException): self.sensor.execute(None) From 53d99fd41072d4d46873ba72936a80e31fa941c2 Mon Sep 17 00:00:00 2001 From: Syed Hussain Date: Wed, 12 Apr 2023 14:56:26 -0700 Subject: [PATCH 10/17] Mock with context to work with python 3.7 --- .../amazon/aws/hooks/test_base_aws.py | 52 ++++++++++--------- 1 file changed, 27 insertions(+), 25 deletions(-) diff --git a/tests/providers/amazon/aws/hooks/test_base_aws.py b/tests/providers/amazon/aws/hooks/test_base_aws.py index 9cb67a4cc013a..b6c39e1ec9f0b 100644 --- a/tests/providers/amazon/aws/hooks/test_base_aws.py +++ b/tests/providers/amazon/aws/hooks/test_base_aws.py @@ -283,32 +283,34 @@ def test_get_credentials_from_role_arn(self, conn_id, conn_extra, region_name): config_for_credentials_test, ) @pytest.mark.parametrize("region_name", ["ap-southeast-2", "sa-east-1"]) - @mock.patch("airflow.providers.amazon.aws.hooks.base_aws.BaseSessionFactory._refresh_credentials") - async def test_async_get_credentials_from_role_arn(self, mock_refresh, conn_id, conn_extra, region_name): - """Test RefreshableCredentials with assume_role for async_conn""" - - def side_effect(): - return { - "access_key": "mock-AccessKeyId", - "secret_key": "mock-SecretAccessKey", - "token": "mock-SessionToken", - "expiry_time": datetime.now(timezone.utc).isoformat(), - } + async def test_async_get_credentials_from_role_arn(self, conn_id, conn_extra, region_name): + """Test RefreshableCredentials with assume_role for async_conn.""" + with mock.patch( + "airflow.providers.amazon.aws.hooks.base_aws.BaseSessionFactory._refresh_credentials" + ) as mock_refresh: - mock_refresh.side_effect = side_effect - extra = { - **conn_extra, - "role_arn": "arn:aws:iam::123456:role/role_arn", - "region_name": region_name, - } - conn = AwsConnectionWrapper.from_connection_metadata(conn_id=conn_id, extra=extra) - sf = BaseSessionFactory(conn=conn) - session = sf.create_session(deferrable=True) - assert session.region_name == region_name - # Validate method of botocore credentials provider. - # It shouldn't be 'explicit' which refers in this case to initial credentials. - credentials = await session.get_credentials() - assert credentials.method == "sts-assume-role" + def side_effect(): + return { + "access_key": "mock-AccessKeyId", + "secret_key": "mock-SecretAccessKey", + "token": "mock-SessionToken", + "expiry_time": datetime.now(timezone.utc).isoformat(), + } + + mock_refresh.side_effect = side_effect + extra = { + **conn_extra, + "role_arn": "arn:aws:iam::123456:role/role_arn", + "region_name": region_name, + } + conn = AwsConnectionWrapper.from_connection_metadata(conn_id=conn_id, extra=extra) + sf = BaseSessionFactory(conn=conn) + session = sf.create_session(deferrable=True) + assert session.region_name == region_name + # Validate method of botocore credentials provider. + # It shouldn't be 'explicit' which refers in this case to initial credentials. + credentials = await session.get_credentials() + assert credentials.method == "sts-assume-role" class TestAwsBaseHook: From 441a6c7f9bf56a553d2cd0c4114e502d97778fe1 Mon Sep 17 00:00:00 2001 From: Syed Hussain Date: Thu, 13 Apr 2023 11:28:55 -0700 Subject: [PATCH 11/17] add importorskip to base_aws.py --- tests/providers/amazon/aws/hooks/test_base_aws.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/providers/amazon/aws/hooks/test_base_aws.py b/tests/providers/amazon/aws/hooks/test_base_aws.py index b6c39e1ec9f0b..daf2c5783024e 100644 --- a/tests/providers/amazon/aws/hooks/test_base_aws.py +++ b/tests/providers/amazon/aws/hooks/test_base_aws.py @@ -47,6 +47,8 @@ from airflow.providers.amazon.aws.utils.connection_wrapper import AwsConnectionWrapper from tests.test_utils.config import conf_vars +pytest.importorskip("aiobotocore") + MOCK_AWS_CONN_ID = "mock-conn-id" MOCK_CONN_TYPE = "aws" MOCK_BOTO3_SESSION = mock.MagicMock(return_value="Mock boto3.session.Session") From 49f376a1f681eb507494ace8a652ef34a41db4d1 Mon Sep 17 00:00:00 2001 From: Syed Hussain Date: Thu, 13 Apr 2023 18:13:32 -0700 Subject: [PATCH 12/17] Add note about async connections in base_aws. Remove typo in system test --- airflow/providers/amazon/aws/hooks/base_aws.py | 3 +++ .../amazon/aws/example_eks_with_nodegroup_in_one_step.py | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/airflow/providers/amazon/aws/hooks/base_aws.py b/airflow/providers/amazon/aws/hooks/base_aws.py index 3814ea668a8e4..4a258e7aa4316 100644 --- a/airflow/providers/amazon/aws/hooks/base_aws.py +++ b/airflow/providers/amazon/aws/hooks/base_aws.py @@ -77,6 +77,9 @@ class BaseSessionFactory(LoggingMixin): User can also derive from this class to have full control of boto3 session creation or to support custom federation. + Note: Not all features implemented for synchronous sessions are available for async + sessions. + .. seealso:: - :ref:`howto/connection:aws:session-factory` """ diff --git a/tests/system/providers/amazon/aws/example_eks_with_nodegroup_in_one_step.py b/tests/system/providers/amazon/aws/example_eks_with_nodegroup_in_one_step.py index 00c6dc1286476..3dd2b90649389 100644 --- a/tests/system/providers/amazon/aws/example_eks_with_nodegroup_in_one_step.py +++ b/tests/system/providers/amazon/aws/example_eks_with_nodegroup_in_one_step.py @@ -20,7 +20,7 @@ import boto3 -from airflow.adecorators import task +from airflow.decorators import task from airflow.models.baseoperator import chain from airflow.models.dag import DAG from airflow.operators.bash import BashOperator From c292de392e4708a9fcfb1670fceb174625ab07e9 Mon Sep 17 00:00:00 2001 From: Syed Hussain Date: Mon, 17 Apr 2023 12:28:54 -0700 Subject: [PATCH 13/17] Remove cached_property on async_conn to prevent awaiting a coro twice --- airflow/providers/amazon/aws/hooks/base_aws.py | 1 - airflow/providers/amazon/aws/triggers/redshift_cluster.py | 2 +- tests/providers/amazon/aws/triggers/test_redshift_cluster.py | 2 +- 3 files changed, 2 insertions(+), 3 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/base_aws.py b/airflow/providers/amazon/aws/hooks/base_aws.py index 4a258e7aa4316..e9013bb181dcf 100644 --- a/airflow/providers/amazon/aws/hooks/base_aws.py +++ b/airflow/providers/amazon/aws/hooks/base_aws.py @@ -684,7 +684,6 @@ def conn(self) -> BaseAwsConnection: else: return self.get_resource_type(region_name=self.region_name) - @cached_property def async_conn(self): """Get an aiobotocore client to use for async operations (cached).""" if not self.client_type: diff --git a/airflow/providers/amazon/aws/triggers/redshift_cluster.py b/airflow/providers/amazon/aws/triggers/redshift_cluster.py index 2f831fa14c2f1..516a9cab44254 100644 --- a/airflow/providers/amazon/aws/triggers/redshift_cluster.py +++ b/airflow/providers/amazon/aws/triggers/redshift_cluster.py @@ -128,7 +128,7 @@ def hook(self) -> RedshiftHook: return RedshiftHook(aws_conn_id=self.aws_conn_id) async def run(self): - async with self.hook.async_conn as client: + async with self.hook.async_conn() as client: await client.get_waiter("cluster_available").wait( ClusterIdentifier=self.cluster_identifier, WaiterConfig={ diff --git a/tests/providers/amazon/aws/triggers/test_redshift_cluster.py b/tests/providers/amazon/aws/triggers/test_redshift_cluster.py index 941258659e9ae..c972c5efad3a1 100644 --- a/tests/providers/amazon/aws/triggers/test_redshift_cluster.py +++ b/tests/providers/amazon/aws/triggers/test_redshift_cluster.py @@ -58,7 +58,7 @@ def test_redshift_create_cluster_trigger_serialize(self): @async_mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.async_conn") async def test_redshift_create_cluster_trigger_run(self, mock_async_conn): mock = async_mock.MagicMock() - mock_async_conn.__aenter__.return_value = mock + mock_async_conn().__aenter__.return_value = mock mock.get_waiter().wait = AsyncMock() redshift_create_cluster_trigger = RedshiftCreateClusterTrigger( From 8e83f4c435658f8e33c255ea002c35176c6775a2 Mon Sep 17 00:00:00 2001 From: Syed Hussain Date: Mon, 17 Apr 2023 16:43:40 -0700 Subject: [PATCH 14/17] Update doc string --- airflow/providers/amazon/aws/hooks/base_aws.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/amazon/aws/hooks/base_aws.py b/airflow/providers/amazon/aws/hooks/base_aws.py index 1f7244d344c72..4e6ea0e02bcde 100644 --- a/airflow/providers/amazon/aws/hooks/base_aws.py +++ b/airflow/providers/amazon/aws/hooks/base_aws.py @@ -685,7 +685,7 @@ def conn(self) -> BaseAwsConnection: return self.get_resource_type(region_name=self.region_name) def async_conn(self): - """Get an aiobotocore client to use for async operations (cached).""" + """Get an aiobotocore client to use for async operations.""" if not self.client_type: raise ValueError("client_type must be specified.") From 93dbf17b19bbdb54c0bae7561a00cfea73caf491 Mon Sep 17 00:00:00 2001 From: Syed Hussain Date: Wed, 19 Apr 2023 05:42:19 -0700 Subject: [PATCH 15/17] Make async_conn a property --- airflow/providers/amazon/aws/hooks/base_aws.py | 1 + airflow/providers/amazon/aws/triggers/redshift_cluster.py | 2 +- tests/providers/amazon/aws/triggers/test_redshift_cluster.py | 2 +- 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/airflow/providers/amazon/aws/hooks/base_aws.py b/airflow/providers/amazon/aws/hooks/base_aws.py index 4e6ea0e02bcde..363d365637323 100644 --- a/airflow/providers/amazon/aws/hooks/base_aws.py +++ b/airflow/providers/amazon/aws/hooks/base_aws.py @@ -684,6 +684,7 @@ def conn(self) -> BaseAwsConnection: else: return self.get_resource_type(region_name=self.region_name) + @property def async_conn(self): """Get an aiobotocore client to use for async operations.""" if not self.client_type: diff --git a/airflow/providers/amazon/aws/triggers/redshift_cluster.py b/airflow/providers/amazon/aws/triggers/redshift_cluster.py index 516a9cab44254..2f831fa14c2f1 100644 --- a/airflow/providers/amazon/aws/triggers/redshift_cluster.py +++ b/airflow/providers/amazon/aws/triggers/redshift_cluster.py @@ -128,7 +128,7 @@ def hook(self) -> RedshiftHook: return RedshiftHook(aws_conn_id=self.aws_conn_id) async def run(self): - async with self.hook.async_conn() as client: + async with self.hook.async_conn as client: await client.get_waiter("cluster_available").wait( ClusterIdentifier=self.cluster_identifier, WaiterConfig={ diff --git a/tests/providers/amazon/aws/triggers/test_redshift_cluster.py b/tests/providers/amazon/aws/triggers/test_redshift_cluster.py index c972c5efad3a1..941258659e9ae 100644 --- a/tests/providers/amazon/aws/triggers/test_redshift_cluster.py +++ b/tests/providers/amazon/aws/triggers/test_redshift_cluster.py @@ -58,7 +58,7 @@ def test_redshift_create_cluster_trigger_serialize(self): @async_mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.async_conn") async def test_redshift_create_cluster_trigger_run(self, mock_async_conn): mock = async_mock.MagicMock() - mock_async_conn().__aenter__.return_value = mock + mock_async_conn.__aenter__.return_value = mock mock.get_waiter().wait = AsyncMock() redshift_create_cluster_trigger = RedshiftCreateClusterTrigger( From bc2f87ee61fd477c3329433f90d2b77eef695fe8 Mon Sep 17 00:00:00 2001 From: Syed Hussain Date: Fri, 21 Apr 2023 00:01:58 -0700 Subject: [PATCH 16/17] Fixed formatting in trigger README.md. Add unit test for deferrable in test_redshift_cluster.py operator Add docstring for deferrable param --- .../amazon/aws/operators/redshift_cluster.py | 1 + airflow/providers/amazon/aws/triggers/README.md | 12 ++++++------ .../aws/operators/test_redshift_cluster.py | 16 ++++++++++++++++ 3 files changed, 23 insertions(+), 6 deletions(-) diff --git a/airflow/providers/amazon/aws/operators/redshift_cluster.py b/airflow/providers/amazon/aws/operators/redshift_cluster.py index dc6835434ff23..a56aeb6a053b1 100644 --- a/airflow/providers/amazon/aws/operators/redshift_cluster.py +++ b/airflow/providers/amazon/aws/operators/redshift_cluster.py @@ -91,6 +91,7 @@ class RedshiftCreateClusterOperator(BaseOperator): :param wait_for_completion: Whether wait for the cluster to be in ``available`` state :param max_attempt: The maximum number of attempts to be made. Default: 5 :param poll_interval: The amount of time in seconds to wait between attempts. Default: 60 + :param deferrable: If True, the operator will run in deferrable mode. """ template_fields: Sequence[str] = ( diff --git a/airflow/providers/amazon/aws/triggers/README.md b/airflow/providers/amazon/aws/triggers/README.md index 7ebc5fc2a3c57..cd0c0baae5d53 100644 --- a/airflow/providers/amazon/aws/triggers/README.md +++ b/airflow/providers/amazon/aws/triggers/README.md @@ -78,7 +78,7 @@ async with self.redshift_hook.async_conn as client: ) ``` -In this case, we are using the built-in cluster_available waiter. If we wanted to use a custom waiter, we would change the code slightly to use the get_waiter function from the hook, rather than the aiobotocore client: +In this case, we are using the built-in cluster_available waiter. If we wanted to use a custom waiter, we would change the code slightly to use the `get_waiter` function from the hook, rather than the aiobotocore client: ```python async with self.redshift_hook.async_conn as client: @@ -92,7 +92,7 @@ async with self.redshift_hook.async_conn as client: ) ``` -Here, we are calling the get_waiter function defined in base_aws.py which takes an optional argument of deferrable (set to True), and the aiobotocore client. cluster_paused is a custom boto waiter defined in redshift.json in the airflow/providers/amazon/aws/waiters folder. In general, the config file for a custom waiter should be named as .json. The config for cluster_paused is shown below: +Here, we are calling the `get_waiter` function defined in `base_aws.py` which takes an optional argument of `deferrable` (set to `True`), and the `aiobotocore` client. `cluster_paused` is a custom boto waiter defined in `redshift.json` in the `airflow/providers/amazon/aws/waiters` folder. In general, the config file for a custom waiter should be named as `.json`. The config for `cluster_paused` is shown below: ```json { @@ -128,22 +128,22 @@ Here, we are calling the get_waiter function defined in base_aws.py which takes For more information about writing custom waiter, see the [README.md](https://github.com/apache/airflow/blob/main/airflow/providers/amazon/aws/waiters/README.md) for custom waiters. -In some cases, a built-in or custom waiter may not be able to solve the problem. In such cases, the asynchronous method used to poll the boto3 API would need to be defined in the hook of the service being used. This method is essentially the same as the synchronous version of the method, except that it will use the aiobotocore client, and will be awaited. For the Redshift example, the async describe_clusters method would look as follows: +In some cases, a built-in or custom waiter may not be able to solve the problem. In such cases, the asynchronous method used to poll the boto3 API would need to be defined in the hook of the service being used. This method is essentially the same as the synchronous version of the method, except that it will use the aiobotocore client, and will be awaited. For the Redshift example, the async `describe_clusters` method would look as follows: ```python async with self.async_conn as client: response = client.describe_clusters(ClusterIdentifier=self.cluster_identifier) ``` -This async method can be used in the Trigger to poll the boto3 API. The polling logic will need to be implemented manually, taking care to use asyncio.sleep() rather than time.sleep(). +This async method can be used in the Trigger to poll the boto3 API. The polling logic will need to be implemented manually, taking care to use `asyncio.sleep()` rather than `time.sleep()`. -The last step in the Trigger is to yield a TriggerEvent that will be used to alert the Triggerer that the Trigger has finished execution. The TriggerEvent can pass information from the trigger to the method_name method named in the self.defer call in the operator. In the Redshift example, the TriggerEvent would look as follows: +The last step in the Trigger is to yield a `TriggerEvent` that will be used to alert the `Triggerer` that the Trigger has finished execution. The `TriggerEvent` can pass information from the trigger to the `method_name` method named in the `self.defer` call in the operator. In the Redshift example, the `TriggerEvent` would look as follows: ``` yield TriggerEvent({"status": "success", "message": "Cluster Created"}) ``` -The object passed through the TrigggerEvent can be captured in the method_name method through an event parameter. This can be used to determine what needs to be done based on the outcome of the Trigger execution. In the Redshift case, we can simply check the status of the event, and raise an Exception if something went wrong. +The object passed through the `TriggerEvent` can be captured in the `method_name` method through an `event` parameter. This can be used to determine what needs to be done based on the outcome of the Trigger execution. In the Redshift case, we can simply check the status of the event, and raise an Exception if something went wrong. ```python def execute_complete(self, context, event=None): diff --git a/tests/providers/amazon/aws/operators/test_redshift_cluster.py b/tests/providers/amazon/aws/operators/test_redshift_cluster.py index 8ec6f67eef152..2f8f8d834097e 100644 --- a/tests/providers/amazon/aws/operators/test_redshift_cluster.py +++ b/tests/providers/amazon/aws/operators/test_redshift_cluster.py @@ -115,6 +115,22 @@ def test_create_multi_node_cluster(self, mock_get_conn): # wait_for_completion is False so check waiter is not called mock_get_conn.return_value.get_waiter.assert_not_called() + + @mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.get_conn") + def test_create_cluster_deferrable(self, mock_get_conn): + redshift_operator = RedshiftCreateClusterOperator( + task_id="task_test", + cluster_identifier="test-cluster", + node_type="dc2.large", + master_username="adminuser", + master_user_password="Test123$", + cluster_type="single-node", + wait_for_completion=True, + deferrable=True, + ) + + with pytest.raises(TaskDeferred): + redshift_operator.execute(None) class TestRedshiftCreateClusterSnapshotOperator: From f8bb4a4e3c96af3095aec7c33d857dcb05a9af6d Mon Sep 17 00:00:00 2001 From: Syed Hussain Date: Fri, 21 Apr 2023 11:45:40 -0700 Subject: [PATCH 17/17] Fix static checks --- airflow/providers/amazon/aws/operators/redshift_cluster.py | 2 +- tests/providers/amazon/aws/operators/test_redshift_cluster.py | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/airflow/providers/amazon/aws/operators/redshift_cluster.py b/airflow/providers/amazon/aws/operators/redshift_cluster.py index a56aeb6a053b1..77ac521c9baf6 100644 --- a/airflow/providers/amazon/aws/operators/redshift_cluster.py +++ b/airflow/providers/amazon/aws/operators/redshift_cluster.py @@ -91,7 +91,7 @@ class RedshiftCreateClusterOperator(BaseOperator): :param wait_for_completion: Whether wait for the cluster to be in ``available`` state :param max_attempt: The maximum number of attempts to be made. Default: 5 :param poll_interval: The amount of time in seconds to wait between attempts. Default: 60 - :param deferrable: If True, the operator will run in deferrable mode. + :param deferrable: If True, the operator will run in deferrable mode """ template_fields: Sequence[str] = ( diff --git a/tests/providers/amazon/aws/operators/test_redshift_cluster.py b/tests/providers/amazon/aws/operators/test_redshift_cluster.py index 2f8f8d834097e..64a276f14d02c 100644 --- a/tests/providers/amazon/aws/operators/test_redshift_cluster.py +++ b/tests/providers/amazon/aws/operators/test_redshift_cluster.py @@ -115,7 +115,7 @@ def test_create_multi_node_cluster(self, mock_get_conn): # wait_for_completion is False so check waiter is not called mock_get_conn.return_value.get_waiter.assert_not_called() - + @mock.patch("airflow.providers.amazon.aws.hooks.redshift_cluster.RedshiftHook.get_conn") def test_create_cluster_deferrable(self, mock_get_conn): redshift_operator = RedshiftCreateClusterOperator( @@ -125,7 +125,6 @@ def test_create_cluster_deferrable(self, mock_get_conn): master_username="adminuser", master_user_password="Test123$", cluster_type="single-node", - wait_for_completion=True, deferrable=True, )