-
Notifications
You must be signed in to change notification settings - Fork 78
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
* feat: add async tests for AsyncClient * feat: add AsyncClient implementation * feat: add AsyncDocument implementation * feat: add AsyncDocument support to AsyncClient * feat: add AsyncDocument tests Note: tests relying on Collection will fail in this commit * feat: add AsyncCollectionReference class * feat: integrate AsyncCollectionReference * feat: add async_collection tests * fix: swap coroutine/function declaration in async_collection * feat: add async_batch implementation * feat: integrate async_batch * feat: add async_batch tests * feat: add async_query implementation * feat: add async_query integration * feat: add async_query tests * fix: AsyncQuery.get async_generator nesting * feat: add async_transaction integration and tests * fix: linter errors * feat: refactor async tests to use aiounittest and pytest-asyncio * feat: remove duplicate code from async_client * feat: remove duplicate code from async_batch * feat: remove duplicate code from async_collection * feat: remove duplicate code from async_document * fix: remove unused imports * fix: remove duplicate test * feat: remove duplicate code from async_transaction * fix: remove unused Python2 compatibility * fix: resolve async generator tests * fix: create mock async generator to get full coverage * fix: copyright date * feat: create Client/AsyncClient superclass * fix: base client test class * feat: create WriteBatch/AsyncWriteBatch superclass * feat: create CollectionReference/AsyncCollectionReference superclass * feat: create DocumentReference/AsyncDocumentReference superclass * fix: base document test class name * feat: create Query/AsyncQuery superclass * refactor: generalize collection tests with mocks * feat: create Transaction/AsyncTransaction superclass * feat: add microgen support to async interface * fix: async client copyright date * fix: standardize assert syntax * fix: incorrect copyright date * fix: incorrect copyright date * fix: clarify _sleep assertions in transaction * fix: clarify error in context manager tests * fix: clarify error in context manager tests
- Loading branch information
Showing
17 changed files
with
4,531 additions
and
14 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,64 @@ | ||
# Copyright 2020 Google LLC All rights reserved. | ||
# | ||
# Licensed under the Apache License, Version 2.0 (the "License"); | ||
# you may not use this file except in compliance with the License. | ||
# You may obtain a copy of the License at | ||
# | ||
# http://www.apache.org/licenses/LICENSE-2.0 | ||
# | ||
# Unless required by applicable law or agreed to in writing, software | ||
# distributed under the License is distributed on an "AS IS" BASIS, | ||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
# See the License for the specific language governing permissions and | ||
# limitations under the License. | ||
|
||
"""Helpers for batch requests to the Google Cloud Firestore API.""" | ||
|
||
|
||
from google.cloud.firestore_v1.base_batch import BaseWriteBatch | ||
|
||
|
||
class AsyncWriteBatch(BaseWriteBatch): | ||
"""Accumulate write operations to be sent in a batch. | ||
This has the same set of methods for write operations that | ||
:class:`~google.cloud.firestore_v1.async_document.AsyncDocumentReference` does, | ||
e.g. :meth:`~google.cloud.firestore_v1.async_document.AsyncDocumentReference.create`. | ||
Args: | ||
client (:class:`~google.cloud.firestore_v1.async_client.AsyncClient`): | ||
The client that created this batch. | ||
""" | ||
|
||
def __init__(self, client): | ||
super(AsyncWriteBatch, self).__init__(client=client) | ||
|
||
async def commit(self): | ||
"""Commit the changes accumulated in this batch. | ||
Returns: | ||
List[:class:`google.cloud.proto.firestore.v1.write.WriteResult`, ...]: | ||
The write results corresponding to the changes committed, returned | ||
in the same order as the changes were applied to this batch. A | ||
write result contains an ``update_time`` field. | ||
""" | ||
commit_response = self._client._firestore_api.commit( | ||
request={ | ||
"database": self._client._database_string, | ||
"writes": self._write_pbs, | ||
"transaction": None, | ||
}, | ||
metadata=self._client._rpc_metadata, | ||
) | ||
|
||
self._write_pbs = [] | ||
self.write_results = results = list(commit_response.write_results) | ||
self.commit_time = commit_response.commit_time | ||
return results | ||
|
||
async def __aenter__(self): | ||
return self | ||
|
||
async def __aexit__(self, exc_type, exc_value, traceback): | ||
if exc_type is None: | ||
await self.commit() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,288 @@ | ||
# Copyright 2020 Google LLC All rights reserved. | ||
# | ||
# Licensed under the Apache License, Version 2.0 (the "License"); | ||
# you may not use this file except in compliance with the License. | ||
# You may obtain a copy of the License at | ||
# | ||
# http://www.apache.org/licenses/LICENSE-2.0 | ||
# | ||
# Unless required by applicable law or agreed to in writing, software | ||
# distributed under the License is distributed on an "AS IS" BASIS, | ||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
# See the License for the specific language governing permissions and | ||
# limitations under the License. | ||
|
||
"""Client for interacting with the Google Cloud Firestore API. | ||
This is the base from which all interactions with the API occur. | ||
In the hierarchy of API concepts | ||
* a :class:`~google.cloud.firestore_v1.client.Client` owns a | ||
:class:`~google.cloud.firestore_v1.async_collection.AsyncCollectionReference` | ||
* a :class:`~google.cloud.firestore_v1.client.Client` owns a | ||
:class:`~google.cloud.firestore_v1.async_document.AsyncDocumentReference` | ||
""" | ||
|
||
from google.cloud.firestore_v1.base_client import ( | ||
BaseClient, | ||
DEFAULT_DATABASE, | ||
_CLIENT_INFO, | ||
_reference_info, | ||
_parse_batch_get, | ||
_get_doc_mask, | ||
_path_helper, | ||
) | ||
|
||
from google.cloud.firestore_v1 import _helpers | ||
from google.cloud.firestore_v1.async_query import AsyncQuery | ||
from google.cloud.firestore_v1.async_batch import AsyncWriteBatch | ||
from google.cloud.firestore_v1.async_collection import AsyncCollectionReference | ||
from google.cloud.firestore_v1.async_document import AsyncDocumentReference | ||
from google.cloud.firestore_v1.async_transaction import AsyncTransaction | ||
|
||
|
||
class AsyncClient(BaseClient): | ||
"""Client for interacting with Google Cloud Firestore API. | ||
.. note:: | ||
Since the Cloud Firestore API requires the gRPC transport, no | ||
``_http`` argument is accepted by this class. | ||
Args: | ||
project (Optional[str]): The project which the client acts on behalf | ||
of. If not passed, falls back to the default inferred | ||
from the environment. | ||
credentials (Optional[~google.auth.credentials.Credentials]): The | ||
OAuth2 Credentials to use for this client. If not passed, falls | ||
back to the default inferred from the environment. | ||
database (Optional[str]): The database name that the client targets. | ||
For now, :attr:`DEFAULT_DATABASE` (the default value) is the | ||
only valid database. | ||
client_info (Optional[google.api_core.gapic_v1.client_info.ClientInfo]): | ||
The client info used to send a user-agent string along with API | ||
requests. If ``None``, then default info will be used. Generally, | ||
you only need to set this if you're developing your own library | ||
or partner tool. | ||
client_options (Union[dict, google.api_core.client_options.ClientOptions]): | ||
Client options used to set user options on the client. API Endpoint | ||
should be set through client_options. | ||
""" | ||
|
||
def __init__( | ||
self, | ||
project=None, | ||
credentials=None, | ||
database=DEFAULT_DATABASE, | ||
client_info=_CLIENT_INFO, | ||
client_options=None, | ||
): | ||
super(AsyncClient, self).__init__( | ||
project=project, | ||
credentials=credentials, | ||
database=database, | ||
client_info=client_info, | ||
client_options=client_options, | ||
) | ||
|
||
def collection(self, *collection_path): | ||
"""Get a reference to a collection. | ||
For a top-level collection: | ||
.. code-block:: python | ||
>>> client.collection('top') | ||
For a sub-collection: | ||
.. code-block:: python | ||
>>> client.collection('mydocs/doc/subcol') | ||
>>> # is the same as | ||
>>> client.collection('mydocs', 'doc', 'subcol') | ||
Sub-collections can be nested deeper in a similar fashion. | ||
Args: | ||
collection_path (Tuple[str, ...]): Can either be | ||
* A single ``/``-delimited path to a collection | ||
* A tuple of collection path segments | ||
Returns: | ||
:class:`~google.cloud.firestore_v1.async_collection.AsyncCollectionReference`: | ||
A reference to a collection in the Firestore database. | ||
""" | ||
return AsyncCollectionReference(*_path_helper(collection_path), client=self) | ||
|
||
def collection_group(self, collection_id): | ||
""" | ||
Creates and returns a new AsyncQuery that includes all documents in the | ||
database that are contained in a collection or subcollection with the | ||
given collection_id. | ||
.. code-block:: python | ||
>>> query = client.collection_group('mygroup') | ||
Args: | ||
collection_id (str) Identifies the collections to query over. | ||
Every collection or subcollection with this ID as the last segment of its | ||
path will be included. Cannot contain a slash. | ||
Returns: | ||
:class:`~google.cloud.firestore_v1.async_query.AsyncQuery`: | ||
The created AsyncQuery. | ||
""" | ||
return AsyncQuery( | ||
self._get_collection_reference(collection_id), all_descendants=True | ||
) | ||
|
||
def document(self, *document_path): | ||
"""Get a reference to a document in a collection. | ||
For a top-level document: | ||
.. code-block:: python | ||
>>> client.document('collek/shun') | ||
>>> # is the same as | ||
>>> client.document('collek', 'shun') | ||
For a document in a sub-collection: | ||
.. code-block:: python | ||
>>> client.document('mydocs/doc/subcol/child') | ||
>>> # is the same as | ||
>>> client.document('mydocs', 'doc', 'subcol', 'child') | ||
Documents in sub-collections can be nested deeper in a similar fashion. | ||
Args: | ||
document_path (Tuple[str, ...]): Can either be | ||
* A single ``/``-delimited path to a document | ||
* A tuple of document path segments | ||
Returns: | ||
:class:`~google.cloud.firestore_v1.document.AsyncDocumentReference`: | ||
A reference to a document in a collection. | ||
""" | ||
return AsyncDocumentReference( | ||
*self._document_path_helper(*document_path), client=self | ||
) | ||
|
||
async def get_all(self, references, field_paths=None, transaction=None): | ||
"""Retrieve a batch of documents. | ||
.. note:: | ||
Documents returned by this method are not guaranteed to be | ||
returned in the same order that they are given in ``references``. | ||
.. note:: | ||
If multiple ``references`` refer to the same document, the server | ||
will only return one result. | ||
See :meth:`~google.cloud.firestore_v1.client.Client.field_path` for | ||
more information on **field paths**. | ||
If a ``transaction`` is used and it already has write operations | ||
added, this method cannot be used (i.e. read-after-write is not | ||
allowed). | ||
Args: | ||
references (List[.AsyncDocumentReference, ...]): Iterable of document | ||
references to be retrieved. | ||
field_paths (Optional[Iterable[str, ...]]): An iterable of field | ||
paths (``.``-delimited list of field names) to use as a | ||
projection of document fields in the returned results. If | ||
no value is provided, all fields will be returned. | ||
transaction (Optional[:class:`~google.cloud.firestore_v1.async_transaction.AsyncTransaction`]): | ||
An existing transaction that these ``references`` will be | ||
retrieved in. | ||
Yields: | ||
.DocumentSnapshot: The next document snapshot that fulfills the | ||
query, or :data:`None` if the document does not exist. | ||
""" | ||
document_paths, reference_map = _reference_info(references) | ||
mask = _get_doc_mask(field_paths) | ||
response_iterator = self._firestore_api.batch_get_documents( | ||
request={ | ||
"database": self._database_string, | ||
"documents": document_paths, | ||
"mask": mask, | ||
"transaction": _helpers.get_transaction_id(transaction), | ||
}, | ||
metadata=self._rpc_metadata, | ||
) | ||
|
||
for get_doc_response in response_iterator: | ||
yield _parse_batch_get(get_doc_response, reference_map, self) | ||
|
||
async def collections(self): | ||
"""List top-level collections of the client's database. | ||
Returns: | ||
Sequence[:class:`~google.cloud.firestore_v1.async_collection.AsyncCollectionReference`]: | ||
iterator of subcollections of the current document. | ||
""" | ||
iterator = self._firestore_api.list_collection_ids( | ||
request={"parent": "{}/documents".format(self._database_string)}, | ||
metadata=self._rpc_metadata, | ||
) | ||
|
||
while True: | ||
for i in iterator.collection_ids: | ||
yield self.collection(i) | ||
if iterator.next_page_token: | ||
iterator = self._firestore_api.list_collection_ids( | ||
request={ | ||
"parent": "{}/documents".format(self._database_string), | ||
"page_token": iterator.next_page_token, | ||
}, | ||
metadata=self._rpc_metadata, | ||
) | ||
else: | ||
return | ||
|
||
# TODO(microgen): currently this method is rewritten to iterate/page itself. | ||
# https://github.com/googleapis/gapic-generator-python/issues/516 | ||
# it seems the generator ought to be able to do this itself. | ||
# iterator.client = self | ||
# iterator.item_to_value = _item_to_collection_ref | ||
# return iterator | ||
|
||
def batch(self): | ||
"""Get a batch instance from this client. | ||
Returns: | ||
:class:`~google.cloud.firestore_v1.async_batch.AsyncWriteBatch`: | ||
A "write" batch to be used for accumulating document changes and | ||
sending the changes all at once. | ||
""" | ||
return AsyncWriteBatch(self) | ||
|
||
def transaction(self, **kwargs): | ||
"""Get a transaction that uses this client. | ||
See :class:`~google.cloud.firestore_v1.async_transaction.AsyncTransaction` for | ||
more information on transactions and the constructor arguments. | ||
Args: | ||
kwargs (Dict[str, Any]): The keyword arguments (other than | ||
``client``) to pass along to the | ||
:class:`~google.cloud.firestore_v1.async_transaction.AsyncTransaction` | ||
constructor. | ||
Returns: | ||
:class:`~google.cloud.firestore_v1.async_transaction.AsyncTransaction`: | ||
A transaction attached to this client. | ||
""" | ||
return AsyncTransaction(self, **kwargs) |
Oops, something went wrong.