Install sqlalchemy-spanner package into Google provider (#31925)

MaksYermak · web-flow · commit 978adb309aee · 2023-07-21T11:57:04.000+02:00
diff --git a/airflow/providers/google/cloud/hooks/spanner.py b/airflow/providers/google/cloud/hooks/spanner.py
@@ -18,28 +18,43 @@
 """This module contains a Google Cloud Spanner Hook."""
 from __future__ import annotations
 
-from typing import Callable, Sequence
+from typing import Callable, NamedTuple, Sequence
 
 from google.api_core.exceptions import AlreadyExists, GoogleAPICallError
 from google.cloud.spanner_v1.client import Client
 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
+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
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook, get_field
 
 
-class SpannerHook(GoogleBaseHook):
+class SpannerConnectionParams(NamedTuple):
+    """Information about Google Spanner connection parameters."""
+
+    project_id: str | None
+    instance_id: str | None
+    database_id: str | None
+
+
+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.
     """
 
+    conn_name_attr = "gcp_conn_id"
+    default_conn_name = "google_cloud_spanner_default"
+    conn_type = "gcpspanner"
+    hook_name = "Google Cloud Spanner"
+
     def __init__(
         self,
         gcp_conn_id: str = "google_cloud_default",
@@ -70,6 +85,34 @@ def _get_client(self, project_id: str) -> Client:
             )
         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)
+
+    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}"
+
+    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
     def get_instance(
         self,
diff --git a/airflow/providers/google/provider.yaml b/airflow/providers/google/provider.yaml
@@ -126,6 +126,7 @@ dependencies:
   - proto-plus>=1.19.6
   - PyOpenSSL
   - sqlalchemy-bigquery>=1.2.1
+  - sqlalchemy-spanner>=1.6.2
 
 integrations:
   - integration-name: Google Analytics360
diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json
@@ -440,7 +440,8 @@
       "pandas-gbq",
       "pandas>=0.17.1",
       "proto-plus>=1.19.6",
-      "sqlalchemy-bigquery>=1.2.1"
+      "sqlalchemy-bigquery>=1.2.1",
+      "sqlalchemy-spanner>=1.6.2"
     ],
     "cross-providers-deps": [
       "amazon",
diff --git a/tests/providers/google/cloud/hooks/test_spanner.py b/tests/providers/google/cloud/hooks/test_spanner.py
@@ -18,9 +18,10 @@
 from __future__ import annotations
 
 from unittest import mock
-from unittest.mock import PropertyMock
+from unittest.mock import MagicMock, PropertyMock
 
 import pytest
+import sqlalchemy
 
 from airflow.providers.google.cloud.hooks.spanner import SpannerHook
 from airflow.providers.google.common.consts import CLIENT_INFO
@@ -33,6 +34,8 @@
 SPANNER_INSTANCE = "instance"
 SPANNER_CONFIGURATION = "configuration"
 SPANNER_DATABASE = "database-name"
+SPANNER_PROJECT_ID = "test_project_id"
+SPANNER_CONN_PARAMS = (SPANNER_PROJECT_ID, SPANNER_INSTANCE, SPANNER_DATABASE)
 
 
 class TestGcpSpannerHookDefaultProjectId:
@@ -431,6 +434,21 @@ def test_execute_dml_overridden_project_id(self, get_client):
         run_in_transaction_method.assert_called_once_with(mock.ANY)
         assert res is None
 
+    def test_get_uri(self):
+        self.spanner_hook_default_project_id._get_conn_params = MagicMock(return_value=SPANNER_CONN_PARAMS)
+        uri = self.spanner_hook_default_project_id.get_uri()
+        assert (
+            uri
+            == f"spanner+spanner:///projects/{SPANNER_PROJECT_ID}/instances/{SPANNER_INSTANCE}/databases/{SPANNER_DATABASE}"
+        )
+
+    @mock.patch("airflow.providers.google.cloud.hooks.spanner.SpannerHook._get_client")
+    def test_get_sqlalchemy_engine(self, get_client):
+        self.spanner_hook_default_project_id._get_conn_params = MagicMock(return_value=SPANNER_CONN_PARAMS)
+        engine = self.spanner_hook_default_project_id.get_sqlalchemy_engine()
+        assert isinstance(engine, sqlalchemy.engine.Engine)
+        assert engine.name == "spanner+spanner"
+
 
 class TestGcpSpannerHookNoDefaultProjectID:
     def setup_method(self):
@@ -675,3 +693,18 @@ def test_execute_dml_overridden_project_id(self, get_client):
         database_method.assert_called_once_with(database_id="database-name")
         run_in_transaction_method.assert_called_once_with(mock.ANY)
         assert res is None
+
+    def test_get_uri(self):
+        self.spanner_hook_no_default_project_id._get_conn_params = MagicMock(return_value=SPANNER_CONN_PARAMS)
+        uri = self.spanner_hook_no_default_project_id.get_uri()
+        assert (
+            uri
+            == f"spanner+spanner:///projects/{SPANNER_PROJECT_ID}/instances/{SPANNER_INSTANCE}/databases/{SPANNER_DATABASE}"
+        )
+
+    @mock.patch("airflow.providers.google.cloud.hooks.spanner.SpannerHook._get_client")
+    def test_get_sqlalchemy_engine(self, get_client):
+        self.spanner_hook_no_default_project_id._get_conn_params = MagicMock(return_value=SPANNER_CONN_PARAMS)
+        engine = self.spanner_hook_no_default_project_id.get_sqlalchemy_engine()
+        assert isinstance(engine, sqlalchemy.engine.Engine)
+        assert engine.name == "spanner+spanner"
