feat: asyncio system tests

google/cloud/firestore.py

@@ -18,6 +18,13 @@
 from google.cloud.firestore_v1 import __version__
 from google.cloud.firestore_v1 import ArrayRemove
 from google.cloud.firestore_v1 import ArrayUnion
+from google.cloud.firestore_v1 import AsyncClient
+from google.cloud.firestore_v1 import AsyncCollectionReference
+from google.cloud.firestore_v1 import AsyncDocumentReference
+from google.cloud.firestore_v1 import AsyncQuery
+from google.cloud.firestore_v1 import AsyncTransaction
+from google.cloud.firestore_v1 import async_transactional
+from google.cloud.firestore_v1 import AsyncWriteBatch
 from google.cloud.firestore_v1 import Client
 from google.cloud.firestore_v1 import CollectionReference
 from google.cloud.firestore_v1 import DELETE_FIELD
@@ -45,6 +52,13 @@
     "__version__",
     "ArrayRemove",
     "ArrayUnion",
+    "AsyncClient",
+    "AsyncCollectionReference",
+    "AsyncDocumentReference",
+    "AsyncQuery",
+    "AsyncTransaction",
+    "async_transactional",
+    "AsyncWriteBatch",
     "Client",
     "CollectionReference",
     "DELETE_FIELD",

google/cloud/firestore_v1/__init__.py

@@ -29,22 +29,31 @@
 from google.cloud.firestore_v1._helpers import LastUpdateOption
 from google.cloud.firestore_v1._helpers import ReadAfterWriteError
 from google.cloud.firestore_v1._helpers import WriteOption
-from google.cloud.firestore_v1.batch import WriteBatch
-from google.cloud.firestore_v1.client import Client
-from google.cloud.firestore_v1.collection import CollectionReference
+from google.cloud.firestore_v1.base_document import DocumentSnapshot
 from google.cloud.firestore_v1.transforms import ArrayRemove
 from google.cloud.firestore_v1.transforms import ArrayUnion
 from google.cloud.firestore_v1.transforms import DELETE_FIELD
 from google.cloud.firestore_v1.transforms import Increment
 from google.cloud.firestore_v1.transforms import Maximum
 from google.cloud.firestore_v1.transforms import Minimum
 from google.cloud.firestore_v1.transforms import SERVER_TIMESTAMP
+from google.cloud.firestore_v1.watch import Watch
+
+from google.cloud.firestore_v1.async_batch import AsyncWriteBatch
+from google.cloud.firestore_v1.async_client import AsyncClient
+from google.cloud.firestore_v1.async_collection import AsyncCollectionReference
+from google.cloud.firestore_v1.async_document import AsyncDocumentReference
+from google.cloud.firestore_v1.async_query import AsyncQuery
+from google.cloud.firestore_v1.async_transaction import AsyncTransaction
+from google.cloud.firestore_v1.async_transaction import async_transactional
+
+from google.cloud.firestore_v1.batch import WriteBatch
+from google.cloud.firestore_v1.client import Client
+from google.cloud.firestore_v1.collection import CollectionReference
 from google.cloud.firestore_v1.document import DocumentReference
-from google.cloud.firestore_v1.document import DocumentSnapshot
 from google.cloud.firestore_v1.query import Query
 from google.cloud.firestore_v1.transaction import Transaction
 from google.cloud.firestore_v1.transaction import transactional
-from google.cloud.firestore_v1.watch import Watch
 
 
 # TODO(https://github.com/googleapis/python-firestore/issues/93): this is all on the generated surface. We require this to match
@@ -100,6 +109,13 @@
     "__version__",
     "ArrayRemove",
     "ArrayUnion",
+    "AsyncClient",
+    "AsyncCollectionReference",
+    "AsyncDocumentReference",
+    "AsyncQuery",
+    "AsyncTransaction",
+    "async_transactional",
+    "AsyncWriteBatch",
     "Client",
     "CollectionReference",
     "DELETE_FIELD",

google/cloud/firestore_v1/async_client.py

@@ -99,11 +99,15 @@ def _firestore_api(self):
             :class:`~google.cloud.gapic.firestore.v1`.async_firestore_client.FirestoreAsyncClient:
             The GAPIC client with the credentials of the current client.
         """
-        return self._firestore_api_helper(
+        firestore_api = self._firestore_api_helper(
             firestore_grpc_transport.FirestoreGrpcAsyncIOTransport,
             firestore_client.FirestoreAsyncClient,
             firestore_client,
         )
+        # Add _transport to FirestoreAsyncClient interface for use in Watch
+        if hasattr(firestore_api, "_client"):
+            firestore_api._transport = firestore_api._client._transport
+        return firestore_api
 
     @property
     def _target(self):
@@ -242,7 +246,7 @@ async def get_all(self, references, field_paths=None, transaction=None):
         """
         document_paths, reference_map = _reference_info(references)
         mask = _get_doc_mask(field_paths)
-        response_iterator = self._firestore_api.batch_get_documents(
+        response_iterator = await self._firestore_api.batch_get_documents(
             request={
                 "database": self._database_string,
                 "documents": document_paths,

google/cloud/firestore_v1/async_collection.py

@@ -22,8 +22,6 @@
     _item_to_document_ref,
 )
 from google.cloud.firestore_v1 import async_query
-from google.cloud.firestore_v1.watch import Watch
-from google.cloud.firestore_v1 import async_document
 
 
 class AsyncCollectionReference(BaseCollectionReference):
@@ -119,7 +117,8 @@ async def list_documents(self, page_size=None):
             },
             metadata=self._client._rpc_metadata,
         )
-        return (_item_to_document_ref(self, i) for i in iterator)
+        async for i in iterator:
+            yield _item_to_document_ref(self, i)
 
     async def get(self, transaction=None):
         """Deprecated alias for :meth:`stream`."""
@@ -161,36 +160,3 @@ async def stream(self, transaction=None):
         query = async_query.AsyncQuery(self)
         async for d in query.stream(transaction=transaction):
             yield d
-
-    def on_snapshot(self, callback):
-        """Monitor the documents in this collection.
-
-        This starts a watch on this collection using a background thread. The
-        provided callback is run on the snapshot of the documents.
-
-        Args:
-            callback (Callable[[:class:`~google.cloud.firestore.collection.CollectionSnapshot`], NoneType]):
-                a callback to run when a change occurs.
-
-        Example:
-            from google.cloud import firestore_v1
-
-            db = firestore_v1.Client()
-            collection_ref = db.collection(u'users')
-
-            def on_snapshot(collection_snapshot, changes, read_time):
-                for doc in collection_snapshot.documents:
-                    print(u'{} => {}'.format(doc.id, doc.to_dict()))
-
-            # Watch this collection
-            collection_watch = collection_ref.on_snapshot(on_snapshot)
-
-            # Terminate this watch
-            collection_watch.unsubscribe()
-        """
-        return Watch.for_query(
-            self._query(),
-            callback,
-            async_document.DocumentSnapshot,
-            async_document.AsyncDocumentReference,
-        )

google/cloud/firestore_v1/async_document.py

@@ -23,7 +23,6 @@
 from google.api_core import exceptions
 from google.cloud.firestore_v1 import _helpers
 from google.cloud.firestore_v1.types import common
-from google.cloud.firestore_v1.watch import Watch
 
 
 class AsyncDocumentReference(BaseDocumentReference):
@@ -385,39 +384,3 @@ async def collections(self, page_size=None):
         # iterator.document = self
         # iterator.item_to_value = _item_to_collection_ref
         # return iterator
-
-    def on_snapshot(self, callback):
-        """Watch this document.
-
-        This starts a watch on this document using a background thread. The
-        provided callback is run on the snapshot.
-
-        Args:
-            callback(Callable[[:class:`~google.cloud.firestore.document.DocumentSnapshot`], NoneType]):
-                a callback to run when a change occurs
-
-        Example:
-
-        .. code-block:: python
-
-            from google.cloud import firestore_v1
-
-            db = firestore_v1.Client()
-            collection_ref = db.collection(u'users')
-
-            def on_snapshot(document_snapshot, changes, read_time):
-                doc = document_snapshot
-                print(u'{} => {}'.format(doc.id, doc.to_dict()))
-
-            doc_ref = db.collection(u'users').document(
-                u'alovelace' + unique_resource_id())
-
-            # Watch this document
-            doc_watch = doc_ref.on_snapshot(on_snapshot)
-
-            # Terminate this watch
-            doc_watch.unsubscribe()
-        """
-        return Watch.for_document(
-            self, callback, DocumentSnapshot, AsyncDocumentReference
-        )

google/cloud/firestore_v1/async_query.py

@@ -27,8 +27,6 @@
 )
 
 from google.cloud.firestore_v1 import _helpers
-from google.cloud.firestore_v1 import async_document
-from google.cloud.firestore_v1.watch import Watch
 
 
 class AsyncQuery(BaseQuery):
@@ -149,7 +147,7 @@ async def stream(self, transaction=None):
             The next document that fulfills the query.
         """
         parent_path, expected_prefix = self._parent._parent_info()
-        response_iterator = self._client._firestore_api.run_query(
+        response_iterator = await self._client._firestore_api.run_query(
             request={
                 "parent": parent_path,
                 "structured_query": self._to_protobuf(),
@@ -169,39 +167,3 @@ async def stream(self, transaction=None):
                 )
             if snapshot is not None:
                 yield snapshot
-
-    def on_snapshot(self, callback):
-        """Monitor the documents in this collection that match this query.
-
-        This starts a watch on this query using a background thread. The
-        provided callback is run on the snapshot of the documents.
-
-        Args:
-            callback(Callable[[:class:`~google.cloud.firestore.query.QuerySnapshot`], NoneType]):
-                a callback to run when a change occurs.
-
-        Example:
-
-        .. code-block:: python
-
-            from google.cloud import firestore_v1
-
-            db = firestore_v1.Client()
-            query_ref = db.collection(u'users').where("user", "==", u'Ada')
-
-            def on_snapshot(docs, changes, read_time):
-                for doc in docs:
-                    print(u'{} => {}'.format(doc.id, doc.to_dict()))
-
-            # Watch this query
-            query_watch = query_ref.on_snapshot(on_snapshot)
-
-            # Terminate this watch
-            query_watch.unsubscribe()
-        """
-        return Watch.for_query(
-            self,
-            callback,
-            async_document.DocumentSnapshot,
-            async_document.AsyncDocumentReference,
-        )

google/cloud/firestore_v1/async_transaction.py

@@ -287,7 +287,7 @@ async def __call__(self, transaction, *args, **kwargs):
         raise ValueError(msg)
 
 
-def transactional(to_wrap):
+def async_transactional(to_wrap):
     """Decorate a callable so that it runs in a transaction.
 
     Args:

noxfile.py

@@ -124,7 +124,7 @@ def system(session):
     # Install all test dependencies, then install this package into the
     # virtualenv's dist-packages.
     session.install(
-        "mock", "pytest", "google-cloud-testutils",
+        "mock", "pytest", "pytest-asyncio", "google-cloud-testutils",
     )
     session.install("-e", ".")