From 05424bf2efd65d8ba9a84e402a8c0f77d018c11d Mon Sep 17 00:00:00 2001 From: Linchin Date: Tue, 6 Aug 2024 17:15:17 -0700 Subject: [PATCH] refactor QueryResultsList --- google/cloud/firestore_v1/aggregation.py | 2 +- google/cloud/firestore_v1/base_aggregation.py | 2 +- google/cloud/firestore_v1/collection.py | 2 +- google/cloud/firestore_v1/query.py | 3 +- google/cloud/firestore_v1/query_results.py | 50 +++++++++ google/cloud/firestore_v1/stream_generator.py | 42 +------- google/cloud/firestore_v1/vector_query.py | 2 +- tests/unit/v1/test_base_document.py | 2 +- tests/unit/v1/test_query_results.py | 102 ++++++++++++++++++ 9 files changed, 160 insertions(+), 47 deletions(-) create mode 100644 google/cloud/firestore_v1/query_results.py create mode 100644 tests/unit/v1/test_query_results.py diff --git a/google/cloud/firestore_v1/aggregation.py b/google/cloud/firestore_v1/aggregation.py index d3182d21a0..8a0d439f94 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.stream_generator import QueryResultsList +from google.cloud.firestore_v1.query_results 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_aggregation.py b/google/cloud/firestore_v1/base_aggregation.py index e414a153e0..4e1b6cf88a 100644 --- a/google/cloud/firestore_v1/base_aggregation.py +++ b/google/cloud/firestore_v1/base_aggregation.py @@ -38,6 +38,7 @@ from google.api_core import retry as retries from google.cloud.firestore_v1 import _helpers +from google.cloud.firestore_v1.query_results import QueryResultsList from google.cloud.firestore_v1.field_path import FieldPath from google.cloud.firestore_v1.types import ( RunAggregationQueryResponse, @@ -50,7 +51,6 @@ from google.cloud.firestore_v1.async_stream_generator import AsyncStreamGenerator from google.cloud.firestore_v1.query_profile import ExplainOptions from google.cloud.firestore_v1.stream_generator import ( - QueryResultsList, StreamGenerator, ) diff --git a/google/cloud/firestore_v1/collection.py b/google/cloud/firestore_v1/collection.py index 4eeb231e4c..5e2f23811e 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.stream_generator import QueryResultsList +from google.cloud.firestore_v1.query_results 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 a7459423db..129238e3c2 100644 --- a/google/cloud/firestore_v1/query.py +++ b/google/cloud/firestore_v1/query.py @@ -27,6 +27,7 @@ from google.cloud import firestore_v1 from google.cloud.firestore_v1 import aggregation, transaction +from google.cloud.firestore_v1.query_results import QueryResultsList from google.cloud.firestore_v1.base_document import ( DocumentSnapshot, ) @@ -38,7 +39,7 @@ _enum_from_direction, _query_response_to_snapshot, ) -from google.cloud.firestore_v1.stream_generator import QueryResultsList, StreamGenerator +from google.cloud.firestore_v1.stream_generator import 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/query_results.py b/google/cloud/firestore_v1/query_results.py new file mode 100644 index 0000000000..3a5c902158 --- /dev/null +++ b/google/cloud/firestore_v1/query_results.py @@ -0,0 +1,50 @@ +from google.cloud.firestore_v1.query_profile import ( + ExplainMetrics, + ExplainOptions, + QueryExplainError, +) + + +from typing import List, Optional, TypeVar + + +T = TypeVar("T") + + +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/stream_generator.py b/google/cloud/firestore_v1/stream_generator.py index 7280f08230..f1899c6771 100644 --- a/google/cloud/firestore_v1/stream_generator.py +++ b/google/cloud/firestore_v1/stream_generator.py @@ -17,11 +17,10 @@ from __future__ import annotations from collections import abc -from typing import TYPE_CHECKING, Generator, List, Optional, TypeVar +from typing import TYPE_CHECKING, Generator, Optional, TypeVar from google.cloud.firestore_v1.query_profile import ( ExplainMetrics, - ExplainOptions, QueryExplainError, ) @@ -115,42 +114,3 @@ 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 cd809d56cd..d57009fd04 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.stream_generator import QueryResultsList +from google.cloud.firestore_v1.query_results 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 0e84188ee3..a1f2a85b4e 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.stream_generator import QueryResultsList + from google.cloud.firestore_v1.query_results import QueryResultsList return QueryResultsList(*args, **kwargs) diff --git a/tests/unit/v1/test_query_results.py b/tests/unit/v1/test_query_results.py new file mode 100644 index 0000000000..b01a3c0297 --- /dev/null +++ b/tests/unit/v1/test_query_results.py @@ -0,0 +1,102 @@ +# Copyright 2020 Google LLC All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import mock + + +def _make_base_document_reference(*args, **kwargs): + from google.cloud.firestore_v1.base_document import BaseDocumentReference + + return BaseDocumentReference(*args, **kwargs) + + +def _make_document_snapshot(*args, **kwargs): + from google.cloud.firestore_v1.document import DocumentSnapshot + + return DocumentSnapshot(*args, **kwargs) + + +def _make_query_results_list(*args, **kwargs): + from google.cloud.firestore_v1.query_results import QueryResultsList + + return QueryResultsList(*args, **kwargs) + + +def _make_explain_metrics(): + from google.cloud.firestore_v1.query_profile import ExplainMetrics, PlanSummary + + plan_summary = PlanSummary( + indexes_used=[{"properties": "(__name__ ASC)", "query_scope": "Collection"}], + ) + return ExplainMetrics(plan_summary=plan_summary) + + +def test_query_results_list_constructor(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + client = mock.sentinel.client + reference = _make_base_document_reference("hi", "bye", client=client) + data_1 = {"zoop": 83} + data_2 = {"zoop": 30} + snapshot_1 = _make_document_snapshot( + reference, + data_1, + True, + mock.sentinel.read_time, + mock.sentinel.create_time, + mock.sentinel.update_time, + ) + snapshot_2 = _make_document_snapshot( + reference, + data_2, + True, + mock.sentinel.read_time, + mock.sentinel.create_time, + mock.sentinel.update_time, + ) + explain_metrics = _make_explain_metrics() + explain_options = ExplainOptions(analyze=True) + snapshot_list = _make_query_results_list( + [snapshot_1, snapshot_2], + explain_options=explain_options, + explain_metrics=explain_metrics, + ) + assert len(snapshot_list) == 2 + assert snapshot_list[0] == snapshot_1 + assert snapshot_list[1] == snapshot_2 + assert snapshot_list._explain_options == explain_options + assert snapshot_list._explain_metrics == explain_metrics + + +def test_query_results_list_explain_options(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + explain_options = ExplainOptions(analyze=True) + snapshot_list = _make_query_results_list([], explain_options=explain_options) + + assert snapshot_list.explain_options == explain_options + + +def test_query_results_list_explain_metrics_w_explain_options(): + from google.cloud.firestore_v1.query_profile import ExplainOptions + + explain_metrics = _make_explain_metrics() + snapshot_list = _make_query_results_list( + [], + explain_options=ExplainOptions(analyze=True), + explain_metrics=explain_metrics, + ) + + assert snapshot_list.explain_metrics == explain_metrics