From 6802d4136fab9b8a02faf22f48cb9f6d259fd523 Mon Sep 17 00:00:00 2001 From: Hussein Awala Date: Sun, 27 Aug 2023 19:05:30 +0200 Subject: [PATCH] Improve modules import in google provider by move some of them into a type-checking block (#33783) --- airflow/providers/google/ads/hooks/ads.py | 12 ++++++---- .../_internal_client/secret_manager_client.py | 5 +++- .../providers/google/cloud/hooks/automl.py | 20 +++++++++------- .../providers/google/cloud/hooks/bigquery.py | 4 ++-- .../google/cloud/hooks/bigquery_dts.py | 8 ++++--- .../providers/google/cloud/hooks/bigtable.py | 9 ++++--- .../google/cloud/hooks/cloud_batch.py | 8 ++++--- .../google/cloud/hooks/cloud_build.py | 10 ++++---- .../google/cloud/hooks/cloud_composer.py | 24 +++++++++++-------- .../google/cloud/hooks/cloud_memorystore.py | 8 ++++--- .../providers/google/cloud/hooks/cloud_run.py | 8 ++++--- .../providers/google/cloud/hooks/cloud_sql.py | 6 +++-- .../hooks/cloud_storage_transfer_service.py | 12 ++++++---- .../providers/google/cloud/hooks/compute.py | 8 ++++--- .../google/cloud/hooks/datacatalog.py | 8 ++++--- .../providers/google/cloud/hooks/dataform.py | 6 +++-- .../providers/google/cloud/hooks/dataplex.py | 10 ++++---- .../providers/google/cloud/hooks/dataproc.py | 16 +++++++------ .../google/cloud/hooks/dataproc_metastore.py | 14 ++++++----- airflow/providers/google/cloud/hooks/dlp.py | 6 +++-- airflow/providers/google/cloud/hooks/gcs.py | 11 +++++---- airflow/providers/google/cloud/hooks/kms.py | 6 +++-- .../google/cloud/hooks/kubernetes_engine.py | 10 ++++---- .../providers/google/cloud/hooks/looker.py | 5 +++- .../providers/google/cloud/hooks/mlengine.py | 8 ++++--- .../google/cloud/hooks/natural_language.py | 6 +++-- .../providers/google/cloud/hooks/os_login.py | 6 +++-- .../providers/google/cloud/hooks/pubsub.py | 24 ++++++++++--------- .../providers/google/cloud/hooks/spanner.py | 12 ++++++---- .../google/cloud/hooks/speech_to_text.py | 6 +++-- .../google/cloud/hooks/stackdriver.py | 6 +++-- airflow/providers/google/cloud/hooks/tasks.py | 8 ++++--- .../google/cloud/hooks/text_to_speech.py | 6 +++-- .../google/cloud/hooks/vertex_ai/auto_ml.py | 12 ++++++---- .../hooks/vertex_ai/batch_prediction_job.py | 10 ++++---- .../cloud/hooks/vertex_ai/custom_job.py | 20 +++++++++------- .../google/cloud/hooks/vertex_ai/dataset.py | 22 +++++++++-------- .../cloud/hooks/vertex_ai/endpoint_service.py | 14 ++++++----- .../vertex_ai/hyperparameter_tuning_job.py | 10 ++++---- .../cloud/hooks/vertex_ai/model_service.py | 12 ++++++---- .../google/cloud/hooks/video_intelligence.py | 8 ++++--- .../providers/google/cloud/hooks/vision.py | 8 ++++--- .../providers/google/cloud/hooks/workflows.py | 14 ++++++----- .../providers/google/cloud/links/bigquery.py | 2 +- .../google/cloud/links/bigquery_dts.py | 2 +- .../google/cloud/links/cloud_functions.py | 2 +- .../google/cloud/links/cloud_memorystore.py | 2 +- .../providers/google/cloud/links/cloud_sql.py | 2 +- .../google/cloud/links/cloud_tasks.py | 2 +- .../providers/google/cloud/links/compute.py | 2 +- .../google/cloud/links/datacatalog.py | 2 +- .../providers/google/cloud/links/dataflow.py | 2 +- .../providers/google/cloud/links/dataform.py | 2 +- .../providers/google/cloud/links/pubsub.py | 2 +- .../providers/google/cloud/links/spanner.py | 2 +- .../google/cloud/links/stackdriver.py | 2 +- .../providers/google/cloud/links/workflows.py | 2 +- .../cloud/log/stackdriver_task_handler.py | 12 ++++++---- .../google/cloud/operators/automl.py | 3 ++- .../google/cloud/operators/bigquery.py | 2 +- .../google/cloud/operators/bigquery_dts.py | 3 ++- .../google/cloud/operators/bigtable.py | 8 ++++--- .../google/cloud/operators/cloud_batch.py | 9 ++++--- .../google/cloud/operators/cloud_build.py | 3 ++- .../google/cloud/operators/cloud_composer.py | 5 ++-- .../cloud/operators/cloud_memorystore.py | 5 ++-- .../google/cloud/operators/cloud_run.py | 5 ++-- .../google/cloud/operators/cloud_sql.py | 6 ++--- .../google/cloud/operators/compute.py | 3 ++- .../google/cloud/operators/datacatalog.py | 5 ++-- .../google/cloud/operators/dataform.py | 3 ++- .../google/cloud/operators/dataplex.py | 3 ++- .../google/cloud/operators/dataproc.py | 7 +++--- .../cloud/operators/dataproc_metastore.py | 3 ++- .../providers/google/cloud/operators/dlp.py | 5 ++-- .../cloud/operators/kubernetes_engine.py | 3 ++- .../cloud/operators/natural_language.py | 5 ++-- .../google/cloud/operators/pubsub.py | 3 ++- .../google/cloud/operators/speech_to_text.py | 5 ++-- .../google/cloud/operators/stackdriver.py | 3 ++- .../providers/google/cloud/operators/tasks.py | 5 ++-- .../google/cloud/operators/text_to_speech.py | 5 ++-- .../cloud/operators/translate_speech.py | 3 ++- .../cloud/operators/vertex_ai/auto_ml.py | 3 ++- .../vertex_ai/batch_prediction_job.py | 5 ++-- .../cloud/operators/vertex_ai/custom_job.py | 3 ++- .../cloud/operators/vertex_ai/dataset.py | 5 ++-- .../operators/vertex_ai/endpoint_service.py | 5 ++-- .../vertex_ai/hyperparameter_tuning_job.py | 5 ++-- .../operators/vertex_ai/model_service.py | 3 ++- .../cloud/operators/video_intelligence.py | 3 ++- .../google/cloud/operators/vision.py | 5 ++-- .../google/cloud/operators/workflows.py | 5 ++-- .../google/cloud/sensors/bigquery_dts.py | 3 ++- .../google/cloud/sensors/dataplex.py | 3 ++- .../cloud/sensors/dataproc_metastore.py | 4 ++-- airflow/providers/google/cloud/sensors/gcs.py | 3 ++- .../google/cloud/sensors/workflows.py | 3 ++- .../google/cloud/transfers/bigquery_to_gcs.py | 3 ++- .../google/cloud/transfers/bigquery_to_sql.py | 2 +- .../google/cloud/transfers/calendar_to_gcs.py | 6 +++-- .../cloud/transfers/facebook_ads_to_gcs.py | 4 ++-- .../google/cloud/transfers/gcs_to_bigquery.py | 3 ++- .../google/cloud/transfers/presto_to_gcs.py | 9 +++---- .../google/cloud/transfers/trino_to_gcs.py | 9 +++---- .../google/cloud/triggers/cloud_run.py | 7 +++--- .../cloud/triggers/kubernetes_engine.py | 6 +++-- .../providers/google/cloud/triggers/pubsub.py | 4 ++-- .../google/cloud/utils/bigquery_get_data.py | 9 ++++--- .../cloud/utils/mlengine_operator_utils.py | 6 +++-- .../google/common/hooks/base_google.py | 6 +++-- .../providers/google/common/links/storage.py | 2 +- .../common/utils/id_token_credentials.py | 5 +++- .../hooks/campaign_manager.py | 6 +++-- .../providers/google/suite/hooks/calendar.py | 6 +++-- .../google/cloud/hooks/test_datacatalog.py | 6 +++-- .../google/cloud/hooks/test_os_login.py | 8 ++++--- .../cloud/operators/test_datacatalog.py | 6 +++-- .../operators/test_dataproc_metastore.py | 5 +++- .../google/cloud/utils/airflow_util.py | 5 +++- 120 files changed, 473 insertions(+), 304 deletions(-) diff --git a/airflow/providers/google/ads/hooks/ads.py b/airflow/providers/google/ads/hooks/ads.py index dbeb42f1daa9..a3ef9b93f787 100644 --- a/airflow/providers/google/ads/hooks/ads.py +++ b/airflow/providers/google/ads/hooks/ads.py @@ -20,20 +20,22 @@ from functools import cached_property from tempfile import NamedTemporaryFile -from typing import IO, Any +from typing import IO, TYPE_CHECKING, Any from google.ads.googleads.client import GoogleAdsClient from google.ads.googleads.errors import GoogleAdsException -from google.ads.googleads.v14.services.services.customer_service import CustomerServiceClient -from google.ads.googleads.v14.services.services.google_ads_service import GoogleAdsServiceClient -from google.ads.googleads.v14.services.types.google_ads_service import GoogleAdsRow -from google.api_core.page_iterator import GRPCIterator from google.auth.exceptions import GoogleAuthError from airflow import AirflowException from airflow.hooks.base import BaseHook from airflow.providers.google.common.hooks.base_google import get_field +if TYPE_CHECKING: + from google.ads.googleads.v14.services.services.customer_service import CustomerServiceClient + from google.ads.googleads.v14.services.services.google_ads_service import GoogleAdsServiceClient + from google.ads.googleads.v14.services.types.google_ads_service import GoogleAdsRow + from google.api_core.page_iterator import GRPCIterator + class GoogleAdsHook(BaseHook): """Interact with Google Ads API. diff --git a/airflow/providers/google/cloud/_internal_client/secret_manager_client.py b/airflow/providers/google/cloud/_internal_client/secret_manager_client.py index 0d346e029420..cc43f11f5245 100644 --- a/airflow/providers/google/cloud/_internal_client/secret_manager_client.py +++ b/airflow/providers/google/cloud/_internal_client/secret_manager_client.py @@ -18,14 +18,17 @@ import re from functools import cached_property +from typing import TYPE_CHECKING -import google from google.api_core.exceptions import InvalidArgument, NotFound, PermissionDenied from google.cloud.secretmanager_v1 import SecretManagerServiceClient from airflow.providers.google.common.consts import CLIENT_INFO from airflow.utils.log.logging_mixin import LoggingMixin +if TYPE_CHECKING: + import google + SECRET_ID_PATTERN = r"^[a-zA-Z0-9-_]*$" diff --git a/airflow/providers/google/cloud/hooks/automl.py b/airflow/providers/google/cloud/hooks/automl.py index cada4eb51129..bda0d49ceaa2 100644 --- a/airflow/providers/google/cloud/hooks/automl.py +++ b/airflow/providers/google/cloud/hooks/automl.py @@ -25,11 +25,9 @@ from __future__ import annotations from functools import cached_property -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.operation import Operation -from google.api_core.retry import Retry from google.cloud.automl_v1beta1 import ( AutoMlClient, BatchPredictInputConfig, @@ -42,17 +40,21 @@ PredictionServiceClient, PredictResponse, ) -from google.cloud.automl_v1beta1.services.auto_ml.pagers import ( - ListColumnSpecsPager, - ListDatasetsPager, - ListTableSpecsPager, -) -from google.protobuf.field_mask_pb2 import FieldMask from airflow import AirflowException from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.operation import Operation + from google.api_core.retry import Retry + from google.cloud.automl_v1beta1.services.auto_ml.pagers import ( + ListColumnSpecsPager, + ListDatasetsPager, + ListTableSpecsPager, + ) + from google.protobuf.field_mask_pb2 import FieldMask + class CloudAutoMLHook(GoogleBaseHook): """ diff --git a/airflow/providers/google/cloud/hooks/bigquery.py b/airflow/providers/google/cloud/hooks/bigquery.py index 9a79465bf4a5..43987aac4c90 100644 --- a/airflow/providers/google/cloud/hooks/bigquery.py +++ b/airflow/providers/google/cloud/hooks/bigquery.py @@ -32,8 +32,6 @@ from aiohttp import ClientSession as ClientSession from gcloud.aio.bigquery import Job, Table as Table_async -from google.api_core.page_iterator import HTTPIterator -from google.api_core.retry import Retry from google.cloud.bigquery import ( DEFAULT_RETRY, Client, @@ -70,6 +68,8 @@ if TYPE_CHECKING: import pandas as pd + from google.api_core.page_iterator import HTTPIterator + from google.api_core.retry import Retry log = logging.getLogger(__name__) diff --git a/airflow/providers/google/cloud/hooks/bigquery_dts.py b/airflow/providers/google/cloud/hooks/bigquery_dts.py index 2d233591fde1..c3f5780ea8b2 100644 --- a/airflow/providers/google/cloud/hooks/bigquery_dts.py +++ b/airflow/providers/google/cloud/hooks/bigquery_dts.py @@ -19,17 +19,15 @@ from __future__ import annotations from copy import copy -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.bigquery_datatransfer_v1 import DataTransferServiceAsyncClient, DataTransferServiceClient from google.cloud.bigquery_datatransfer_v1.types import ( StartManualTransferRunsResponse, TransferConfig, TransferRun, ) -from googleapiclient.discovery import Resource from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import ( @@ -38,6 +36,10 @@ GoogleBaseHook, ) +if TYPE_CHECKING: + from google.api_core.retry import Retry + from googleapiclient.discovery import Resource + def get_object_id(obj: dict) -> str: """Returns unique id of the object.""" diff --git a/airflow/providers/google/cloud/hooks/bigtable.py b/airflow/providers/google/cloud/hooks/bigtable.py index 888d8d023611..786d02b4961c 100644 --- a/airflow/providers/google/cloud/hooks/bigtable.py +++ b/airflow/providers/google/cloud/hooks/bigtable.py @@ -18,18 +18,21 @@ """This module contains a Google Cloud Bigtable Hook.""" from __future__ import annotations -import enum -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.cloud.bigtable import Client, enums from google.cloud.bigtable.cluster import Cluster -from google.cloud.bigtable.column_family import ColumnFamily, GarbageCollectionRule from google.cloud.bigtable.instance import Instance from google.cloud.bigtable.table import ClusterState, Table from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import GoogleBaseHook +if TYPE_CHECKING: + import enum + + from google.cloud.bigtable.column_family import ColumnFamily, GarbageCollectionRule + class BigtableHook(GoogleBaseHook): """ diff --git a/airflow/providers/google/cloud/hooks/cloud_batch.py b/airflow/providers/google/cloud/hooks/cloud_batch.py index f85283047da4..ef642e4ec9c3 100644 --- a/airflow/providers/google/cloud/hooks/cloud_batch.py +++ b/airflow/providers/google/cloud/hooks/cloud_batch.py @@ -20,9 +20,8 @@ import itertools import json from time import sleep -from typing import Iterable, Sequence +from typing import TYPE_CHECKING, Iterable, Sequence -from google.api_core import operation # type: ignore from google.cloud.batch import ListJobsRequest, ListTasksRequest from google.cloud.batch_v1 import ( BatchServiceAsyncClient, @@ -32,12 +31,15 @@ JobStatus, Task, ) -from google.cloud.batch_v1.services.batch_service import pagers from airflow.exceptions import AirflowException from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook +if TYPE_CHECKING: + from google.api_core import operation + from google.cloud.batch_v1.services.batch_service import pagers + class CloudBatchHook(GoogleBaseHook): """ diff --git a/airflow/providers/google/cloud/hooks/cloud_build.py b/airflow/providers/google/cloud/hooks/cloud_build.py index 1b7815d9035a..5cd9b798ea1e 100644 --- a/airflow/providers/google/cloud/hooks/cloud_build.py +++ b/airflow/providers/google/cloud/hooks/cloud_build.py @@ -19,20 +19,22 @@ from __future__ import annotations import warnings -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core.client_options import ClientOptions from google.api_core.exceptions import AlreadyExists from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.operation import Operation -from google.api_core.retry import Retry from google.cloud.devtools.cloudbuild_v1 import CloudBuildAsyncClient, CloudBuildClient, GetBuildRequest -from google.cloud.devtools.cloudbuild_v1.types import Build, BuildTrigger, RepoSource from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.operation import Operation + from google.api_core.retry import Retry + from google.cloud.devtools.cloudbuild_v1.types import Build, BuildTrigger, RepoSource + # Time to sleep between active checks of the operation results TIME_TO_SLEEP_IN_SECONDS = 5 diff --git a/airflow/providers/google/cloud/hooks/cloud_composer.py b/airflow/providers/google/cloud/hooks/cloud_composer.py index 3ad0ba1098ac..d0543811c1ba 100644 --- a/airflow/providers/google/cloud/hooks/cloud_composer.py +++ b/airflow/providers/google/cloud/hooks/cloud_composer.py @@ -17,29 +17,33 @@ # under the License. from __future__ import annotations -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core.client_options import ClientOptions from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.operation import Operation -from google.api_core.operation_async import AsyncOperation -from google.api_core.retry import Retry from google.cloud.orchestration.airflow.service_v1 import ( EnvironmentsAsyncClient, EnvironmentsClient, ImageVersionsClient, ) -from google.cloud.orchestration.airflow.service_v1.services.environments.pagers import ListEnvironmentsPager -from google.cloud.orchestration.airflow.service_v1.services.image_versions.pagers import ( - ListImageVersionsPager, -) -from google.cloud.orchestration.airflow.service_v1.types import Environment -from google.protobuf.field_mask_pb2 import FieldMask from airflow import AirflowException from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.operation import Operation + from google.api_core.operation_async import AsyncOperation + from google.api_core.retry import Retry + from google.cloud.orchestration.airflow.service_v1.services.environments.pagers import ( + ListEnvironmentsPager, + ) + from google.cloud.orchestration.airflow.service_v1.services.image_versions.pagers import ( + ListImageVersionsPager, + ) + from google.cloud.orchestration.airflow.service_v1.types import Environment + from google.protobuf.field_mask_pb2 import FieldMask + class CloudComposerHook(GoogleBaseHook): """Hook for Google Cloud Composer APIs.""" diff --git a/airflow/providers/google/cloud/hooks/cloud_memorystore.py b/airflow/providers/google/cloud/hooks/cloud_memorystore.py index a6b63db7534e..4ce501506660 100644 --- a/airflow/providers/google/cloud/hooks/cloud_memorystore.py +++ b/airflow/providers/google/cloud/hooks/cloud_memorystore.py @@ -27,12 +27,11 @@ """ from __future__ import annotations -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core import path_template from google.api_core.exceptions import NotFound from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.memcache_v1beta2 import CloudMemcacheClient from google.cloud.memcache_v1beta2.types import cloud_memcache from google.cloud.redis_v1 import ( @@ -42,12 +41,15 @@ Instance, OutputConfig, ) -from google.protobuf.field_mask_pb2 import FieldMask from airflow import version from airflow.exceptions import AirflowException from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.retry import Retry + from google.protobuf.field_mask_pb2 import FieldMask + class CloudMemorystoreHook(GoogleBaseHook): """ diff --git a/airflow/providers/google/cloud/hooks/cloud_run.py b/airflow/providers/google/cloud/hooks/cloud_run.py index f33b5fa3afd5..d2acbc4462ea 100644 --- a/airflow/providers/google/cloud/hooks/cloud_run.py +++ b/airflow/providers/google/cloud/hooks/cloud_run.py @@ -18,9 +18,8 @@ from __future__ import annotations import itertools -from typing import Iterable, Sequence +from typing import TYPE_CHECKING, Iterable, Sequence -from google.api_core import operation from google.cloud.run_v2 import ( CreateJobRequest, DeleteJobRequest, @@ -32,13 +31,16 @@ RunJobRequest, UpdateJobRequest, ) -from google.cloud.run_v2.services.jobs import pagers from google.longrunning import operations_pb2 from airflow.exceptions import AirflowException from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook +if TYPE_CHECKING: + from google.api_core import operation + from google.cloud.run_v2.services.jobs import pagers + class CloudRunHook(GoogleBaseHook): """ diff --git a/airflow/providers/google/cloud/hooks/cloud_sql.py b/airflow/providers/google/cloud/hooks/cloud_sql.py index f89b460c0e38..f3fa10bb97ab 100644 --- a/airflow/providers/google/cloud/hooks/cloud_sql.py +++ b/airflow/providers/google/cloud/hooks/cloud_sql.py @@ -35,7 +35,7 @@ from pathlib import Path from subprocess import PIPE, Popen from tempfile import gettempdir -from typing import Any, Sequence +from typing import TYPE_CHECKING, Any, Sequence from urllib.parse import quote_plus import httpx @@ -43,7 +43,6 @@ from gcloud.aio.auth import AioSession, Token from googleapiclient.discovery import Resource, build from googleapiclient.errors import HttpError -from requests import Session # Number of retries - used by googleapiclient method calls to perform retries # For requests that are "retriable" @@ -55,6 +54,9 @@ from airflow.providers.postgres.hooks.postgres import PostgresHook from airflow.utils.log.logging_mixin import LoggingMixin +if TYPE_CHECKING: + from requests import Session + UNIX_PATH_MAX = 108 # Time to sleep between active checks of the operation results diff --git a/airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py b/airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py index ec452153f142..b8c1215b8214 100644 --- a/airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py +++ b/airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py @@ -33,7 +33,7 @@ import warnings from copy import deepcopy from datetime import timedelta -from typing import Any, Sequence +from typing import TYPE_CHECKING, Any, Sequence from google.cloud.storage_transfer_v1 import ( ListTransferJobsRequest, @@ -41,16 +41,18 @@ TransferJob, TransferOperation, ) -from google.cloud.storage_transfer_v1.services.storage_transfer_service.pagers import ( - ListTransferJobsAsyncPager, -) from googleapiclient.discovery import Resource, build from googleapiclient.errors import HttpError -from proto import Message from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook +if TYPE_CHECKING: + from google.cloud.storage_transfer_v1.services.storage_transfer_service.pagers import ( + ListTransferJobsAsyncPager, + ) + from proto import Message + log = logging.getLogger(__name__) # Time to sleep between active checks of the operation results diff --git a/airflow/providers/google/cloud/hooks/compute.py b/airflow/providers/google/cloud/hooks/compute.py index 1a95e94c317f..d3ecdc97e15e 100644 --- a/airflow/providers/google/cloud/hooks/compute.py +++ b/airflow/providers/google/cloud/hooks/compute.py @@ -19,18 +19,20 @@ from __future__ import annotations import time -from typing import Any, Sequence +from typing import TYPE_CHECKING, Any, Sequence -from google.api_core.retry import Retry from google.cloud.compute_v1.services.instance_group_managers import InstanceGroupManagersClient from google.cloud.compute_v1.services.instance_templates import InstanceTemplatesClient from google.cloud.compute_v1.services.instances import InstancesClient -from google.cloud.compute_v1.types import Instance, InstanceGroupManager, InstanceTemplate from googleapiclient.discovery import build from airflow.exceptions import AirflowException from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.retry import Retry + from google.cloud.compute_v1.types import Instance, InstanceGroupManager, InstanceTemplate + # Time to sleep between active checks of the operation results TIME_TO_SLEEP_IN_SECONDS = 1 diff --git a/airflow/providers/google/cloud/hooks/datacatalog.py b/airflow/providers/google/cloud/hooks/datacatalog.py index 2dd60931fe42..91f9d69d9b9f 100644 --- a/airflow/providers/google/cloud/hooks/datacatalog.py +++ b/airflow/providers/google/cloud/hooks/datacatalog.py @@ -16,10 +16,9 @@ # under the License. from __future__ import annotations -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud import datacatalog from google.cloud.datacatalog import ( CreateTagRequest, @@ -31,12 +30,15 @@ TagTemplate, TagTemplateField, ) -from google.protobuf.field_mask_pb2 import FieldMask from airflow import AirflowException from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.retry import Retry + from google.protobuf.field_mask_pb2 import FieldMask + class CloudDataCatalogHook(GoogleBaseHook): """ diff --git a/airflow/providers/google/cloud/hooks/dataform.py b/airflow/providers/google/cloud/hooks/dataform.py index d7f86157f581..5f26591e0914 100644 --- a/airflow/providers/google/cloud/hooks/dataform.py +++ b/airflow/providers/google/cloud/hooks/dataform.py @@ -17,10 +17,9 @@ from __future__ import annotations import time -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.dataform_v1beta1 import DataformClient from google.cloud.dataform_v1beta1.types import ( CompilationResult, @@ -34,6 +33,9 @@ from airflow import AirflowException from airflow.providers.google.common.hooks.base_google import GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.retry import Retry + class DataformHook(GoogleBaseHook): """Hook for Google Cloud DataForm APIs.""" diff --git a/airflow/providers/google/cloud/hooks/dataplex.py b/airflow/providers/google/cloud/hooks/dataplex.py index 5fbfb0bd6efe..d3c958820a0a 100644 --- a/airflow/providers/google/cloud/hooks/dataplex.py +++ b/airflow/providers/google/cloud/hooks/dataplex.py @@ -18,12 +18,10 @@ from __future__ import annotations import time -from typing import Any, Sequence +from typing import TYPE_CHECKING, Any, Sequence from google.api_core.client_options import ClientOptions from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.operation import Operation -from google.api_core.retry import Retry from google.cloud.dataplex_v1 import DataplexServiceClient, DataScanServiceClient from google.cloud.dataplex_v1.types import ( Asset, @@ -34,12 +32,16 @@ Zone, ) from google.protobuf.field_mask_pb2 import FieldMask -from googleapiclient.discovery import Resource from airflow.exceptions import AirflowException from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.operation import Operation + from google.api_core.retry import Retry + from googleapiclient.discovery import Resource + PATH_DATA_SCAN = "projects/{project_id}/locations/{region}/dataScans/{data_scan_id}" diff --git a/airflow/providers/google/cloud/hooks/dataproc.py b/airflow/providers/google/cloud/hooks/dataproc.py index e7eb48bef105..349a8488bc87 100644 --- a/airflow/providers/google/cloud/hooks/dataproc.py +++ b/airflow/providers/google/cloud/hooks/dataproc.py @@ -20,15 +20,11 @@ import time import uuid -from typing import Any, Sequence +from typing import TYPE_CHECKING, Any, Sequence from google.api_core.client_options import ClientOptions from google.api_core.exceptions import ServerError from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.operation import Operation -from google.api_core.operation_async import AsyncOperation -from google.api_core.operations_v1.operations_client import OperationsClient -from google.api_core.retry import Retry from google.cloud.dataproc_v1 import ( Batch, BatchControllerAsyncClient, @@ -44,14 +40,20 @@ WorkflowTemplateServiceAsyncClient, WorkflowTemplateServiceClient, ) -from google.protobuf.duration_pb2 import Duration -from google.protobuf.field_mask_pb2 import FieldMask from airflow.exceptions import AirflowException from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import GoogleBaseHook from airflow.version import version as airflow_version +if TYPE_CHECKING: + from google.api_core.operation import Operation + from google.api_core.operation_async import AsyncOperation + from google.api_core.operations_v1.operations_client import OperationsClient + from google.api_core.retry import Retry + from google.protobuf.duration_pb2 import Duration + from google.protobuf.field_mask_pb2 import FieldMask + class DataProcJobBuilder: """A helper class for building Dataproc job.""" diff --git a/airflow/providers/google/cloud/hooks/dataproc_metastore.py b/airflow/providers/google/cloud/hooks/dataproc_metastore.py index 55814970f80c..8066258fdb85 100644 --- a/airflow/providers/google/cloud/hooks/dataproc_metastore.py +++ b/airflow/providers/google/cloud/hooks/dataproc_metastore.py @@ -18,21 +18,23 @@ """This module contains a Google Cloud Dataproc Metastore hook.""" from __future__ import annotations -from typing import Any, Sequence +from typing import TYPE_CHECKING, Any, Sequence from google.api_core.client_options import ClientOptions from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.operation import Operation -from google.api_core.retry import Retry from google.cloud.metastore_v1 import DataprocMetastoreClient -from google.cloud.metastore_v1.types import Backup, MetadataImport, Service -from google.cloud.metastore_v1.types.metastore import DatabaseDumpSpec, Restore -from google.protobuf.field_mask_pb2 import FieldMask from airflow.exceptions import AirflowException from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.operation import Operation + from google.api_core.retry import Retry + from google.cloud.metastore_v1.types import Backup, MetadataImport, Service + from google.cloud.metastore_v1.types.metastore import DatabaseDumpSpec, Restore + from google.protobuf.field_mask_pb2 import FieldMask + class DataprocMetastoreHook(GoogleBaseHook): """Hook for Google Cloud Dataproc Metastore APIs.""" diff --git a/airflow/providers/google/cloud/hooks/dlp.py b/airflow/providers/google/cloud/hooks/dlp.py index d65d8845582b..1f96ec0cd255 100644 --- a/airflow/providers/google/cloud/hooks/dlp.py +++ b/airflow/providers/google/cloud/hooks/dlp.py @@ -27,10 +27,9 @@ import re import time -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.dlp import DlpServiceClient from google.cloud.dlp_v2.types import ( ByteContentItem, @@ -58,6 +57,9 @@ from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.retry import Retry + DLP_JOB_PATH_PATTERN = "^projects/[^/]+/dlpJobs/(?P.*?)$" diff --git a/airflow/providers/google/cloud/hooks/gcs.py b/airflow/providers/google/cloud/hooks/gcs.py index 72c555bbda16..b27bedafa61b 100644 --- a/airflow/providers/google/cloud/hooks/gcs.py +++ b/airflow/providers/google/cloud/hooks/gcs.py @@ -26,18 +26,15 @@ import time import warnings from contextlib import contextmanager -from datetime import datetime from functools import partial from io import BytesIO from os import path from tempfile import NamedTemporaryFile -from typing import IO, Any, Callable, Generator, Sequence, TypeVar, cast, overload +from typing import IO, TYPE_CHECKING, Any, Callable, Generator, Sequence, TypeVar, cast, overload from urllib.parse import urlsplit -from aiohttp import ClientSession from gcloud.aio.storage import Storage from google.api_core.exceptions import GoogleAPICallError, NotFound -from google.api_core.retry import Retry # not sure why but mypy complains on missing `storage` but it is clearly there and is importable from google.cloud import storage # type: ignore[attr-defined] @@ -52,6 +49,12 @@ from airflow.utils import timezone from airflow.version import version +if TYPE_CHECKING: + from datetime import datetime + + from aiohttp import ClientSession + from google.api_core.retry import Retry + try: # Airflow 2.3 doesn't have this yet from airflow.typing_compat import ParamSpec diff --git a/airflow/providers/google/cloud/hooks/kms.py b/airflow/providers/google/cloud/hooks/kms.py index 3be5f09887ed..b132bf8631b8 100644 --- a/airflow/providers/google/cloud/hooks/kms.py +++ b/airflow/providers/google/cloud/hooks/kms.py @@ -19,15 +19,17 @@ from __future__ import annotations import base64 -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.kms_v1 import KeyManagementServiceClient from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.retry import Retry + def _b64encode(s: bytes) -> str: """Base 64 encodes a bytes object to a string.""" diff --git a/airflow/providers/google/cloud/hooks/kubernetes_engine.py b/airflow/providers/google/cloud/hooks/kubernetes_engine.py index 1837613ea065..c29e3106492a 100644 --- a/airflow/providers/google/cloud/hooks/kubernetes_engine.py +++ b/airflow/providers/google/cloud/hooks/kubernetes_engine.py @@ -30,13 +30,11 @@ import time import warnings from functools import cached_property -from typing import Sequence +from typing import TYPE_CHECKING, Sequence -import google.auth.credentials from gcloud.aio.auth import Token from google.api_core.exceptions import NotFound from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.auth.transport import requests as google_requests # not sure why but mypy complains on missing `container_v1` but it is clearly there and is importable @@ -45,7 +43,6 @@ from google.cloud.container_v1.types import Cluster, Operation from kubernetes import client from kubernetes_asyncio import client as async_client -from kubernetes_asyncio.client.models import V1Pod from kubernetes_asyncio.config.kube_config import FileOrData from urllib3.exceptions import HTTPError @@ -59,6 +56,11 @@ GoogleBaseHook, ) +if TYPE_CHECKING: + import google.auth.credentials + from google.api_core.retry import Retry + from kubernetes_asyncio.client.models import V1Pod + OPERATIONAL_POLL_INTERVAL = 15 diff --git a/airflow/providers/google/cloud/hooks/looker.py b/airflow/providers/google/cloud/hooks/looker.py index d1179d707cae..616d583221b5 100644 --- a/airflow/providers/google/cloud/hooks/looker.py +++ b/airflow/providers/google/cloud/hooks/looker.py @@ -21,6 +21,7 @@ import json import time from enum import Enum +from typing import TYPE_CHECKING from looker_sdk.rtl import api_settings, auth_session, requests_transport, serialize from looker_sdk.sdk.api40 import methods as methods40 @@ -28,9 +29,11 @@ from airflow.exceptions import AirflowException from airflow.hooks.base import BaseHook -from airflow.models.connection import Connection from airflow.version import version +if TYPE_CHECKING: + from airflow.models.connection import Connection + class LookerHook(BaseHook): """Hook for Looker APIs.""" diff --git a/airflow/providers/google/cloud/hooks/mlengine.py b/airflow/providers/google/cloud/hooks/mlengine.py index aabe2b88ae02..fdc339299de7 100644 --- a/airflow/providers/google/cloud/hooks/mlengine.py +++ b/airflow/providers/google/cloud/hooks/mlengine.py @@ -21,19 +21,21 @@ import logging import random import time -from typing import Callable +from typing import TYPE_CHECKING, Callable from aiohttp import ClientSession from gcloud.aio.auth import AioSession, Token from googleapiclient.discovery import Resource, build from googleapiclient.errors import HttpError -from httplib2 import Response -from requests import Session from airflow.exceptions import AirflowException from airflow.providers.google.common.hooks.base_google import GoogleBaseAsyncHook, GoogleBaseHook from airflow.version import version as airflow_version +if TYPE_CHECKING: + from httplib2 import Response + from requests import Session + log = logging.getLogger(__name__) _AIRFLOW_VERSION = "v" + airflow_version.replace(".", "-").replace("+", "-") diff --git a/airflow/providers/google/cloud/hooks/natural_language.py b/airflow/providers/google/cloud/hooks/natural_language.py index 673df34b2328..6c8ddbcb2654 100644 --- a/airflow/providers/google/cloud/hooks/natural_language.py +++ b/airflow/providers/google/cloud/hooks/natural_language.py @@ -18,10 +18,9 @@ """This module contains a Google Cloud Natural Language Hook.""" from __future__ import annotations -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.language_v1 import EncodingType, LanguageServiceClient from google.cloud.language_v1.types import ( AnalyzeEntitiesResponse, @@ -37,6 +36,9 @@ from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.retry import Retry + class CloudNaturalLanguageHook(GoogleBaseHook): """ diff --git a/airflow/providers/google/cloud/hooks/os_login.py b/airflow/providers/google/cloud/hooks/os_login.py index ceb8ff74384e..75945a004e05 100644 --- a/airflow/providers/google/cloud/hooks/os_login.py +++ b/airflow/providers/google/cloud/hooks/os_login.py @@ -22,15 +22,17 @@ """ from __future__ import annotations -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.oslogin_v1 import ImportSshPublicKeyResponse, OsLoginServiceClient from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.retry import Retry + class OSLoginHook(GoogleBaseHook): """ diff --git a/airflow/providers/google/cloud/hooks/pubsub.py b/airflow/providers/google/cloud/hooks/pubsub.py index 6e867c2bdb1f..603749a55b02 100644 --- a/airflow/providers/google/cloud/hooks/pubsub.py +++ b/airflow/providers/google/cloud/hooks/pubsub.py @@ -28,23 +28,13 @@ import warnings from base64 import b64decode from functools import cached_property -from typing import Any, Sequence +from typing import TYPE_CHECKING, Any, Sequence from uuid import uuid4 from google.api_core.exceptions import AlreadyExists, GoogleAPICallError from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.exceptions import NotFound from google.cloud.pubsub_v1 import PublisherClient, SubscriberClient -from google.cloud.pubsub_v1.types import ( - DeadLetterPolicy, - Duration, - ExpirationPolicy, - MessageStoragePolicy, - PushConfig, - ReceivedMessage, - RetryPolicy, -) from google.pubsub_v1.services.subscriber.async_client import SubscriberAsyncClient from googleapiclient.errors import HttpError @@ -57,6 +47,18 @@ ) from airflow.version import version +if TYPE_CHECKING: + from google.api_core.retry import Retry + from google.cloud.pubsub_v1.types import ( + DeadLetterPolicy, + Duration, + ExpirationPolicy, + MessageStoragePolicy, + PushConfig, + ReceivedMessage, + RetryPolicy, + ) + class PubSubException(Exception): """Alias for Exception.""" diff --git a/airflow/providers/google/cloud/hooks/spanner.py b/airflow/providers/google/cloud/hooks/spanner.py index a3e652dcbd88..a2dab7b858cb 100644 --- a/airflow/providers/google/cloud/hooks/spanner.py +++ b/airflow/providers/google/cloud/hooks/spanner.py @@ -18,14 +18,10 @@ """This module contains a Google Cloud Spanner Hook.""" from __future__ import annotations -from typing import Callable, NamedTuple, Sequence +from typing import TYPE_CHECKING, Callable, NamedTuple, Sequence from google.api_core.exceptions import AlreadyExists, GoogleAPICallError from google.cloud.spanner_v1.client import Client -from google.cloud.spanner_v1.database import Database -from google.cloud.spanner_v1.instance import Instance -from google.cloud.spanner_v1.transaction import Transaction -from google.longrunning.operations_grpc_pb2 import Operation from sqlalchemy import create_engine from airflow.exceptions import AirflowException @@ -33,6 +29,12 @@ from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import GoogleBaseHook, get_field +if TYPE_CHECKING: + from google.cloud.spanner_v1.database import Database + from google.cloud.spanner_v1.instance import Instance + from google.cloud.spanner_v1.transaction import Transaction + from google.longrunning.operations_grpc_pb2 import Operation + class SpannerConnectionParams(NamedTuple): """Information about Google Spanner connection parameters.""" diff --git a/airflow/providers/google/cloud/hooks/speech_to_text.py b/airflow/providers/google/cloud/hooks/speech_to_text.py index 33ddc653c37c..765befdee55f 100644 --- a/airflow/providers/google/cloud/hooks/speech_to_text.py +++ b/airflow/providers/google/cloud/hooks/speech_to_text.py @@ -18,16 +18,18 @@ """This module contains a Google Cloud Speech Hook.""" from __future__ import annotations -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.speech_v1 import SpeechClient from google.cloud.speech_v1.types import RecognitionAudio, RecognitionConfig from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.retry import Retry + class CloudSpeechToTextHook(GoogleBaseHook): """ diff --git a/airflow/providers/google/cloud/hooks/stackdriver.py b/airflow/providers/google/cloud/hooks/stackdriver.py index 28e0b8830357..b16a2233ee59 100644 --- a/airflow/providers/google/cloud/hooks/stackdriver.py +++ b/airflow/providers/google/cloud/hooks/stackdriver.py @@ -19,11 +19,10 @@ from __future__ import annotations import json -from typing import Any, Sequence +from typing import TYPE_CHECKING, Any, Sequence from google.api_core.exceptions import InvalidArgument from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud import monitoring_v3 from google.cloud.monitoring_v3 import AlertPolicy, NotificationChannel from google.protobuf.field_mask_pb2 import FieldMask @@ -32,6 +31,9 @@ from airflow.exceptions import AirflowException from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.retry import Retry + class StackdriverHook(GoogleBaseHook): """Stackdriver Hook for connecting with Google Cloud Stackdriver.""" diff --git a/airflow/providers/google/cloud/hooks/tasks.py b/airflow/providers/google/cloud/hooks/tasks.py index 4807d0fde77c..e9bca0a4564e 100644 --- a/airflow/providers/google/cloud/hooks/tasks.py +++ b/airflow/providers/google/cloud/hooks/tasks.py @@ -19,18 +19,20 @@ from __future__ import annotations -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.tasks_v2 import CloudTasksClient from google.cloud.tasks_v2.types import Queue, Task -from google.protobuf.field_mask_pb2 import FieldMask from airflow.exceptions import AirflowException from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.retry import Retry + from google.protobuf.field_mask_pb2 import FieldMask + class CloudTasksHook(GoogleBaseHook): """ diff --git a/airflow/providers/google/cloud/hooks/text_to_speech.py b/airflow/providers/google/cloud/hooks/text_to_speech.py index 9adf4bb0695b..cb56b3c55743 100644 --- a/airflow/providers/google/cloud/hooks/text_to_speech.py +++ b/airflow/providers/google/cloud/hooks/text_to_speech.py @@ -18,10 +18,9 @@ """This module contains a Google Cloud Text to Speech Hook.""" from __future__ import annotations -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.texttospeech_v1 import TextToSpeechClient from google.cloud.texttospeech_v1.types import ( AudioConfig, @@ -33,6 +32,9 @@ from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.retry import Retry + class CloudTextToSpeechHook(GoogleBaseHook): """ diff --git a/airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py b/airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py index f4446d98310e..16e6b1c87c33 100644 --- a/airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py +++ b/airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py @@ -45,12 +45,10 @@ from __future__ import annotations import warnings -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core.client_options import ClientOptions from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.operation import Operation -from google.api_core.retry import Retry from google.cloud.aiplatform import ( AutoMLForecastingTrainingJob, AutoMLImageTrainingJob, @@ -61,12 +59,16 @@ models, ) from google.cloud.aiplatform_v1 import JobServiceClient, PipelineServiceClient -from google.cloud.aiplatform_v1.services.pipeline_service.pagers import ListTrainingPipelinesPager -from google.cloud.aiplatform_v1.types import TrainingPipeline from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning from airflow.providers.google.common.hooks.base_google import GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.operation import Operation + from google.api_core.retry import Retry + from google.cloud.aiplatform_v1.services.pipeline_service.pagers import ListTrainingPipelinesPager + from google.cloud.aiplatform_v1.types import TrainingPipeline + class AutoMLHook(GoogleBaseHook): """Hook for Google Cloud Vertex AI Auto ML APIs.""" diff --git a/airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py b/airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py index 249001989fd1..31552451669e 100644 --- a/airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py +++ b/airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py @@ -26,19 +26,21 @@ """ from __future__ import annotations -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core.client_options import ClientOptions from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.operation import Operation -from google.api_core.retry import Retry from google.cloud.aiplatform import BatchPredictionJob, Model, explain from google.cloud.aiplatform_v1 import JobServiceClient -from google.cloud.aiplatform_v1.services.job_service.pagers import ListBatchPredictionJobsPager from airflow import AirflowException from airflow.providers.google.common.hooks.base_google import GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.operation import Operation + from google.api_core.retry import Retry + from google.cloud.aiplatform_v1.services.job_service.pagers import ListBatchPredictionJobsPager + class BatchPredictionJobHook(GoogleBaseHook): """Hook for Google Cloud Vertex AI Batch Prediction Job APIs.""" diff --git a/airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py b/airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py index b5f43d643f1e..8e3ccf228610 100644 --- a/airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py +++ b/airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py @@ -18,12 +18,10 @@ """This module contains a Google Cloud Vertex AI hook.""" from __future__ import annotations -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core.client_options import ClientOptions from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.operation import Operation -from google.api_core.retry import Retry from google.cloud.aiplatform import ( CustomContainerTrainingJob, CustomPythonPackageTrainingJob, @@ -32,17 +30,21 @@ models, ) from google.cloud.aiplatform_v1 import JobServiceClient, PipelineServiceClient -from google.cloud.aiplatform_v1.services.job_service.pagers import ListCustomJobsPager -from google.cloud.aiplatform_v1.services.pipeline_service.pagers import ( - ListPipelineJobsPager, - ListTrainingPipelinesPager, -) -from google.cloud.aiplatform_v1.types import CustomJob, PipelineJob, TrainingPipeline from airflow import AirflowException from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.operation import Operation + from google.api_core.retry import Retry + from google.cloud.aiplatform_v1.services.job_service.pagers import ListCustomJobsPager + from google.cloud.aiplatform_v1.services.pipeline_service.pagers import ( + ListPipelineJobsPager, + ListTrainingPipelinesPager, + ) + from google.cloud.aiplatform_v1.types import CustomJob, PipelineJob, TrainingPipeline + class CustomJobHook(GoogleBaseHook): """Hook for Google Cloud Vertex AI Custom Job APIs.""" diff --git a/airflow/providers/google/cloud/hooks/vertex_ai/dataset.py b/airflow/providers/google/cloud/hooks/vertex_ai/dataset.py index 9b69d1edd9a6..d195bdce5cf2 100644 --- a/airflow/providers/google/cloud/hooks/vertex_ai/dataset.py +++ b/airflow/providers/google/cloud/hooks/vertex_ai/dataset.py @@ -18,25 +18,27 @@ """This module contains a Google Cloud Vertex AI hook.""" from __future__ import annotations -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core.client_options import ClientOptions from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.operation import Operation -from google.api_core.retry import Retry from google.cloud.aiplatform_v1 import DatasetServiceClient -from google.cloud.aiplatform_v1.services.dataset_service.pagers import ( - ListAnnotationsPager, - ListDataItemsPager, - ListDatasetsPager, -) -from google.cloud.aiplatform_v1.types import AnnotationSpec, Dataset, ExportDataConfig, ImportDataConfig -from google.protobuf.field_mask_pb2 import FieldMask from airflow import AirflowException from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.operation import Operation + from google.api_core.retry import Retry + from google.cloud.aiplatform_v1.services.dataset_service.pagers import ( + ListAnnotationsPager, + ListDataItemsPager, + ListDatasetsPager, + ) + from google.cloud.aiplatform_v1.types import AnnotationSpec, Dataset, ExportDataConfig, ImportDataConfig + from google.protobuf.field_mask_pb2 import FieldMask + class DatasetHook(GoogleBaseHook): """Hook for Google Cloud Vertex AI Dataset APIs.""" diff --git a/airflow/providers/google/cloud/hooks/vertex_ai/endpoint_service.py b/airflow/providers/google/cloud/hooks/vertex_ai/endpoint_service.py index 3c1fb9b3d019..31204490b3b1 100644 --- a/airflow/providers/google/cloud/hooks/vertex_ai/endpoint_service.py +++ b/airflow/providers/google/cloud/hooks/vertex_ai/endpoint_service.py @@ -28,20 +28,22 @@ """ from __future__ import annotations -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core.client_options import ClientOptions from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.operation import Operation -from google.api_core.retry import Retry from google.cloud.aiplatform_v1 import EndpointServiceClient -from google.cloud.aiplatform_v1.services.endpoint_service.pagers import ListEndpointsPager -from google.cloud.aiplatform_v1.types import DeployedModel, Endpoint -from google.protobuf.field_mask_pb2 import FieldMask from airflow import AirflowException from airflow.providers.google.common.hooks.base_google import GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.operation import Operation + from google.api_core.retry import Retry + from google.cloud.aiplatform_v1.services.endpoint_service.pagers import ListEndpointsPager + from google.cloud.aiplatform_v1.types import DeployedModel, Endpoint + from google.protobuf.field_mask_pb2 import FieldMask + class EndpointServiceHook(GoogleBaseHook): """Hook for Google Cloud Vertex AI Endpoint Service APIs.""" diff --git a/airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py b/airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py index 35b7991e5f47..6ca741c7b217 100644 --- a/airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py +++ b/airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py @@ -27,19 +27,21 @@ """ from __future__ import annotations -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core.client_options import ClientOptions from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.operation import Operation -from google.api_core.retry import Retry from google.cloud.aiplatform import CustomJob, HyperparameterTuningJob, gapic, hyperparameter_tuning from google.cloud.aiplatform_v1 import JobServiceClient, types -from google.cloud.aiplatform_v1.services.job_service.pagers import ListHyperparameterTuningJobsPager from airflow import AirflowException from airflow.providers.google.common.hooks.base_google import GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.operation import Operation + from google.api_core.retry import Retry + from google.cloud.aiplatform_v1.services.job_service.pagers import ListHyperparameterTuningJobsPager + class HyperparameterTuningJobHook(GoogleBaseHook): """Hook for Google Cloud Vertex AI Hyperparameter Tuning Job APIs.""" diff --git a/airflow/providers/google/cloud/hooks/vertex_ai/model_service.py b/airflow/providers/google/cloud/hooks/vertex_ai/model_service.py index db550bc9b9f3..f12960eef4f1 100644 --- a/airflow/providers/google/cloud/hooks/vertex_ai/model_service.py +++ b/airflow/providers/google/cloud/hooks/vertex_ai/model_service.py @@ -24,19 +24,21 @@ """ from __future__ import annotations -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core.client_options import ClientOptions from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.operation import Operation -from google.api_core.retry import Retry from google.cloud.aiplatform_v1 import ModelServiceClient -from google.cloud.aiplatform_v1.services.model_service.pagers import ListModelsPager -from google.cloud.aiplatform_v1.types import Model, model_service from airflow import AirflowException from airflow.providers.google.common.hooks.base_google import GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.operation import Operation + from google.api_core.retry import Retry + from google.cloud.aiplatform_v1.services.model_service.pagers import ListModelsPager + from google.cloud.aiplatform_v1.types import Model, model_service + class ModelServiceHook(GoogleBaseHook): """Hook for Google Cloud Vertex AI Endpoint Service APIs.""" diff --git a/airflow/providers/google/cloud/hooks/video_intelligence.py b/airflow/providers/google/cloud/hooks/video_intelligence.py index 498c04cc91b1..a7874edae550 100644 --- a/airflow/providers/google/cloud/hooks/video_intelligence.py +++ b/airflow/providers/google/cloud/hooks/video_intelligence.py @@ -18,11 +18,9 @@ """This module contains a Google Cloud Video Intelligence Hook.""" from __future__ import annotations -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.operation import Operation -from google.api_core.retry import Retry from google.cloud.videointelligence_v1 import ( Feature, VideoContext, @@ -32,6 +30,10 @@ from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.operation import Operation + from google.api_core.retry import Retry + class CloudVideoIntelligenceHook(GoogleBaseHook): """ diff --git a/airflow/providers/google/cloud/hooks/vision.py b/airflow/providers/google/cloud/hooks/vision.py index a6e504ab4fcb..4e17d5054d40 100644 --- a/airflow/providers/google/cloud/hooks/vision.py +++ b/airflow/providers/google/cloud/hooks/vision.py @@ -20,10 +20,9 @@ from copy import deepcopy from functools import cached_property -from typing import Any, Callable, Sequence +from typing import TYPE_CHECKING, Any, Callable, Sequence from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.vision_v1 import ( AnnotateImageRequest, Image, @@ -33,13 +32,16 @@ ProductSet, ReferenceImage, ) -from google.protobuf import field_mask_pb2 from google.protobuf.json_format import MessageToDict from airflow.exceptions import AirflowException from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.retry import Retry + from google.protobuf import field_mask_pb2 + ERR_DIFF_NAMES = """The {label} name provided in the object ({explicit_name}) is different than the name created from the input parameters ({constructed_name}). Please either: 1) Remove the {label} name, diff --git a/airflow/providers/google/cloud/hooks/workflows.py b/airflow/providers/google/cloud/hooks/workflows.py index 6e3942e3b965..9e87e79c7891 100644 --- a/airflow/providers/google/cloud/hooks/workflows.py +++ b/airflow/providers/google/cloud/hooks/workflows.py @@ -16,20 +16,22 @@ # under the License. from __future__ import annotations -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.operation import Operation -from google.api_core.retry import Retry from google.cloud.workflows.executions_v1beta import Execution, ExecutionsClient -from google.cloud.workflows.executions_v1beta.services.executions.pagers import ListExecutionsPager from google.cloud.workflows_v1beta import Workflow, WorkflowsClient -from google.cloud.workflows_v1beta.services.workflows.pagers import ListWorkflowsPager -from google.protobuf.field_mask_pb2 import FieldMask from airflow.providers.google.common.consts import CLIENT_INFO from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID, GoogleBaseHook +if TYPE_CHECKING: + from google.api_core.operation import Operation + from google.api_core.retry import Retry + from google.cloud.workflows.executions_v1beta.services.executions.pagers import ListExecutionsPager + from google.cloud.workflows_v1beta.services.workflows.pagers import ListWorkflowsPager + from google.protobuf.field_mask_pb2 import FieldMask + class WorkflowsHook(GoogleBaseHook): """ diff --git a/airflow/providers/google/cloud/links/bigquery.py b/airflow/providers/google/cloud/links/bigquery.py index 0c22cfb9a9b4..2e6eb8b24cb4 100644 --- a/airflow/providers/google/cloud/links/bigquery.py +++ b/airflow/providers/google/cloud/links/bigquery.py @@ -20,10 +20,10 @@ from typing import TYPE_CHECKING -from airflow.models import BaseOperator from airflow.providers.google.cloud.links.base import BaseGoogleLink if TYPE_CHECKING: + from airflow.models import BaseOperator from airflow.utils.context import Context BIGQUERY_BASE_LINK = "/bigquery" diff --git a/airflow/providers/google/cloud/links/bigquery_dts.py b/airflow/providers/google/cloud/links/bigquery_dts.py index 134982ea9cad..4f8a5b32a5f8 100644 --- a/airflow/providers/google/cloud/links/bigquery_dts.py +++ b/airflow/providers/google/cloud/links/bigquery_dts.py @@ -20,10 +20,10 @@ from typing import TYPE_CHECKING -from airflow.models import BaseOperator from airflow.providers.google.cloud.links.base import BaseGoogleLink if TYPE_CHECKING: + from airflow.models import BaseOperator from airflow.utils.context import Context BIGQUERY_BASE_LINK = "/bigquery/transfers" diff --git a/airflow/providers/google/cloud/links/cloud_functions.py b/airflow/providers/google/cloud/links/cloud_functions.py index 078b2b9bb514..c81aa0cdd0c4 100644 --- a/airflow/providers/google/cloud/links/cloud_functions.py +++ b/airflow/providers/google/cloud/links/cloud_functions.py @@ -21,10 +21,10 @@ from typing import TYPE_CHECKING -from airflow.models import BaseOperator from airflow.providers.google.cloud.links.base import BaseGoogleLink if TYPE_CHECKING: + from airflow.models import BaseOperator from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/links/cloud_memorystore.py b/airflow/providers/google/cloud/links/cloud_memorystore.py index 8c43312fc559..1ab253f65ebf 100644 --- a/airflow/providers/google/cloud/links/cloud_memorystore.py +++ b/airflow/providers/google/cloud/links/cloud_memorystore.py @@ -20,10 +20,10 @@ from typing import TYPE_CHECKING -from airflow.models import BaseOperator from airflow.providers.google.cloud.links.base import BaseGoogleLink if TYPE_CHECKING: + from airflow.models import BaseOperator from airflow.utils.context import Context BASE_LINK = "/memorystore" diff --git a/airflow/providers/google/cloud/links/cloud_sql.py b/airflow/providers/google/cloud/links/cloud_sql.py index 6b21b0885d5f..c3dd81b90370 100644 --- a/airflow/providers/google/cloud/links/cloud_sql.py +++ b/airflow/providers/google/cloud/links/cloud_sql.py @@ -20,10 +20,10 @@ from typing import TYPE_CHECKING -from airflow.models import BaseOperator from airflow.providers.google.cloud.links.base import BaseGoogleLink if TYPE_CHECKING: + from airflow.models import BaseOperator from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/links/cloud_tasks.py b/airflow/providers/google/cloud/links/cloud_tasks.py index 55f3979bb5b9..07bb959cf2f7 100644 --- a/airflow/providers/google/cloud/links/cloud_tasks.py +++ b/airflow/providers/google/cloud/links/cloud_tasks.py @@ -20,10 +20,10 @@ from typing import TYPE_CHECKING -from airflow.models import BaseOperator from airflow.providers.google.cloud.links.base import BaseGoogleLink if TYPE_CHECKING: + from airflow.models import BaseOperator from airflow.utils.context import Context CLOUD_TASKS_BASE_LINK = "/cloudtasks" diff --git a/airflow/providers/google/cloud/links/compute.py b/airflow/providers/google/cloud/links/compute.py index e754f7085825..3c27cd18649d 100644 --- a/airflow/providers/google/cloud/links/compute.py +++ b/airflow/providers/google/cloud/links/compute.py @@ -20,10 +20,10 @@ from typing import TYPE_CHECKING -from airflow.models import BaseOperator from airflow.providers.google.cloud.links.base import BaseGoogleLink if TYPE_CHECKING: + from airflow.models import BaseOperator from airflow.utils.context import Context COMPUTE_BASE_LINK = "https://console.cloud.google.com/compute" diff --git a/airflow/providers/google/cloud/links/datacatalog.py b/airflow/providers/google/cloud/links/datacatalog.py index 7547d6152adf..695b0a643a2b 100644 --- a/airflow/providers/google/cloud/links/datacatalog.py +++ b/airflow/providers/google/cloud/links/datacatalog.py @@ -20,10 +20,10 @@ from typing import TYPE_CHECKING -from airflow.models import BaseOperator from airflow.providers.google.cloud.links.base import BaseGoogleLink if TYPE_CHECKING: + from airflow.models import BaseOperator from airflow.utils.context import Context DATACATALOG_BASE_LINK = "/datacatalog" diff --git a/airflow/providers/google/cloud/links/dataflow.py b/airflow/providers/google/cloud/links/dataflow.py index e2422b2653c6..3c8b8e92d870 100644 --- a/airflow/providers/google/cloud/links/dataflow.py +++ b/airflow/providers/google/cloud/links/dataflow.py @@ -20,10 +20,10 @@ from typing import TYPE_CHECKING -from airflow.models import BaseOperator from airflow.providers.google.cloud.links.base import BaseGoogleLink if TYPE_CHECKING: + from airflow.models import BaseOperator from airflow.utils.context import Context DATAFLOW_BASE_LINK = "/dataflow/jobs" diff --git a/airflow/providers/google/cloud/links/dataform.py b/airflow/providers/google/cloud/links/dataform.py index 5744f22e1d5f..99b316a325e0 100644 --- a/airflow/providers/google/cloud/links/dataform.py +++ b/airflow/providers/google/cloud/links/dataform.py @@ -20,10 +20,10 @@ from typing import TYPE_CHECKING -from airflow.models import BaseOperator from airflow.providers.google.cloud.links.base import BaseGoogleLink if TYPE_CHECKING: + from airflow.models import BaseOperator from airflow.utils.context import Context DATAFORM_BASE_LINK = "/bigquery/dataform" diff --git a/airflow/providers/google/cloud/links/pubsub.py b/airflow/providers/google/cloud/links/pubsub.py index 6ec47e7ea701..a8aed4fc8e23 100644 --- a/airflow/providers/google/cloud/links/pubsub.py +++ b/airflow/providers/google/cloud/links/pubsub.py @@ -20,10 +20,10 @@ from typing import TYPE_CHECKING -from airflow.models import BaseOperator from airflow.providers.google.cloud.links.base import BaseGoogleLink if TYPE_CHECKING: + from airflow.models import BaseOperator from airflow.utils.context import Context PUBSUB_BASE_LINK = "/cloudpubsub" diff --git a/airflow/providers/google/cloud/links/spanner.py b/airflow/providers/google/cloud/links/spanner.py index c83facb14075..85c669ba556c 100644 --- a/airflow/providers/google/cloud/links/spanner.py +++ b/airflow/providers/google/cloud/links/spanner.py @@ -20,10 +20,10 @@ from typing import TYPE_CHECKING -from airflow.models import BaseOperator from airflow.providers.google.cloud.links.base import BaseGoogleLink if TYPE_CHECKING: + from airflow.models import BaseOperator from airflow.utils.context import Context SPANNER_BASE_LINK = "/spanner/instances" diff --git a/airflow/providers/google/cloud/links/stackdriver.py b/airflow/providers/google/cloud/links/stackdriver.py index 4ca8580e99a8..2c516f932f76 100644 --- a/airflow/providers/google/cloud/links/stackdriver.py +++ b/airflow/providers/google/cloud/links/stackdriver.py @@ -20,10 +20,10 @@ from typing import TYPE_CHECKING -from airflow.models import BaseOperator from airflow.providers.google.cloud.links.base import BaseGoogleLink if TYPE_CHECKING: + from airflow.models import BaseOperator from airflow.utils.context import Context STACKDRIVER_BASE_LINK = "/monitoring/alerting" diff --git a/airflow/providers/google/cloud/links/workflows.py b/airflow/providers/google/cloud/links/workflows.py index 160539579cb2..1d944a8e8f34 100644 --- a/airflow/providers/google/cloud/links/workflows.py +++ b/airflow/providers/google/cloud/links/workflows.py @@ -20,10 +20,10 @@ from typing import TYPE_CHECKING -from airflow.models import BaseOperator from airflow.providers.google.cloud.links.base import BaseGoogleLink if TYPE_CHECKING: + from airflow.models import BaseOperator from airflow.utils.context import Context WORKFLOWS_BASE_LINK = "workflows" diff --git a/airflow/providers/google/cloud/log/stackdriver_task_handler.py b/airflow/providers/google/cloud/log/stackdriver_task_handler.py index d1e95bc7b0ad..bf0e3bf4fb44 100644 --- a/airflow/providers/google/cloud/log/stackdriver_task_handler.py +++ b/airflow/providers/google/cloud/log/stackdriver_task_handler.py @@ -18,22 +18,26 @@ from __future__ import annotations import logging -from contextvars import ContextVar from functools import cached_property -from typing import Collection +from typing import TYPE_CHECKING, Collection from urllib.parse import urlencode -from google.auth.credentials import Credentials from google.cloud import logging as gcp_logging from google.cloud.logging import Resource from google.cloud.logging.handlers.transports import BackgroundThreadTransport, Transport from google.cloud.logging_v2.services.logging_service_v2 import LoggingServiceV2Client from google.cloud.logging_v2.types import ListLogEntriesRequest, ListLogEntriesResponse -from airflow.models import TaskInstance from airflow.providers.google.cloud.utils.credentials_provider import get_credentials_and_project_id from airflow.providers.google.common.consts import CLIENT_INFO +if TYPE_CHECKING: + from contextvars import ContextVar + + from google.auth.credentials import Credentials + + from airflow.models import TaskInstance + try: # todo: remove this conditional import when min airflow version >= 2.6 ctx_indiv_trigger: ContextVar | None diff --git a/airflow/providers/google/cloud/operators/automl.py b/airflow/providers/google/cloud/operators/automl.py index e590dc3d0620..6fefa3081ac5 100644 --- a/airflow/providers/google/cloud/operators/automl.py +++ b/airflow/providers/google/cloud/operators/automl.py @@ -22,7 +22,6 @@ from typing import TYPE_CHECKING, Sequence, Tuple from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.automl_v1beta1 import ( BatchPredictResult, ColumnSpec, @@ -43,6 +42,8 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator if TYPE_CHECKING: + from google.api_core.retry import Retry + from airflow.utils.context import Context MetaData = Sequence[Tuple[str, str]] diff --git a/airflow/providers/google/cloud/operators/bigquery.py b/airflow/providers/google/cloud/operators/bigquery.py index 75ae73b6f9ff..ae7e48eed797 100644 --- a/airflow/providers/google/cloud/operators/bigquery.py +++ b/airflow/providers/google/cloud/operators/bigquery.py @@ -26,7 +26,6 @@ import attr from google.api_core.exceptions import Conflict -from google.api_core.retry import Retry from google.cloud.bigquery import DEFAULT_RETRY, CopyJob, ExtractJob, LoadJob, QueryJob from google.cloud.bigquery.table import RowIterator @@ -56,6 +55,7 @@ from airflow.providers.google.cloud.utils.bigquery import convert_job_id if TYPE_CHECKING: + from google.api_core.retry import Retry from google.cloud.bigquery import UnknownJob from airflow.models.taskinstancekey import TaskInstanceKey diff --git a/airflow/providers/google/cloud/operators/bigquery_dts.py b/airflow/providers/google/cloud/operators/bigquery_dts.py index 9edae0b65978..2b942d1f4366 100644 --- a/airflow/providers/google/cloud/operators/bigquery_dts.py +++ b/airflow/providers/google/cloud/operators/bigquery_dts.py @@ -23,7 +23,6 @@ from typing import TYPE_CHECKING, Sequence from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.bigquery_datatransfer_v1 import ( StartManualTransferRunsResponse, TransferConfig, @@ -39,6 +38,8 @@ from airflow.providers.google.cloud.triggers.bigquery_dts import BigQueryDataTransferRunTrigger if TYPE_CHECKING: + from google.api_core.retry import Retry + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/bigtable.py b/airflow/providers/google/cloud/operators/bigtable.py index c6af84ddd1a4..ffdf96c36c13 100644 --- a/airflow/providers/google/cloud/operators/bigtable.py +++ b/airflow/providers/google/cloud/operators/bigtable.py @@ -18,12 +18,9 @@ """This module contains Google Cloud Bigtable operators.""" from __future__ import annotations -import enum from typing import TYPE_CHECKING, Iterable, Sequence import google.api_core.exceptions -from google.cloud.bigtable import enums -from google.cloud.bigtable.column_family import GarbageCollectionRule from airflow.exceptions import AirflowException from airflow.providers.google.cloud.hooks.bigtable import BigtableHook @@ -35,6 +32,11 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator if TYPE_CHECKING: + import enum + + from google.cloud.bigtable import enums + from google.cloud.bigtable.column_family import GarbageCollectionRule + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/cloud_batch.py b/airflow/providers/google/cloud/operators/cloud_batch.py index 26c0af06c1d9..60687f3c1da4 100644 --- a/airflow/providers/google/cloud/operators/cloud_batch.py +++ b/airflow/providers/google/cloud/operators/cloud_batch.py @@ -17,9 +17,8 @@ # under the License. from __future__ import annotations -from typing import Sequence +from typing import TYPE_CHECKING, Sequence -from google.api_core import operation # type: ignore from google.cloud.batch_v1 import Job, Task from airflow.configuration import conf @@ -27,7 +26,11 @@ from airflow.providers.google.cloud.hooks.cloud_batch import CloudBatchHook from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator from airflow.providers.google.cloud.triggers.cloud_batch import CloudBatchJobFinishedTrigger -from airflow.utils.context import Context + +if TYPE_CHECKING: + from google.api_core import operation + + from airflow.utils.context import Context class CloudBatchSubmitJobOperator(GoogleCloudBaseOperator): diff --git a/airflow/providers/google/cloud/operators/cloud_build.py b/airflow/providers/google/cloud/operators/cloud_build.py index db1f44db3715..9daacefa7213 100644 --- a/airflow/providers/google/cloud/operators/cloud_build.py +++ b/airflow/providers/google/cloud/operators/cloud_build.py @@ -25,7 +25,6 @@ from urllib.parse import unquote, urlsplit from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.devtools.cloudbuild_v1.types import Build, BuildTrigger, RepoSource from airflow.configuration import conf @@ -44,6 +43,8 @@ from airflow.utils.helpers import exactly_one if TYPE_CHECKING: + from google.api_core.retry import Retry + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/cloud_composer.py b/airflow/providers/google/cloud/operators/cloud_composer.py index c9b52d855915..3aaf58ef04b0 100644 --- a/airflow/providers/google/cloud/operators/cloud_composer.py +++ b/airflow/providers/google/cloud/operators/cloud_composer.py @@ -21,10 +21,8 @@ from google.api_core.exceptions import AlreadyExists from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.orchestration.airflow.service_v1 import ImageVersion from google.cloud.orchestration.airflow.service_v1.types import Environment -from google.protobuf.field_mask_pb2 import FieldMask from airflow import AirflowException from airflow.configuration import conf @@ -35,6 +33,9 @@ from airflow.providers.google.common.consts import GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME if TYPE_CHECKING: + from google.api_core.retry import Retry + from google.protobuf.field_mask_pb2 import FieldMask + from airflow.utils.context import Context CLOUD_COMPOSER_BASE_LINK = "https://console.cloud.google.com/composer/environments" diff --git a/airflow/providers/google/cloud/operators/cloud_memorystore.py b/airflow/providers/google/cloud/operators/cloud_memorystore.py index 4c27b8982f70..54e9ff7ef79c 100644 --- a/airflow/providers/google/cloud/operators/cloud_memorystore.py +++ b/airflow/providers/google/cloud/operators/cloud_memorystore.py @@ -28,10 +28,8 @@ from typing import TYPE_CHECKING, Sequence from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.memcache_v1beta2.types import cloud_memcache from google.cloud.redis_v1 import FailoverInstanceRequest, InputConfig, Instance, OutputConfig -from google.protobuf.field_mask_pb2 import FieldMask from airflow.providers.google.cloud.hooks.cloud_memorystore import ( CloudMemorystoreHook, @@ -46,6 +44,9 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator if TYPE_CHECKING: + from google.api_core.retry import Retry + from google.protobuf.field_mask_pb2 import FieldMask + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/cloud_run.py b/airflow/providers/google/cloud/operators/cloud_run.py index d27b17973dff..c0b1610e4bf6 100644 --- a/airflow/providers/google/cloud/operators/cloud_run.py +++ b/airflow/providers/google/cloud/operators/cloud_run.py @@ -19,9 +19,7 @@ from typing import TYPE_CHECKING, Sequence -from google.api_core import operation from google.cloud.run_v2 import Job -from google.cloud.run_v2.types import Execution from airflow.configuration import conf from airflow.exceptions import AirflowException @@ -30,6 +28,9 @@ from airflow.providers.google.cloud.triggers.cloud_run import CloudRunJobFinishedTrigger, RunJobStatus if TYPE_CHECKING: + from google.api_core import operation + from google.cloud.run_v2.types import Execution + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/cloud_sql.py b/airflow/providers/google/cloud/operators/cloud_sql.py index 3ddf366217c6..c371e7f11221 100644 --- a/airflow/providers/google/cloud/operators/cloud_sql.py +++ b/airflow/providers/google/cloud/operators/cloud_sql.py @@ -25,7 +25,6 @@ from airflow.configuration import conf from airflow.exceptions import AirflowException from airflow.hooks.base import BaseHook -from airflow.models import Connection from airflow.providers.google.cloud.hooks.cloud_sql import CloudSQLDatabaseHook, CloudSQLHook from airflow.providers.google.cloud.links.cloud_sql import CloudSQLInstanceDatabaseLink, CloudSQLInstanceLink from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator @@ -33,10 +32,11 @@ from airflow.providers.google.cloud.utils.field_validator import GcpBodyFieldValidator from airflow.providers.google.common.hooks.base_google import get_field from airflow.providers.google.common.links.storage import FileDetailsLink -from airflow.providers.mysql.hooks.mysql import MySqlHook -from airflow.providers.postgres.hooks.postgres import PostgresHook if TYPE_CHECKING: + from airflow.models import Connection + from airflow.providers.mysql.hooks.mysql import MySqlHook + from airflow.providers.postgres.hooks.postgres import PostgresHook from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/compute.py b/airflow/providers/google/cloud/operators/compute.py index 4386798e9e61..2abc46d35a03 100644 --- a/airflow/providers/google/cloud/operators/compute.py +++ b/airflow/providers/google/cloud/operators/compute.py @@ -22,7 +22,6 @@ from typing import TYPE_CHECKING, Any, Sequence from google.api_core import exceptions -from google.api_core.retry import Retry from google.cloud.compute_v1.types import Instance, InstanceGroupManager, InstanceTemplate from json_merge_patch import merge @@ -38,6 +37,8 @@ from airflow.providers.google.cloud.utils.field_validator import GcpBodyFieldValidator if TYPE_CHECKING: + from google.api_core.retry import Retry + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/datacatalog.py b/airflow/providers/google/cloud/operators/datacatalog.py index c96a668ad372..f3ec08441d16 100644 --- a/airflow/providers/google/cloud/operators/datacatalog.py +++ b/airflow/providers/google/cloud/operators/datacatalog.py @@ -20,7 +20,6 @@ from google.api_core.exceptions import AlreadyExists, NotFound from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.datacatalog import ( DataCatalogClient, Entry, @@ -31,7 +30,6 @@ TagTemplate, TagTemplateField, ) -from google.protobuf.field_mask_pb2 import FieldMask from airflow.providers.google.cloud.hooks.datacatalog import CloudDataCatalogHook from airflow.providers.google.cloud.links.datacatalog import ( @@ -42,6 +40,9 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator if TYPE_CHECKING: + from google.api_core.retry import Retry + from google.protobuf.field_mask_pb2 import FieldMask + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/dataform.py b/airflow/providers/google/cloud/operators/dataform.py index aeb6e05c0fbc..c035ccaaecd0 100644 --- a/airflow/providers/google/cloud/operators/dataform.py +++ b/airflow/providers/google/cloud/operators/dataform.py @@ -25,10 +25,11 @@ ) if TYPE_CHECKING: + from google.api_core.retry import Retry + from airflow.utils.context import Context from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.dataform_v1beta1.types import ( CompilationResult, InstallNpmPackagesResponse, diff --git a/airflow/providers/google/cloud/operators/dataplex.py b/airflow/providers/google/cloud/operators/dataplex.py index 2b27cb70dba9..ffbb75a7b608 100644 --- a/airflow/providers/google/cloud/operators/dataplex.py +++ b/airflow/providers/google/cloud/operators/dataplex.py @@ -24,13 +24,14 @@ from airflow import AirflowException if TYPE_CHECKING: + from google.protobuf.field_mask_pb2 import FieldMask + from airflow.utils.context import Context from google.api_core.exceptions import AlreadyExists, GoogleAPICallError from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault from google.api_core.retry import Retry, exponential_sleep_generator from google.cloud.dataplex_v1.types import Asset, DataScan, DataScanJob, Lake, Task, Zone -from google.protobuf.field_mask_pb2 import FieldMask from googleapiclient.errors import HttpError from airflow.providers.google.cloud.hooks.dataplex import AirflowDataQualityScanException, DataplexHook diff --git a/airflow/providers/google/cloud/operators/dataproc.py b/airflow/providers/google/cloud/operators/dataproc.py index eaa65ba7c15d..aee0b6fa7226 100644 --- a/airflow/providers/google/cloud/operators/dataproc.py +++ b/airflow/providers/google/cloud/operators/dataproc.py @@ -29,13 +29,10 @@ from enum import Enum from typing import TYPE_CHECKING, Any, Sequence -from google.api_core import operation # type: ignore from google.api_core.exceptions import AlreadyExists, NotFound from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault from google.api_core.retry import Retry, exponential_sleep_generator from google.cloud.dataproc_v1 import Batch, Cluster, ClusterStatus, JobStatus -from google.protobuf.duration_pb2 import Duration -from google.protobuf.field_mask_pb2 import FieldMask from airflow.configuration import conf from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning @@ -64,6 +61,10 @@ from airflow.utils import timezone if TYPE_CHECKING: + from google.api_core import operation + from google.protobuf.duration_pb2 import Duration + from google.protobuf.field_mask_pb2 import FieldMask + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/dataproc_metastore.py b/airflow/providers/google/cloud/operators/dataproc_metastore.py index 25ef15318f85..ff802b6801a2 100644 --- a/airflow/providers/google/cloud/operators/dataproc_metastore.py +++ b/airflow/providers/google/cloud/operators/dataproc_metastore.py @@ -26,7 +26,6 @@ from google.cloud.metastore_v1 import MetadataExport, MetadataManagementActivity from google.cloud.metastore_v1.types import Backup, MetadataImport, Service from google.cloud.metastore_v1.types.metastore import DatabaseDumpSpec, Restore -from google.protobuf.field_mask_pb2 import FieldMask from googleapiclient.errors import HttpError from airflow import AirflowException @@ -37,6 +36,8 @@ from airflow.providers.google.common.links.storage import StorageLink if TYPE_CHECKING: + from google.protobuf.field_mask_pb2 import FieldMask + from airflow.models.taskinstancekey import TaskInstanceKey from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/dlp.py b/airflow/providers/google/cloud/operators/dlp.py index 60b52a4751c8..db705b5cd504 100644 --- a/airflow/providers/google/cloud/operators/dlp.py +++ b/airflow/providers/google/cloud/operators/dlp.py @@ -23,7 +23,6 @@ from google.api_core.exceptions import AlreadyExists, InvalidArgument, NotFound from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.dlp_v2.types import ( ByteContentItem, ContentItem, @@ -44,7 +43,6 @@ StoredInfoType, StoredInfoTypeConfig, ) -from google.protobuf.field_mask_pb2 import FieldMask from airflow.providers.google.cloud.hooks.dlp import CloudDLPHook from airflow.providers.google.cloud.links.data_loss_prevention import ( @@ -63,6 +61,9 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator if TYPE_CHECKING: + from google.api_core.retry import Retry + from google.protobuf.field_mask_pb2 import FieldMask + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/kubernetes_engine.py b/airflow/providers/google/cloud/operators/kubernetes_engine.py index 642c4bf27955..4f1d68c29be4 100644 --- a/airflow/providers/google/cloud/operators/kubernetes_engine.py +++ b/airflow/providers/google/cloud/operators/kubernetes_engine.py @@ -24,7 +24,6 @@ from google.api_core.exceptions import AlreadyExists from google.cloud.container_v1.types import Cluster -from kubernetes.client.models import V1Pod from airflow.configuration import conf from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning @@ -46,6 +45,8 @@ from airflow.utils.timezone import utcnow if TYPE_CHECKING: + from kubernetes.client.models import V1Pod + from airflow.utils.context import Context KUBE_CONFIG_ENV_VAR = "KUBECONFIG" diff --git a/airflow/providers/google/cloud/operators/natural_language.py b/airflow/providers/google/cloud/operators/natural_language.py index 21ddb01f8a89..0389350acc6d 100644 --- a/airflow/providers/google/cloud/operators/natural_language.py +++ b/airflow/providers/google/cloud/operators/natural_language.py @@ -21,14 +21,15 @@ from typing import TYPE_CHECKING, Sequence, Tuple from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry -from google.cloud.language_v1.types import Document, EncodingType from google.protobuf.json_format import MessageToDict from airflow.providers.google.cloud.hooks.natural_language import CloudNaturalLanguageHook from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator if TYPE_CHECKING: + from google.api_core.retry import Retry + from google.cloud.language_v1.types import Document, EncodingType + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/pubsub.py b/airflow/providers/google/cloud/operators/pubsub.py index f91e9ea8c18d..530f232af84f 100644 --- a/airflow/providers/google/cloud/operators/pubsub.py +++ b/airflow/providers/google/cloud/operators/pubsub.py @@ -27,7 +27,6 @@ from typing import TYPE_CHECKING, Any, Callable, Sequence from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.pubsub_v1.types import ( DeadLetterPolicy, Duration, @@ -43,6 +42,8 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator if TYPE_CHECKING: + from google.api_core.retry import Retry + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/speech_to_text.py b/airflow/providers/google/cloud/operators/speech_to_text.py index 8d4d92b30b50..867ad0b1dde8 100644 --- a/airflow/providers/google/cloud/operators/speech_to_text.py +++ b/airflow/providers/google/cloud/operators/speech_to_text.py @@ -21,8 +21,6 @@ from typing import TYPE_CHECKING, Sequence from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry -from google.cloud.speech_v1.types import RecognitionConfig from google.protobuf.json_format import MessageToDict from airflow.exceptions import AirflowException @@ -31,6 +29,9 @@ from airflow.providers.google.common.links.storage import FileDetailsLink if TYPE_CHECKING: + from google.api_core.retry import Retry + from google.cloud.speech_v1.types import RecognitionConfig + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/stackdriver.py b/airflow/providers/google/cloud/operators/stackdriver.py index aa2b3d368c67..66fed54ebe00 100644 --- a/airflow/providers/google/cloud/operators/stackdriver.py +++ b/airflow/providers/google/cloud/operators/stackdriver.py @@ -20,7 +20,6 @@ from typing import TYPE_CHECKING, Sequence from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.monitoring_v3 import AlertPolicy, NotificationChannel from airflow.providers.google.cloud.hooks.stackdriver import StackdriverHook @@ -31,6 +30,8 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator if TYPE_CHECKING: + from google.api_core.retry import Retry + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/tasks.py b/airflow/providers/google/cloud/operators/tasks.py index f78399d76673..921610d2956c 100644 --- a/airflow/providers/google/cloud/operators/tasks.py +++ b/airflow/providers/google/cloud/operators/tasks.py @@ -23,15 +23,16 @@ from google.api_core.exceptions import AlreadyExists from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.tasks_v2.types import Queue, Task -from google.protobuf.field_mask_pb2 import FieldMask from airflow.providers.google.cloud.hooks.tasks import CloudTasksHook from airflow.providers.google.cloud.links.cloud_tasks import CloudTasksLink, CloudTasksQueueLink from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator if TYPE_CHECKING: + from google.api_core.retry import Retry + from google.protobuf.field_mask_pb2 import FieldMask + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/text_to_speech.py b/airflow/providers/google/cloud/operators/text_to_speech.py index 90ca92576392..398b4950e5bc 100644 --- a/airflow/providers/google/cloud/operators/text_to_speech.py +++ b/airflow/providers/google/cloud/operators/text_to_speech.py @@ -22,8 +22,6 @@ from typing import TYPE_CHECKING, Sequence from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry -from google.cloud.texttospeech_v1.types import AudioConfig, SynthesisInput, VoiceSelectionParams from airflow.exceptions import AirflowException from airflow.providers.google.cloud.hooks.gcs import GCSHook @@ -32,6 +30,9 @@ from airflow.providers.google.common.links.storage import FileDetailsLink if TYPE_CHECKING: + from google.api_core.retry import Retry + from google.cloud.texttospeech_v1.types import AudioConfig, SynthesisInput, VoiceSelectionParams + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/translate_speech.py b/airflow/providers/google/cloud/operators/translate_speech.py index 9295d77f1ddc..1edc9b6b4dad 100644 --- a/airflow/providers/google/cloud/operators/translate_speech.py +++ b/airflow/providers/google/cloud/operators/translate_speech.py @@ -20,7 +20,6 @@ from typing import TYPE_CHECKING, Sequence -from google.cloud.speech_v1.types import RecognitionAudio, RecognitionConfig from google.protobuf.json_format import MessageToDict from airflow.exceptions import AirflowException @@ -30,6 +29,8 @@ from airflow.providers.google.common.links.storage import FileDetailsLink if TYPE_CHECKING: + from google.cloud.speech_v1.types import RecognitionAudio, RecognitionConfig + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py b/airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py index 88a8560e85dc..d9b042669510 100644 --- a/airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py +++ b/airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py @@ -22,7 +22,6 @@ from google.api_core.exceptions import NotFound from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.aiplatform import datasets from google.cloud.aiplatform.models import Model from google.cloud.aiplatform_v1.types.training_pipeline import TrainingPipeline @@ -36,6 +35,8 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator if TYPE_CHECKING: + from google.api_core.retry import Retry + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/vertex_ai/batch_prediction_job.py b/airflow/providers/google/cloud/operators/vertex_ai/batch_prediction_job.py index 684cfdd8f2e8..f6d9229dfe8e 100644 --- a/airflow/providers/google/cloud/operators/vertex_ai/batch_prediction_job.py +++ b/airflow/providers/google/cloud/operators/vertex_ai/batch_prediction_job.py @@ -30,8 +30,6 @@ from google.api_core.exceptions import NotFound from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry -from google.cloud.aiplatform import Model, explain from google.cloud.aiplatform_v1.types import BatchPredictionJob from airflow.providers.google.cloud.hooks.vertex_ai.batch_prediction_job import BatchPredictionJobHook @@ -42,6 +40,9 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator if TYPE_CHECKING: + from google.api_core.retry import Retry + from google.cloud.aiplatform import Model, explain + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py b/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py index 9896cdcccd69..87611c10a658 100644 --- a/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py +++ b/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py @@ -22,7 +22,6 @@ from google.api_core.exceptions import NotFound from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.aiplatform.models import Model from google.cloud.aiplatform_v1.types.dataset import Dataset from google.cloud.aiplatform_v1.types.training_pipeline import TrainingPipeline @@ -36,6 +35,8 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator if TYPE_CHECKING: + from google.api_core.retry import Retry + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/vertex_ai/dataset.py b/airflow/providers/google/cloud/operators/vertex_ai/dataset.py index 222cc659727f..900d926472de 100644 --- a/airflow/providers/google/cloud/operators/vertex_ai/dataset.py +++ b/airflow/providers/google/cloud/operators/vertex_ai/dataset.py @@ -22,15 +22,16 @@ from google.api_core.exceptions import NotFound from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.aiplatform_v1.types import Dataset, ExportDataConfig, ImportDataConfig -from google.protobuf.field_mask_pb2 import FieldMask from airflow.providers.google.cloud.hooks.vertex_ai.dataset import DatasetHook from airflow.providers.google.cloud.links.vertex_ai import VertexAIDatasetLink, VertexAIDatasetListLink from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator if TYPE_CHECKING: + from google.api_core.retry import Retry + from google.protobuf.field_mask_pb2 import FieldMask + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/vertex_ai/endpoint_service.py b/airflow/providers/google/cloud/operators/vertex_ai/endpoint_service.py index b2bf7c421339..eb8eec39363a 100644 --- a/airflow/providers/google/cloud/operators/vertex_ai/endpoint_service.py +++ b/airflow/providers/google/cloud/operators/vertex_ai/endpoint_service.py @@ -32,9 +32,7 @@ from google.api_core.exceptions import NotFound from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.aiplatform_v1.types import DeployedModel, Endpoint, endpoint_service -from google.protobuf.field_mask_pb2 import FieldMask from airflow.providers.google.cloud.hooks.vertex_ai.endpoint_service import EndpointServiceHook from airflow.providers.google.cloud.links.vertex_ai import ( @@ -45,6 +43,9 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator if TYPE_CHECKING: + from google.api_core.retry import Retry + from google.protobuf.field_mask_pb2 import FieldMask + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/vertex_ai/hyperparameter_tuning_job.py b/airflow/providers/google/cloud/operators/vertex_ai/hyperparameter_tuning_job.py index 69c7e76016d3..70cbddc701f5 100644 --- a/airflow/providers/google/cloud/operators/vertex_ai/hyperparameter_tuning_job.py +++ b/airflow/providers/google/cloud/operators/vertex_ai/hyperparameter_tuning_job.py @@ -31,8 +31,6 @@ from google.api_core.exceptions import NotFound from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry -from google.cloud.aiplatform import gapic, hyperparameter_tuning from google.cloud.aiplatform_v1.types import HyperparameterTuningJob from airflow.providers.google.cloud.hooks.vertex_ai.hyperparameter_tuning_job import ( @@ -45,6 +43,9 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator if TYPE_CHECKING: + from google.api_core.retry import Retry + from google.cloud.aiplatform import gapic, hyperparameter_tuning + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/vertex_ai/model_service.py b/airflow/providers/google/cloud/operators/vertex_ai/model_service.py index 0e179e86ef07..379557fa8e6f 100644 --- a/airflow/providers/google/cloud/operators/vertex_ai/model_service.py +++ b/airflow/providers/google/cloud/operators/vertex_ai/model_service.py @@ -28,7 +28,6 @@ from google.api_core.exceptions import NotFound from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.aiplatform_v1.types import Model, model_service from airflow.providers.google.cloud.hooks.vertex_ai.model_service import ModelServiceHook @@ -40,6 +39,8 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator if TYPE_CHECKING: + from google.api_core.retry import Retry + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/video_intelligence.py b/airflow/providers/google/cloud/operators/video_intelligence.py index 2069513d063d..8bf696fa2f0d 100644 --- a/airflow/providers/google/cloud/operators/video_intelligence.py +++ b/airflow/providers/google/cloud/operators/video_intelligence.py @@ -21,7 +21,6 @@ from typing import TYPE_CHECKING, Sequence from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.videointelligence_v1 import Feature, VideoContext from google.protobuf.json_format import MessageToDict @@ -29,6 +28,8 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator if TYPE_CHECKING: + from google.api_core.retry import Retry + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/vision.py b/airflow/providers/google/cloud/operators/vision.py index f774b693cd6d..7abada3928ee 100644 --- a/airflow/providers/google/cloud/operators/vision.py +++ b/airflow/providers/google/cloud/operators/vision.py @@ -23,7 +23,6 @@ from google.api_core.exceptions import AlreadyExists from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.vision_v1 import ( AnnotateImageRequest, Image, @@ -31,12 +30,14 @@ ProductSet, ReferenceImage, ) -from google.protobuf.field_mask_pb2 import FieldMask # type: ignore from airflow.providers.google.cloud.hooks.vision import CloudVisionHook from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator if TYPE_CHECKING: + from google.api_core.retry import Retry + from google.protobuf.field_mask_pb2 import FieldMask + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/operators/workflows.py b/airflow/providers/google/cloud/operators/workflows.py index 2887bd135846..0bfefe5df458 100644 --- a/airflow/providers/google/cloud/operators/workflows.py +++ b/airflow/providers/google/cloud/operators/workflows.py @@ -24,10 +24,8 @@ from google.api_core.exceptions import AlreadyExists from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.workflows.executions_v1beta import Execution from google.cloud.workflows_v1beta import Workflow -from google.protobuf.field_mask_pb2 import FieldMask from airflow.providers.google.cloud.hooks.workflows import WorkflowsHook from airflow.providers.google.cloud.links.workflows import ( @@ -38,6 +36,9 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator if TYPE_CHECKING: + from google.api_core.retry import Retry + from google.protobuf.field_mask_pb2 import FieldMask + from airflow.utils.context import Context try: from airflow.utils.hashlib_wrapper import md5 diff --git a/airflow/providers/google/cloud/sensors/bigquery_dts.py b/airflow/providers/google/cloud/sensors/bigquery_dts.py index 9dc2623783c1..34198d2819be 100644 --- a/airflow/providers/google/cloud/sensors/bigquery_dts.py +++ b/airflow/providers/google/cloud/sensors/bigquery_dts.py @@ -21,7 +21,6 @@ from typing import TYPE_CHECKING, Sequence from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.bigquery_datatransfer_v1 import TransferState from airflow.exceptions import AirflowException @@ -29,6 +28,8 @@ from airflow.sensors.base import BaseSensorOperator if TYPE_CHECKING: + from google.api_core.retry import Retry + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/sensors/dataplex.py b/airflow/providers/google/cloud/sensors/dataplex.py index e1581fe1dac9..c00373f9476a 100644 --- a/airflow/providers/google/cloud/sensors/dataplex.py +++ b/airflow/providers/google/cloud/sensors/dataplex.py @@ -21,10 +21,11 @@ from typing import TYPE_CHECKING, Sequence if TYPE_CHECKING: + from google.api_core.retry import Retry + from airflow.utils.context import Context from google.api_core.exceptions import GoogleAPICallError from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.dataplex_v1.types import DataScanJob from airflow.exceptions import AirflowException diff --git a/airflow/providers/google/cloud/sensors/dataproc_metastore.py b/airflow/providers/google/cloud/sensors/dataproc_metastore.py index fb895abba33d..4a03ac6d06e6 100644 --- a/airflow/providers/google/cloud/sensors/dataproc_metastore.py +++ b/airflow/providers/google/cloud/sensors/dataproc_metastore.py @@ -19,14 +19,14 @@ from typing import TYPE_CHECKING, Sequence -from google.api_core.operation import Operation - from airflow import AirflowException from airflow.providers.google.cloud.hooks.dataproc_metastore import DataprocMetastoreHook from airflow.providers.google.cloud.hooks.gcs import parse_json_from_gcs from airflow.sensors.base import BaseSensorOperator if TYPE_CHECKING: + from google.api_core.operation import Operation + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/sensors/gcs.py b/airflow/providers/google/cloud/sensors/gcs.py index e427d1c6bef0..fe3572573a90 100644 --- a/airflow/providers/google/cloud/sensors/gcs.py +++ b/airflow/providers/google/cloud/sensors/gcs.py @@ -24,7 +24,6 @@ from datetime import datetime, timedelta from typing import TYPE_CHECKING, Any, Callable, Sequence -from google.api_core.retry import Retry from google.cloud.storage.retry import DEFAULT_RETRY from airflow.configuration import conf @@ -39,6 +38,8 @@ from airflow.sensors.base import BaseSensorOperator, poke_mode_only if TYPE_CHECKING: + from google.api_core.retry import Retry + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/sensors/workflows.py b/airflow/providers/google/cloud/sensors/workflows.py index e352fcc4a05e..712e328bdd67 100644 --- a/airflow/providers/google/cloud/sensors/workflows.py +++ b/airflow/providers/google/cloud/sensors/workflows.py @@ -19,7 +19,6 @@ from typing import TYPE_CHECKING, Sequence from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault -from google.api_core.retry import Retry from google.cloud.workflows.executions_v1beta import Execution from airflow.exceptions import AirflowException @@ -27,6 +26,8 @@ from airflow.sensors.base import BaseSensorOperator if TYPE_CHECKING: + from google.api_core.retry import Retry + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/transfers/bigquery_to_gcs.py b/airflow/providers/google/cloud/transfers/bigquery_to_gcs.py index a1ec7a1b5141..2fd968bf741e 100644 --- a/airflow/providers/google/cloud/transfers/bigquery_to_gcs.py +++ b/airflow/providers/google/cloud/transfers/bigquery_to_gcs.py @@ -21,7 +21,6 @@ from typing import TYPE_CHECKING, Any, Sequence from google.api_core.exceptions import Conflict -from google.api_core.retry import Retry from google.cloud.bigquery import DEFAULT_RETRY, UnknownJob from airflow import AirflowException @@ -32,6 +31,8 @@ from airflow.providers.google.cloud.triggers.bigquery import BigQueryInsertJobTrigger if TYPE_CHECKING: + from google.api_core.retry import Retry + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/transfers/bigquery_to_sql.py b/airflow/providers/google/cloud/transfers/bigquery_to_sql.py index 36acb9363811..736880726c94 100644 --- a/airflow/providers/google/cloud/transfers/bigquery_to_sql.py +++ b/airflow/providers/google/cloud/transfers/bigquery_to_sql.py @@ -22,11 +22,11 @@ from typing import TYPE_CHECKING, Sequence from airflow.models import BaseOperator -from airflow.providers.common.sql.hooks.sql import DbApiHook from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook from airflow.providers.google.cloud.utils.bigquery_get_data import bigquery_get_data if TYPE_CHECKING: + from airflow.providers.common.sql.hooks.sql import DbApiHook from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/transfers/calendar_to_gcs.py b/airflow/providers/google/cloud/transfers/calendar_to_gcs.py index 0fec0b36981c..c9c3f585859f 100644 --- a/airflow/providers/google/cloud/transfers/calendar_to_gcs.py +++ b/airflow/providers/google/cloud/transfers/calendar_to_gcs.py @@ -17,14 +17,16 @@ from __future__ import annotations import json -from datetime import datetime from tempfile import NamedTemporaryFile -from typing import Any, Sequence +from typing import TYPE_CHECKING, Any, Sequence from airflow.models import BaseOperator from airflow.providers.google.cloud.hooks.gcs import GCSHook from airflow.providers.google.suite.hooks.calendar import GoogleCalendarHook +if TYPE_CHECKING: + from datetime import datetime + class GoogleCalendarToGCSOperator(BaseOperator): """ diff --git a/airflow/providers/google/cloud/transfers/facebook_ads_to_gcs.py b/airflow/providers/google/cloud/transfers/facebook_ads_to_gcs.py index 9273553e1bc1..bc0dae153a2d 100644 --- a/airflow/providers/google/cloud/transfers/facebook_ads_to_gcs.py +++ b/airflow/providers/google/cloud/transfers/facebook_ads_to_gcs.py @@ -23,14 +23,14 @@ from enum import Enum from typing import TYPE_CHECKING, Any, Sequence -from facebook_business.adobjects.adsinsights import AdsInsights - from airflow.exceptions import AirflowException from airflow.models import BaseOperator from airflow.providers.facebook.ads.hooks.ads import FacebookAdsReportingHook from airflow.providers.google.cloud.hooks.gcs import GCSHook if TYPE_CHECKING: + from facebook_business.adobjects.adsinsights import AdsInsights + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/transfers/gcs_to_bigquery.py b/airflow/providers/google/cloud/transfers/gcs_to_bigquery.py index 7b8b99e8b4b2..97b2943fb555 100644 --- a/airflow/providers/google/cloud/transfers/gcs_to_bigquery.py +++ b/airflow/providers/google/cloud/transfers/gcs_to_bigquery.py @@ -22,7 +22,6 @@ from typing import TYPE_CHECKING, Any, Sequence from google.api_core.exceptions import BadRequest, Conflict -from google.api_core.retry import Retry from google.cloud.bigquery import ( DEFAULT_RETRY, CopyJob, @@ -44,6 +43,8 @@ from airflow.providers.google.cloud.triggers.bigquery import BigQueryInsertJobTrigger if TYPE_CHECKING: + from google.api_core.retry import Retry + from airflow.utils.context import Context ALLOWED_FORMATS = [ diff --git a/airflow/providers/google/cloud/transfers/presto_to_gcs.py b/airflow/providers/google/cloud/transfers/presto_to_gcs.py index b9385876de6e..34e412802c2f 100644 --- a/airflow/providers/google/cloud/transfers/presto_to_gcs.py +++ b/airflow/providers/google/cloud/transfers/presto_to_gcs.py @@ -17,14 +17,15 @@ # under the License. from __future__ import annotations -from typing import Any - -from prestodb.client import PrestoResult -from prestodb.dbapi import Cursor as PrestoCursor +from typing import TYPE_CHECKING, Any from airflow.providers.google.cloud.transfers.sql_to_gcs import BaseSQLToGCSOperator from airflow.providers.presto.hooks.presto import PrestoHook +if TYPE_CHECKING: + from prestodb.client import PrestoResult + from prestodb.dbapi import Cursor as PrestoCursor + class _PrestoToGCSPrestoCursorAdapter: """ diff --git a/airflow/providers/google/cloud/transfers/trino_to_gcs.py b/airflow/providers/google/cloud/transfers/trino_to_gcs.py index 138f69acf1c0..d30827722e76 100644 --- a/airflow/providers/google/cloud/transfers/trino_to_gcs.py +++ b/airflow/providers/google/cloud/transfers/trino_to_gcs.py @@ -17,14 +17,15 @@ # under the License. from __future__ import annotations -from typing import Any - -from trino.client import TrinoResult -from trino.dbapi import Cursor as TrinoCursor +from typing import TYPE_CHECKING, Any from airflow.providers.google.cloud.transfers.sql_to_gcs import BaseSQLToGCSOperator from airflow.providers.trino.hooks.trino import TrinoHook +if TYPE_CHECKING: + from trino.client import TrinoResult + from trino.dbapi import Cursor as TrinoCursor + class _TrinoToGCSTrinoCursorAdapter: """ diff --git a/airflow/providers/google/cloud/triggers/cloud_run.py b/airflow/providers/google/cloud/triggers/cloud_run.py index ddbd74864c58..9506245d20cd 100644 --- a/airflow/providers/google/cloud/triggers/cloud_run.py +++ b/airflow/providers/google/cloud/triggers/cloud_run.py @@ -18,14 +18,15 @@ import asyncio from enum import Enum -from typing import Any, AsyncIterator, Sequence - -from google.longrunning import operations_pb2 +from typing import TYPE_CHECKING, Any, AsyncIterator, Sequence from airflow.exceptions import AirflowException from airflow.providers.google.cloud.hooks.cloud_run import CloudRunAsyncHook from airflow.triggers.base import BaseTrigger, TriggerEvent +if TYPE_CHECKING: + from google.longrunning import operations_pb2 + DEFAULT_BATCH_LOCATION = "us-central1" diff --git a/airflow/providers/google/cloud/triggers/kubernetes_engine.py b/airflow/providers/google/cloud/triggers/kubernetes_engine.py index 1e0780fbc66c..d47538a05f0c 100644 --- a/airflow/providers/google/cloud/triggers/kubernetes_engine.py +++ b/airflow/providers/google/cloud/triggers/kubernetes_engine.py @@ -19,8 +19,7 @@ import asyncio import warnings -from datetime import datetime -from typing import Any, AsyncIterator, Sequence +from typing import TYPE_CHECKING, Any, AsyncIterator, Sequence from google.cloud.container_v1.types import Operation @@ -35,6 +34,9 @@ from airflow.providers.google.cloud.hooks.kubernetes_engine import GKEAsyncHook, GKEPodAsyncHook from airflow.triggers.base import BaseTrigger, TriggerEvent +if TYPE_CHECKING: + from datetime import datetime + class GKEStartPodTrigger(KubernetesPodTrigger): """ diff --git a/airflow/providers/google/cloud/triggers/pubsub.py b/airflow/providers/google/cloud/triggers/pubsub.py index 27dbd8285cac..b39d491fc6f7 100644 --- a/airflow/providers/google/cloud/triggers/pubsub.py +++ b/airflow/providers/google/cloud/triggers/pubsub.py @@ -20,12 +20,12 @@ import asyncio from typing import TYPE_CHECKING, Any, AsyncIterator, Callable, Sequence -from google.cloud.pubsub_v1.types import ReceivedMessage - from airflow.providers.google.cloud.hooks.pubsub import PubSubAsyncHook from airflow.triggers.base import BaseTrigger, TriggerEvent if TYPE_CHECKING: + from google.cloud.pubsub_v1.types import ReceivedMessage + from airflow.utils.context import Context diff --git a/airflow/providers/google/cloud/utils/bigquery_get_data.py b/airflow/providers/google/cloud/utils/bigquery_get_data.py index 8c5c38add754..8fb61fc52ccc 100644 --- a/airflow/providers/google/cloud/utils/bigquery_get_data.py +++ b/airflow/providers/google/cloud/utils/bigquery_get_data.py @@ -16,12 +16,15 @@ # under the License. from __future__ import annotations -from collections.abc import Iterator -from logging import Logger +from typing import TYPE_CHECKING from google.cloud.bigquery.table import Row, RowIterator -from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook +if TYPE_CHECKING: + from collections.abc import Iterator + from logging import Logger + + from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook def bigquery_get_data( diff --git a/airflow/providers/google/cloud/utils/mlengine_operator_utils.py b/airflow/providers/google/cloud/utils/mlengine_operator_utils.py index e4845bd7f6ba..d58e9190d6e1 100644 --- a/airflow/providers/google/cloud/utils/mlengine_operator_utils.py +++ b/airflow/providers/google/cloud/utils/mlengine_operator_utils.py @@ -21,12 +21,11 @@ import json import os import re -from typing import Callable, Iterable, TypeVar +from typing import TYPE_CHECKING, Callable, Iterable, TypeVar from urllib.parse import urlsplit import dill -from airflow import DAG from airflow.exceptions import AirflowException from airflow.operators.python import PythonOperator from airflow.providers.apache.beam.hooks.beam import BeamRunnerType @@ -34,6 +33,9 @@ from airflow.providers.google.cloud.hooks.gcs import GCSHook from airflow.providers.google.cloud.operators.mlengine import MLEngineStartBatchPredictionJobOperator +if TYPE_CHECKING: + from airflow import DAG + T = TypeVar("T", bound=Callable) diff --git a/airflow/providers/google/common/hooks/base_google.py b/airflow/providers/google/common/hooks/base_google.py index 30287c27049f..94cc420bae34 100644 --- a/airflow/providers/google/common/hooks/base_google.py +++ b/airflow/providers/google/common/hooks/base_google.py @@ -26,7 +26,7 @@ import warnings from contextlib import ExitStack, contextmanager from subprocess import check_output -from typing import Any, Callable, Generator, Sequence, TypeVar, cast +from typing import TYPE_CHECKING, Any, Callable, Generator, Sequence, TypeVar, cast import google.auth import google.auth.credentials @@ -36,7 +36,6 @@ import tenacity from asgiref.sync import sync_to_async from google.api_core.exceptions import Forbidden, ResourceExhausted, TooManyRequests -from google.api_core.gapic_v1.client_info import ClientInfo from google.auth import _cloud_sdk, compute_engine from google.auth.environment_vars import CLOUD_SDK_CONFIG_DIR, CREDENTIALS from google.auth.exceptions import RefreshError @@ -56,6 +55,9 @@ from airflow.providers.google.common.consts import CLIENT_INFO from airflow.utils.process_utils import patch_environ +if TYPE_CHECKING: + from google.api_core.gapic_v1.client_info import ClientInfo + log = logging.getLogger(__name__) # Constants used by the mechanism of repeating requests in reaction to exceeding the temporary quota. diff --git a/airflow/providers/google/common/links/storage.py b/airflow/providers/google/common/links/storage.py index 3252e552443b..6b4cb94690c2 100644 --- a/airflow/providers/google/common/links/storage.py +++ b/airflow/providers/google/common/links/storage.py @@ -19,7 +19,6 @@ from typing import TYPE_CHECKING -from airflow.models import BaseOperator from airflow.providers.google.cloud.links.base import BaseGoogleLink BASE_LINK = "https://console.cloud.google.com" @@ -27,6 +26,7 @@ GCS_FILE_DETAILS_LINK = BASE_LINK + "/storage/browser/_details/{uri};tab=live_object?project={project_id}" if TYPE_CHECKING: + from airflow.models import BaseOperator from airflow.utils.context import Context diff --git a/airflow/providers/google/common/utils/id_token_credentials.py b/airflow/providers/google/common/utils/id_token_credentials.py index 6978edc4d2eb..6a41438a22a0 100644 --- a/airflow/providers/google/common/utils/id_token_credentials.py +++ b/airflow/providers/google/common/utils/id_token_credentials.py @@ -32,12 +32,15 @@ import json import os +from typing import TYPE_CHECKING import google.auth.transport -import google.oauth2 from google.auth import credentials as google_auth_credentials, environment_vars, exceptions from google.oauth2 import credentials as oauth2_credentials, service_account +if TYPE_CHECKING: + import google.oauth2 + # Valid types accepted for file-based credentials. # They are taken from "google.auth._default" and since they are all "protected" and the imports might # change any time and fail the whole Google provider functionality - we should inline them diff --git a/airflow/providers/google/marketing_platform/hooks/campaign_manager.py b/airflow/providers/google/marketing_platform/hooks/campaign_manager.py index f0a02ddd255e..42e976db1607 100644 --- a/airflow/providers/google/marketing_platform/hooks/campaign_manager.py +++ b/airflow/providers/google/marketing_platform/hooks/campaign_manager.py @@ -18,14 +18,16 @@ """This module contains Google Campaign Manager hook.""" from __future__ import annotations -from typing import Any, Sequence +from typing import TYPE_CHECKING, Any, Sequence -from googleapiclient import http from googleapiclient.discovery import Resource, build from airflow.exceptions import AirflowException from airflow.providers.google.common.hooks.base_google import GoogleBaseHook +if TYPE_CHECKING: + from googleapiclient import http + class GoogleCampaignManagerHook(GoogleBaseHook): """Hook for Google Campaign Manager.""" diff --git a/airflow/providers/google/suite/hooks/calendar.py b/airflow/providers/google/suite/hooks/calendar.py index b5b595067507..22b9f127612c 100644 --- a/airflow/providers/google/suite/hooks/calendar.py +++ b/airflow/providers/google/suite/hooks/calendar.py @@ -18,14 +18,16 @@ """This module contains a Google Calendar API hook.""" from __future__ import annotations -from datetime import datetime -from typing import Any, Sequence +from typing import TYPE_CHECKING, Any, Sequence from googleapiclient.discovery import build from airflow.exceptions import AirflowException from airflow.providers.google.common.hooks.base_google import GoogleBaseHook +if TYPE_CHECKING: + from datetime import datetime + class GoogleCalendarHook(GoogleBaseHook): """ diff --git a/tests/providers/google/cloud/hooks/test_datacatalog.py b/tests/providers/google/cloud/hooks/test_datacatalog.py index c82f29d28e45..4b109f8b9ac2 100644 --- a/tests/providers/google/cloud/hooks/test_datacatalog.py +++ b/tests/providers/google/cloud/hooks/test_datacatalog.py @@ -18,11 +18,10 @@ import re from copy import deepcopy -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from unittest import mock import pytest -from google.api_core.gapic_v1.method import _MethodDefault from google.api_core.retry import Retry from google.cloud.datacatalog import CreateTagRequest, CreateTagTemplateRequest, Entry, Tag, TagTemplate from google.protobuf.field_mask_pb2 import FieldMask @@ -34,6 +33,9 @@ mock_base_gcp_hook_no_default_project_id, ) +if TYPE_CHECKING: + from google.api_core.gapic_v1.method import _MethodDefault + TEST_GCP_CONN_ID: str = "test-gcp-conn-id" TEST_LOCATION: str = "europe-west-3b" TEST_ENTRY_ID: str = "test-entry-id" diff --git a/tests/providers/google/cloud/hooks/test_os_login.py b/tests/providers/google/cloud/hooks/test_os_login.py index b5605f4be0e5..49452d109d17 100644 --- a/tests/providers/google/cloud/hooks/test_os_login.py +++ b/tests/providers/google/cloud/hooks/test_os_login.py @@ -17,12 +17,10 @@ from __future__ import annotations import re -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from unittest import mock import pytest -from google.api_core.gapic_v1.method import _MethodDefault -from google.api_core.retry import Retry from airflow import AirflowException from airflow.providers.google.cloud.hooks.os_login import OSLoginHook @@ -31,6 +29,10 @@ mock_base_gcp_hook_no_default_project_id, ) +if TYPE_CHECKING: + from google.api_core.gapic_v1.method import _MethodDefault + from google.api_core.retry import Retry + TEST_GCP_CONN_ID: str = "test-gcp-conn-id" TEST_PROJECT_ID: str = "test-project-id" TEST_PROJECT_ID_2: str = "test-project-id-2" diff --git a/tests/providers/google/cloud/operators/test_datacatalog.py b/tests/providers/google/cloud/operators/test_datacatalog.py index 7ae4940a0ff5..424783193333 100644 --- a/tests/providers/google/cloud/operators/test_datacatalog.py +++ b/tests/providers/google/cloud/operators/test_datacatalog.py @@ -16,11 +16,10 @@ # under the License. from __future__ import annotations -from typing import Sequence +from typing import TYPE_CHECKING, Sequence from unittest import mock from google.api_core.exceptions import AlreadyExists -from google.api_core.gapic_v1.method import _MethodDefault from google.api_core.retry import Retry from google.cloud.datacatalog import Entry, EntryGroup, Tag, TagTemplate, TagTemplateField from google.protobuf.field_mask_pb2 import FieldMask @@ -49,6 +48,9 @@ CloudDataCatalogUpdateTagTemplateOperator, ) +if TYPE_CHECKING: + from google.api_core.gapic_v1.method import _MethodDefault + BASE_PATH = "airflow.providers.google.cloud.operators.datacatalog.{}" TEST_PROJECT_ID: str = "example_id" TEST_LOCATION: str = "en-west-3" diff --git a/tests/providers/google/cloud/operators/test_dataproc_metastore.py b/tests/providers/google/cloud/operators/test_dataproc_metastore.py index afaed9355b94..3fffd892ada8 100644 --- a/tests/providers/google/cloud/operators/test_dataproc_metastore.py +++ b/tests/providers/google/cloud/operators/test_dataproc_metastore.py @@ -16,9 +16,9 @@ # under the License. from __future__ import annotations +from typing import TYPE_CHECKING from unittest import mock -from google.api_core.gapic_v1.method import _MethodDefault from google.api_core.retry import Retry from google.protobuf.field_mask_pb2 import FieldMask @@ -35,6 +35,9 @@ DataprocMetastoreUpdateServiceOperator, ) +if TYPE_CHECKING: + from google.api_core.gapic_v1.method import _MethodDefault + TASK_ID: str = "task_id" GCP_LOCATION: str = "test-location" GCP_PROJECT_ID: str = "test-project-id" diff --git a/tests/providers/google/cloud/utils/airflow_util.py b/tests/providers/google/cloud/utils/airflow_util.py index da3faa6fcf49..b856081dd10c 100644 --- a/tests/providers/google/cloud/utils/airflow_util.py +++ b/tests/providers/google/cloud/utils/airflow_util.py @@ -16,17 +16,20 @@ # under the License. from __future__ import annotations +from typing import TYPE_CHECKING from unittest import mock import pendulum from airflow.models import DAG, Connection -from airflow.models.baseoperator import BaseOperator from airflow.models.dagrun import DagRun from airflow.models.taskinstance import TaskInstance from airflow.utils import timezone from airflow.utils.types import DagRunType +if TYPE_CHECKING: + from airflow.models.baseoperator import BaseOperator + def get_dag_run(dag_id: str = "test_dag_id", run_id: str = "test_dag_id") -> DagRun: dag_run = DagRun(