refactor QueryResultsList

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

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,
     )
 

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

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

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

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

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,

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)
 

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