refactor: removing usage of GetDocument proto in favor of BatchGetDocuments (#316)

* Began refactoring away from GetDocument proto

which contains no `read_time` field and in general is a shortcut around `BatchGetDocuments`

* Correctly instantiate snapshots for missing documents

* Removed stale NotFound exception

* Removed unnecessary empty list check

* Linting fix

* Expanded batch_get change to async classes

* Updated variable name

* Added get_batch to async classes

* Improved consumption of async generators

* Fixed test coverage

* Fixed broken mock in test

* Linting

* Reverted the move of AsyncIter

Co-authored-by: Craig Labenz <craiglabenz@google.com>

Author: craiglabenz

packages/google-cloud-firestore/google/cloud/firestore_v1/async_client.py

@@ -49,7 +49,7 @@
 from google.cloud.firestore_v1.services.firestore.transports import (
     grpc_asyncio as firestore_grpc_transport,
 )
-from typing import Any, AsyncGenerator, Iterable
+from typing import Any, AsyncGenerator, Iterable, List
 
 
 class AsyncClient(BaseClient):
@@ -209,7 +209,7 @@ def document(self, *document_path: str) -> AsyncDocumentReference:
 
     async def get_all(
         self,
-        references: list,
+        references: List[AsyncDocumentReference],
         field_paths: Iterable[str] = None,
         transaction=None,
         retry: retries.Retry = gapic_v1.method.DEFAULT,

packages/google-cloud-firestore/google/cloud/firestore_v1/async_document.py

@@ -13,23 +13,27 @@
 # limitations under the License.
 
 """Classes for representing documents for the Google Cloud Firestore API."""
+import datetime
+import logging
 
 from google.api_core import gapic_v1  # type: ignore
 from google.api_core import retry as retries  # type: ignore
+from google.cloud._helpers import _datetime_to_pb_timestamp  # type: ignore
 
 from google.cloud.firestore_v1.base_document import (
     BaseDocumentReference,
     DocumentSnapshot,
     _first_write_result,
 )
-
-from google.api_core import exceptions  # type: ignore
 from google.cloud.firestore_v1 import _helpers
 from google.cloud.firestore_v1.types import write
-from google.protobuf import timestamp_pb2
+from google.protobuf.timestamp_pb2 import Timestamp
 from typing import Any, AsyncGenerator, Coroutine, Iterable, Union
 
 
+logger = logging.getLogger(__name__)
+
+
 class AsyncDocumentReference(BaseDocumentReference):
     """A reference to a document in a Firestore database.
 
@@ -289,7 +293,7 @@ async def delete(
         option: _helpers.WriteOption = None,
         retry: retries.Retry = gapic_v1.method.DEFAULT,
         timeout: float = None,
-    ) -> timestamp_pb2.Timestamp:
+    ) -> Timestamp:
         """Delete the current document in the Firestore database.
 
         Args:
@@ -353,31 +357,34 @@ async def get(
                 :attr:`create_time` attributes will all be ``None`` and
                 its :attr:`exists` attribute will be ``False``.
         """
-        request, kwargs = self._prep_get(field_paths, transaction, retry, timeout)
+        from google.cloud.firestore_v1.base_client import _parse_batch_get
 
-        firestore_api = self._client._firestore_api
-        try:
-            document_pb = await firestore_api.get_document(
-                request=request, metadata=self._client._rpc_metadata, **kwargs,
+        request, kwargs = self._prep_batch_get(field_paths, transaction, retry, timeout)
+
+        response_iter = await self._client._firestore_api.batch_get_documents(
+            request=request, metadata=self._client._rpc_metadata, **kwargs,
+        )
+
+        async for resp in response_iter:
+            # Immediate return as the iterator should only ever have one item.
+            return _parse_batch_get(
+                get_doc_response=resp,
+                reference_map={self._document_path: self},
+                client=self._client,
             )
-        except exceptions.NotFound:
-            data = None
-            exists = False
-            create_time = None
-            update_time = None
-        else:
-            data = _helpers.decode_dict(document_pb.fields, self._client)
-            exists = True
-            create_time = document_pb.create_time
-            update_time = document_pb.update_time
+
+        logger.warning(
+            "`batch_get_documents` unexpectedly returned empty "
+            "stream. Expected one object.",
+        )
 
         return DocumentSnapshot(
-            reference=self,
-            data=data,
-            exists=exists,
-            read_time=None,  # No server read_time available
-            create_time=create_time,
-            update_time=update_time,
+            self,
+            None,
+            exists=False,
+            read_time=_datetime_to_pb_timestamp(datetime.datetime.now()),
+            create_time=None,
+            update_time=None,
         )
 
     async def collections(

packages/google-cloud-firestore/google/cloud/firestore_v1/base_document.py

@@ -268,7 +268,7 @@ def delete(
     ) -> NoReturn:
         raise NotImplementedError
 
-    def _prep_get(
+    def _prep_batch_get(
         self,
         field_paths: Iterable[str] = None,
         transaction=None,
@@ -285,7 +285,8 @@ def _prep_get(
             mask = None
 
         request = {
-            "name": self._document_path,
+            "database": self._client._database_string,
+            "documents": [self._document_path],
             "mask": mask,
             "transaction": _helpers.get_transaction_id(transaction),
         }

packages/google-cloud-firestore/google/cloud/firestore_v1/document.py

@@ -13,24 +13,28 @@
 # limitations under the License.
 
 """Classes for representing documents for the Google Cloud Firestore API."""
+import datetime
+import logging
 
 from google.api_core import gapic_v1  # type: ignore
 from google.api_core import retry as retries  # type: ignore
+from google.cloud._helpers import _datetime_to_pb_timestamp  # type: ignore
 
 from google.cloud.firestore_v1.base_document import (
     BaseDocumentReference,
     DocumentSnapshot,
     _first_write_result,
 )
-
-from google.api_core import exceptions  # type: ignore
 from google.cloud.firestore_v1 import _helpers
 from google.cloud.firestore_v1.types import write
 from google.cloud.firestore_v1.watch import Watch
-from google.protobuf import timestamp_pb2
+from google.protobuf.timestamp_pb2 import Timestamp
 from typing import Any, Callable, Generator, Iterable
 
 
+logger = logging.getLogger(__name__)
+
+
 class DocumentReference(BaseDocumentReference):
     """A reference to a document in a Firestore database.
 
@@ -325,7 +329,7 @@ def delete(
         option: _helpers.WriteOption = None,
         retry: retries.Retry = gapic_v1.method.DEFAULT,
         timeout: float = None,
-    ) -> timestamp_pb2.Timestamp:
+    ) -> Timestamp:
         """Delete the current document in the Firestore database.
 
         Args:
@@ -389,31 +393,35 @@ def get(
                 :attr:`create_time` attributes will all be ``None`` and
                 its :attr:`exists` attribute will be ``False``.
         """
-        request, kwargs = self._prep_get(field_paths, transaction, retry, timeout)
+        from google.cloud.firestore_v1.base_client import _parse_batch_get
+
+        request, kwargs = self._prep_batch_get(field_paths, transaction, retry, timeout)
+
+        response_iter = self._client._firestore_api.batch_get_documents(
+            request=request, metadata=self._client._rpc_metadata, **kwargs,
+        )
+
+        get_doc_response = next(response_iter, None)
 
-        firestore_api = self._client._firestore_api
-        try:
-            document_pb = firestore_api.get_document(
-                request=request, metadata=self._client._rpc_metadata, **kwargs,
+        if get_doc_response is not None:
+            return _parse_batch_get(
+                get_doc_response=get_doc_response,
+                reference_map={self._document_path: self},
+                client=self._client,
             )
-        except exceptions.NotFound:
-            data = None
-            exists = False
-            create_time = None
-            update_time = None
-        else:
-            data = _helpers.decode_dict(document_pb.fields, self._client)
-            exists = True
-            create_time = document_pb.create_time
-            update_time = document_pb.update_time
+
+        logger.warning(
+            "`batch_get_documents` unexpectedly returned empty "
+            "stream. Expected one object.",
+        )
 
         return DocumentSnapshot(
-            reference=self,
-            data=data,
-            exists=exists,
-            read_time=None,  # No server read_time available
-            create_time=create_time,
-            update_time=update_time,
+            self,
+            None,
+            exists=False,
+            read_time=_datetime_to_pb_timestamp(datetime.datetime.now()),
+            create_time=None,
+            update_time=None,
         )
 
     def collections(

packages/google-cloud-firestore/tests/unit/v1/test__helpers.py

@@ -13,11 +13,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import aiounittest
 import datetime
 import sys
 import unittest
 
 import mock
+import pytest
+from typing import List
 
 
 class AsyncMock(mock.MagicMock):
@@ -26,10 +29,14 @@ async def __call__(self, *args, **kwargs):
 
 
 class AsyncIter:
+    """Utility to help recreate the effect of an async generator. Useful when
+    you need to mock a system that requires `async for`.
+    """
+
     def __init__(self, items):
         self.items = items
 
-    async def __aiter__(self, **_):
+    async def __aiter__(self):
         for i in self.items:
             yield i
 
@@ -2424,6 +2431,15 @@ def test_retry_and_timeout(self):
         self.assertEqual(kwargs, expected)
 
 
+class TestAsyncGenerator(aiounittest.AsyncTestCase):
+    @pytest.mark.asyncio
+    async def test_async_iter(self):
+        consumed: List[int] = []
+        async for el in AsyncIter([1, 2, 3]):
+            consumed.append(el)
+        self.assertEqual(consumed, [1, 2, 3])
+
+
 def _value_pb(**kwargs):
     from google.cloud.firestore_v1.types.document import Value
 

packages/google-cloud-firestore/tests/unit/v1/test_async_client.py

@@ -18,7 +18,7 @@
 import aiounittest
 
 import mock
-from tests.unit.v1.test__helpers import AsyncMock, AsyncIter
+from tests.unit.v1.test__helpers import AsyncIter, AsyncMock
 
 
 class TestAsyncClient(aiounittest.AsyncTestCase):

packages/google-cloud-firestore/tests/unit/v1/test_async_collection.py

@@ -17,7 +17,7 @@
 import aiounittest
 
 import mock
-from tests.unit.v1.test__helpers import AsyncMock, AsyncIter
+from tests.unit.v1.test__helpers import AsyncIter, AsyncMock
 
 
 class TestAsyncCollectionReference(aiounittest.AsyncTestCase):