diff --git a/google/cloud/firestore_v1/aggregation.py b/google/cloud/firestore_v1/aggregation.py index 763202ff8..183c869a7 100644 --- a/google/cloud/firestore_v1/aggregation.py +++ b/google/cloud/firestore_v1/aggregation.py @@ -30,7 +30,7 @@ BaseAggregationQuery, _query_response_to_result, ) -from google.cloud.firestore_v1.base_document import QueryResultsList +from google.cloud.firestore_v1.stream_generator import QueryResultsList from google.cloud.firestore_v1.stream_generator import StreamGenerator # Types needed only for Type Hints diff --git a/google/cloud/firestore_v1/base_document.py b/google/cloud/firestore_v1/base_document.py index f82b760ec..e2f8238ea 100644 --- a/google/cloud/firestore_v1/base_document.py +++ b/google/cloud/firestore_v1/base_document.py @@ -21,7 +21,6 @@ Any, Dict, Iterable, - List, NoReturn, Optional, Tuple, @@ -33,12 +32,10 @@ from google.cloud.firestore_v1 import _helpers from google.cloud.firestore_v1 import field_path as field_path_module -from google.cloud.firestore_v1.query_profile import QueryExplainError from google.cloud.firestore_v1.types import common # Types needed only for Type Hints if TYPE_CHECKING: # pragma: NO COVER - from google.cloud.firestore_v1.query_profile import ExplainMetrics, ExplainOptions from google.cloud.firestore_v1.types import Document, firestore, write @@ -522,45 +519,6 @@ def _to_protobuf(self) -> Optional[Document]: return _helpers.document_snapshot_to_protobuf(self) -class QueryResultsList(list): - """A list of received query results from the query call. - - This is a subclass of the built-in list. A new property `explain_metrics` - is added to return the query profile results. - - Args: - docs (list[T]): - The list of query results. - explain_options - (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. - explain_metrics (Optional[ExplainMetrics]): - Query profile results. - """ - - def __init__( - self, - docs: List[T], - explain_options: Optional[ExplainOptions] = None, - explain_metrics: Optional[ExplainMetrics] = None, - ): - super().__init__(docs) - self._explain_options = explain_options - self._explain_metrics = explain_metrics - - @property - def explain_options(self): - return self._explain_options - - @property - def explain_metrics(self): - if self._explain_options is None: - raise QueryExplainError("explain_options not set on query.") - else: - return self._explain_metrics - - def _get_document_path(client, path: Tuple[str]) -> str: """Convert a path tuple into a full path string. diff --git a/google/cloud/firestore_v1/collection.py b/google/cloud/firestore_v1/collection.py index 66b3a6b47..4eeb231e4 100644 --- a/google/cloud/firestore_v1/collection.py +++ b/google/cloud/firestore_v1/collection.py @@ -27,7 +27,7 @@ BaseCollectionReference, _item_to_document_ref, ) -from google.cloud.firestore_v1.base_document import QueryResultsList +from google.cloud.firestore_v1.stream_generator import QueryResultsList from google.cloud.firestore_v1.watch import Watch if TYPE_CHECKING: # pragma: NO COVER diff --git a/google/cloud/firestore_v1/query.py b/google/cloud/firestore_v1/query.py index 51bbbb922..a7459423d 100644 --- a/google/cloud/firestore_v1/query.py +++ b/google/cloud/firestore_v1/query.py @@ -29,7 +29,6 @@ from google.cloud.firestore_v1 import aggregation, transaction from google.cloud.firestore_v1.base_document import ( DocumentSnapshot, - QueryResultsList, ) from google.cloud.firestore_v1.base_query import ( BaseCollectionGroup, @@ -39,7 +38,7 @@ _enum_from_direction, _query_response_to_snapshot, ) -from google.cloud.firestore_v1.stream_generator import StreamGenerator +from google.cloud.firestore_v1.stream_generator import QueryResultsList, StreamGenerator from google.cloud.firestore_v1.vector import Vector from google.cloud.firestore_v1.vector_query import VectorQuery from google.cloud.firestore_v1.watch import Watch diff --git a/google/cloud/firestore_v1/stream_generator.py b/google/cloud/firestore_v1/stream_generator.py index f6b07a7f9..7280f0823 100644 --- a/google/cloud/firestore_v1/stream_generator.py +++ b/google/cloud/firestore_v1/stream_generator.py @@ -17,9 +17,13 @@ from __future__ import annotations from collections import abc -from typing import TYPE_CHECKING, Generator, Optional, TypeVar +from typing import TYPE_CHECKING, Generator, List, Optional, TypeVar -from google.cloud.firestore_v1.query_profile import ExplainMetrics, QueryExplainError +from google.cloud.firestore_v1.query_profile import ( + ExplainMetrics, + ExplainOptions, + QueryExplainError, +) if TYPE_CHECKING: # pragma: NO COVER from google.cloud.firestore_v1.query_profile import ExplainOptions @@ -111,3 +115,42 @@ def explain_metrics(self) -> ExplainMetrics: raise QueryExplainError( "explain_metrics not available until query is complete." ) + + +class QueryResultsList(list): + """A list of received query results from the query call. + + This is a subclass of the built-in list. A new property `explain_metrics` + is added to return the query profile results. + + Args: + docs (list[T]): + The list of query results. + explain_options + (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. + explain_metrics (Optional[ExplainMetrics]): + Query profile results. + """ + + def __init__( + self, + docs: List[T], + explain_options: Optional[ExplainOptions] = None, + explain_metrics: Optional[ExplainMetrics] = None, + ): + super().__init__(docs) + self._explain_options = explain_options + self._explain_metrics = explain_metrics + + @property + def explain_options(self): + return self._explain_options + + @property + def explain_metrics(self): + if self._explain_options is None: + raise QueryExplainError("explain_options not set on query.") + else: + return self._explain_metrics diff --git a/google/cloud/firestore_v1/vector_query.py b/google/cloud/firestore_v1/vector_query.py index ee289f62c..cd809d56c 100644 --- a/google/cloud/firestore_v1/vector_query.py +++ b/google/cloud/firestore_v1/vector_query.py @@ -21,7 +21,7 @@ from google.api_core import gapic_v1 from google.api_core import retry as retries -from google.cloud.firestore_v1.base_document import QueryResultsList +from google.cloud.firestore_v1.stream_generator import QueryResultsList from google.cloud.firestore_v1.base_query import ( BaseQuery, _collection_group_query_response_to_snapshot, diff --git a/tests/unit/v1/test_base_document.py b/tests/unit/v1/test_base_document.py index 18c714d05..0e84188ee 100644 --- a/tests/unit/v1/test_base_document.py +++ b/tests/unit/v1/test_base_document.py @@ -363,7 +363,7 @@ def test_documentsnapshot_non_existent(): def _make_query_results_list(*args, **kwargs): - from google.cloud.firestore_v1.base_document import QueryResultsList + from google.cloud.firestore_v1.stream_generator import QueryResultsList return QueryResultsList(*args, **kwargs)