feat: Add support for Directed Reads

google/cloud/spanner_v1/__init__.py

@@ -38,6 +38,7 @@
 from .types.spanner import CommitRequest
 from .types.spanner import CreateSessionRequest
 from .types.spanner import DeleteSessionRequest
+from .types.spanner import DirectedReadOptions
 from .types.spanner import ExecuteBatchDmlRequest
 from .types.spanner import ExecuteBatchDmlResponse
 from .types.spanner import ExecuteSqlRequest
@@ -77,6 +78,7 @@
 This value can only be used for timestamp columns that have set the option
 ``(allow_commit_timestamp=true)`` in the schema.
 """
+TransactionTypes = DirectedReadOptions.ReplicaSelection.Type
 
 
 __all__ = (
@@ -96,6 +98,7 @@
     "TransactionPingingPool",
     # local
     "COMMIT_TIMESTAMP",
+    "TransactionTypes",
     # google.cloud.spanner_v1.types
     "BatchCreateSessionsRequest",
     "BatchCreateSessionsResponse",
@@ -104,6 +107,7 @@
     "CommitResponse",
     "CreateSessionRequest",
     "DeleteSessionRequest",
+    "DirectedReadOptions",
     "ExecuteBatchDmlRequest",
     "ExecuteBatchDmlResponse",
     "ExecuteSqlRequest",

google/cloud/spanner_v1/_helpers.py

@@ -28,6 +28,8 @@
 from google.cloud.spanner_v1 import TypeCode
 from google.cloud.spanner_v1 import ExecuteSqlRequest
 from google.cloud.spanner_v1 import JsonObject
+from google.cloud.spanner_v1 import DirectedReadOptions
+from google.api_core.exceptions import InvalidArgument
 
 # Validation error messages
 NUMERIC_MAX_SCALE_ERR_MSG = (
@@ -358,3 +360,51 @@ def _metadata_with_leader_aware_routing(value, **kw):
         List[Tuple[str, str]]: RPC metadata with leader aware routing header
     """
     return ("x-goog-spanner-route-to-leader", str(value).lower())
+
+
+def verify_directed_read_options(directed_read_options):
+    """Verifies if value of directed_read_options is correct.
+    :type directed_read_options: :class:`~googlecloud.spanner_v1.types.DirectedReadOptions`
+            or :class:`dict`
+    :param directed_read_options: directed_read_options for ReadRequests and ExecuteSqlRequests.
+
+    :raises InvalidArguement: if directed_read_options is incorrect.
+    """
+    if type(directed_read_options) == dict:
+        if (
+            "include_replicas" in directed_read_options.keys()
+            and "exclude_replicas" in directed_read_options.keys()
+        ):
+            raise InvalidArgument(
+                "Only one of include_replicas or exclude_replicas can be set"
+            )
+        if (
+            "include_replicas" in directed_read_options.keys()
+            and "replica_selections" in directed_read_options["include_replicas"].keys()
+            and len(directed_read_options["include_replicas"]["replica_selections"])
+            > 10
+        ) or (
+            "exclude_replicas" in directed_read_options.keys()
+            and "replica_selections" in directed_read_options["exclude_replicas"].keys()
+            and len(directed_read_options["exclude_replicas"]["replica_selections"])
+            > 10
+        ):
+            raise InvalidArgument("Maximum length of replica selection allowed is 10")
+    elif isinstance(directed_read_options, DirectedReadOptions):
+        if (
+            directed_read_options.include_replicas is not None
+            and directed_read_options.exclude_replicas is not None
+        ):
+            raise InvalidArgument(
+                "Only one of include_replicas or exclude_replicas can be set"
+            )
+        if (
+            directed_read_options.include_replicas is not None
+            and directed_read_options.include_replicas.replica_selections is not None
+            and len(directed_read_options.include_replicas.replica_selections) > 10
+        ) or (
+            directed_read_options.include_replicas is not None
+            and directed_read_options.exclude_replicas.replica_selections is not None
+            and len(directed_read_options.exclude_replicas.replica_selections) > 10
+        ):
+            raise InvalidArgument("Maximum length of replica selection allowed is 10")

google/cloud/spanner_v1/client.py

@@ -47,6 +47,7 @@
 from google.cloud.spanner_v1 import ExecuteSqlRequest
 from google.cloud.spanner_v1._helpers import _merge_query_options
 from google.cloud.spanner_v1._helpers import _metadata_with_prefix
+from google.cloud.spanner_v1._helpers import verify_directed_read_options
 from google.cloud.spanner_v1.instance import Instance
 
 _CLIENT_INFO = client_info.ClientInfo(client_library_version=__version__)
@@ -120,6 +121,12 @@ class Client(ClientWithProject):
         disable leader aware routing. Disabling leader aware routing would
         route all requests in RW/PDML transactions to the closest region.
 
+    :type directed_read_options: :class:`~googlecloud.spanner_v1.types.DirectedReadOptions`
+        or :class:`dict`
+    :param directed_read_options: (Optional) Client options used to set the directed_read_options
+        for all ReadRequests and ExecuteSqlRequests for the Client which indicate which replicas
+        or regions should be used for non-transactional reads or queries.
+
     :raises: :class:`ValueError <exceptions.ValueError>` if both ``read_only``
              and ``admin`` are :data:`True`
     """
@@ -139,6 +146,7 @@ def __init__(
         client_options=None,
         query_options=None,
         route_to_leader_enabled=True,
+        directed_read_options=None,
     ):
         self._emulator_host = _get_spanner_emulator_host()
 
@@ -179,6 +187,8 @@ def __init__(
             warnings.warn(_EMULATOR_HOST_HTTP_SCHEME)
 
         self._route_to_leader_enabled = route_to_leader_enabled
+        verify_directed_read_options(directed_read_options)
+        self._directed_read_options = directed_read_options
 
     @property
     def credentials(self):
@@ -260,6 +270,17 @@ def route_to_leader_enabled(self):
         """
         return self._route_to_leader_enabled
 
+
+    def directed_read_options(self):
+        """Getter for directed_read_options.
+
+        :rtype:
+            :class:`~googlecloud.spanner_v1.types.DirectedReadOptions`
+            or :class:`dict`
+        :returns: The directed_read_options for the client.
+        """
+        return self._directed_read_options
+
     def copy(self):
         """Make a copy of this client.
 
@@ -383,3 +404,13 @@ def list_instances(self, filter_="", page_size=None):
             request=request, metadata=metadata
         )
         return page_iter
+
+    def set_directed_read_options(self, directed_read_options):
+        """Sets directed_read_options for the client
+        :type directed_read_options: :class:`~googlecloud.spanner_v1.types.DirectedReadOptions`
+            or :class:`dict`
+        :param directed_read_options: Client options used to set the directed_read_options
+            for all ReadRequests and ExecuteSqlRequests for the Client which indicate which replicas
+            or regions should be used for non-transactional reads or queries.
+        """
+        self._directed_read_options = directed_read_options

google/cloud/spanner_v1/database.py

@@ -166,6 +166,7 @@ def __init__(
         self._route_to_leader_enabled = self._instance._client.route_to_leader_enabled
         self._enable_drop_protection = enable_drop_protection
         self._reconciling = False
+        self._directed_read_options = self._instance._client.directed_read_options
 
         if pool is None:
             pool = BurstyPool(database_role=database_role)

google/cloud/spanner_v1/snapshot.py

@@ -28,6 +28,7 @@
 from google.api_core.exceptions import InternalServerError
 from google.api_core.exceptions import ServiceUnavailable
 from google.api_core.exceptions import InvalidArgument
+from google.api_core.exceptions import BadRequest
 from google.api_core import gapic_v1
 from google.cloud.spanner_v1._helpers import (
     _make_value_pb,
@@ -41,6 +42,7 @@
 from google.cloud.spanner_v1._opentelemetry_tracing import trace_call
 from google.cloud.spanner_v1.streamed import StreamedResultSet
 from google.cloud.spanner_v1 import RequestOptions
+from google.cloud.spanner_v1._helpers import verify_directed_read_options
 
 _STREAM_RESUMPTION_INTERNAL_ERROR_MESSAGES = (
     "RST_STREAM",
@@ -176,6 +178,7 @@ def read(
         *,
         retry=gapic_v1.method.DEFAULT,
         timeout=gapic_v1.method.DEFAULT,
+        directed_read_options=None,
     ):
         """Perform a ``StreamingRead`` API request for rows in a table.
 
@@ -224,12 +227,22 @@ def read(
                 ``partition_token``, the API will return an
                 ``INVALID_ARGUMENT`` error.
 
+        :type directed_read_options: :class:`~googlecloud.spanner_v1.types.DirectedReadOptions`
+            or :class:`dict`
+        :param directed_read_options: (Optional) Client options used to set the directed_read_options
+            for all ReadRequests and ExecuteSqlRequests for the Client which indicate which replicas
+            or regions should be used for non-transactional reads or queries.
+
         :rtype: :class:`~google.cloud.spanner_v1.streamed.StreamedResultSet`
         :returns: a result set instance which can be used to consume rows.
 
         :raises ValueError:
             for reuse of single-use snapshots, or if a transaction ID is
             already pending for multiple-use snapshots.
+
+        :raises InvalidArguement: if directed_read_options is incorrect.
+
+        :raises BadRequest: if directed_read_options are set for READ-WRITE Transaction or PDML.
         """
         if self._read_request_count > 0:
             if not self._multi_use:
@@ -255,6 +268,19 @@ def read(
             request_options.transaction_tag = None
         elif self.transaction_tag is not None:
             request_options.transaction_tag = self.transaction_tag
+            if (
+                directed_read_options is None
+                and database._directed_read_options is not None
+            ):
+                directed_read_options = database._directed_read_options
+        else:
+            if self.transaction_tag is not None:
+                request_options.transaction_tag = self.transaction_tag
+            if directed_read_options is not None:
+                raise BadRequest(
+                    "directed_read_options can't be set for readWrite transactions or partitioned dml requests"
+                )
+        verify_directed_read_options(directed_read_options)
 
         request = ReadRequest(
             session=self._session.name,
@@ -266,6 +292,7 @@ def read(
             partition_token=partition,
             request_options=request_options,
             data_boost_enabled=data_boost_enabled,
+            directed_read_options=directed_read_options,
         )
         restart = functools.partial(
             api.streaming_read,
@@ -322,6 +349,7 @@ def execute_sql(
         retry=gapic_v1.method.DEFAULT,
         timeout=gapic_v1.method.DEFAULT,
         data_boost_enabled=False,
+        directed_read_options=None,
     ):
         """Perform an ``ExecuteStreamingSql`` API request.
 
@@ -379,9 +407,19 @@ def execute_sql(
                 ``partition_token``, the API will return an
                 ``INVALID_ARGUMENT`` error.
 
+        :type directed_read_options: :class:`~googlecloud.spanner_v1.types.DirectedReadOptions`
+            or :class:`dict`
+        :param directed_read_options: (Optional) Client options used to set the directed_read_options
+            for all ReadRequests and ExecuteSqlRequests for the Client which indicate which replicas
+            or regions should be used for non-transactional reads or queries.
+
         :raises ValueError:
             for reuse of single-use snapshots, or if a transaction ID is
             already pending for multiple-use snapshots.
+
+        :raises InvalidArguement: if directed_read_options is incorrect.
+
+        :raises BadRequest: if directed_read_options are set for READ-WRITE Transaction or PDML.
         """
         if self._read_request_count > 0:
             if not self._multi_use:
@@ -421,6 +459,19 @@ def execute_sql(
             request_options.transaction_tag = None
         elif self.transaction_tag is not None:
             request_options.transaction_tag = self.transaction_tag
+            if (
+                directed_read_options is None
+                and database._directed_read_options is not None
+            ):
+                directed_read_options = database._directed_read_options
+        else:
+            if self.transaction_tag is not None:
+                request_options.transaction_tag = self.transaction_tag
+            if directed_read_options is not None:
+                raise BadRequest(
+                    "directed_read_options can't be set for readWrite transactions or partitioned dml requests"
+                )
+        verify_directed_read_options(directed_read_options)
 
         request = ExecuteSqlRequest(
             session=self._session.name,
@@ -433,6 +484,7 @@ def execute_sql(
             query_options=query_options,
             request_options=request_options,
             data_boost_enabled=data_boost_enabled,
+            directed_read_options=directed_read_options,
         )
         restart = functools.partial(
             api.execute_streaming_sql,