From 5a1ef50a018873a00d2cf41b9d7fdbc57acdcd65 Mon Sep 17 00:00:00 2001 From: Tres Seaver Date: Tue, 13 Oct 2020 18:26:01 -0400 Subject: [PATCH] feat: add retry/timeout to base class signatures Toward #221 --- google/cloud/firestore_v1/base_client.py | 5 ++- google/cloud/firestore_v1/base_collection.py | 20 ++++++++-- google/cloud/firestore_v1/base_document.py | 37 ++++++++++++++++--- google/cloud/firestore_v1/base_query.py | 14 ++++++- google/cloud/firestore_v1/base_transaction.py | 9 ++++- 5 files changed, 70 insertions(+), 15 deletions(-) diff --git a/google/cloud/firestore_v1/base_client.py b/google/cloud/firestore_v1/base_client.py index b2a422291..28d63bb85 100644 --- a/google/cloud/firestore_v1/base_client.py +++ b/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], diff --git a/google/cloud/firestore_v1/base_collection.py b/google/cloud/firestore_v1/base_collection.py index 72480a911..167cac9d4 100644 --- a/google/cloud/firestore_v1/base_collection.py +++ b/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 diff --git a/google/cloud/firestore_v1/base_document.py b/google/cloud/firestore_v1/base_document.py index 68534c471..6207e458e 100644 --- a/google/cloud/firestore_v1/base_document.py +++ b/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: diff --git a/google/cloud/firestore_v1/base_query.py b/google/cloud/firestore_v1/base_query.py index 188c15b6a..78afd5750 100644 --- a/google/cloud/firestore_v1/base_query.py +++ b/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. diff --git a/google/cloud/firestore_v1/base_transaction.py b/google/cloud/firestore_v1/base_transaction.py index c676d3d7a..5eac1d7fe 100644 --- a/google/cloud/firestore_v1/base_transaction.py +++ b/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