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
