def _prep_batch_get( self, field_paths: Iterable[str] = None, transaction=None, retry: retries.Retry = None, timeout: float = None, ) -> Tuple[dict, dict]: """Shared setup for async/sync :meth:`get`.""" if isinstance(field_paths, str): raise ValueError( "'field_paths' must be a sequence of paths, not a string.") if field_paths is not None: mask = common.DocumentMask(field_paths=sorted(field_paths)) else: mask = None request = { "database": self._client._database_string, "documents": [self._document_path], "mask": mask, "transaction": _helpers.get_transaction_id(transaction), } kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) return request, kwargs
async def get( self, ref_or_query, retry: retries.Retry = gapic_v1.method.DEFAULT, timeout: float = None, ) -> AsyncGenerator[DocumentSnapshot, Any]: """ Retrieve a document or a query result from the database. Args: ref_or_query The document references or query object to return. retry (google.api_core.retry.Retry): Designation of what errors, if any, should be retried. Defaults to a system-specified policy. timeout (float): The timeout for this request. Defaults to a system-specified value. Yields: .DocumentSnapshot: The next document snapshot that fulfills the query, or :data:`None` if the document does not exist. """ kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) if isinstance(ref_or_query, AsyncDocumentReference): return await self._client.get_all([ref_or_query], transaction=self, **kwargs) elif isinstance(ref_or_query, AsyncQuery): return await ref_or_query.stream(transaction=self, **kwargs) else: raise ValueError( 'Value for argument "ref_or_query" must be a AsyncDocumentReference or a AsyncQuery.' )
def _create_helper(self, retry=None, timeout=None): from google.cloud.firestore_v1 import _helpers # Create a minimal fake GAPIC with a dummy response. firestore_api = mock.Mock() firestore_api.commit.mock_add_spec(spec=["commit"]) firestore_api.commit.return_value = self._make_commit_repsonse() # Attach the fake GAPIC to a real client. client = _make_client("dignity") client._firestore_api_internal = firestore_api # Actually make a document and call create(). document = self._make_one("foo", "twelve", client=client) document_data = {"hello": "goodbye", "count": 99} kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) write_result = document.create(document_data, **kwargs) # Verify the response and the mocks. self.assertIs(write_result, mock.sentinel.write_result) write_pb = self._write_pb_for_create(document._document_path, document_data) firestore_api.commit.assert_called_once_with( request={ "database": client._database_string, "writes": [write_pb], "transaction": None, }, metadata=client._rpc_metadata, **kwargs, )
def _set_helper(merge=False, retry=None, timeout=None, **option_kwargs): from google.cloud.firestore_v1 import _helpers # Create a minimal fake GAPIC with a dummy response. firestore_api = mock.Mock(spec=["commit"]) firestore_api.commit.return_value = _make_commit_repsonse() # Attach the fake GAPIC to a real client. client = _make_client("db-dee-bee") client._firestore_api_internal = firestore_api # Actually make a document and call create(). document = _make_document_reference("User", "Interface", client=client) document_data = {"And": 500, "Now": b"\xba\xaa\xaa \xba\xaa\xaa"} kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) write_result = document.set(document_data, merge, **kwargs) # Verify the response and the mocks. assert write_result is mock.sentinel.write_result write_pb = _write_pb_for_set(document._document_path, document_data, merge) firestore_api.commit.assert_called_once_with( request={ "database": client._database_string, "writes": [write_pb], "transaction": None, }, metadata=client._rpc_metadata, **kwargs, )
async def _collections_helper(self, retry=None, timeout=None): from google.cloud.firestore_v1.async_collection import AsyncCollectionReference from google.cloud.firestore_v1 import _helpers collection_ids = ["users", "projects"] class Pager(object): async def __aiter__(self, **_): for collection_id in collection_ids: yield collection_id firestore_api = AsyncMock() firestore_api.mock_add_spec(spec=["list_collection_ids"]) firestore_api.list_collection_ids.return_value = Pager() client = self._make_default_one() client._firestore_api_internal = firestore_api kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) collections = [c async for c in client.collections(**kwargs)] self.assertEqual(len(collections), len(collection_ids)) for collection, collection_id in zip(collections, collection_ids): self.assertIsInstance(collection, AsyncCollectionReference) self.assertEqual(collection.parent, None) self.assertEqual(collection.id, collection_id) base_path = client._database_string + "/documents" firestore_api.list_collection_ids.assert_called_once_with( request={"parent": base_path}, metadata=client._rpc_metadata, **kwargs, )
def _collections_helper(retry=None, timeout=None): from google.cloud.firestore_v1 import _helpers from google.cloud.firestore_v1.collection import CollectionReference collection_ids = ["users", "projects"] class Pager(object): def __iter__(self): yield from collection_ids firestore_api = mock.Mock(spec=["list_collection_ids"]) firestore_api.list_collection_ids.return_value = Pager() client = _make_default_client() client._firestore_api_internal = firestore_api kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) collections = list(client.collections(**kwargs)) assert len(collections) == len(collection_ids) for collection, collection_id in zip(collections, collection_ids): assert isinstance(collection, CollectionReference) assert collection.parent is None assert collection.id == collection_id base_path = client._database_string + "/documents" firestore_api.list_collection_ids.assert_called_once_with( request={"parent": base_path}, metadata=client._rpc_metadata, **kwargs, )
def _list_documents_helper(self, page_size=None, retry=None, timeout=None): from google.cloud.firestore_v1 import _helpers as _fs_v1_helpers from google.api_core.page_iterator import Iterator from google.api_core.page_iterator import Page from google.cloud.firestore_v1.document import DocumentReference from google.cloud.firestore_v1.services.firestore.client import FirestoreClient from google.cloud.firestore_v1.types.document import Document class _Iterator(Iterator): def __init__(self, pages): super(_Iterator, self).__init__(client=None) self._pages = pages def _next_page(self): if self._pages: page, self._pages = self._pages[0], self._pages[1:] return Page(self, page, self.item_to_value) client = _test_helpers.make_client() template = client._database_string + "/documents/{}" document_ids = ["doc-1", "doc-2"] documents = [ Document(name=template.format(document_id)) for document_id in document_ids ] iterator = _Iterator(pages=[documents]) api_client = mock.create_autospec(FirestoreClient) api_client.list_documents.return_value = iterator client._firestore_api_internal = api_client collection = self._make_one("collection", client=client) kwargs = _fs_v1_helpers.make_retry_timeout_kwargs(retry, timeout) if page_size is not None: documents = list( collection.list_documents(page_size=page_size, **kwargs)) else: documents = list(collection.list_documents(**kwargs)) # Verify the response and the mocks. self.assertEqual(len(documents), len(document_ids)) for document, document_id in zip(documents, document_ids): self.assertIsInstance(document, DocumentReference) self.assertEqual(document.parent, collection) self.assertEqual(document.id, document_id) parent, _ = collection._parent_info() api_client.list_documents.assert_called_once_with( request={ "parent": parent, "collection_id": collection.id, "page_size": page_size, "show_missing": True, "mask": { "field_paths": None }, }, metadata=client._rpc_metadata, **kwargs, )
def _prep_collections( self, retry: retries.Retry = None, timeout: float = None, ) -> Tuple[dict, dict]: """Shared setup for async/sync :meth:`collections`.""" request = {"parent": "{}/documents".format(self._database_string)} kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) return request, kwargs
def _prep_get_or_stream( self, retry: retries.Retry = None, timeout: float = None, ) -> Tuple[Any, dict]: """Shared setup for async / sync :meth:`get` / :meth:`stream`""" query = self._query() kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) return query, kwargs
def _prep_create( self, document_data: dict, retry: retries.Retry = None, timeout: float = None, ) -> Tuple[Any, dict]: batch = self._client.batch() batch.create(self, document_data) kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) return batch, kwargs
def _prep_commit(self, retry: retries.Retry, timeout: float): request = { "database": self._client._database_string, "writes": self._write_pbs, "labels": None, } kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) return request, kwargs
def _prep_collections( self, page_size: int = None, retry: retries.Retry = None, timeout: float = None, ) -> Tuple[dict, dict]: """Shared setup for async/sync :meth:`collections`.""" request = {"parent": self._document_path, "page_size": page_size} kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) return request, kwargs
async def _list_documents_helper(page_size=None, retry=None, timeout=None): from google.cloud.firestore_v1 import _helpers from google.api_core.page_iterator_async import AsyncIterator from google.api_core.page_iterator import Page from google.cloud.firestore_v1.async_document import AsyncDocumentReference from google.cloud.firestore_v1.types.document import Document class _AsyncIterator(AsyncIterator): def __init__(self, pages): super(_AsyncIterator, self).__init__(client=None) self._pages = pages async def _next_page(self): if self._pages: page, self._pages = self._pages[0], self._pages[1:] return Page(self, page, self.item_to_value) client = _make_client() template = client._database_string + "/documents/{}" document_ids = ["doc-1", "doc-2"] documents = [ Document(name=template.format(document_id)) for document_id in document_ids ] iterator = _AsyncIterator(pages=[documents]) firestore_api = AsyncMock() firestore_api.mock_add_spec(spec=["list_documents"]) firestore_api.list_documents.return_value = iterator client._firestore_api_internal = firestore_api collection = _make_async_collection_reference("collection", client=client) kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) if page_size is not None: documents = [ i async for i in collection.list_documents(page_size=page_size, **kwargs,) ] else: documents = [i async for i in collection.list_documents(**kwargs)] # Verify the response and the mocks. assert len(documents) == len(document_ids) for document, document_id in zip(documents, document_ids): assert isinstance(document, AsyncDocumentReference) assert document.parent == collection assert document.id == document_id parent, _ = collection._parent_info() firestore_api.list_documents.assert_called_once_with( request={ "parent": parent, "collection_id": collection.id, "page_size": page_size, "show_missing": True, "mask": {"field_paths": None}, }, metadata=client._rpc_metadata, **kwargs, )
def _prep_commit(self, retry: retries.Retry, timeout: float): """Shared setup for async/sync :meth:`commit`.""" request = { "database": self._client._database_string, "writes": self._write_pbs, "transaction": None, } kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) return request, kwargs
def _prep_set( self, document_data: dict, merge: bool = False, retry: retries.Retry = None, timeout: float = None, ) -> Tuple[Any, dict]: batch = self._client.batch() batch.set(self, document_data, merge=merge) kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) return batch, kwargs
def _prep_update( self, field_updates: dict, option: _helpers.WriteOption = None, retry: retries.Retry = None, timeout: float = None, ) -> Tuple[Any, dict]: batch = self._client.batch() batch.update(self, field_updates, option=option) kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) return batch, kwargs
async def _get_w_document_ref_helper(retry=None, timeout=None): from google.cloud.firestore_v1.async_document import AsyncDocumentReference from google.cloud.firestore_v1 import _helpers client = AsyncMock(spec=["get_all"]) transaction = _make_async_transaction(client) ref = AsyncDocumentReference("documents", "doc-id") kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) result = await transaction.get(ref, **kwargs) client.get_all.assert_called_once_with([ref], transaction=transaction, **kwargs) assert result is client.get_all.return_value
def _update_helper(self, retry=None, timeout=None, **option_kwargs): from google.cloud.firestore_v1 import _helpers from google.cloud.firestore_v1.transforms import DELETE_FIELD # Create a minimal fake GAPIC with a dummy response. firestore_api = mock.Mock(spec=["commit"]) firestore_api.commit.return_value = self._make_commit_repsonse() # Attach the fake GAPIC to a real client. client = _make_client("potato-chip") client._firestore_api_internal = firestore_api # Actually make a document and call create(). document = self._make_one("baked", "Alaska", client=client) # "Cheat" and use OrderedDict-s so that iteritems() is deterministic. field_updates = collections.OrderedDict( (("hello", 1), ("then.do", False), ("goodbye", DELETE_FIELD))) kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) if option_kwargs: option = client.write_option(**option_kwargs) write_result = document.update(field_updates, option=option, **kwargs) else: option = None write_result = document.update(field_updates, **kwargs) # Verify the response and the mocks. self.assertIs(write_result, mock.sentinel.write_result) update_values = { "hello": field_updates["hello"], "then": { "do": field_updates["then.do"] }, } field_paths = list(field_updates.keys()) write_pb = self._write_pb_for_update(document._document_path, update_values, sorted(field_paths)) if option is not None: option.modify_write(write_pb) firestore_api.commit.assert_called_once_with( request={ "database": client._database_string, "writes": [write_pb], "transaction": None, }, metadata=client._rpc_metadata, **kwargs, )
async def _get_all_helper(retry=None, timeout=None): from google.cloud.firestore_v1 import _helpers client = AsyncMock(spec=["get_all"]) transaction = _make_async_transaction(client) ref1, ref2 = mock.Mock(), mock.Mock() kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) result = await transaction.get_all([ref1, ref2], **kwargs) client.get_all.assert_called_once_with( [ref1, ref2], transaction=transaction, **kwargs, ) assert result is client.get_all.return_value
def _transaction_get_w_query_helper(retry=None, timeout=None): from google.cloud.firestore_v1 import _helpers from google.cloud.firestore_v1.query import Query client = mock.Mock(spec=[]) transaction = _make_transaction(client) query = Query(parent=mock.Mock(spec=[])) query.stream = mock.MagicMock() kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) result = transaction.get(query, **kwargs) assert result is query.stream.return_value query.stream.assert_called_once_with(transaction=transaction, **kwargs)
def _get_w_document_ref_helper(self, retry=None, timeout=None): from google.cloud.firestore_v1.document import DocumentReference from google.cloud.firestore_v1 import _helpers client = mock.Mock(spec=["get_all"]) transaction = self._make_one(client) ref = DocumentReference("documents", "doc-id") kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) result = transaction.get(ref, **kwargs) self.assertIs(result, client.get_all.return_value) client.get_all.assert_called_once_with([ref], transaction=transaction, **kwargs)
def _prep_add( self, document_data: dict, document_id: str = None, retry: retries.Retry = None, timeout: float = None, ) -> Tuple[DocumentReference, dict]: """Shared setup for async / sync :method:`add`""" if document_id is None: document_id = _auto_id() document_ref = self.document(document_id) kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) return document_ref, kwargs
async def _add_helper(self, retry=None, timeout=None): from google.cloud.firestore_v1.async_document import AsyncDocumentReference from google.cloud.firestore_v1 import _helpers # Create a minimal fake GAPIC with a dummy response. firestore_api = AsyncMock(spec=["commit"]) write_result = mock.Mock(update_time=mock.sentinel.update_time, spec=["update_time"]) commit_response = mock.Mock( write_results=[write_result], spec=["write_results", "commit_time"], commit_time=mock.sentinel.commit_time, ) firestore_api.commit.return_value = commit_response # Attach the fake GAPIC to a real client. client = _make_client() client._firestore_api_internal = firestore_api # Actually make a collection and call add(). collection = self._make_one("parent", client=client) document_data = {"zorp": 208.75, "i-did-not": b"know that"} doc_id = "child" kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) update_time, document_ref = await collection.add( document_data, document_id=doc_id, **kwargs, ) # Verify the response and the mocks. self.assertIs(update_time, mock.sentinel.update_time) self.assertIsInstance(document_ref, AsyncDocumentReference) self.assertIs(document_ref._client, client) self.assertEqual(document_ref._path, (collection.id, doc_id)) write_pb = self._write_pb_for_create(document_ref._document_path, document_data) firestore_api.commit.assert_called_once_with( request={ "database": client._database_string, "writes": [write_pb], "transaction": None, }, metadata=client._rpc_metadata, **kwargs, )
def _get_partitions_helper(self, retry=None, timeout=None): from google.cloud.firestore_v1 import _helpers # Create a minimal fake GAPIC. firestore_api = mock.Mock(spec=["partition_query"]) # Attach the fake GAPIC to a real client. client = _make_client() client._firestore_api_internal = firestore_api # Make a **real** collection reference as parent. parent = client.collection("charles") # Make two **real** document references to use as cursors document1 = parent.document("one") document2 = parent.document("two") # Add cursor pb's to the minimal fake GAPIC. cursor_pb1 = _make_cursor_pb(([document1], False)) cursor_pb2 = _make_cursor_pb(([document2], False)) firestore_api.partition_query.return_value = iter( [cursor_pb1, cursor_pb2]) kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) # Execute the query and check the response. query = self._make_one(parent) get_response = query.get_partitions(2, **kwargs) self.assertIsInstance(get_response, types.GeneratorType) returned = list(get_response) self.assertEqual(len(returned), 3) # Verify the mock call. parent_path, _ = parent._parent_info() partition_query = self._make_one( parent, orders=(query._make_order("__name__", query.ASCENDING), ), ) firestore_api.partition_query.assert_called_once_with( request={ "parent": parent_path, "structured_query": partition_query._to_protobuf(), "partition_count": 2, }, metadata=client._rpc_metadata, **kwargs, )
async def _get_w_query_helper(retry=None, timeout=None): from google.cloud.firestore_v1.async_query import AsyncQuery from google.cloud.firestore_v1 import _helpers client = AsyncMock(spec=[]) transaction = _make_async_transaction(client) query = AsyncQuery(parent=AsyncMock(spec=[])) query.stream = AsyncMock() kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) result = await transaction.get(query, **kwargs,) query.stream.assert_called_once_with( transaction=transaction, **kwargs, ) assert result is query.stream.return_value
def _get_all_helper(self, retry=None, timeout=None): from google.cloud.firestore_v1 import _helpers client = mock.Mock(spec=["get_all"]) transaction = self._make_one(client) ref1, ref2 = mock.Mock(), mock.Mock() kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) result = transaction.get_all([ref1, ref2], **kwargs) client.get_all.assert_called_once_with( [ref1, ref2], transaction=transaction, **kwargs, ) self.assertIs(result, client.get_all.return_value)
def _prep_delete( self, option: _helpers.WriteOption = None, retry: retries.Retry = None, timeout: float = None, ) -> Tuple[dict, dict]: """Shared setup for async/sync :meth:`delete`.""" write_pb = _helpers.pb_for_delete(self._document_path, option) request = { "database": self._client._database_string, "writes": [write_pb], "transaction": None, } kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) return request, kwargs
async def _collections_helper(self, page_size=None, retry=None, timeout=None): from google.cloud.firestore_v1 import _helpers from google.cloud.firestore_v1.async_collection import AsyncCollectionReference collection_ids = ["coll-1", "coll-2"] class Pager(object): async def __aiter__(self, **_): for collection_id in collection_ids: yield collection_id firestore_api = AsyncMock() firestore_api.mock_add_spec(spec=["list_collection_ids"]) firestore_api.list_collection_ids.return_value = Pager() client = _make_client() client._firestore_api_internal = firestore_api kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) # Actually make a document and call delete(). document = self._make_one("where", "we-are", client=client) if page_size is not None: collections = [ c async for c in document.collections(page_size=page_size, **kwargs) ] else: collections = [c async for c in document.collections(**kwargs)] # Verify the response and the mocks. self.assertEqual(len(collections), len(collection_ids)) for collection, collection_id in zip(collections, collection_ids): self.assertIsInstance(collection, AsyncCollectionReference) self.assertEqual(collection.parent, document) self.assertEqual(collection.id, collection_id) firestore_api.list_collection_ids.assert_called_once_with( request={ "parent": document._document_path, "page_size": page_size }, metadata=client._rpc_metadata, **kwargs, )
async def _commit_helper(self, retry=None, timeout=None): from google.protobuf import timestamp_pb2 from google.cloud.firestore_v1 import _helpers from google.cloud.firestore_v1.types import firestore from google.cloud.firestore_v1.types import write # Create a minimal fake GAPIC with a dummy result. firestore_api = AsyncMock(spec=["commit"]) timestamp = timestamp_pb2.Timestamp(seconds=1234567, nanos=123456798) commit_response = firestore.CommitResponse( write_results=[write.WriteResult(), write.WriteResult()], commit_time=timestamp, ) firestore_api.commit.return_value = commit_response kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) # Attach the fake GAPIC to a real client. client = _make_client("grand") client._firestore_api_internal = firestore_api # Actually make a batch with some mutations and call commit(). batch = self._make_one(client) document1 = client.document("a", "b") batch.create(document1, {"ten": 10, "buck": "ets"}) document2 = client.document("c", "d", "e", "f") batch.delete(document2) write_pbs = batch._write_pbs[::] write_results = await batch.commit(**kwargs) self.assertEqual(write_results, list(commit_response.write_results)) self.assertEqual(batch.write_results, write_results) self.assertEqual(batch.commit_time.timestamp_pb(), timestamp) # Make sure batch has no more "changes". self.assertEqual(batch._write_pbs, []) # Verify the mocks. firestore_api.commit.assert_called_once_with( request={ "database": client._database_string, "writes": write_pbs, "transaction": None, }, metadata=client._rpc_metadata, **kwargs, )
def _prep_list_documents( self, page_size: int = None, retry: retries.Retry = None, timeout: float = None, ) -> Tuple[dict, dict]: """Shared setup for async / sync :method:`list_documents`""" parent, _ = self._parent_info() request = { "parent": parent, "collection_id": self.id, "page_size": page_size, "show_missing": True, } kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) return request, kwargs