Source code for airflow.providers.google.cloud.hooks.spanner

#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""This module contains a Google Cloud Spanner Hook."""
from __future__ import annotations

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 sqlalchemy import create_engine

from airflow.exceptions import AirflowException
from airflow.providers.common.sql.hooks.sql import DbApiHook
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


[docs]class SpannerConnectionParams(NamedTuple): """Information about Google Spanner connection parameters."""
[docs] project_id: str | None
[docs] instance_id: str | None
[docs] database_id: str | None
[docs]class SpannerHook(GoogleBaseHook, DbApiHook): """ Hook for Google Cloud Spanner APIs. All the methods in the hook where project_id is used must be called with keyword arguments rather than positional. """
[docs] conn_name_attr = "gcp_conn_id"
[docs] default_conn_name = "google_cloud_spanner_default"
[docs] conn_type = "gcpspanner"
[docs] hook_name = "Google Cloud Spanner"
def __init__( self, gcp_conn_id: str = "google_cloud_default", impersonation_chain: str | Sequence[str] | None = None, **kwargs, ) -> None: if kwargs.get("delegate_to") is not None: raise RuntimeError( "The `delegate_to` parameter has been deprecated before and finally removed in this version" " of Google Provider. You MUST convert it to `impersonate_chain`" ) super().__init__( gcp_conn_id=gcp_conn_id, impersonation_chain=impersonation_chain, ) self._client: Client | None = None def _get_client(self, project_id: str) -> Client: """ Provide a client for interacting with the Cloud Spanner API. :param project_id: The ID of the Google Cloud project. :return: Client """ if not self._client: self._client = Client( project=project_id, credentials=self.get_credentials(), client_info=CLIENT_INFO ) return self._client def _get_conn_params(self) -> SpannerConnectionParams: """Extract spanner database connection parameters.""" extras = self.get_connection(self.gcp_conn_id).extra_dejson project_id = get_field(extras, "project_id") or self.project_id instance_id = get_field(extras, "instance_id") database_id = get_field(extras, "database_id") return SpannerConnectionParams(project_id, instance_id, database_id)
[docs] def get_uri(self) -> str: """Override DbApiHook get_uri method for get_sqlalchemy_engine().""" project_id, instance_id, database_id = self._get_conn_params() if not all([instance_id, database_id]): raise AirflowException("The instance_id or database_id were not specified") return f"spanner+spanner:///projects/{project_id}/instances/{instance_id}/databases/{database_id}"
[docs] def get_sqlalchemy_engine(self, engine_kwargs=None): """ Get an sqlalchemy_engine object. :param engine_kwargs: Kwargs used in :func:`~sqlalchemy.create_engine`. :return: the created engine. """ if engine_kwargs is None: engine_kwargs = {} project_id, _, _ = self._get_conn_params() spanner_client = self._get_client(project_id=project_id) return create_engine(self.get_uri(), connect_args={"client": spanner_client}, **engine_kwargs)
@GoogleBaseHook.fallback_to_default_project_id
[docs] def get_instance( self, instance_id: str, project_id: str, ) -> Instance | None: """ Get information about a particular instance. :param project_id: Optional, The ID of the Google Cloud project that owns the Cloud Spanner database. If set to None or missing, the default project_id from the Google Cloud connection is used. :param instance_id: The ID of the Cloud Spanner instance. :return: Spanner instance """ instance = self._get_client(project_id=project_id).instance(instance_id=instance_id) if not instance.exists(): return None return instance
def _apply_to_instance( self, project_id: str, instance_id: str, configuration_name: str, node_count: int, display_name: str, func: Callable[[Instance], Operation], ) -> None: """ Invoke a method on a given instance by applying a specified Callable. :param project_id: The ID of the Google Cloud project that owns the Cloud Spanner database. :param instance_id: The ID of the instance. :param configuration_name: Name of the instance configuration defining how the instance will be created. Required for instances which do not yet exist. :param node_count: (Optional) Number of nodes allocated to the instance. :param display_name: (Optional) The display name for the instance in the Cloud Console UI. (Must be between 4 and 30 characters.) If this value is not set in the constructor, will fall back to the instance ID. :param func: Method of the instance to be called. """ instance = self._get_client(project_id=project_id).instance( instance_id=instance_id, configuration_name=configuration_name, node_count=node_count, display_name=display_name, ) try: operation: Operation = func(instance) except GoogleAPICallError as e: self.log.error("An error occurred: %s. Exiting.", e.message) raise e if operation: result = operation.result() self.log.info(result) @GoogleBaseHook.fallback_to_default_project_id
[docs] def create_instance( self, instance_id: str, configuration_name: str, node_count: int, display_name: str, project_id: str, ) -> None: """ Create a new Cloud Spanner instance. :param instance_id: The ID of the Cloud Spanner instance. :param configuration_name: The name of the instance configuration defining how the instance will be created. Possible configuration values can be retrieved via https://cloud.google.com/spanner/docs/reference/rest/v1/projects.instanceConfigs/list :param node_count: (Optional) The number of nodes allocated to the Cloud Spanner instance. :param display_name: (Optional) The display name for the instance in the Google Cloud Console. Must be between 4 and 30 characters. If this value is not passed, the name falls back to the instance ID. :param project_id: Optional, the ID of the Google Cloud project that owns the Cloud Spanner database. If set to None or missing, the default project_id from the Google Cloud connection is used. :return: None """ self._apply_to_instance( project_id, instance_id, configuration_name, node_count, display_name, lambda x: x.create() )
@GoogleBaseHook.fallback_to_default_project_id
[docs] def update_instance( self, instance_id: str, configuration_name: str, node_count: int, display_name: str, project_id: str, ) -> None: """ Update an existing Cloud Spanner instance. :param instance_id: The ID of the Cloud Spanner instance. :param configuration_name: The name of the instance configuration defining how the instance will be created. Possible configuration values can be retrieved via https://cloud.google.com/spanner/docs/reference/rest/v1/projects.instanceConfigs/list :param node_count: (Optional) The number of nodes allocated to the Cloud Spanner instance. :param display_name: (Optional) The display name for the instance in the Google Cloud Console. Must be between 4 and 30 characters. If this value is not set in the constructor, the name falls back to the instance ID. :param project_id: Optional, the ID of the Google Cloud project that owns the Cloud Spanner database. If set to None or missing, the default project_id from the Google Cloud connection is used. :return: None """ self._apply_to_instance( project_id, instance_id, configuration_name, node_count, display_name, lambda x: x.update() )
@GoogleBaseHook.fallback_to_default_project_id
[docs] def delete_instance(self, instance_id: str, project_id: str) -> None: """ Delete an existing Cloud Spanner instance. :param instance_id: The ID of the Cloud Spanner instance. :param project_id: Optional, the ID of the Google Cloud project that owns the Cloud Spanner database. If set to None or missing, the default project_id from the Google Cloud connection is used. :return: None """ instance = self._get_client(project_id=project_id).instance(instance_id) try: instance.delete() return except GoogleAPICallError as e: self.log.error("An error occurred: %s. Exiting.", e.message) raise e
@GoogleBaseHook.fallback_to_default_project_id
[docs] def get_database( self, instance_id: str, database_id: str, project_id: str, ) -> Database | None: """ Retrieve a database in Cloud Spanner; return None if the database does not exist in the instance. :param instance_id: The ID of the Cloud Spanner instance. :param database_id: The ID of the database in Cloud Spanner. :param project_id: Optional, the ID of the Google Cloud project that owns the Cloud Spanner database. If set to None or missing, the default project_id from the Google Cloud connection is used. :return: Database object or None if database does not exist """ instance = self._get_client(project_id=project_id).instance(instance_id=instance_id) if not instance.exists(): raise AirflowException(f"The instance {instance_id} does not exist in project {project_id} !") database = instance.database(database_id=database_id) if not database.exists(): return None return database
@GoogleBaseHook.fallback_to_default_project_id
[docs] def create_database( self, instance_id: str, database_id: str, ddl_statements: list[str], project_id: str, ) -> None: """ Create a new database in Cloud Spanner. :param instance_id: The ID of the Cloud Spanner instance. :param database_id: The ID of the database to create in Cloud Spanner. :param ddl_statements: The string list containing DDL for the new database. :param project_id: Optional, the ID of the Google Cloud project that owns the Cloud Spanner database. If set to None or missing, the default project_id from the Google Cloud connection is used. :return: None """ instance = self._get_client(project_id=project_id).instance(instance_id=instance_id) if not instance.exists(): raise AirflowException(f"The instance {instance_id} does not exist in project {project_id} !") database = instance.database(database_id=database_id, ddl_statements=ddl_statements) try: operation: Operation = database.create() except GoogleAPICallError as e: self.log.error("An error occurred: %s. Exiting.", e.message) raise e if operation: result = operation.result() self.log.info(result)
@GoogleBaseHook.fallback_to_default_project_id
[docs] def update_database( self, instance_id: str, database_id: str, ddl_statements: list[str], project_id: str, operation_id: str | None = None, ) -> None: """ Update DDL of a database in Cloud Spanner. :param instance_id: The ID of the Cloud Spanner instance. :param database_id: The ID of the database in Cloud Spanner. :param ddl_statements: The string list containing DDL for the new database. :param project_id: Optional, the ID of the Google Cloud project that owns the Cloud Spanner database. If set to None or missing, the default project_id from the Google Cloud connection is used. :param operation_id: (Optional) The unique per database operation ID that can be specified to implement idempotency check. :return: None """ instance = self._get_client(project_id=project_id).instance(instance_id=instance_id) if not instance.exists(): raise AirflowException(f"The instance {instance_id} does not exist in project {project_id} !") database = instance.database(database_id=database_id) try: operation = database.update_ddl(ddl_statements=ddl_statements, operation_id=operation_id) if operation: result = operation.result() self.log.info(result) return except AlreadyExists as e: if e.code == 409 and operation_id in e.message: self.log.info( "Replayed update_ddl message - the operation id %s was already done before.", operation_id, ) return except GoogleAPICallError as e: self.log.error("An error occurred: %s. Exiting.", e.message) raise e
@GoogleBaseHook.fallback_to_default_project_id
[docs] def delete_database(self, instance_id: str, database_id, project_id: str) -> bool: """ Drop a database in Cloud Spanner. :param instance_id: The ID of the Cloud Spanner instance. :param database_id: The ID of the database in Cloud Spanner. :param project_id: Optional, the ID of the Google Cloud project that owns the Cloud Spanner database. If set to None or missing, the default project_id from the Google Cloud connection is used. :return: True if everything succeeded """ instance = self._get_client(project_id=project_id).instance(instance_id=instance_id) if not instance.exists(): raise AirflowException(f"The instance {instance_id} does not exist in project {project_id} !") database = instance.database(database_id=database_id) if not database.exists(): self.log.info( "The database %s is already deleted from instance %s. Exiting.", database_id, instance_id ) return False try: database.drop() except GoogleAPICallError as e: self.log.error("An error occurred: %s. Exiting.", e.message) raise e return True
@GoogleBaseHook.fallback_to_default_project_id
[docs] def execute_dml( self, instance_id: str, database_id: str, queries: list[str], project_id: str, ) -> None: """ Execute an arbitrary DML query (INSERT, UPDATE, DELETE). :param instance_id: The ID of the Cloud Spanner instance. :param database_id: The ID of the database in Cloud Spanner. :param queries: The queries to execute. :param project_id: Optional, the ID of the Google Cloud project that owns the Cloud Spanner database. If set to None or missing, the default project_id from the Google Cloud connection is used. """ self._get_client(project_id=project_id).instance(instance_id=instance_id).database( database_id=database_id ).run_in_transaction(lambda transaction: self._execute_sql_in_transaction(transaction, queries))
@staticmethod def _execute_sql_in_transaction(transaction: Transaction, queries: list[str]): for sql in queries: transaction.execute_update(sql)

Was this entry helpful?