Make GoogleBaseHook credentials functions public (#25785)

apache · Aug 19, 2022 · 6e66dd7 · 6e66dd7
1 parent 728a3ce
commit 6e66dd7
Show file tree

Hide file tree

Showing 64 changed files with 218 additions and 218 deletions.
diff --git a/airflow/providers/google/cloud/hooks/automl.py b/airflow/providers/google/cloud/hooks/automl.py
@@ -86,7 +86,7 @@ def get_conn(self) -> AutoMlClient:
  :rtype: google.cloud.automl_v1beta1.AutoMlClient
  """
  if self._client is None:
- self._client = AutoMlClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
+ self._client = AutoMlClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
  return self._client
 
  @cached_property
@@ -97,7 +97,7 @@ def prediction_client(self) -> PredictionServiceClient:
  :return: Google Cloud AutoML PredictionServiceClient client object.
  :rtype: google.cloud.automl_v1beta1.PredictionServiceClient
  """
- return PredictionServiceClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
+ return PredictionServiceClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
 
  @GoogleBaseHook.fallback_to_default_project_id
  def create_model(

diff --git a/airflow/providers/google/cloud/hooks/bigquery.py b/airflow/providers/google/cloud/hooks/bigquery.py
@@ -136,7 +136,7 @@ def get_client(self, project_id: Optional[str] = None, location: Optional[str] =
  client_info=CLIENT_INFO,
  project=project_id,
  location=location,
- credentials=self._get_credentials(),
+ credentials=self.get_credentials(),
  )
 
  def get_uri(self) -> str:
@@ -247,7 +247,7 @@ def get_pandas_df(
  if dialect is None:
  dialect = 'legacy' if self.use_legacy_sql else 'standard'
 
- credentials, project_id = self._get_credentials_and_project_id()
+ credentials, project_id = self.get_credentials_and_project_id()
 
  return read_gbq(
  sql, project_id=project_id, dialect=dialect, verbose=False, credentials=credentials, **kwargs

diff --git a/airflow/providers/google/cloud/hooks/bigquery_dts.py b/airflow/providers/google/cloud/hooks/bigquery_dts.py
@@ -94,7 +94,7 @@ def get_conn(self) -> DataTransferServiceClient:
  """
  if not self._conn:
  self._conn = DataTransferServiceClient(
- credentials=self._get_credentials(), client_info=CLIENT_INFO
+ credentials=self.get_credentials(), client_info=CLIENT_INFO
  )
  return self._conn
 

diff --git a/airflow/providers/google/cloud/hooks/bigtable.py b/airflow/providers/google/cloud/hooks/bigtable.py
@@ -55,7 +55,7 @@ def _get_client(self, project_id: str):
  if not self._client:
  self._client = Client(
  project=project_id,
- credentials=self._get_credentials(),
+ credentials=self.get_credentials(),
  client_info=CLIENT_INFO,
  admin=True,
  )

diff --git a/airflow/providers/google/cloud/hooks/cloud_build.py b/airflow/providers/google/cloud/hooks/cloud_build.py
@@ -86,7 +86,7 @@ def get_conn(self) -> CloudBuildClient:
  :rtype: `google.cloud.devtools.cloudbuild_v1.CloudBuildClient`
  """
  if not self._client:
- self._client = CloudBuildClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
+ self._client = CloudBuildClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
  return self._client
 
  @GoogleBaseHook.fallback_to_default_project_id

diff --git a/airflow/providers/google/cloud/hooks/cloud_composer.py b/airflow/providers/google/cloud/hooks/cloud_composer.py
@@ -43,7 +43,7 @@ class CloudComposerHook(GoogleBaseHook):
  def get_environment_client(self) -> EnvironmentsClient:
  """Retrieves client library object that allow access Environments service."""
  return EnvironmentsClient(
- credentials=self._get_credentials(),
+ credentials=self.get_credentials(),
  client_info=CLIENT_INFO,
  client_options=self.client_options,
  )
@@ -53,7 +53,7 @@ def get_image_versions_client(
  ) -> ImageVersionsClient:
  """Retrieves client library object that allow access Image Versions service."""
  return ImageVersionsClient(
- credentials=self._get_credentials(),
+ credentials=self.get_credentials(),
  client_info=CLIENT_INFO,
  client_options=self.client_options,
  )

diff --git a/airflow/providers/google/cloud/hooks/cloud_memorystore.py b/airflow/providers/google/cloud/hooks/cloud_memorystore.py
@@ -84,7 +84,7 @@ def __init__(
  def get_conn(self) -> CloudRedisClient:
  """Retrieves client library object that allow access to Cloud Memorystore service."""
  if not self._client:
- self._client = CloudRedisClient(credentials=self._get_credentials())
+ self._client = CloudRedisClient(credentials=self.get_credentials())
  return self._client
 
  @staticmethod
@@ -522,7 +522,7 @@ def get_conn(
  ):
  """Retrieves client library object that allow access to Cloud Memorystore Memcached service."""
  if not self._client:
- self._client = CloudMemcacheClient(credentials=self._get_credentials())
+ self._client = CloudMemcacheClient(credentials=self.get_credentials())
  return self._client
 
  @staticmethod

diff --git a/airflow/providers/google/cloud/hooks/datacatalog.py b/airflow/providers/google/cloud/hooks/datacatalog.py
@@ -71,7 +71,7 @@ def __init__(
  def get_conn(self) -> DataCatalogClient:
  """Retrieves client library object that allow access to Cloud Data Catalog service."""
  if not self._client:
- self._client = DataCatalogClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
+ self._client = DataCatalogClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
  return self._client
 
  @GoogleBaseHook.fallback_to_default_project_id

diff --git a/airflow/providers/google/cloud/hooks/dataform.py b/airflow/providers/google/cloud/hooks/dataform.py
@@ -34,7 +34,7 @@ def get_dataform_client(
  self,
  ) -> DataformClient:
  """Retrieves client library object that allow access to Cloud Dataform service."""
- return DataformClient(credentials=self._get_credentials())
+ return DataformClient(credentials=self.get_credentials())
 
  @GoogleBaseHook.fallback_to_default_project_id
  def wait_for_workflow_invocation(

diff --git a/airflow/providers/google/cloud/hooks/datafusion.py b/airflow/providers/google/cloud/hooks/datafusion.py
@@ -143,7 +143,7 @@ def _cdap_request(
  headers: Dict[str, str] = {"Content-Type": "application/json"}
  request = google.auth.transport.requests.Request()
 
- credentials = self._get_credentials()
+ credentials = self.get_credentials()
  credentials.before_request(request=request, method=method, url=url, headers=headers)
 
  payload = json.dumps(body) if body else None

diff --git a/airflow/providers/google/cloud/hooks/dataplex.py b/airflow/providers/google/cloud/hooks/dataplex.py
@@ -69,7 +69,7 @@ def get_dataplex_client(self) -> DataplexServiceClient:
  client_options = ClientOptions(api_endpoint='dataplex.googleapis.com:443')
 
  return DataplexServiceClient(
- credentials=self._get_credentials(), client_info=self.client_info, client_options=client_options
+ credentials=self.get_credentials(), client_info=self.client_info, client_options=client_options
  )
 
  def wait_for_operation(self, timeout: Optional[float], operation: Operation):

diff --git a/airflow/providers/google/cloud/hooks/dataproc.py b/airflow/providers/google/cloud/hooks/dataproc.py
@@ -207,7 +207,7 @@ def get_cluster_client(self, region: Optional[str] = None) -> ClusterControllerC
  client_options = ClientOptions(api_endpoint=f'{region}-dataproc.googleapis.com:443')
 
  return ClusterControllerClient(
- credentials=self._get_credentials(), client_info=CLIENT_INFO, client_options=client_options
+ credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options
  )
 
  def get_template_client(self, region: Optional[str] = None) -> WorkflowTemplateServiceClient:
@@ -217,7 +217,7 @@ def get_template_client(self, region: Optional[str] = None) -> WorkflowTemplateS
  client_options = ClientOptions(api_endpoint=f'{region}-dataproc.googleapis.com:443')
 
  return WorkflowTemplateServiceClient(
- credentials=self._get_credentials(), client_info=CLIENT_INFO, client_options=client_options
+ credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options
  )
 
  def get_job_client(self, region: Optional[str] = None) -> JobControllerClient:
@@ -227,7 +227,7 @@ def get_job_client(self, region: Optional[str] = None) -> JobControllerClient:
  client_options = ClientOptions(api_endpoint=f'{region}-dataproc.googleapis.com:443')
 
  return JobControllerClient(
- credentials=self._get_credentials(), client_info=CLIENT_INFO, client_options=client_options
+ credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options
  )
 
  def get_batch_client(self, region: Optional[str] = None) -> BatchControllerClient:
@@ -237,7 +237,7 @@ def get_batch_client(self, region: Optional[str] = None) -> BatchControllerClien
  client_options = ClientOptions(api_endpoint=f'{region}-dataproc.googleapis.com:443')
 
  return BatchControllerClient(
- credentials=self._get_credentials(), client_info=CLIENT_INFO, client_options=client_options
+ credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options
  )
 
  def wait_for_operation(self, operation: Operation, timeout: Optional[float] = None):

diff --git a/airflow/providers/google/cloud/hooks/dataproc_metastore.py b/airflow/providers/google/cloud/hooks/dataproc_metastore.py
@@ -42,7 +42,7 @@ def get_dataproc_metastore_client(self) -> DataprocMetastoreClient:
  client_options = ClientOptions(api_endpoint='metastore.googleapis.com:443')
 
  return DataprocMetastoreClient(
- credentials=self._get_credentials(), client_info=CLIENT_INFO, client_options=client_options
+ credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options
  )
 
  def wait_for_operation(self, timeout: Optional[float], operation: Operation):

diff --git a/airflow/providers/google/cloud/hooks/dlp.py b/airflow/providers/google/cloud/hooks/dlp.py
@@ -106,7 +106,7 @@ def get_conn(self) -> DlpServiceClient:
  :rtype: google.cloud.dlp_v2.DlpServiceClient
  """
  if not self._client:
- self._client = DlpServiceClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
+ self._client = DlpServiceClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
  return self._client
 
  @GoogleBaseHook.fallback_to_default_project_id

diff --git a/airflow/providers/google/cloud/hooks/gcs.py b/airflow/providers/google/cloud/hooks/gcs.py
@@ -157,7 +157,7 @@ def get_conn(self) -> storage.Client:
  """Returns a Google Cloud Storage service object."""
  if not self._conn:
  self._conn = storage.Client(
- credentials=self._get_credentials(), client_info=CLIENT_INFO, project=self.project_id
+ credentials=self.get_credentials(), client_info=CLIENT_INFO, project=self.project_id
  )
 
  return self._conn

diff --git a/airflow/providers/google/cloud/hooks/kms.py b/airflow/providers/google/cloud/hooks/kms.py
@@ -80,7 +80,7 @@ def get_conn(self) -> KeyManagementServiceClient:
  """
  if not self._conn:
  self._conn = KeyManagementServiceClient(
- credentials=self._get_credentials(), client_info=CLIENT_INFO
+ credentials=self.get_credentials(), client_info=CLIENT_INFO
  )
  return self._conn
 

diff --git a/airflow/providers/google/cloud/hooks/kubernetes_engine.py b/airflow/providers/google/cloud/hooks/kubernetes_engine.py
@@ -73,7 +73,7 @@ def __init__(
  def get_cluster_manager_client(self) -> ClusterManagerClient:
  """Returns ClusterManagerClient."""
  if self._client is None:
- self._client = ClusterManagerClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
+ self._client = ClusterManagerClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
  return self._client
 
  # To preserve backward compatibility

diff --git a/airflow/providers/google/cloud/hooks/natural_language.py b/airflow/providers/google/cloud/hooks/natural_language.py
@@ -75,7 +75,7 @@ def get_conn(self) -> LanguageServiceClient:
  :rtype: google.cloud.language_v1.LanguageServiceClient
  """
  if not self._conn:
- self._conn = LanguageServiceClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
+ self._conn = LanguageServiceClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
  return self._conn
 
  @GoogleBaseHook.quota_retry()

diff --git a/airflow/providers/google/cloud/hooks/os_login.py b/airflow/providers/google/cloud/hooks/os_login.py
@@ -58,7 +58,7 @@ def get_conn(self) -> OsLoginServiceClient:
  if self._conn:
  return self._conn
 
- self._conn = OsLoginServiceClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
+ self._conn = OsLoginServiceClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
  return self._conn
 
  @GoogleBaseHook.fallback_to_default_project_id

diff --git a/airflow/providers/google/cloud/hooks/pubsub.py b/airflow/providers/google/cloud/hooks/pubsub.py
@@ -83,7 +83,7 @@ def get_conn(self) -> PublisherClient:
  :rtype: google.cloud.pubsub_v1.PublisherClient
  """
  if not self._client:
- self._client = PublisherClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
+ self._client = PublisherClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
  return self._client
 
  @cached_property
@@ -94,7 +94,7 @@ def subscriber_client(self) -> SubscriberClient:
  :return: Google Cloud Pub/Sub client object.
  :rtype: google.cloud.pubsub_v1.SubscriberClient
  """
- return SubscriberClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
+ return SubscriberClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
 
  @GoogleBaseHook.fallback_to_default_project_id
  def publish(

diff --git a/airflow/providers/google/cloud/hooks/secret_manager.py b/airflow/providers/google/cloud/hooks/secret_manager.py
@@ -56,7 +56,7 @@ def __init__(
  delegate_to=delegate_to,
  impersonation_chain=impersonation_chain,
  )
- self.client = _SecretManagerClient(credentials=self._get_credentials())
+ self.client = _SecretManagerClient(credentials=self.get_credentials())
 
  def get_conn(self) -> _SecretManagerClient:
  """

diff --git a/airflow/providers/google/cloud/hooks/spanner.py b/airflow/providers/google/cloud/hooks/spanner.py
@@ -61,7 +61,7 @@ def _get_client(self, project_id: str) -> Client:
  """
  if not self._client:
  self._client = Client(
- project=project_id, credentials=self._get_credentials(), client_info=CLIENT_INFO
+ project=project_id, credentials=self.get_credentials(), client_info=CLIENT_INFO
  )
  return self._client
 

diff --git a/airflow/providers/google/cloud/hooks/speech_to_text.py b/airflow/providers/google/cloud/hooks/speech_to_text.py
@@ -66,7 +66,7 @@ def get_conn(self) -> SpeechClient:
  :rtype: google.cloud.speech_v1.SpeechClient
  """
  if not self._client:
- self._client = SpeechClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
+ self._client = SpeechClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
  return self._client
 
  @GoogleBaseHook.quota_retry()

diff --git a/airflow/providers/google/cloud/hooks/tasks.py b/airflow/providers/google/cloud/hooks/tasks.py
@@ -78,7 +78,7 @@ def get_conn(self) -> CloudTasksClient:
  :rtype: google.cloud.tasks_v2.CloudTasksClient
  """
  if self._client is None:
- self._client = CloudTasksClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
+ self._client = CloudTasksClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
  return self._client
 
  @GoogleBaseHook.fallback_to_default_project_id

diff --git a/airflow/providers/google/cloud/hooks/text_to_speech.py b/airflow/providers/google/cloud/hooks/text_to_speech.py
@@ -75,7 +75,7 @@ def get_conn(self) -> TextToSpeechClient:
  """
  if not self._client:
 
- self._client = TextToSpeechClient(credentials=self._get_credentials(), client_info=CLIENT_INFO)
+ self._client = TextToSpeechClient(credentials=self.get_credentials(), client_info=CLIENT_INFO)
 
  return self._client
 

diff --git a/airflow/providers/google/cloud/hooks/translate.py b/airflow/providers/google/cloud/hooks/translate.py
@@ -53,7 +53,7 @@ def get_conn(self) -> Client:
  :rtype: google.cloud.translate_v2.Client
  """
  if not self._client:
- self._client = Client(credentials=self._get_credentials(), client_info=CLIENT_INFO)
+ self._client = Client(credentials=self.get_credentials(), client_info=CLIENT_INFO)
  return self._client
 
  @GoogleBaseHook.quota_retry()

diff --git a/airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py b/airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py
@@ -103,7 +103,7 @@ def get_pipeline_service_client(
  client_options = ClientOptions()
 
  return PipelineServiceClient(
- credentials=self._get_credentials(), client_info=self.client_info, client_options=client_options
+ credentials=self.get_credentials(), client_info=self.client_info, client_options=client_options
  )
 
  def get_job_service_client(
@@ -117,7 +117,7 @@ def get_job_service_client(
  client_options = ClientOptions()
 
  return JobServiceClient(
- credentials=self._get_credentials(), client_info=self.client_info, client_options=client_options
+ credentials=self.get_credentials(), client_info=self.client_info, client_options=client_options
  )
 
  def get_auto_ml_tabular_training_job(
@@ -146,7 +146,7 @@ def get_auto_ml_tabular_training_job(
  optimization_objective_precision_value=optimization_objective_precision_value,
  project=project,
  location=location,
- credentials=self._get_credentials(),
+ credentials=self.get_credentials(),
  labels=labels,
  training_encryption_spec_key_name=training_encryption_spec_key_name,
  model_encryption_spec_key_name=model_encryption_spec_key_name,
@@ -172,7 +172,7 @@ def get_auto_ml_forecasting_training_job(
  column_transformations=column_transformations,
  project=project,
  location=location,
- credentials=self._get_credentials(),
+ credentials=self.get_credentials(),
  labels=labels,
  training_encryption_spec_key_name=training_encryption_spec_key_name,
  model_encryption_spec_key_name=model_encryption_spec_key_name,
@@ -200,7 +200,7 @@ def get_auto_ml_image_training_job(
  base_model=base_model,
  project=project,
  location=location,
- credentials=self._get_credentials(),
+ credentials=self.get_credentials(),
  labels=labels,
  training_encryption_spec_key_name=training_encryption_spec_key_name,
  model_encryption_spec_key_name=model_encryption_spec_key_name,
@@ -226,7 +226,7 @@ def get_auto_ml_text_training_job(
  sentiment_max=sentiment_max,
  project=project,
  location=location,
- credentials=self._get_credentials(),
+ credentials=self.get_credentials(),
  labels=labels,
  training_encryption_spec_key_name=training_encryption_spec_key_name,
  model_encryption_spec_key_name=model_encryption_spec_key_name,
@@ -250,7 +250,7 @@ def get_auto_ml_video_training_job(
  model_type=model_type,
  project=project,
  location=location,
- credentials=self._get_credentials(),
+ credentials=self.get_credentials(),
  labels=labels,
  training_encryption_spec_key_name=training_encryption_spec_key_name,
  model_encryption_spec_key_name=model_encryption_spec_key_name,

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
@@ -64,7 +64,7 @@ def get_job_service_client(self, region: Optional[str] = None) -> JobServiceClie
  client_options = ClientOptions()
 
  return JobServiceClient(
- credentials=self._get_credentials(), client_info=self.client_info, client_options=client_options
+ credentials=self.get_credentials(), client_info=self.client_info, client_options=client_options
  )
 
  def wait_for_operation(self, operation: Operation, timeout: Optional[float] = None):
@@ -225,7 +225,7 @@ def create_batch_prediction_job(
  labels=labels,
  project=project_id,
  location=region,
- credentials=self._get_credentials(),
+ credentials=self.get_credentials(),
  encryption_spec_key_name=encryption_spec_key_name,
  sync=sync,
  )