address comments

google/cloud/firestore_v1/async_aggregation.py

@@ -137,8 +137,6 @@ async def _make_stream(
             returned explain metrics, yield `query_profile_pb.ExplainMetrics`
             individually.
         """
-        metrics: query_profile_pb.ExplainMetrics | None = None
-
         request, kwargs = self._prep_stream(
             transaction,
             retry,
@@ -157,7 +155,7 @@ async def _make_stream(
             if result:
                 yield result
 
-            if metrics is None and response.explain_metrics:
+            if response.explain_metrics:
                 metrics = response.explain_metrics
                 yield metrics
 

google/cloud/firestore_v1/async_query.py

@@ -373,8 +373,6 @@ async def _make_stream(
             yielded as `DocumentSnapshot`. When the result contains returned
             explain metrics, yield `query_profile_pb.ExplainMetrics` individually.
         """
-        metrics: query_profile_pb.ExplainMetrics | None = None
-
         request, expected_prefix, kwargs = self._prep_stream(
             transaction,
             retry,
@@ -400,7 +398,7 @@ async def _make_stream(
             if snapshot is not None:
                 yield snapshot
 
-            if metrics is None and response.explain_metrics:
+            if response.explain_metrics:
                 metrics = response.explain_metrics
                 yield metrics
 

google/cloud/firestore_v1/async_transaction.py

@@ -16,7 +16,6 @@
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, AsyncGenerator, Callable, Coroutine, Optional
-import warnings
 
 from google.api_core import exceptions, gapic_v1
 from google.api_core import retry_async as retries
@@ -195,26 +194,29 @@ async def get(
                 explain_metrics will be available on the returned generator.
                 Can only be used when running a query, not a document reference.
 
-
         Yields:
             DocumentSnapshot: The next document snapshot that fulfills the query,
             or :data:`None` if the document does not exist.
+
+        Raises:
+            ValueError: if `ref_or_query` is not one of the supported types, or
+            explain_options is provided when `ref_or_query` is a document
+            reference.
         """
         kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
         if isinstance(ref_or_query, AsyncDocumentReference):
             if explain_options is not None:
-                warnings.warn(
-                    "explain_options not supported in transanction with "
-                    "document references and will be ignored. To use "
-                    "explain_options, use transaction with query instead."
+                raise ValueError(
+                    "When type of `ref_or_query` is `AsyncDocumentReference`, "
+                    "`explain_options` cannot be provided."
                 )
             return await self._client.get_all(
                 [ref_or_query], transaction=self, **kwargs
             )
         elif isinstance(ref_or_query, AsyncQuery):
             if explain_options is not None:
                 kwargs["explain_options"] = explain_options
-            return await ref_or_query.stream(transaction=self, **kwargs)
+            return ref_or_query.stream(transaction=self, **kwargs)
         else:
             raise ValueError(
                 'Value for argument "ref_or_query" must be a AsyncDocumentReference or a AsyncQuery.'

google/cloud/firestore_v1/async_vector_query.py

@@ -139,8 +139,6 @@ async def _make_stream(
             yielded as `DocumentSnapshot`. When the result contains returned
             explain metrics, yield `query_profile_pb.ExplainMetrics` individually.
         """
-        metrics: query_profile_pb.ExplainMetrics | None = None
-
         request, expected_prefix, kwargs = self._prep_stream(
             transaction,
             retry,
@@ -166,7 +164,7 @@ async def _make_stream(
             if snapshot is not None:
                 yield snapshot
 
-            if metrics is None and response.explain_metrics:
+            if response.explain_metrics:
                 metrics = response.explain_metrics
                 yield metrics
 

google/cloud/firestore_v1/transaction.py

@@ -16,7 +16,6 @@
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, Callable, Generator, Optional
-import warnings
 
 from google.api_core import exceptions, gapic_v1
 from google.api_core import retry as retries
@@ -197,14 +196,18 @@ def get(
         Yields:
             .DocumentSnapshot: The next document snapshot that fulfills the
             query, or :data:`None` if the document does not exist.
+
+        Raises:
+            ValueError: if `ref_or_query` is not one of the supported types, or
+            explain_options is provided when `ref_or_query` is a document
+            reference.
         """
         kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)
         if isinstance(ref_or_query, DocumentReference):
             if explain_options is not None:
-                warnings.warn(
-                    "explain_options not supported in transanction with "
-                    "document references and will be ignored. To use "
-                    "explain_options, use transaction with query instead."
+                raise ValueError(
+                    "When type of `ref_or_query` is `AsyncDocumentReference`, "
+                    "`explain_options` cannot be provided."
                 )
             return self._client.get_all([ref_or_query], transaction=self, **kwargs)
         elif isinstance(ref_or_query, Query):

tests/system/test_system_async.py

@@ -2946,21 +2946,134 @@ async def _make_transaction_query(client, cleanup):
     FIRESTORE_EMULATOR, reason="Query profile not supported in emulator."
 )
 @pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True)
-async def test_query_in_transaction_with_no_explain_options(client, cleanup, database):
+async def test_transaction_w_query_w_no_explain_options(client, cleanup, database):
     from google.cloud.firestore_v1.query_profile import QueryExplainError
 
+    inner_fn_ran = False
     query = await _make_transaction_query(client, cleanup)
+    transaction = client.transaction()
+
+    # should work when transaction is initiated through transactional decorator
+    @firestore.async_transactional
+    async def in_transaction(transaction):
+        nonlocal inner_fn_ran
+
+        # When no explain_options value is passed, an exception shoud be raised
+        # when accessing explain_metrics.
+        returned_generator = await transaction.get(query)
+
+        with pytest.raises(
+            QueryExplainError, match="explain_options not set on query."
+        ):
+            await returned_generator.get_explain_metrics()
+
+        inner_fn_ran = True
+
+    await in_transaction(transaction)
+
+    # make sure we didn't skip assertions in inner function
+    assert inner_fn_ran is True
+
+
+@pytest.mark.skipif(
+    FIRESTORE_EMULATOR, reason="Query profile not supported in emulator."
+)
+@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True)
+async def test_transaction_w_query_w_explain_options_analyze_true(
+    client, cleanup, database
+):
+    from google.cloud.firestore_v1.query_profile import ExplainOptions
 
+    inner_fn_ran = False
+    query = await _make_transaction_query(client, cleanup)
     transaction = client.transaction()
 
     # should work when transaction is initiated through transactional decorator
     @firestore.async_transactional
     async def in_transaction(transaction):
-        global inner_fn_ran
+        nonlocal inner_fn_ran
+
+        returned_generator = await transaction.get(
+            query,
+            explain_options=ExplainOptions(analyze=True),
+        )
+
+        # explain_metrics should not be available before reading all results.
+        with pytest.raises(
+            QueryExplainError,
+            match="explain_metrics not available until query is complete",
+        ):
+            await returned_generator.get_explain_metrics()
 
-        # When no explain_options value is passed,  an exception shoud be
-        # raised when accessing explain_metrics.
+        result = [x async for x in returned_generator]
+        explain_metrics = await returned_generator.get_explain_metrics()
+        _verify_explain_metrics_analyze_true(explain_metrics, len(result))
+
+        inner_fn_ran = True
+
+    await in_transaction(transaction)
+
+    # make sure we didn't skip assertions in inner function
+    assert inner_fn_ran is True
+
+
+@pytest.mark.skipif(
+    FIRESTORE_EMULATOR, reason="Query profile not supported in emulator."
+)
+@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True)
+async def test_transaction_w_query_w_explain_options_analyze_false(
+    client, cleanup, database
+):
+    from google.cloud.firestore_v1.query_profile import ExplainOptions
+
+    inner_fn_ran = False
+    query = await _make_transaction_query(client, cleanup)
+    transaction = client.transaction()
+
+    # should work when transaction is initiated through transactional decorator
+    @firestore.async_transactional
+    async def in_transaction(transaction):
+        nonlocal inner_fn_ran
+
+        returned_generator = await transaction.get(
+            query,
+            explain_options=ExplainOptions(analyze=False),
+        )
+        explain_metrics = await returned_generator.get_explain_metrics()
+        _verify_explain_metrics_analyze_false(explain_metrics)
+
+        # When analyze == False, result should be empty.
+        result = [x async for x in returned_generator]
+        assert not result
+
+        inner_fn_ran = True
+
+    await in_transaction(transaction)
+
+    # make sure we didn't skip assertions in inner function
+    assert inner_fn_ran is True
+
+
+@pytest.mark.skipif(
+    FIRESTORE_EMULATOR, reason="Query profile not supported in emulator."
+)
+@pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True)
+async def test_query_in_transaction_w_no_explain_options(client, cleanup, database):
+    from google.cloud.firestore_v1.query_profile import QueryExplainError
+
+    inner_fn_ran = False
+    query = await _make_transaction_query(client, cleanup)
+    transaction = client.transaction()
+
+    # should work when transaction is initiated through transactional decorator
+    @firestore.async_transactional
+    async def in_transaction(transaction):
+        nonlocal inner_fn_ran
+
+        # When no explain_options value is passed, an exception shoud be raised
+        # when accessing explain_metrics.
         result = await query.get(transaction=transaction)
+
         with pytest.raises(
             QueryExplainError, match="explain_options not set on query."
         ):
@@ -2978,24 +3091,25 @@ async def in_transaction(transaction):
     FIRESTORE_EMULATOR, reason="Query profile not supported in emulator."
 )
 @pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True)
-async def test_query_in_transaction_with_explain_options_analyze_true(
+async def test_query_in_transaction_w_explain_options_analyze_true(
     client, cleanup, database
 ):
     from google.cloud.firestore_v1.query_profile import ExplainOptions
 
+    inner_fn_ran = False
     query = await _make_transaction_query(client, cleanup)
-
     transaction = client.transaction()
 
     # should work when transaction is initiated through transactional decorator
     @firestore.async_transactional
     async def in_transaction(transaction):
-        global inner_fn_ran
+        nonlocal inner_fn_ran
 
         result = await query.get(
             transaction=transaction,
             explain_options=ExplainOptions(analyze=True),
         )
+
         explain_metrics = result.get_explain_metrics()
         _verify_explain_metrics_analyze_true(explain_metrics, len(result))
 
@@ -3011,19 +3125,19 @@ async def in_transaction(transaction):
     FIRESTORE_EMULATOR, reason="Query profile not supported in emulator."
 )
 @pytest.mark.parametrize("database", [None, FIRESTORE_OTHER_DB], indirect=True)
-async def test_query_in_transaction_with_explain_options_analyze_false(
+async def test_query_in_transaction_w_explain_options_analyze_false(
     client, cleanup, database
 ):
     from google.cloud.firestore_v1.query_profile import ExplainOptions
 
+    inner_fn_ran = False
     query = await _make_transaction_query(client, cleanup)
-
     transaction = client.transaction()
 
     # should work when transaction is initiated through transactional decorator
     @firestore.async_transactional
     async def in_transaction(transaction):
-        global inner_fn_ran
+        nonlocal inner_fn_ran
 
         result = await query.get(
             transaction=transaction,
@@ -3032,6 +3146,9 @@ async def in_transaction(transaction):
         explain_metrics = result.get_explain_metrics()
         _verify_explain_metrics_analyze_false(explain_metrics)
 
+        # When analyze == False, result should be empty.
+        assert not result
+
         inner_fn_ran = True
 
     await in_transaction(transaction)