feat: add retry/timeout to base class signatures

Toward #221

google/cloud/firestore_v1/base_client.py

@@ -28,6 +28,7 @@
 
 import google.api_core.client_options  # type: ignore
 import google.api_core.path_template  # type: ignore
+from google.api_core import retry as retries  # type: ignore
 from google.api_core.gapic_v1 import client_info  # type: ignore
 from google.cloud.client import ClientWithProject  # type: ignore
 
@@ -358,13 +359,15 @@ def get_all(
         references: list,
         field_paths: Iterable[str] = None,
         transaction: BaseTransaction = None,
+        retry: retries.Retry = None,
+        timeout: float = None,
     ) -> Union[
         AsyncGenerator[DocumentSnapshot, Any], Generator[DocumentSnapshot, Any, Any]
     ]:
         raise NotImplementedError
 
     def collections(
-        self,
+        self, retry: retries.Retry = None, timeout: float = None,
     ) -> Union[
         AsyncGenerator[BaseCollectionReference, Any],
         Generator[BaseCollectionReference, Any, Any],

google/cloud/firestore_v1/base_collection.py

@@ -15,6 +15,8 @@
 """Classes for representing collections for the Google Cloud Firestore API."""
 import random
 
+from google.api_core import retry as retries  # type: ignore
+
 from google.cloud.firestore_v1 import _helpers
 from google.cloud.firestore_v1.document import DocumentReference
 from typing import (
@@ -147,12 +149,16 @@ def _parent_info(self) -> Tuple[Any, str]:
         return parent_path, expected_prefix
 
     def add(
-        self, document_data: dict, document_id: str = None
+        self,
+        document_data: dict,
+        document_id: str = None,
+        retry: retries.Retry = None,
+        timeout: float = None,
     ) -> Union[Tuple[Any, Any], Coroutine[Any, Any, Tuple[Any, Any]]]:
         raise NotImplementedError
 
     def list_documents(
-        self, page_size: int = None
+        self, page_size: int = None, retry: retries.Retry = None, timeout: float = None,
     ) -> Union[
         Generator[DocumentReference, Any, Any], AsyncGenerator[DocumentReference, Any]
     ]:
@@ -374,14 +380,20 @@ def end_at(
         return query.end_at(document_fields)
 
     def get(
-        self, transaction: Transaction = None
+        self,
+        transaction: Transaction = None,
+        retry: retries.Retry = None,
+        timeout: float = None,
     ) -> Union[
         Generator[DocumentSnapshot, Any, Any], AsyncGenerator[DocumentSnapshot, Any]
     ]:
         raise NotImplementedError
 
     def stream(
-        self, transaction: Transaction = None
+        self,
+        transaction: Transaction = None,
+        retry: retries.Retry = None,
+        timeout: float = None,
     ) -> Union[Iterator[DocumentSnapshot], AsyncIterator[DocumentSnapshot]]:
         raise NotImplementedError
 

google/cloud/firestore_v1/base_document.py

@@ -16,6 +16,8 @@
 
 import copy
 
+from google.api_core import retry as retries  # type: ignore
+
 from google.cloud.firestore_v1 import _helpers
 from google.cloud.firestore_v1 import field_path as field_path_module
 from typing import Any, Iterable, NoReturn, Tuple
@@ -178,26 +180,49 @@ def collection(self, collection_id: str) -> Any:
         child_path = self._path + (collection_id,)
         return self._client.collection(*child_path)
 
-    def create(self, document_data: dict) -> NoReturn:
+    def create(
+        self, document_data: dict, retry: retries.Retry = None, timeout: float = None,
+    ) -> NoReturn:
         raise NotImplementedError
 
-    def set(self, document_data: dict, merge: bool = False) -> NoReturn:
+    def set(
+        self,
+        document_data: dict,
+        merge: bool = False,
+        retry: retries.Retry = None,
+        timeout: float = None,
+    ) -> NoReturn:
         raise NotImplementedError
 
     def update(
-        self, field_updates: dict, option: _helpers.WriteOption = None
+        self,
+        field_updates: dict,
+        option: _helpers.WriteOption = None,
+        retry: retries.Retry = None,
+        timeout: float = None,
     ) -> NoReturn:
         raise NotImplementedError
 
-    def delete(self, option: _helpers.WriteOption = None) -> NoReturn:
+    def delete(
+        self,
+        option: _helpers.WriteOption = None,
+        retry: retries.Retry = None,
+        timeout: float = None,
+    ) -> NoReturn:
         raise NotImplementedError
 
     def get(
-        self, field_paths: Iterable[str] = None, transaction=None
+        self,
+        field_paths: Iterable[str] = None,
+        transaction=None,
+        retry: retries.Retry = None,
+        timeout: float = None,
     ) -> "DocumentSnapshot":
         raise NotImplementedError
 
-    def collections(self, page_size: int = None) -> NoReturn:
+    def collections(
+        self, page_size: int = None, retry: retries.Retry = None, timeout: float = None,
+    ) -> NoReturn:
         raise NotImplementedError
 
     def on_snapshot(self, callback) -> NoReturn:

google/cloud/firestore_v1/base_query.py

@@ -21,6 +21,7 @@
 import copy
 import math
 
+from google.api_core import retry as retries  # type: ignore
 from google.protobuf import wrappers_pb2
 
 from google.cloud.firestore_v1 import _helpers
@@ -800,10 +801,14 @@ def _to_protobuf(self) -> StructuredQuery:
 
         return query.StructuredQuery(**query_kwargs)
 
-    def get(self, transaction=None) -> NoReturn:
+    def get(
+        self, transaction=None, retry: retries.Retry = None, timeout: float = None,
+    ) -> NoReturn:
         raise NotImplementedError
 
-    def stream(self, transaction=None) -> NoReturn:
+    def stream(
+        self, transaction=None, retry: retries.Retry = None, timeout: float = None,
+    ) -> NoReturn:
         raise NotImplementedError
 
     def on_snapshot(self, callback) -> NoReturn:
@@ -1099,6 +1104,11 @@ def _validate_partition_query(self):
         if self._offset:
             raise ValueError("Can't partition query with offset.")
 
+    def get_partitions(
+        self, partition_count, retry: retries.Retry = None, timeout: float = None,
+    ) -> NoReturn:
+        raise NotImplementedError
+
 
 class QueryPartition:
     """Represents a bounded partition of a collection group query.

google/cloud/firestore_v1/base_transaction.py

@@ -14,6 +14,7 @@
 
 """Helpers for applying Google Cloud Firestore changes in a transaction."""
 
+from google.api_core import retry as retries  # type: ignore
 
 from google.cloud.firestore_v1 import types
 from typing import Any, Coroutine, NoReturn, Optional, Union
@@ -141,10 +142,14 @@ def _rollback(self) -> NoReturn:
     def _commit(self) -> Union[list, Coroutine[Any, Any, list]]:
         raise NotImplementedError
 
-    def get_all(self, references: list) -> NoReturn:
+    def get_all(
+        self, references: list, retry: retries.Retry = None, timeout: float = None,
+    ) -> NoReturn:
         raise NotImplementedError
 
-    def get(self, ref_or_query) -> NoReturn:
+    def get(
+        self, ref_or_query, retry: retries.Retry = None, timeout: float = None,
+    ) -> NoReturn:
         raise NotImplementedError