Skip to content

Commit

Permalink
Improve modules import in google provider by move some of them into a…
Browse files Browse the repository at this point in the history
… type-checking block (#33783)
  • Loading branch information
hussein-awala authored Aug 27, 2023
1 parent 452a978 commit 6802d41
Show file tree
Hide file tree
Showing 120 changed files with 473 additions and 304 deletions.
12 changes: 7 additions & 5 deletions airflow/providers/google/ads/hooks/ads.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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-_]*$"


Expand Down
20 changes: 11 additions & 9 deletions airflow/providers/google/cloud/hooks/automl.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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):
"""
Expand Down
4 changes: 2 additions & 2 deletions airflow/providers/google/cloud/hooks/bigquery.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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__)

Expand Down
8 changes: 5 additions & 3 deletions airflow/providers/google/cloud/hooks/bigquery_dts.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand All @@ -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."""
Expand Down
9 changes: 6 additions & 3 deletions airflow/providers/google/cloud/hooks/bigtable.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down
8 changes: 5 additions & 3 deletions airflow/providers/google/cloud/hooks/cloud_batch.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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):
"""
Expand Down
10 changes: 6 additions & 4 deletions airflow/providers/google/cloud/hooks/cloud_build.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
24 changes: 14 additions & 10 deletions airflow/providers/google/cloud/hooks/cloud_composer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
8 changes: 5 additions & 3 deletions airflow/providers/google/cloud/hooks/cloud_memorystore.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand All @@ -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):
"""
Expand Down
8 changes: 5 additions & 3 deletions airflow/providers/google/cloud/hooks/cloud_run.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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):
"""
Expand Down
6 changes: 4 additions & 2 deletions airflow/providers/google/cloud/hooks/cloud_sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,15 +35,14 @@
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
from aiohttp import ClientSession
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"
Expand All @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -33,24 +33,26 @@
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,
StorageTransferServiceAsyncClient,
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
Expand Down
Loading

0 comments on commit 6802d41

Please sign in to comment.
  翻译: