diff --git a/.circleci/integration-tests/astronomer_migration_dag.py b/.circleci/integration-tests/astronomer_migration_dag.py index 064739c88..016637067 100644 --- a/.circleci/integration-tests/astronomer_migration_dag.py +++ b/.circleci/integration-tests/astronomer_migration_dag.py @@ -1,4 +1,5 @@ """Astronomer migration DAG to transform metadata from source deployment to target Astro Cloud deployment.""" + from datetime import datetime from airflow import DAG diff --git a/.circleci/integration-tests/master_dag.py b/.circleci/integration-tests/master_dag.py index eb2eb6aab..445ca265a 100644 --- a/.circleci/integration-tests/master_dag.py +++ b/.circleci/integration-tests/master_dag.py @@ -1,4 +1,5 @@ """Master Dag to run all the example dags.""" + import logging import os import time diff --git a/.github/scripts/get_latest_runtime_image_tag.py b/.github/scripts/get_latest_runtime_image_tag.py index 74ad191a1..3ca879c63 100644 --- a/.github/scripts/get_latest_runtime_image_tag.py +++ b/.github/scripts/get_latest_runtime_image_tag.py @@ -1,4 +1,5 @@ """This script fetches the latest runtime image tag from the provided Quay.io repository URL.""" + from __future__ import annotations import sys diff --git a/astronomer/providers/amazon/aws/example_dags/example_aws_nuke.py b/astronomer/providers/amazon/aws/example_dags/example_aws_nuke.py index 0bb3601ea..9e58e31fd 100644 --- a/astronomer/providers/amazon/aws/example_dags/example_aws_nuke.py +++ b/astronomer/providers/amazon/aws/example_dags/example_aws_nuke.py @@ -1,4 +1,5 @@ """DAG to nuke AWS resources.""" + import logging import os from datetime import datetime, timedelta diff --git a/astronomer/providers/amazon/aws/example_dags/example_batch.py b/astronomer/providers/amazon/aws/example_dags/example_batch.py index 1776ea791..5bfa4c382 100644 --- a/astronomer/providers/amazon/aws/example_dags/example_batch.py +++ b/astronomer/providers/amazon/aws/example_dags/example_batch.py @@ -1,4 +1,5 @@ """Example DAG for the AWS Batch Operator Async""" + import logging import os import time diff --git a/astronomer/providers/amazon/aws/example_dags/example_redshift_cluster_management.py b/astronomer/providers/amazon/aws/example_dags/example_redshift_cluster_management.py index f6506d159..5ed666a92 100644 --- a/astronomer/providers/amazon/aws/example_dags/example_redshift_cluster_management.py +++ b/astronomer/providers/amazon/aws/example_dags/example_redshift_cluster_management.py @@ -1,4 +1,5 @@ """Airflow operators example to manage AWS Redshift cluster.""" + import os from datetime import datetime, timedelta from typing import Any diff --git a/astronomer/providers/amazon/aws/hooks/base_aws.py b/astronomer/providers/amazon/aws/hooks/base_aws.py index aa38b0ec3..7ecedacec 100644 --- a/astronomer/providers/amazon/aws/hooks/base_aws.py +++ b/astronomer/providers/amazon/aws/hooks/base_aws.py @@ -1,4 +1,5 @@ """This module contains async AWS Base Hook for deferrable operators and sensors.""" + from __future__ import annotations from typing import Any diff --git a/astronomer/providers/amazon/aws/operators/batch.py b/astronomer/providers/amazon/aws/operators/batch.py index 5370ab75e..1eda818ea 100644 --- a/astronomer/providers/amazon/aws/operators/batch.py +++ b/astronomer/providers/amazon/aws/operators/batch.py @@ -7,6 +7,7 @@ - `Batch `_ - `Welcome `_ """ + from __future__ import annotations import warnings diff --git a/astronomer/providers/apache/hive/example_dags/example_hive.py b/astronomer/providers/apache/hive/example_dags/example_hive.py index 6e4e18052..1e91b4e0d 100644 --- a/astronomer/providers/apache/hive/example_dags/example_hive.py +++ b/astronomer/providers/apache/hive/example_dags/example_hive.py @@ -1,4 +1,5 @@ """This is an example dag for hive partition sensors.""" + import logging import os import time diff --git a/astronomer/providers/apache/hive/hooks/hive.py b/astronomer/providers/apache/hive/hooks/hive.py index 11d0bc3fd..2eccbee59 100644 --- a/astronomer/providers/apache/hive/hooks/hive.py +++ b/astronomer/providers/apache/hive/hooks/hive.py @@ -1,4 +1,5 @@ """This module contains the Apache HiveCli hook async.""" + import asyncio from typing import Tuple diff --git a/astronomer/providers/apache/livy/example_dags/example_livy.py b/astronomer/providers/apache/livy/example_dags/example_livy.py index e32aae1af..f2cc834d1 100644 --- a/astronomer/providers/apache/livy/example_dags/example_livy.py +++ b/astronomer/providers/apache/livy/example_dags/example_livy.py @@ -3,6 +3,7 @@ The tasks below trigger the computation of pi on the Spark instance using the Java and Python executables provided in the example library. """ + import logging import os import time diff --git a/astronomer/providers/apache/livy/hooks/livy.py b/astronomer/providers/apache/livy/hooks/livy.py index cbe2e2a4b..b3d55c504 100644 --- a/astronomer/providers/apache/livy/hooks/livy.py +++ b/astronomer/providers/apache/livy/hooks/livy.py @@ -1,4 +1,5 @@ """This module contains the Apache Livy hook async.""" + import asyncio import re import warnings diff --git a/astronomer/providers/apache/livy/operators/livy.py b/astronomer/providers/apache/livy/operators/livy.py index 94d9c0eea..a2b8a0dad 100644 --- a/astronomer/providers/apache/livy/operators/livy.py +++ b/astronomer/providers/apache/livy/operators/livy.py @@ -1,4 +1,5 @@ """This module contains the Apache Livy operator async.""" + from __future__ import annotations import warnings diff --git a/astronomer/providers/apache/livy/triggers/livy.py b/astronomer/providers/apache/livy/triggers/livy.py index 7d48cc6aa..610c37c97 100644 --- a/astronomer/providers/apache/livy/triggers/livy.py +++ b/astronomer/providers/apache/livy/triggers/livy.py @@ -1,4 +1,5 @@ """This module contains the Apache Livy Trigger.""" + import asyncio import warnings from typing import Any, AsyncIterator, Dict, Optional, Tuple, Union diff --git a/astronomer/providers/cncf/kubernetes/example_dags/example_kubernetes_pod_operator.py b/astronomer/providers/cncf/kubernetes/example_dags/example_kubernetes_pod_operator.py index 3344a5660..cafb204d8 100644 --- a/astronomer/providers/cncf/kubernetes/example_dags/example_kubernetes_pod_operator.py +++ b/astronomer/providers/cncf/kubernetes/example_dags/example_kubernetes_pod_operator.py @@ -1,4 +1,5 @@ """Example DAG demonstrating the usage of the KubernetesPodOperatorAsync.""" + import os from datetime import datetime, timedelta diff --git a/astronomer/providers/google/cloud/example_dags/example_bigquery_impersonation_chain.py b/astronomer/providers/google/cloud/example_dags/example_bigquery_impersonation_chain.py index 600c0db99..7770b8d59 100644 --- a/astronomer/providers/google/cloud/example_dags/example_bigquery_impersonation_chain.py +++ b/astronomer/providers/google/cloud/example_dags/example_bigquery_impersonation_chain.py @@ -1,4 +1,5 @@ """Example Airflow DAG which uses impersonation parameters for authenticating with Google BigQuery service.""" + import os from datetime import datetime, timedelta from typing import Any diff --git a/astronomer/providers/google/cloud/example_dags/example_bigquery_queries.py b/astronomer/providers/google/cloud/example_dags/example_bigquery_queries.py index 0bbba4234..e2139e62a 100644 --- a/astronomer/providers/google/cloud/example_dags/example_bigquery_queries.py +++ b/astronomer/providers/google/cloud/example_dags/example_bigquery_queries.py @@ -2,6 +2,7 @@ Example Airflow DAG for Google BigQuery service. Uses Async version of BigQueryInsertJobOperator and BigQueryCheckOperator. """ + import os from datetime import datetime, timedelta from typing import Any diff --git a/astronomer/providers/google/cloud/example_dags/example_bigquery_sensors.py b/astronomer/providers/google/cloud/example_dags/example_bigquery_sensors.py index 65a0661e2..77d4b7823 100644 --- a/astronomer/providers/google/cloud/example_dags/example_bigquery_sensors.py +++ b/astronomer/providers/google/cloud/example_dags/example_bigquery_sensors.py @@ -1,4 +1,5 @@ """Example Airflow DAG for Google BigQuery Sensors.""" + import os from datetime import datetime, timedelta from typing import Any diff --git a/astronomer/providers/google/cloud/example_dags/example_gcp_nuke.py b/astronomer/providers/google/cloud/example_dags/example_gcp_nuke.py index b7ec46004..509698f87 100644 --- a/astronomer/providers/google/cloud/example_dags/example_gcp_nuke.py +++ b/astronomer/providers/google/cloud/example_dags/example_gcp_nuke.py @@ -1,4 +1,5 @@ """DAG to delete stale GCP resources.""" + import os from datetime import timedelta diff --git a/astronomer/providers/google/cloud/example_dags/example_gcs.py b/astronomer/providers/google/cloud/example_dags/example_gcs.py index 212592630..eafb034a7 100644 --- a/astronomer/providers/google/cloud/example_dags/example_gcs.py +++ b/astronomer/providers/google/cloud/example_dags/example_gcs.py @@ -1,4 +1,5 @@ """Example Airflow DAG for Google Cloud Storage operators.""" + import os from datetime import datetime, timedelta diff --git a/astronomer/providers/google/cloud/gke_utils.py b/astronomer/providers/google/cloud/gke_utils.py index 2ae5018e6..f02f2c83f 100644 --- a/astronomer/providers/google/cloud/gke_utils.py +++ b/astronomer/providers/google/cloud/gke_utils.py @@ -76,9 +76,11 @@ def _get_gke_config_file( # Write config to a temp file and set the environment variable to point to it. # This is to avoid race conditions of reading/writing a single file - with tempfile.NamedTemporaryFile() as conf_file, patch_environ( - {KUBE_CONFIG_ENV_VAR: conf_file.name} - ), hook.provide_authorized_gcloud(): + with ( + tempfile.NamedTemporaryFile() as conf_file, + patch_environ({KUBE_CONFIG_ENV_VAR: conf_file.name}), + hook.provide_authorized_gcloud(), + ): # Attempt to get/update credentials # We call gcloud directly instead of using google-cloud-python api # because there is no way to write kubernetes config to a file, which is diff --git a/astronomer/providers/google/cloud/hooks/gcs.py b/astronomer/providers/google/cloud/hooks/gcs.py index 8a2bc848e..c60fbf1bc 100644 --- a/astronomer/providers/google/cloud/hooks/gcs.py +++ b/astronomer/providers/google/cloud/hooks/gcs.py @@ -1,4 +1,5 @@ """This module contains a Google Cloud Storage hook.""" + import warnings from typing import Any, cast diff --git a/astronomer/providers/google/cloud/operators/bigquery.py b/astronomer/providers/google/cloud/operators/bigquery.py index 0d9a130de..3d07efc2e 100644 --- a/astronomer/providers/google/cloud/operators/bigquery.py +++ b/astronomer/providers/google/cloud/operators/bigquery.py @@ -1,4 +1,5 @@ """This module contains Google BigQueryAsync providers.""" + from __future__ import annotations import warnings diff --git a/astronomer/providers/google/cloud/operators/dataproc.py b/astronomer/providers/google/cloud/operators/dataproc.py index fd455aa4b..dff462020 100644 --- a/astronomer/providers/google/cloud/operators/dataproc.py +++ b/astronomer/providers/google/cloud/operators/dataproc.py @@ -1,4 +1,5 @@ """This module contains Google Dataproc operators.""" + from __future__ import annotations import warnings diff --git a/astronomer/providers/google/cloud/operators/kubernetes_engine.py b/astronomer/providers/google/cloud/operators/kubernetes_engine.py index b500d4993..904e87e8a 100644 --- a/astronomer/providers/google/cloud/operators/kubernetes_engine.py +++ b/astronomer/providers/google/cloud/operators/kubernetes_engine.py @@ -1,4 +1,5 @@ """This module contains Google GKE operators.""" + from __future__ import annotations from typing import Any, Sequence diff --git a/astronomer/providers/google/cloud/sensors/bigquery.py b/astronomer/providers/google/cloud/sensors/bigquery.py index e4c27eb35..030857b15 100644 --- a/astronomer/providers/google/cloud/sensors/bigquery.py +++ b/astronomer/providers/google/cloud/sensors/bigquery.py @@ -1,4 +1,5 @@ """This module contains Google Big Query sensors.""" + from __future__ import annotations import warnings diff --git a/astronomer/providers/google/cloud/sensors/gcs.py b/astronomer/providers/google/cloud/sensors/gcs.py index 5f9925014..0465e2de7 100644 --- a/astronomer/providers/google/cloud/sensors/gcs.py +++ b/astronomer/providers/google/cloud/sensors/gcs.py @@ -1,4 +1,5 @@ """This module contains Google Cloud Storage sensors.""" + from __future__ import annotations import warnings diff --git a/astronomer/providers/microsoft/azure/hooks/data_factory.py b/astronomer/providers/microsoft/azure/hooks/data_factory.py index 64622015b..f5354cea1 100644 --- a/astronomer/providers/microsoft/azure/hooks/data_factory.py +++ b/astronomer/providers/microsoft/azure/hooks/data_factory.py @@ -1,4 +1,5 @@ """This module contains the Azure Data Factory hook's asynchronous implementation.""" + from __future__ import annotations import inspect diff --git a/astronomer/providers/microsoft/azure/hooks/wasb.py b/astronomer/providers/microsoft/azure/hooks/wasb.py index b50dd610e..d837d47dd 100644 --- a/astronomer/providers/microsoft/azure/hooks/wasb.py +++ b/astronomer/providers/microsoft/azure/hooks/wasb.py @@ -1,4 +1,5 @@ """This module contains the Azure WASB hook's asynchronous implementation.""" + from __future__ import annotations import warnings diff --git a/dev/integration_test_scripts/replace_dependencies.py b/dev/integration_test_scripts/replace_dependencies.py index eb9dfbbab..dfd4716fa 100644 --- a/dev/integration_test_scripts/replace_dependencies.py +++ b/dev/integration_test_scripts/replace_dependencies.py @@ -65,9 +65,11 @@ def parse_providers_release_testing_gh_issue(issue_url: str) -> list[str]: h2_titles = first_comment.find_all("h2") package_urls = [_parse_pypi_url_from_h2_title(h2_title) for h2_title in h2_titles] pinned_packages = [ - _parse_pinned_package_from_pypi_url(url) - if "apache-airflow-providers-cncf-kubernetes" not in url - else _parse_pinned_pacakge_from_cncf(url) + ( + _parse_pinned_package_from_pypi_url(url) + if "apache-airflow-providers-cncf-kubernetes" not in url + else _parse_pinned_pacakge_from_cncf(url) + ) for url in package_urls ] return pinned_packages diff --git a/tests/http/hooks/test_http.py b/tests/http/hooks/test_http.py index dae4d4364..2841fcd64 100644 --- a/tests/http/hooks/test_http.py +++ b/tests/http/hooks/test_http.py @@ -15,9 +15,12 @@ async def test_do_api_call_async_non_retryable_error(self, aioresponse): hook = HttpHookAsync(method="GET") aioresponse.get("http://httpbin.org/non_existent_endpoint", status=400) - with pytest.raises(AirflowException) as exc, mock.patch.dict( - "os.environ", - AIRFLOW_CONN_HTTP_DEFAULT="http://httpbin.org/", + with ( + pytest.raises(AirflowException) as exc, + mock.patch.dict( + "os.environ", + AIRFLOW_CONN_HTTP_DEFAULT="http://httpbin.org/", + ), ): await hook.run(endpoint="non_existent_endpoint") @@ -29,9 +32,12 @@ async def test_do_api_call_async_retryable_error(self, caplog, aioresponse): hook = HttpHookAsync(method="GET") aioresponse.get("http://httpbin.org/non_existent_endpoint", status=500, repeat=True) - with pytest.raises(AirflowException) as exc, mock.patch.dict( - "os.environ", - AIRFLOW_CONN_HTTP_DEFAULT="http://httpbin.org/", + with ( + pytest.raises(AirflowException) as exc, + mock.patch.dict( + "os.environ", + AIRFLOW_CONN_HTTP_DEFAULT="http://httpbin.org/", + ), ): await hook.run(endpoint="non_existent_endpoint") diff --git a/tests/snowflake/hooks/test_snowflake_sql_api.py b/tests/snowflake/hooks/test_snowflake_sql_api.py index 58785d6d3..519b8dd33 100644 --- a/tests/snowflake/hooks/test_snowflake_sql_api.py +++ b/tests/snowflake/hooks/test_snowflake_sql_api.py @@ -322,12 +322,15 @@ def test_get_private_key_raise_exception(self, encrypted_temporary_private_key: }, } hook = SnowflakeSqlApiHookAsync(snowflake_conn_id="test_conn") - with unittest.mock.patch.dict( - "os.environ", AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri() - ), pytest.raises( - AirflowException, - match="The private_key_file and private_key_content extra fields are mutually " - "exclusive. Please remove one.", + with ( + unittest.mock.patch.dict( + "os.environ", AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri() + ), + pytest.raises( + AirflowException, + match="The private_key_file and private_key_content extra fields are mutually " + "exclusive. Please remove one.", + ), ): hook.get_private_key() @@ -384,9 +387,12 @@ def test_get_private_key_should_support_private_auth_with_unencrypted_key( hook.get_private_key() assert hook.private_key is not None connection_kwargs["password"] = _PASSWORD - with unittest.mock.patch.dict( - "os.environ", AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri() - ), pytest.raises(TypeError, match="Password was given but private key is not encrypted."): + with ( + unittest.mock.patch.dict( + "os.environ", AIRFLOW_CONN_TEST_CONN=Connection(**connection_kwargs).get_uri() + ), + pytest.raises(TypeError, match="Password was given but private key is not encrypted."), + ): SnowflakeSqlApiHookAsync(snowflake_conn_id="test_conn").get_private_key() @pytest.mark.asyncio