feat: Added read_time as a parameter to various calls (async classes) (#1059)

* feat: Added read_time as a parameter to various calls (async classes)

* used TYPE_CHECKING; fixed unit tests

* linting + fixing cover

* final linting

Author: gkevinzheng

google/cloud/firestore_v1/async_aggregation.py

@@ -37,6 +37,7 @@
     from google.cloud.firestore_v1.base_aggregation import AggregationResult
     from google.cloud.firestore_v1.query_profile import ExplainMetrics, ExplainOptions
     import google.cloud.firestore_v1.types.query_profile as query_profile_pb
+    import datetime
 
 
 class AsyncAggregationQuery(BaseAggregationQuery):
@@ -55,6 +56,7 @@ async def get(
         timeout: float | None = None,
         *,
         explain_options: Optional[ExplainOptions] = None,
+        read_time: Optional[datetime.datetime] = None,
     ) -> QueryResultsList[List[AggregationResult]]:
         """Runs the aggregation query.
 
@@ -75,6 +77,10 @@ async def get(
                 (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]):
                 Options to enable query profiling for this query. When set,
                 explain_metrics will be available on the returned generator.
+            read_time (Optional[datetime.datetime]): If set, reads documents as they were at the given
+                time. This must be a timestamp within the past one hour, or if Point-in-Time Recovery
+                is enabled, can additionally be a whole minute timestamp within the past 7 days. If no
+                timezone is specified in the :class:`datetime.datetime` object, it is assumed to be UTC.
 
         Returns:
             QueryResultsList[List[AggregationResult]]: The aggregation query results.
@@ -87,6 +93,7 @@ async def get(
             retry=retry,
             timeout=timeout,
             explain_options=explain_options,
+            read_time=read_time,
         )
         try:
             result = [aggregation async for aggregation in stream_result]
@@ -106,6 +113,7 @@ async def _make_stream(
         retry: retries.AsyncRetry | object | None = gapic_v1.method.DEFAULT,
         timeout: Optional[float] = None,
         explain_options: Optional[ExplainOptions] = None,
+        read_time: Optional[datetime.datetime] = None,
     ) -> AsyncGenerator[List[AggregationResult] | query_profile_pb.ExplainMetrics, Any]:
         """Internal method for stream(). Runs the aggregation query.
 
@@ -130,6 +138,10 @@ async def _make_stream(
                 (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]):
                 Options to enable query profiling for this query. When set,
                 explain_metrics will be available on the returned generator.
+            read_time (Optional[datetime.datetime]): If set, reads documents as they were at the given
+                time. This must be a timestamp within the past one hour, or if Point-in-Time Recovery
+                is enabled, can additionally be a whole minute timestamp within the past 7 days. If no
+                timezone is specified in the :class:`datetime.datetime` object, it is assumed to be UTC.
 
         Yields:
             List[AggregationResult] | query_profile_pb.ExplainMetrics:
@@ -143,6 +155,7 @@ async def _make_stream(
             retry,
             timeout,
             explain_options,
+            read_time,
         )
 
         response_iterator = await self._client._firestore_api.run_aggregation_query(
@@ -167,6 +180,7 @@ def stream(
         timeout: Optional[float] = None,
         *,
         explain_options: Optional[ExplainOptions] = None,
+        read_time: Optional[datetime.datetime] = None,
     ) -> AsyncStreamGenerator[List[AggregationResult]]:
         """Runs the aggregation query.
 
@@ -190,6 +204,10 @@ def stream(
                 (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]):
                 Options to enable query profiling for this query. When set,
                 explain_metrics will be available on the returned generator.
+            read_time (Optional[datetime.datetime]): If set, reads documents as they were at the given
+                time. This must be a timestamp within the past one hour, or if Point-in-Time Recovery
+                is enabled, can additionally be a whole minute timestamp within the past 7 days. If no
+                timezone is specified in the :class:`datetime.datetime` object, it is assumed to be UTC.
 
         Returns:
             `AsyncStreamGenerator[List[AggregationResult]]`:
@@ -201,5 +219,6 @@ def stream(
             retry=retry,
             timeout=timeout,
             explain_options=explain_options,
+            read_time=read_time,
         )
         return AsyncStreamGenerator(inner_generator, explain_options)

google/cloud/firestore_v1/async_client.py

@@ -48,8 +48,10 @@
     grpc_asyncio as firestore_grpc_transport,
 )
 
-if TYPE_CHECKING:
-    from google.cloud.firestore_v1.bulk_writer import BulkWriter  # pragma: NO COVER
+if TYPE_CHECKING:  # pragma: NO COVER
+    import datetime
+
+    from google.cloud.firestore_v1.bulk_writer import BulkWriter
 
 
 class AsyncClient(BaseClient):
@@ -227,6 +229,8 @@ async def get_all(
         transaction: AsyncTransaction | None = None,
         retry: retries.AsyncRetry | object | None = gapic_v1.method.DEFAULT,
         timeout: float | None = None,
+        *,
+        read_time: datetime.datetime | None = None,
     ) -> AsyncGenerator[DocumentSnapshot, Any]:
         """Retrieve a batch of documents.
 
@@ -261,13 +265,17 @@ async def get_all(
                 should be retried.  Defaults to a system-specified policy.
             timeout (float): The timeout for this request.  Defaults to a
                 system-specified value.
+            read_time (Optional[datetime.datetime]): If set, reads documents as they were at the given
+                time. This must be a timestamp within the past one hour, or if Point-in-Time Recovery
+                is enabled, can additionally be a whole minute timestamp within the past 7 days. If no
+                timezone is specified in the :class:`datetime.datetime` object, it is assumed to be UTC.
 
         Yields:
             .DocumentSnapshot: The next document snapshot that fulfills the
             query, or :data:`None` if the document does not exist.
         """
         request, reference_map, kwargs = self._prep_get_all(
-            references, field_paths, transaction, retry, timeout
+            references, field_paths, transaction, retry, timeout, read_time
         )
 
         response_iterator = await self._firestore_api.batch_get_documents(
@@ -283,6 +291,8 @@ async def collections(
         self,
         retry: retries.AsyncRetry | object | None = gapic_v1.method.DEFAULT,
         timeout: float | None = None,
+        *,
+        read_time: datetime.datetime | None = None,
     ) -> AsyncGenerator[AsyncCollectionReference, Any]:
         """List top-level collections of the client's database.
 
@@ -291,12 +301,16 @@ async def collections(
                 should be retried.  Defaults to a system-specified policy.
             timeout (float): The timeout for this request.  Defaults to a
                 system-specified value.
+            read_time (Optional[datetime.datetime]): If set, reads documents as they were at the given
+                time. This must be a timestamp within the past one hour, or if Point-in-Time Recovery
+                is enabled, can additionally be a whole minute timestamp within the past 7 days. If no
+                timezone is specified in the :class:`datetime.datetime` object, it is assumed to be UTC.
 
         Returns:
             Sequence[:class:`~google.cloud.firestore_v1.async_collection.AsyncCollectionReference`]:
                 iterator of subcollections of the current document.
         """
-        request, kwargs = self._prep_collections(retry, timeout)
+        request, kwargs = self._prep_collections(retry, timeout, read_time)
         iterator = await self._firestore_api.list_collection_ids(
             request=request,
             metadata=self._rpc_metadata,

google/cloud/firestore_v1/async_collection.py

@@ -34,6 +34,8 @@
 from google.cloud.firestore_v1.document import DocumentReference
 
 if TYPE_CHECKING:  # pragma: NO COVER
+    import datetime
+
     from google.cloud.firestore_v1.async_stream_generator import AsyncStreamGenerator
     from google.cloud.firestore_v1.base_document import DocumentSnapshot
     from google.cloud.firestore_v1.query_profile import ExplainOptions
@@ -162,6 +164,8 @@ async def list_documents(
         page_size: int | None = None,
         retry: retries.AsyncRetry | object | None = gapic_v1.method.DEFAULT,
         timeout: float | None = None,
+        *,
+        read_time: datetime.datetime | None = None,
     ) -> AsyncGenerator[DocumentReference, None]:
         """List all subdocuments of the current collection.
 
@@ -173,14 +177,20 @@ async def list_documents(
                 should be retried.  Defaults to a system-specified policy.
             timeout (float): The timeout for this request.  Defaults to a
                 system-specified value.
+            read_time (Optional[datetime.datetime]): If set, reads documents as they were at the given
+                time. This must be a timestamp within the past one hour, or if Point-in-Time Recovery
+                is enabled, can additionally be a whole minute timestamp within the past 7 days. If no
+                timezone is specified in the :class:`datetime.datetime` object, it is assumed to be UTC.
 
         Returns:
             Sequence[:class:`~google.cloud.firestore_v1.collection.DocumentReference`]:
                 iterator of subdocuments of the current collection. If the
                 collection does not exist at the time of `snapshot`, the
                 iterator will be empty
         """
-        request, kwargs = self._prep_list_documents(page_size, retry, timeout)
+        request, kwargs = self._prep_list_documents(
+            page_size, retry, timeout, read_time
+        )
 
         iterator = await self._client._firestore_api.list_documents(
             request=request,
@@ -197,6 +207,7 @@ async def get(
         timeout: Optional[float] = None,
         *,
         explain_options: Optional[ExplainOptions] = None,
+        read_time: Optional[datetime.datetime] = None,
     ) -> QueryResultsList[DocumentSnapshot]:
         """Read the documents in this collection.
 
@@ -216,6 +227,10 @@ async def get(
                 (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]):
                 Options to enable query profiling for this query. When set,
                 explain_metrics will be available on the returned generator.
+            read_time (Optional[datetime.datetime]): If set, reads documents as they were at the given
+                time. This must be a timestamp within the past one hour, or if Point-in-Time Recovery
+                is enabled, can additionally be a whole minute timestamp within the past 7 days. If no
+                timezone is specified in the :class:`datetime.datetime` object, it is assumed to be UTC.
 
         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).
@@ -227,6 +242,8 @@ async def get(
         query, kwargs = self._prep_get_or_stream(retry, timeout)
         if explain_options is not None:
             kwargs["explain_options"] = explain_options
+        if read_time is not None:
+            kwargs["read_time"] = read_time
 
         return await query.get(transaction=transaction, **kwargs)
 
@@ -237,6 +254,7 @@ def stream(
         timeout: Optional[float] = None,
         *,
         explain_options: Optional[ExplainOptions] = None,
+        read_time: Optional[datetime.datetime] = None,
     ) -> AsyncStreamGenerator[DocumentSnapshot]:
         """Read the documents in this collection.
 
@@ -268,6 +286,10 @@ def stream(
                 (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]):
                 Options to enable query profiling for this query. When set,
                 explain_metrics will be available on the returned generator.
+            read_time (Optional[datetime.datetime]): If set, reads documents as they were at the given
+                time. This must be a timestamp within the past one hour, or if Point-in-Time Recovery
+                is enabled, can additionally be a whole minute timestamp within the past 7 days. If no
+                timezone is specified in the :class:`datetime.datetime` object, it is assumed to be UTC.
 
         Returns:
             `AsyncStreamGenerator[DocumentSnapshot]`: A generator of the query
@@ -276,5 +298,7 @@ def stream(
         query, kwargs = self._prep_get_or_stream(retry, timeout)
         if explain_options:
             kwargs["explain_options"] = explain_options
+        if read_time is not None:
+            kwargs["read_time"] = read_time
 
         return query.stream(transaction=transaction, **kwargs)

google/cloud/firestore_v1/async_document.py

@@ -329,6 +329,8 @@ async def get(
         transaction=None,
         retry: retries.AsyncRetry | object | None = gapic_v1.method.DEFAULT,
         timeout: float | None = None,
+        *,
+        read_time: datetime.datetime | None = None,
     ) -> DocumentSnapshot:
         """Retrieve a snapshot of the current document.
 
@@ -351,6 +353,10 @@ async def get(
                 should be retried.  Defaults to a system-specified policy.
             timeout (float): The timeout for this request.  Defaults to a
                 system-specified value.
+            read_time (Optional[datetime.datetime]): If set, reads documents as they were at the given
+                time. This must be a timestamp within the past one hour, or if Point-in-Time Recovery
+                is enabled, can additionally be a whole minute timestamp within the past 7 days. If no
+                timezone is specified in the :class:`datetime.datetime` object, it is assumed to be UTC.
 
         Returns:
             :class:`~google.cloud.firestore_v1.base_document.DocumentSnapshot`:
@@ -362,7 +368,9 @@ async def get(
         """
         from google.cloud.firestore_v1.base_client import _parse_batch_get
 
-        request, kwargs = self._prep_batch_get(field_paths, transaction, retry, timeout)
+        request, kwargs = self._prep_batch_get(
+            field_paths, transaction, retry, timeout, read_time
+        )
 
         response_iter = await self._client._firestore_api.batch_get_documents(
             request=request,
@@ -397,6 +405,8 @@ async def collections(
         page_size: int | None = None,
         retry: retries.AsyncRetry | object | None = gapic_v1.method.DEFAULT,
         timeout: float | None = None,
+        *,
+        read_time: datetime.datetime | None = None,
     ) -> AsyncGenerator:
         """List subcollections of the current document.
 
@@ -408,14 +418,18 @@ async def collections(
                 should be retried.  Defaults to a system-specified policy.
             timeout (float): The timeout for this request.  Defaults to a
                 system-specified value.
+            read_time (Optional[datetime.datetime]): If set, reads documents as they were at the given
+                time. This must be a timestamp within the past one hour, or if Point-in-Time Recovery
+                is enabled, can additionally be a whole minute timestamp within the past 7 days. If no
+                timezone is specified in the :class:`datetime.datetime` object, it is assumed to be UTC.
 
         Returns:
             Sequence[:class:`~google.cloud.firestore_v1.async_collection.AsyncCollectionReference`]:
                 iterator of subcollections of the current document. If the
                 document does not exist at the time of `snapshot`, the
                 iterator will be empty
         """
-        request, kwargs = self._prep_collections(page_size, retry, timeout)
+        request, kwargs = self._prep_collections(page_size, retry, timeout, read_time)
 
         iterator = await self._client._firestore_api.list_collection_ids(
             request=request,