Skip to content

Commit

Permalink
[AIRFLOW-6120] Rename GoogleCloudBaseHook (#6734)
Browse files Browse the repository at this point in the history 
* [AIRFLOW-6120] Rename GoogleCloudBaseHook
  • Loading branch information
@michalslowikowski00 @potiuk
michalslowikowski00 authored and potiuk committed Dec 9, 2019
1 parent 572151c commit ed0a14f
Show file tree
Hide file tree
Showing 65 changed files with 639 additions and 568 deletions.
15 changes: 13 additions & 2 deletions airflow/contrib/hooks/gcp_api_base_hook.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,10 +19,21 @@
"""This module is deprecated. Please use `airflow.gcp.hooks.base`."""
import warnings

# pylint: disable=unused-import
from airflow.gcp.hooks.base import GoogleCloudBaseHook # noqa
from airflow.gcp.hooks.base import CloudBaseHook

warnings.warn(
"This module is deprecated. Please use `airflow.gcp.hooks.base`.",
DeprecationWarning, stacklevel=2
)


class GoogleCloudBaseHook(CloudBaseHook):
"""
This class is deprecated. Please use `airflow.gcp.hooks.base.CloudBaseHook`.
"""
def __init__(self, *args, **kwargs):
warnings.warn(
"This class is deprecated. Please use `airflow.gcp.hooks.base.CloudBaseHook`.",
DeprecationWarning, stacklevel=2
)
super().__init__(*args, **kwargs)
4 changes: 2 additions & 2 deletions airflow/contrib/hooks/gdrive_hook.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,11 +22,11 @@
from googleapiclient.discovery import Resource, build
from googleapiclient.http import MediaFileUpload

from airflow.gcp.hooks.base import GoogleCloudBaseHook
from airflow.gcp.hooks.base import CloudBaseHook


# noinspection PyAbstractClass
class GoogleDriveHook(GoogleCloudBaseHook):
class GoogleDriveHook(CloudBaseHook):
"""
Hook for the Google Drive APIs.
Expand Down
52 changes: 26 additions & 26 deletions airflow/gcp/hooks/automl.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,10 +30,10 @@
ImageObjectDetectionModelDeploymentMetadata, InputConfig, Model, Operation, PredictResponse, TableSpec,
)

from airflow.gcp.hooks.base import GoogleCloudBaseHook
from airflow.gcp.hooks.base import CloudBaseHook


class CloudAutoMLHook(GoogleCloudBaseHook):
class CloudAutoMLHook(CloudBaseHook):
"""
Google Cloud AutoML hook.
Expand Down Expand Up @@ -79,8 +79,8 @@ def prediction_client(self) -> PredictionServiceClient:
credentials=self._get_credentials(), client_info=self.client_info
)

@GoogleCloudBaseHook.catch_http_exception
@GoogleCloudBaseHook.fallback_to_default_project_id
@CloudBaseHook.catch_http_exception
@CloudBaseHook.fallback_to_default_project_id
def create_model(
self,
model: Union[dict, Model],
Expand Down Expand Up @@ -122,8 +122,8 @@ def create_model(
parent=parent, model=model, retry=retry, timeout=timeout, metadata=metadata
)

@GoogleCloudBaseHook.catch_http_exception
@GoogleCloudBaseHook.fallback_to_default_project_id
@CloudBaseHook.catch_http_exception
@CloudBaseHook.fallback_to_default_project_id
def batch_predict(
self,
model_id: str,
Expand Down Expand Up @@ -184,8 +184,8 @@ def batch_predict(
)
return result

@GoogleCloudBaseHook.catch_http_exception
@GoogleCloudBaseHook.fallback_to_default_project_id
@CloudBaseHook.catch_http_exception
@CloudBaseHook.fallback_to_default_project_id
def predict(
self,
model_id: str,
Expand Down Expand Up @@ -238,8 +238,8 @@ def predict(
)
return result

@GoogleCloudBaseHook.catch_http_exception
@GoogleCloudBaseHook.fallback_to_default_project_id
@CloudBaseHook.catch_http_exception
@CloudBaseHook.fallback_to_default_project_id
def create_dataset(
self,
dataset: Union[dict, Dataset],
Expand Down Expand Up @@ -283,8 +283,8 @@ def create_dataset(
)
return result

@GoogleCloudBaseHook.catch_http_exception
@GoogleCloudBaseHook.fallback_to_default_project_id
@CloudBaseHook.catch_http_exception
@CloudBaseHook.fallback_to_default_project_id
def import_data(
self,
dataset_id: str,
Expand Down Expand Up @@ -333,8 +333,8 @@ def import_data(
)
return result

@GoogleCloudBaseHook.catch_http_exception
@GoogleCloudBaseHook.fallback_to_default_project_id
@CloudBaseHook.catch_http_exception
@CloudBaseHook.fallback_to_default_project_id
def list_column_specs( # pylint: disable=too-many-arguments
self,
dataset_id: str,
Expand Down Expand Up @@ -401,8 +401,8 @@ def list_column_specs( # pylint: disable=too-many-arguments
)
return result

@GoogleCloudBaseHook.catch_http_exception
@GoogleCloudBaseHook.fallback_to_default_project_id
@CloudBaseHook.catch_http_exception
@CloudBaseHook.fallback_to_default_project_id
def get_model(
self,
model_id: str,
Expand Down Expand Up @@ -441,8 +441,8 @@ def get_model(
)
return result

@GoogleCloudBaseHook.catch_http_exception
@GoogleCloudBaseHook.fallback_to_default_project_id
@CloudBaseHook.catch_http_exception
@CloudBaseHook.fallback_to_default_project_id
def delete_model(
self,
model_id: str,
Expand Down Expand Up @@ -481,8 +481,8 @@ def delete_model(
)
return result

@GoogleCloudBaseHook.catch_http_exception
@GoogleCloudBaseHook.fallback_to_default_project_id
@CloudBaseHook.catch_http_exception
@CloudBaseHook.fallback_to_default_project_id
def update_dataset(
self,
dataset: Union[dict, Dataset],
Expand Down Expand Up @@ -526,8 +526,8 @@ def update_dataset(
)
return result

@GoogleCloudBaseHook.catch_http_exception
@GoogleCloudBaseHook.fallback_to_default_project_id
@CloudBaseHook.catch_http_exception
@CloudBaseHook.fallback_to_default_project_id
def deploy_model(
self,
model_id: str,
Expand Down Expand Up @@ -640,8 +640,8 @@ def list_table_specs(
)
return result

@GoogleCloudBaseHook.catch_http_exception
@GoogleCloudBaseHook.fallback_to_default_project_id
@CloudBaseHook.catch_http_exception
@CloudBaseHook.fallback_to_default_project_id
def list_datasets(
self,
location: str,
Expand Down Expand Up @@ -680,8 +680,8 @@ def list_datasets(
)
return result

@GoogleCloudBaseHook.catch_http_exception
@GoogleCloudBaseHook.fallback_to_default_project_id
@CloudBaseHook.catch_http_exception
@CloudBaseHook.fallback_to_default_project_id
def delete_dataset(
self,
dataset_id: str,
Expand Down
8 changes: 4 additions & 4 deletions airflow/gcp/hooks/base.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,7 @@ def __init__(self):
RT = TypeVar('RT') # pylint: disable=invalid-name


class GoogleCloudBaseHook(BaseHook):
class CloudBaseHook(BaseHook):
"""
A base hook for Google cloud-related hooks. Google cloud has a shared REST
API client that is built in the same way no matter which service you use.
Expand Down Expand Up @@ -302,7 +302,7 @@ def catch_http_exception(func: Callable[..., RT]) -> Callable[..., RT]:
"""

@functools.wraps(func)
def wrapper_decorator(self: GoogleCloudBaseHook, *args, **kwargs) -> RT:
def wrapper_decorator(self: CloudBaseHook, *args, **kwargs) -> RT:
try:
return func(self, *args, **kwargs)
except GoogleAPICallError as e:
Expand Down Expand Up @@ -335,7 +335,7 @@ def fallback_to_default_project_id(func: Callable[..., RT]) -> Callable[..., RT]
:return: result of the function call
"""
@functools.wraps(func)
def inner_wrapper(self: GoogleCloudBaseHook, *args, **kwargs) -> RT:
def inner_wrapper(self: CloudBaseHook, *args, **kwargs) -> RT:
if args:
raise AirflowException(
"You must use keyword arguments in this methods rather than"
Expand All @@ -362,7 +362,7 @@ def provide_gcp_credential_file(func: Callable[..., RT]) -> Callable[..., RT]:
makes it easier to use multiple connection in one function.
"""
@functools.wraps(func)
def wrapper(self: GoogleCloudBaseHook, *args, **kwargs) -> RT:
def wrapper(self: CloudBaseHook, *args, **kwargs) -> RT:
with self.provide_gcp_credential_file_as_context():
return func(self, *args, **kwargs)
return wrapper
Expand Down
6 changes: 3 additions & 3 deletions airflow/gcp/hooks/bigquery.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,12 +37,12 @@
)

from airflow import AirflowException
from airflow.gcp.hooks.base import GoogleCloudBaseHook
from airflow.gcp.hooks.base import CloudBaseHook
from airflow.hooks.dbapi_hook import DbApiHook
from airflow.utils.log.logging_mixin import LoggingMixin


class BigQueryHook(GoogleCloudBaseHook, DbApiHook):
class BigQueryHook(CloudBaseHook, DbApiHook):
"""
Interact with BigQuery. This hook uses the Google Cloud Platform
connection.
Expand Down Expand Up @@ -1846,7 +1846,7 @@ def get_datasets_list(self, project_id: Optional[str] = None) -> List:

return datasets_list

@GoogleCloudBaseHook.catch_http_exception
@CloudBaseHook.catch_http_exception
def get_dataset_tables_list(self, dataset_id, project_id=None, table_prefix=None, max_results=None):
"""
Method returns tables list of a BigQuery dataset. If table prefix is specified,
Expand Down
20 changes: 10 additions & 10 deletions airflow/gcp/hooks/bigquery_dts.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@
from google.protobuf.json_format import MessageToDict, ParseDict
from googleapiclient.discovery import Resource

from airflow.gcp.hooks.base import GoogleCloudBaseHook
from airflow.gcp.hooks.base import CloudBaseHook


def get_object_id(obj: dict) -> str:
Expand All @@ -41,7 +41,7 @@ def get_object_id(obj: dict) -> str:
return obj["name"].rpartition("/")[-1]


class BiqQueryDataTransferServiceHook(GoogleCloudBaseHook):
class BiqQueryDataTransferServiceHook(CloudBaseHook):
"""
Hook for Google Bigquery Transfer API.
Expand Down Expand Up @@ -94,8 +94,8 @@ def get_conn(self) -> DataTransferServiceClient:
)
return self._conn

@GoogleCloudBaseHook.catch_http_exception
@GoogleCloudBaseHook.fallback_to_default_project_id
@CloudBaseHook.catch_http_exception
@CloudBaseHook.fallback_to_default_project_id
def create_transfer_config(
self,
transfer_config: Union[dict, TransferConfig],
Expand Down Expand Up @@ -139,8 +139,8 @@ def create_transfer_config(
metadata=metadata,
)

@GoogleCloudBaseHook.catch_http_exception
@GoogleCloudBaseHook.fallback_to_default_project_id
@CloudBaseHook.catch_http_exception
@CloudBaseHook.fallback_to_default_project_id
def delete_transfer_config(
self,
transfer_config_id: str,
Expand Down Expand Up @@ -177,8 +177,8 @@ def delete_transfer_config(
name=name, retry=retry, timeout=timeout, metadata=metadata
)

@GoogleCloudBaseHook.catch_http_exception
@GoogleCloudBaseHook.fallback_to_default_project_id
@CloudBaseHook.catch_http_exception
@CloudBaseHook.fallback_to_default_project_id
def start_manual_transfer_runs(
self,
transfer_config_id: str,
Expand Down Expand Up @@ -234,8 +234,8 @@ def start_manual_transfer_runs(
metadata=metadata,
)

@GoogleCloudBaseHook.catch_http_exception
@GoogleCloudBaseHook.fallback_to_default_project_id
@CloudBaseHook.catch_http_exception
@CloudBaseHook.fallback_to_default_project_id
def get_transfer_run(
self,
run_id: str,
Expand Down
12 changes: 6 additions & 6 deletions airflow/gcp/hooks/bigtable.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,10 +28,10 @@
from google.cloud.bigtable.table import ClusterState, Table
from google.cloud.bigtable_admin_v2 import enums

from airflow.gcp.hooks.base import GoogleCloudBaseHook
from airflow.gcp.hooks.base import CloudBaseHook


class BigtableHook(GoogleCloudBaseHook):
class BigtableHook(CloudBaseHook):
"""
Hook for Google Cloud Bigtable APIs.
Expand All @@ -54,7 +54,7 @@ def _get_client(self, project_id: str):
)
return self._client

@GoogleCloudBaseHook.fallback_to_default_project_id
@CloudBaseHook.fallback_to_default_project_id
def get_instance(self, instance_id: str, project_id: Optional[str] = None) -> Instance:
"""
Retrieves and returns the specified Cloud Bigtable instance if it exists.
Expand All @@ -74,7 +74,7 @@ def get_instance(self, instance_id: str, project_id: Optional[str] = None) -> In
return None
return instance

@GoogleCloudBaseHook.fallback_to_default_project_id
@CloudBaseHook.fallback_to_default_project_id
def delete_instance(self, instance_id: str, project_id: Optional[str] = None) -> None:
"""
Deletes the specified Cloud Bigtable instance.
Expand All @@ -96,7 +96,7 @@ def delete_instance(self, instance_id: str, project_id: Optional[str] = None) ->
self.log.info("The instance '%s' does not exist in project '%s'. Exiting", instance_id,
project_id)

@GoogleCloudBaseHook.fallback_to_default_project_id
@CloudBaseHook.fallback_to_default_project_id
def create_instance(
self,
instance_id: str,
Expand Down Expand Up @@ -210,7 +210,7 @@ def create_table(
table = Table(table_id, instance)
table.create(initial_split_keys, column_families)

@GoogleCloudBaseHook.fallback_to_default_project_id
@CloudBaseHook.fallback_to_default_project_id
def delete_table(self, instance_id: str, table_id: str, project_id: Optional[str] = None) -> None:
"""
Deletes the specified table in Cloud Bigtable.
Expand Down
6 changes: 3 additions & 3 deletions airflow/gcp/hooks/cloud_build.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,14 +24,14 @@
from googleapiclient.discovery import build

from airflow import AirflowException
from airflow.gcp.hooks.base import GoogleCloudBaseHook
from airflow.gcp.hooks.base import CloudBaseHook

# Time to sleep between active checks of the operation results
TIME_TO_SLEEP_IN_SECONDS = 5


# noinspection PyAbstractClass
class CloudBuildHook(GoogleCloudBaseHook):
class CloudBuildHook(CloudBaseHook):
"""
Hook for the Google Cloud Build APIs.
Expand Down Expand Up @@ -70,7 +70,7 @@ def get_conn(self):
self._conn = build("cloudbuild", self.api_version, http=http_authorized, cache_discovery=False)
return self._conn

@GoogleCloudBaseHook.fallback_to_default_project_id
@CloudBaseHook.fallback_to_default_project_id
def create_build(self, body: Dict, project_id: Optional[str] = None) -> Dict:
"""
Starts a build with the specified configuration.
Expand Down

0 comments on commit ed0a14f

Please sign in to comment.