feat: added retry and timeout params to partition read in database and snapshot class (#278)

* feat: added retry and timeout params to partition read in database and snapshot class

* feat: lint corrections

* feat: added retry and timeout support in process_read_batch and process_query_batch

* feat: added retry and timeout support in process_read_batch and process_query_batch

* feat: changed retry to retry object in tests

Author: vi3k6i5

google/cloud/spanner_v1/database.py

@@ -26,6 +26,7 @@
 from google.api_core.retry import if_exception_type
 from google.cloud.exceptions import NotFound
 from google.api_core.exceptions import Aborted
+from google.api_core import gapic_v1
 import six
 
 # pylint: disable=ungrouped-imports
@@ -915,6 +916,9 @@ def generate_read_batches(
         index="",
         partition_size_bytes=None,
         max_partitions=None,
+        *,
+        retry=gapic_v1.method.DEFAULT,
+        timeout=gapic_v1.method.DEFAULT,
     ):
         """Start a partitioned batch read operation.
 
@@ -946,6 +950,12 @@ def generate_read_batches(
             service uses this as a hint, the actual number of partitions may
             differ.
 
+        :type retry: :class:`~google.api_core.retry.Retry`
+        :param retry: (Optional) The retry settings for this request.
+
+        :type timeout: float
+        :param timeout: (Optional) The timeout for this request.
+
         :rtype: iterable of dict
         :returns:
             mappings of information used perform actual partitioned reads via
@@ -958,6 +968,8 @@ def generate_read_batches(
             index=index,
             partition_size_bytes=partition_size_bytes,
             max_partitions=max_partitions,
+            retry=retry,
+            timeout=timeout,
         )
 
         read_info = {
@@ -969,21 +981,32 @@ def generate_read_batches(
         for partition in partitions:
             yield {"partition": partition, "read": read_info.copy()}
 
-    def process_read_batch(self, batch):
+    def process_read_batch(
+        self, batch, *, retry=gapic_v1.method.DEFAULT, timeout=gapic_v1.method.DEFAULT,
+    ):
         """Process a single, partitioned read.
 
         :type batch: mapping
         :param batch:
             one of the mappings returned from an earlier call to
             :meth:`generate_read_batches`.
 
+        :type retry: :class:`~google.api_core.retry.Retry`
+        :param retry: (Optional) The retry settings for this request.
+
+        :type timeout: float
+        :param timeout: (Optional) The timeout for this request.
+
+
         :rtype: :class:`~google.cloud.spanner_v1.streamed.StreamedResultSet`
         :returns: a result set instance which can be used to consume rows.
         """
         kwargs = copy.deepcopy(batch["read"])
         keyset_dict = kwargs.pop("keyset")
         kwargs["keyset"] = KeySet._from_dict(keyset_dict)
-        return self._get_snapshot().read(partition=batch["partition"], **kwargs)
+        return self._get_snapshot().read(
+            partition=batch["partition"], **kwargs, retry=retry, timeout=timeout
+        )
 
     def generate_query_batches(
         self,
@@ -993,6 +1016,9 @@ def generate_query_batches(
         partition_size_bytes=None,
         max_partitions=None,
         query_options=None,
+        *,
+        retry=gapic_v1.method.DEFAULT,
+        timeout=gapic_v1.method.DEFAULT,
     ):
         """Start a partitioned query operation.
 
@@ -1036,6 +1062,12 @@ def generate_query_batches(
                 If a dict is provided, it must be of the same form as the protobuf
                 message :class:`~google.cloud.spanner_v1.types.QueryOptions`
 
+        :type retry: :class:`~google.api_core.retry.Retry`
+        :param retry: (Optional) The retry settings for this request.
+
+        :type timeout: float
+        :param timeout: (Optional) The timeout for this request.
+
         :rtype: iterable of dict
         :returns:
             mappings of information used perform actual partitioned reads via
@@ -1047,6 +1079,8 @@ def generate_query_batches(
             param_types=param_types,
             partition_size_bytes=partition_size_bytes,
             max_partitions=max_partitions,
+            retry=retry,
+            timeout=timeout,
         )
 
         query_info = {"sql": sql}
@@ -1064,19 +1098,27 @@ def generate_query_batches(
         for partition in partitions:
             yield {"partition": partition, "query": query_info}
 
-    def process_query_batch(self, batch):
+    def process_query_batch(
+        self, batch, *, retry=gapic_v1.method.DEFAULT, timeout=gapic_v1.method.DEFAULT,
+    ):
         """Process a single, partitioned query.
 
         :type batch: mapping
         :param batch:
             one of the mappings returned from an earlier call to
             :meth:`generate_query_batches`.
 
+        :type retry: :class:`~google.api_core.retry.Retry`
+        :param retry: (Optional) The retry settings for this request.
+
+        :type timeout: float
+        :param timeout: (Optional) The timeout for this request.
+
         :rtype: :class:`~google.cloud.spanner_v1.streamed.StreamedResultSet`
         :returns: a result set instance which can be used to consume rows.
         """
         return self._get_snapshot().execute_sql(
-            partition=batch["partition"], **batch["query"]
+            partition=batch["partition"], **batch["query"], retry=retry, timeout=timeout
         )
 
     def process(self, batch):

google/cloud/spanner_v1/session.py

@@ -258,6 +258,12 @@ def execute_sql(
             or :class:`dict`
         :param query_options: (Optional) Options that are provided for query plan stability.
 
+        :type retry: :class:`~google.api_core.retry.Retry`
+        :param retry: (Optional) The retry settings for this request.
+
+        :type timeout: float
+        :param timeout: (Optional) The timeout for this request.
+
         :rtype: :class:`~google.cloud.spanner_v1.streamed.StreamedResultSet`
         :returns: a result set instance which can be used to consume rows.
         """

google/cloud/spanner_v1/snapshot.py

@@ -27,7 +27,7 @@
 
 from google.api_core.exceptions import InternalServerError
 from google.api_core.exceptions import ServiceUnavailable
-import google.api_core.gapic_v1.method
+from google.api_core import gapic_v1
 from google.cloud.spanner_v1._helpers import _make_value_pb
 from google.cloud.spanner_v1._helpers import _merge_query_options
 from google.cloud.spanner_v1._helpers import _metadata_with_prefix
@@ -109,7 +109,18 @@ def _make_txn_selector(self):  # pylint: disable=redundant-returns-doc
         """
         raise NotImplementedError
 
-    def read(self, table, columns, keyset, index="", limit=0, partition=None):
+    def read(
+        self,
+        table,
+        columns,
+        keyset,
+        index="",
+        limit=0,
+        partition=None,
+        *,
+        retry=gapic_v1.method.DEFAULT,
+        timeout=gapic_v1.method.DEFAULT,
+    ):
         """Perform a ``StreamingRead`` API request for rows in a table.
 
         :type table: str
@@ -134,6 +145,12 @@ def read(self, table, columns, keyset, index="", limit=0, partition=None):
                           from :meth:`partition_read`.  Incompatible with
                           ``limit``.
 
+        :type retry: :class:`~google.api_core.retry.Retry`
+        :param retry: (Optional) The retry settings for this request.
+
+        :type timeout: float
+        :param timeout: (Optional) The timeout for this request.
+
         :rtype: :class:`~google.cloud.spanner_v1.streamed.StreamedResultSet`
         :returns: a result set instance which can be used to consume rows.
 
@@ -163,7 +180,11 @@ def read(self, table, columns, keyset, index="", limit=0, partition=None):
             partition_token=partition,
         )
         restart = functools.partial(
-            api.streaming_read, request=request, metadata=metadata,
+            api.streaming_read,
+            request=request,
+            metadata=metadata,
+            retry=retry,
+            timeout=timeout,
         )
 
         trace_attributes = {"table_id": table, "columns": columns}
@@ -186,8 +207,8 @@ def execute_sql(
         query_mode=None,
         query_options=None,
         partition=None,
-        retry=google.api_core.gapic_v1.method.DEFAULT,
-        timeout=google.api_core.gapic_v1.method.DEFAULT,
+        retry=gapic_v1.method.DEFAULT,
+        timeout=gapic_v1.method.DEFAULT,
     ):
         """Perform an ``ExecuteStreamingSql`` API request.
 
@@ -224,6 +245,12 @@ def execute_sql(
         :rtype: :class:`~google.cloud.spanner_v1.streamed.StreamedResultSet`
         :returns: a result set instance which can be used to consume rows.
 
+        :type retry: :class:`~google.api_core.retry.Retry`
+        :param retry: (Optional) The retry settings for this request.
+
+        :type timeout: float
+        :param timeout: (Optional) The timeout for this request.
+
         :raises ValueError:
             for reuse of single-use snapshots, or if a transaction ID is
             already pending for multiple-use snapshots.
@@ -296,6 +323,9 @@ def partition_read(
         index="",
         partition_size_bytes=None,
         max_partitions=None,
+        *,
+        retry=gapic_v1.method.DEFAULT,
+        timeout=gapic_v1.method.DEFAULT,
     ):
         """Perform a ``PartitionRead`` API request for rows in a table.
 
@@ -323,6 +353,12 @@ def partition_read(
             service uses this as a hint, the actual number of partitions may
             differ.
 
+        :type retry: :class:`~google.api_core.retry.Retry`
+        :param retry: (Optional) The retry settings for this request.
+
+        :type timeout: float
+        :param timeout: (Optional) The timeout for this request.
+
         :rtype: iterable of bytes
         :returns: a sequence of partition tokens
 
@@ -357,7 +393,9 @@ def partition_read(
         with trace_call(
             "CloudSpanner.PartitionReadOnlyTransaction", self._session, trace_attributes
         ):
-            response = api.partition_read(request=request, metadata=metadata,)
+            response = api.partition_read(
+                request=request, metadata=metadata, retry=retry, timeout=timeout,
+            )
 
         return [partition.partition_token for partition in response.partitions]
 
@@ -368,6 +406,9 @@ def partition_query(
         param_types=None,
         partition_size_bytes=None,
         max_partitions=None,
+        *,
+        retry=gapic_v1.method.DEFAULT,
+        timeout=gapic_v1.method.DEFAULT,
     ):
         """Perform a ``PartitionQuery`` API request.
 
@@ -394,6 +435,12 @@ def partition_query(
             service uses this as a hint, the actual number of partitions may
             differ.
 
+        :type retry: :class:`~google.api_core.retry.Retry`
+        :param retry: (Optional) The retry settings for this request.
+
+        :type timeout: float
+        :param timeout: (Optional) The timeout for this request.
+
         :rtype: iterable of bytes
         :returns: a sequence of partition tokens
 
@@ -438,7 +485,9 @@ def partition_query(
             self._session,
             trace_attributes,
         ):
-            response = api.partition_query(request=request, metadata=metadata,)
+            response = api.partition_query(
+                request=request, metadata=metadata, retry=retry, timeout=timeout,
+            )
 
         return [partition.partition_token for partition in response.partitions]