feat: add support for experimental host (#1452)

* feat: add support for experimental host

* fix lint issues

* fixed unit tests

* added docmentation for new client option

---------

Co-authored-by: rahul2393 <irahul@google.com>

Author: sagnghos

google/cloud/spanner_dbapi/connection.py

@@ -15,6 +15,7 @@
 """DB-API Connection for the Google Cloud Spanner."""
 import warnings
 
+from google.api_core.client_options import ClientOptions
 from google.api_core.exceptions import Aborted
 from google.api_core.gapic_v1.client_info import ClientInfo
 from google.auth.credentials import AnonymousCredentials
@@ -734,6 +735,7 @@ def connect(
     client=None,
     route_to_leader_enabled=True,
     database_role=None,
+    experimental_host=None,
     **kwargs,
 ):
     """Creates a connection to a Google Cloud Spanner database.
@@ -805,6 +807,10 @@ def connect(
             client_options = None
             if isinstance(credentials, AnonymousCredentials):
                 client_options = kwargs.get("client_options")
+            if experimental_host is not None:
+                project = "default"
+                credentials = AnonymousCredentials()
+                client_options = ClientOptions(api_endpoint=experimental_host)
             client = spanner.Client(
                 project=project,
                 credentials=credentials,

google/cloud/spanner_v1/client.py

@@ -176,6 +176,11 @@ class Client(ClientWithProject):
         or :class:`dict`
     :param default_transaction_options: (Optional) Default options to use for all transactions.
 
+    :type experimental_host: str
+    :param experimental_host: (Optional) The endpoint for a spanner experimental host deployment.
+        This is intended only for experimental host spanner endpoints.
+        If set, this will override the `api_endpoint` in `client_options`.
+
     :raises: :class:`ValueError <exceptions.ValueError>` if both ``read_only``
              and ``admin`` are :data:`True`
     """
@@ -200,8 +205,10 @@ def __init__(
         directed_read_options=None,
         observability_options=None,
         default_transaction_options: Optional[DefaultTransactionOptions] = None,
+        experimental_host=None,
     ):
         self._emulator_host = _get_spanner_emulator_host()
+        self._experimental_host = experimental_host
 
         if client_options and type(client_options) is dict:
             self._client_options = google.api_core.client_options.from_dict(
@@ -212,6 +219,8 @@ def __init__(
 
         if self._emulator_host:
             credentials = AnonymousCredentials()
+        elif self._experimental_host:
+            credentials = AnonymousCredentials()
         elif isinstance(credentials, AnonymousCredentials):
             self._emulator_host = self._client_options.api_endpoint
 
@@ -324,6 +333,15 @@ def instance_admin_api(self):
                     client_options=self._client_options,
                     transport=transport,
                 )
+            elif self._experimental_host:
+                transport = InstanceAdminGrpcTransport(
+                    channel=grpc.insecure_channel(target=self._experimental_host)
+                )
+                self._instance_admin_api = InstanceAdminClient(
+                    client_info=self._client_info,
+                    client_options=self._client_options,
+                    transport=transport,
+                )
             else:
                 self._instance_admin_api = InstanceAdminClient(
                     credentials=self.credentials,
@@ -345,6 +363,15 @@ def database_admin_api(self):
                     client_options=self._client_options,
                     transport=transport,
                 )
+            elif self._experimental_host:
+                transport = DatabaseAdminGrpcTransport(
+                    channel=grpc.insecure_channel(target=self._experimental_host)
+                )
+                self._database_admin_api = DatabaseAdminClient(
+                    client_info=self._client_info,
+                    client_options=self._client_options,
+                    transport=transport,
+                )
             else:
                 self._database_admin_api = DatabaseAdminClient(
                     credentials=self.credentials,
@@ -485,6 +512,7 @@ def instance(
             self._emulator_host,
             labels,
             processing_units,
+            self._experimental_host,
         )
 
     def list_instances(self, filter_="", page_size=None):

google/cloud/spanner_v1/database.py

@@ -203,8 +203,11 @@ def __init__(
 
         self._pool = pool
         pool.bind(self)
+        is_experimental_host = self._instance.experimental_host is not None
 
-        self._sessions_manager = DatabaseSessionsManager(self, pool)
+        self._sessions_manager = DatabaseSessionsManager(
+            self, pool, is_experimental_host
+        )
 
     @classmethod
     def from_pb(cls, database_pb, instance, pool=None):
@@ -449,6 +452,16 @@ def spanner_api(self):
                     client_info=client_info, transport=transport
                 )
                 return self._spanner_api
+            if self._instance.experimental_host is not None:
+                transport = SpannerGrpcTransport(
+                    channel=grpc.insecure_channel(self._instance.experimental_host)
+                )
+                self._spanner_api = SpannerClient(
+                    client_info=client_info,
+                    transport=transport,
+                    client_options=client_options,
+                )
+                return self._spanner_api
             credentials = self._instance._client.credentials
             if isinstance(credentials, google.auth.credentials.Scoped):
                 credentials = credentials.with_scopes((SPANNER_DATA_SCOPE,))

google/cloud/spanner_v1/database_sessions_manager.py

@@ -62,9 +62,10 @@ class DatabaseSessionsManager(object):
     _MAINTENANCE_THREAD_POLLING_INTERVAL = timedelta(minutes=10)
     _MAINTENANCE_THREAD_REFRESH_INTERVAL = timedelta(days=7)
 
-    def __init__(self, database, pool):
+    def __init__(self, database, pool, is_experimental_host: bool = False):
         self._database = database
         self._pool = pool
+        self._is_experimental_host = is_experimental_host
 
         # Declare multiplexed session attributes. When a multiplexed session for the
         # database session manager is created, a maintenance thread is initialized to
@@ -88,7 +89,7 @@ def get_session(self, transaction_type: TransactionType) -> Session:
 
         session = (
             self._get_multiplexed_session()
-            if self._use_multiplexed(transaction_type)
+            if self._use_multiplexed(transaction_type) or self._is_experimental_host
             else self._pool.get()
         )
 

google/cloud/spanner_v1/instance.py

@@ -122,6 +122,7 @@ def __init__(
         emulator_host=None,
         labels=None,
         processing_units=None,
+        experimental_host=None,
     ):
         self.instance_id = instance_id
         self._client = client
@@ -142,6 +143,7 @@ def __init__(
             self._node_count = processing_units // PROCESSING_UNITS_PER_NODE
         self.display_name = display_name or instance_id
         self.emulator_host = emulator_host
+        self.experimental_host = experimental_host
         if labels is None:
             labels = {}
         self.labels = labels

google/cloud/spanner_v1/testing/database_test.py

@@ -86,6 +86,18 @@ def spanner_api(self):
                     transport=transport,
                 )
                 return self._spanner_api
+            if self._instance.experimental_host is not None:
+                channel = grpc.insecure_channel(self._instance.experimental_host)
+                self._x_goog_request_id_interceptor = XGoogRequestIDHeaderInterceptor()
+                self._interceptors.append(self._x_goog_request_id_interceptor)
+                channel = grpc.intercept_channel(channel, *self._interceptors)
+                transport = SpannerGrpcTransport(channel=channel)
+                self._spanner_api = SpannerClient(
+                    client_info=client_info,
+                    transport=transport,
+                    client_options=client_options,
+                )
+                return self._spanner_api
             credentials = client.credentials
             if isinstance(credentials, google.auth.credentials.Scoped):
                 credentials = credentials.with_scopes((SPANNER_DATA_SCOPE,))

tests/system/_helpers.py

@@ -56,6 +56,12 @@
 EMULATOR_PROJECT_DEFAULT = "emulator-test-project"
 EMULATOR_PROJECT = os.getenv(EMULATOR_PROJECT_ENVVAR, EMULATOR_PROJECT_DEFAULT)
 
+USE_EXPERIMENTAL_HOST_ENVVAR = "SPANNER_EXPERIMENTAL_HOST"
+EXPERIMENTAL_HOST = os.getenv(USE_EXPERIMENTAL_HOST_ENVVAR)
+USE_EXPERIMENTAL_HOST = EXPERIMENTAL_HOST is not None
+
+EXPERIMENTAL_HOST_PROJECT = "default"
+EXPERIMENTAL_HOST_INSTANCE = "default"
 
 DDL_STATEMENTS = (
     _fixtures.PG_DDL_STATEMENTS

tests/system/conftest.py

@@ -49,6 +49,12 @@ def not_emulator():
         pytest.skip(f"{_helpers.USE_EMULATOR_ENVVAR} set in environment.")
 
 
+@pytest.fixture(scope="module")
+def not_experimental_host():
+    if _helpers.USE_EXPERIMENTAL_HOST:
+        pytest.skip(f"{_helpers.USE_EXPERIMENTAL_HOST_ENVVAR} set in environment.")
+
+
 @pytest.fixture(scope="session")
 def not_postgres(database_dialect):
     if database_dialect == DatabaseDialect.POSTGRESQL:
@@ -104,6 +110,15 @@ def spanner_client():
             project=_helpers.EMULATOR_PROJECT,
             credentials=credentials,
         )
+    elif _helpers.USE_EXPERIMENTAL_HOST:
+        from google.auth.credentials import AnonymousCredentials
+
+        credentials = AnonymousCredentials()
+        return spanner_v1.Client(
+            project=_helpers.EXPERIMENTAL_HOST_PROJECT,
+            credentials=credentials,
+            experimental_host=_helpers.EXPERIMENTAL_HOST,
+        )
     else:
         client_options = {"api_endpoint": _helpers.API_ENDPOINT}
         return spanner_v1.Client(
@@ -130,15 +145,16 @@ def backup_operation_timeout():
 def shared_instance_id():
     if _helpers.CREATE_INSTANCE:
         return f"{_helpers.unique_id('google-cloud')}"
-
+    if _helpers.USE_EXPERIMENTAL_HOST:
+        return _helpers.EXPERIMENTAL_HOST_INSTANCE
     return _helpers.INSTANCE_ID
 
 
 @pytest.fixture(scope="session")
 def instance_configs(spanner_client):
     configs = list(_helpers.retry_503(spanner_client.list_instance_configs)())
 
-    if not _helpers.USE_EMULATOR:
+    if not _helpers.USE_EMULATOR and not _helpers.USE_EXPERIMENTAL_HOST:
         # Defend against back-end returning configs for regions we aren't
         # actually allowed to use.
         configs = [config for config in configs if "-us-" in config.name]

tests/system/test_backup_api.py

@@ -26,10 +26,16 @@
 Remove {_helpers.SKIP_BACKUP_TESTS_ENVVAR} from environment to run these tests.\
 """
 skip_emulator_reason = "Backup operations not supported by emulator."
+skip_experimental_host_reason = (
+    "Backup operations not supported on experimental host yet."
+)
 
 pytestmark = [
     pytest.mark.skipif(_helpers.SKIP_BACKUP_TESTS, reason=skip_env_reason),
     pytest.mark.skipif(_helpers.USE_EMULATOR, reason=skip_emulator_reason),
+    pytest.mark.skipif(
+        _helpers.USE_EXPERIMENTAL_HOST, reason=skip_experimental_host_reason
+    ),
 ]
 
 

tests/system/test_database_api.py

@@ -47,7 +47,9 @@
 
 
 @pytest.fixture(scope="module")
-def multiregion_instance(spanner_client, instance_operation_timeout, not_postgres):
+def multiregion_instance(
+    spanner_client, instance_operation_timeout, not_postgres, not_experimental_host
+):
     multi_region_instance_id = _helpers.unique_id("multi-region")
     multi_region_config = "nam3"
     config_name = "{}/instanceConfigs/{}".format(
@@ -97,6 +99,7 @@ def test_database_binding_of_fixed_size_pool(
     databases_to_delete,
     not_postgres,
     proto_descriptor_file,
+    not_experimental_host,
 ):
     temp_db_id = _helpers.unique_id("fixed_size_db", separator="_")
     temp_db = shared_instance.database(temp_db_id)
@@ -130,6 +133,7 @@ def test_database_binding_of_pinging_pool(
     databases_to_delete,
     not_postgres,
     proto_descriptor_file,
+    not_experimental_host,
 ):
     temp_db_id = _helpers.unique_id("binding_db", separator="_")
     temp_db = shared_instance.database(temp_db_id)
@@ -217,6 +221,7 @@ def test_create_database_pitr_success(
 def test_create_database_with_default_leader_success(
     not_emulator,  # Default leader setting not supported by the emulator
     not_postgres,
+    not_experimental_host,
     multiregion_instance,
     databases_to_delete,
 ):
@@ -253,6 +258,7 @@ def test_create_database_with_default_leader_success(
 
 def test_iam_policy(
     not_emulator,
+    not_experimental_host,
     shared_instance,
     databases_to_delete,
 ):
@@ -414,6 +420,7 @@ def test_update_ddl_w_pitr_success(
 def test_update_ddl_w_default_leader_success(
     not_emulator,
     not_postgres,
+    not_experimental_host,
     multiregion_instance,
     databases_to_delete,
     proto_descriptor_file,
@@ -448,6 +455,7 @@ def test_update_ddl_w_default_leader_success(
 
 def test_create_role_grant_access_success(
     not_emulator,
+    not_experimental_host,
     shared_instance,
     databases_to_delete,
     database_dialect,
@@ -514,6 +522,7 @@ def test_create_role_grant_access_success(
 
 def test_list_database_role_success(
     not_emulator,
+    not_experimental_host,
     shared_instance,
     databases_to_delete,
     database_dialect,
@@ -757,7 +766,11 @@ def test_information_schema_referential_constraints_fkadc(
 
 
 def test_update_database_success(
-    not_emulator, shared_database, shared_instance, database_operation_timeout
+    not_emulator,
+    not_experimental_host,
+    shared_database,
+    shared_instance,
+    database_operation_timeout,
 ):
     old_protection = shared_database.enable_drop_protection
     new_protection = True