Skip to content

Commit

Permalink
feat: create async interface (#61)
Browse files Browse the repository at this point in the history
* 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
rafilong committed Jul 16, 2020
1 parent cb8606a commit eaba25e
Show file tree
Hide file tree
Showing 17 changed files with 4,531 additions and 14 deletions.
64 changes: 64 additions & 0 deletions google/cloud/firestore_v1/async_batch.py
@@ -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()
288 changes: 288 additions & 0 deletions google/cloud/firestore_v1/async_client.py
@@ -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)

0 comments on commit eaba25e

Please sign in to comment.