Rename CloudBaseHook to GoogleBaseHook and move it to google.common (#…

…8011)
apache · Mar 31, 2020 · 8a02402 · 8a02402
1 parent 8e89780
commit 8a02402
Show file tree

Hide file tree

Showing 90 changed files with 811 additions and 741 deletions.
diff --git a/UPDATING.md b/UPDATING.md
@@ -233,7 +233,7 @@ The fix to `none_failed` trigger rule breaks workflows that depend on the previo
 ### Standardize handling http exception in BigQuery
 
 Since BigQuery is the part of the GCP it was possible to simplify the code by handling the exceptions
-by usage of the `airflow.providers.google.cloud.hooks.base.CloudBaseHook.catch_http_exception` decorator however it changes
+by usage of the `airflow.providers.google.common.hooks.base.GoogleBaseHook.catch_http_exception` decorator however it changes
 exceptions raised by the following methods:
 * `airflow.providers.google.cloud.hooks.bigquery.BigQueryBaseCursor.run_table_delete` raises `AirflowException` instead of `Exception`.
 * `airflow.providers.google.cloud.hooks.bigquery.BigQueryBaseCursor.create_empty_dataset` raises `AirflowException` instead of `ValueError`.
@@ -914,12 +914,12 @@ The following variables were removed from the task instance context:
 - latest_date
 - tables
 
-### Moved provide_gcp_credential_file decorator to GoogleCloudBaseHook
+### Moved provide_gcp_credential_file decorator to GoogleBaseHook
 
 To simplify the code, the decorator has been moved from the inner-class.
 
-Instead of `@GoogleCloudBaseHook._Decorators.provide_gcp_credential_file`,
-you should write `@GoogleCloudBaseHook.provide_gcp_credential_file`
+Instead of `@GoogleBaseHook._Decorators.provide_gcp_credential_file`,
+you should write `@GoogleBaseHook.provide_gcp_credential_file`
 
 ### Changes to S3Hook
 

diff --git a/airflow/contrib/hooks/gcp_api_base_hook.py b/airflow/contrib/hooks/gcp_api_base_hook.py
@@ -15,24 +15,25 @@
 # KIND, either express or implied. See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""This module is deprecated. Please use `airflow.providers.google.cloud.hooks.base`."""
+"""This module is deprecated. Please use `airflow.providers.google.common.hooks.base_google`."""
 import warnings
 
-from airflow.providers.google.cloud.hooks.base import CloudBaseHook
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 
 warnings.warn(
- "This module is deprecated. Please use `airflow.providers.google.cloud.hooks.base`.",
+ "This module is deprecated. Please use `airflow.providers.google.common.hooks.base_google`.",
  DeprecationWarning, stacklevel=2
 )
 
 
-class GoogleCloudBaseHook(CloudBaseHook):
+class GoogleCloudBaseHook(GoogleBaseHook):
  """
- This class is deprecated. Please use `airflow.providers.google.cloud.hooks.base.CloudBaseHook`.
+ This class is deprecated. Please use `airflow.providers.google.common.hooks.base_google.GoogleBaseHook`.
  """
  def __init__(self, *args, **kwargs):
  warnings.warn(
- "This class is deprecated. Please use `airflow.providers.google.cloud.hooks.base.CloudBaseHook`.",
+ "This class is deprecated. Please use "
+ "`airflow.providers.google.common.hooks.base_google.GoogleBaseHook`.",
  DeprecationWarning, stacklevel=2
  )
  super().__init__(*args, **kwargs)
diff --git a/airflow/providers/amazon/aws/operators/google_api_to_s3_transfer.py b/airflow/providers/amazon/aws/operators/google_api_to_s3_transfer.py
@@ -25,7 +25,7 @@
 from airflow.models import BaseOperator
 from airflow.models.xcom import MAX_XCOM_SIZE
 from airflow.providers.amazon.aws.hooks.s3 import S3Hook
-from airflow.providers.google.cloud.hooks.discovery_api import GoogleDiscoveryApiHook
+from airflow.providers.google.common.hooks.discovery_api import GoogleDiscoveryApiHook
 from airflow.utils.decorators import apply_defaults
 
 

diff --git a/airflow/providers/google/cloud/hooks/automl.py b/airflow/providers/google/cloud/hooks/automl.py
@@ -29,10 +29,10 @@
  ImageObjectDetectionModelDeploymentMetadata, InputConfig, Model, Operation, PredictResponse, TableSpec,
 )
 
-from airflow.providers.google.cloud.hooks.base import CloudBaseHook
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 
 
-class CloudAutoMLHook(CloudBaseHook):
+class CloudAutoMLHook(GoogleBaseHook):
  """
  Google Cloud AutoML hook.
 
@@ -78,7 +78,7 @@ def prediction_client(self) -> PredictionServiceClient:
  credentials=self._get_credentials(), client_info=self.client_info
  )
 
- @CloudBaseHook.fallback_to_default_project_id
+ @GoogleBaseHook.fallback_to_default_project_id
  def create_model(
  self,
  model: Union[dict, Model],
@@ -121,7 +121,7 @@ def create_model(
  parent=parent, model=model, retry=retry, timeout=timeout, metadata=metadata
  )
 
- @CloudBaseHook.fallback_to_default_project_id
+ @GoogleBaseHook.fallback_to_default_project_id
  def batch_predict(
  self,
  model_id: str,
@@ -183,7 +183,7 @@ def batch_predict(
  )
  return result
 
- @CloudBaseHook.fallback_to_default_project_id
+ @GoogleBaseHook.fallback_to_default_project_id
  def predict(
  self,
  model_id: str,
@@ -237,7 +237,7 @@ def predict(
  )
  return result
 
- @CloudBaseHook.fallback_to_default_project_id
+ @GoogleBaseHook.fallback_to_default_project_id
  def create_dataset(
  self,
  dataset: Union[dict, Dataset],
@@ -282,7 +282,7 @@ def create_dataset(
  )
  return result
 
- @CloudBaseHook.fallback_to_default_project_id
+ @GoogleBaseHook.fallback_to_default_project_id
  def import_data(
  self,
  dataset_id: str,
@@ -332,7 +332,7 @@ def import_data(
  )
  return result
 
- @CloudBaseHook.fallback_to_default_project_id
+ @GoogleBaseHook.fallback_to_default_project_id
  def list_column_specs( # pylint: disable=too-many-arguments
  self,
  dataset_id: str,
@@ -400,7 +400,7 @@ def list_column_specs( # pylint: disable=too-many-arguments
  )
  return result
 
- @CloudBaseHook.fallback_to_default_project_id
+ @GoogleBaseHook.fallback_to_default_project_id
  def get_model(
  self,
  model_id: str,
@@ -440,7 +440,7 @@ def get_model(
  )
  return result
 
- @CloudBaseHook.fallback_to_default_project_id
+ @GoogleBaseHook.fallback_to_default_project_id
  def delete_model(
  self,
  model_id: str,
@@ -480,7 +480,7 @@ def delete_model(
  )
  return result
 
- @CloudBaseHook.fallback_to_default_project_id
+ @GoogleBaseHook.fallback_to_default_project_id
  def update_dataset(
  self,
  dataset: Union[dict, Dataset],
@@ -525,7 +525,7 @@ def update_dataset(
  )
  return result
 
- @CloudBaseHook.fallback_to_default_project_id
+ @GoogleBaseHook.fallback_to_default_project_id
  def deploy_model(
  self,
  model_id: str,
@@ -640,7 +640,7 @@ def list_table_specs(
  )
  return result
 
- @CloudBaseHook.fallback_to_default_project_id
+ @GoogleBaseHook.fallback_to_default_project_id
  def list_datasets(
  self,
  location: str,
@@ -680,7 +680,7 @@ def list_datasets(
  )
  return result
 
- @CloudBaseHook.fallback_to_default_project_id
+ @GoogleBaseHook.fallback_to_default_project_id
  def delete_dataset(
  self,
  dataset_id: str,

diff --git a/airflow/providers/google/cloud/hooks/bigquery.py b/airflow/providers/google/cloud/hooks/bigquery.py
@@ -37,14 +37,14 @@
 
 from airflow.exceptions import AirflowException
 from airflow.hooks.dbapi_hook import DbApiHook
-from airflow.providers.google.cloud.hooks.base import CloudBaseHook
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from airflow.utils.log.logging_mixin import LoggingMixin
 
 log = logging.getLogger(__name__)
 
 
 # pylint: disable=too-many-public-methods
-class BigQueryHook(CloudBaseHook, DbApiHook):
+class BigQueryHook(GoogleBaseHook, DbApiHook):
  """
  Interact with BigQuery. This hook uses the Google Cloud Platform
  connection.

diff --git a/airflow/providers/google/cloud/hooks/bigquery_dts.py b/airflow/providers/google/cloud/hooks/bigquery_dts.py
@@ -30,7 +30,7 @@
 from google.protobuf.json_format import MessageToDict, ParseDict
 from googleapiclient.discovery import Resource
 
-from airflow.providers.google.cloud.hooks.base import CloudBaseHook
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 
 
 def get_object_id(obj: dict) -> str:
@@ -40,7 +40,7 @@ def get_object_id(obj: dict) -> str:
  return obj["name"].rpartition("/")[-1]
 
 
-class BiqQueryDataTransferServiceHook(CloudBaseHook):
+class BiqQueryDataTransferServiceHook(GoogleBaseHook):
  """
  Hook for Google Bigquery Transfer API.
 
@@ -93,7 +93,7 @@ def get_conn(self) -> DataTransferServiceClient:
  )
  return self._conn
 
- @CloudBaseHook.fallback_to_default_project_id
+ @GoogleBaseHook.fallback_to_default_project_id
  def create_transfer_config(
  self,
  transfer_config: Union[dict, TransferConfig],
@@ -138,7 +138,7 @@ def create_transfer_config(
  metadata=metadata,
  )
 
- @CloudBaseHook.fallback_to_default_project_id
+ @GoogleBaseHook.fallback_to_default_project_id
  def delete_transfer_config(
  self,
  transfer_config_id: str,
@@ -176,7 +176,7 @@ def delete_transfer_config(
  name=name, retry=retry, timeout=timeout, metadata=metadata
  )
 
- @CloudBaseHook.fallback_to_default_project_id
+ @GoogleBaseHook.fallback_to_default_project_id
  def start_manual_transfer_runs(
  self,
  transfer_config_id: str,
@@ -233,7 +233,7 @@ def start_manual_transfer_runs(
  metadata=metadata,
  )
 
- @CloudBaseHook.fallback_to_default_project_id
+ @GoogleBaseHook.fallback_to_default_project_id
  def get_transfer_run(
  self,
  run_id: str,

diff --git a/airflow/providers/google/cloud/hooks/bigtable.py b/airflow/providers/google/cloud/hooks/bigtable.py
@@ -27,10 +27,10 @@
 from google.cloud.bigtable.table import ClusterState, Table
 from google.cloud.bigtable_admin_v2 import enums
 
-from airflow.providers.google.cloud.hooks.base import CloudBaseHook
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 
 
-class BigtableHook(CloudBaseHook):
+class BigtableHook(GoogleBaseHook):
  """
  Hook for Google Cloud Bigtable APIs.
 
@@ -53,7 +53,7 @@ def _get_client(self, project_id: str):
  )
  return self._client
 
- @CloudBaseHook.fallback_to_default_project_id
+ @GoogleBaseHook.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.
@@ -74,7 +74,7 @@ def get_instance(self, instance_id: str, project_id: Optional[str] = None) -> In
  return None
  return instance
 
- @CloudBaseHook.fallback_to_default_project_id
+ @GoogleBaseHook.fallback_to_default_project_id
  def delete_instance(self, instance_id: str, project_id: Optional[str] = None) -> None:
  """
  Deletes the specified Cloud Bigtable instance.
@@ -97,7 +97,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)
 
- @CloudBaseHook.fallback_to_default_project_id
+ @GoogleBaseHook.fallback_to_default_project_id
  def create_instance(
  self,
  instance_id: str,
@@ -212,7 +212,7 @@ def create_table(
  table = Table(table_id, instance)
  table.create(initial_split_keys, column_families)
 
- @CloudBaseHook.fallback_to_default_project_id
+ @GoogleBaseHook.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.

diff --git a/airflow/providers/google/cloud/hooks/cloud_build.py b/airflow/providers/google/cloud/hooks/cloud_build.py
@@ -23,14 +23,14 @@
 from googleapiclient.discovery import build
 
 from airflow.exceptions import AirflowException
-from airflow.providers.google.cloud.hooks.base import CloudBaseHook
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 
 # Time to sleep between active checks of the operation results
 TIME_TO_SLEEP_IN_SECONDS = 5
 
 
 # noinspection PyAbstractClass
-class CloudBuildHook(CloudBaseHook):
+class CloudBuildHook(GoogleBaseHook):
  """
  Hook for the Google Cloud Build APIs.
 
@@ -69,7 +69,7 @@ def get_conn(self):
  self._conn = build("cloudbuild", self.api_version, http=http_authorized, cache_discovery=False)
  return self._conn
 
- @CloudBaseHook.fallback_to_default_project_id
+ @GoogleBaseHook.fallback_to_default_project_id
  def create_build(self, body: Dict, project_id: Optional[str] = None) -> Dict:
  """
  Starts a build with the specified configuration.