From 1fd12c75efc8cca0537edc8dcd3dcf62135c0ed3 Mon Sep 17 00:00:00 2001 From: dirrao Date: Sat, 6 Jul 2024 10:35:19 +0530 Subject: [PATCH] Resolve google deprecations in tests --- airflow/providers/google/cloud/sensors/gcs.py | 10 +++++++++- .../google/cloud/transfers/bigquery_to_postgres.py | 2 +- tests/deprecations_ignore.yml | 6 ------ 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/airflow/providers/google/cloud/sensors/gcs.py b/airflow/providers/google/cloud/sensors/gcs.py index 99e1a2380e570..c140f3b279d23 100644 --- a/airflow/providers/google/cloud/sensors/gcs.py +++ b/airflow/providers/google/cloud/sensors/gcs.py @@ -188,7 +188,15 @@ def ts_function(context): try: return context["data_interval_end"] except KeyError: - return context["dag"].following_schedule(context["execution_date"]) + from airflow.utils import timezone + + data_interval = context["dag"].infer_automated_data_interval( + timezone.coerce_datetime(context["execution_date"]) + ) + next_info = context["dag"].next_dagrun_info(data_interval, restricted=False) + if next_info is None: + return None + return next_info.data_interval.start class GCSObjectUpdateSensor(BaseSensorOperator): diff --git a/airflow/providers/google/cloud/transfers/bigquery_to_postgres.py b/airflow/providers/google/cloud/transfers/bigquery_to_postgres.py index 18b59e2a11744..ad4d44b39c014 100644 --- a/airflow/providers/google/cloud/transfers/bigquery_to_postgres.py +++ b/airflow/providers/google/cloud/transfers/bigquery_to_postgres.py @@ -76,7 +76,7 @@ def __init__( self.replace_index = replace_index def get_sql_hook(self) -> PostgresHook: - return PostgresHook(schema=self.database, postgres_conn_id=self.postgres_conn_id) + return PostgresHook(database=self.database, postgres_conn_id=self.postgres_conn_id) def execute(self, context: Context) -> None: big_query_hook = BigQueryHook( diff --git a/tests/deprecations_ignore.yml b/tests/deprecations_ignore.yml index 35620e6417f4c..03aebdcaa1813 100644 --- a/tests/deprecations_ignore.yml +++ b/tests/deprecations_ignore.yml @@ -158,12 +158,6 @@ - tests/providers/google/cloud/operators/vertex_ai/test_generative_model.py::TestVertexAIGenerateTextEmbeddingsOperator::test_execute - tests/providers/google/cloud/operators/vertex_ai/test_generative_model.py::TestVertexAIPromptMultimodalModelOperator::test_execute - tests/providers/google/cloud/operators/vertex_ai/test_generative_model.py::TestVertexAIPromptMultimodalModelWithMediaOperator::test_execute -- tests/providers/google/cloud/secrets/test_secret_manager.py::TestCloudSecretManagerBackend::test_connections_prefix_none_value -- tests/providers/google/cloud/secrets/test_secret_manager.py::TestCloudSecretManagerBackend::test_get_conn_uri -- tests/providers/google/cloud/secrets/test_secret_manager.py::TestCloudSecretManagerBackend::test_get_conn_uri_non_existent_key -- tests/providers/google/cloud/sensors/test_gcs.py::TestTsFunction::test_should_support_cron -- tests/providers/google/cloud/sensors/test_gcs.py::TestTsFunction::test_should_support_datetime -- tests/providers/google/cloud/transfers/test_bigquery_to_postgres.py::TestBigQueryToPostgresOperator::test_execute_good_request_to_bq - tests/providers/google/cloud/transfers/test_gcs_to_gcs.py::TestGoogleCloudStorageToCloudStorageOperator::test_copy_files_into_a_folder - tests/providers/google/cloud/transfers/test_gcs_to_gcs.py::TestGoogleCloudStorageToCloudStorageOperator::test_execute_last_modified_time - tests/providers/google/cloud/transfers/test_gcs_to_gcs.py::TestGoogleCloudStorageToCloudStorageOperator::test_execute_more_than_1_wildcard