Skip to content

Commit

Permalink
move QueryResultsList to stream_generator.py
Browse files Browse the repository at this point in the history
  • Loading branch information
Linchin committed Aug 6, 2024
1 parent cb693f8 commit 7439e76
Show file tree
Hide file tree
Showing 7 changed files with 50 additions and 50 deletions.
2 changes: 1 addition & 1 deletion google/cloud/firestore_v1/aggregation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
42 changes: 0 additions & 42 deletions google/cloud/firestore_v1/base_document.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@
Any,
Dict,
Iterable,
List,
NoReturn,
Optional,
Tuple,
Expand All @@ -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


Expand Down Expand Up @@ -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.
Expand Down
2 changes: 1 addition & 1 deletion google/cloud/firestore_v1/collection.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 1 addition & 2 deletions google/cloud/firestore_v1/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down
47 changes: 45 additions & 2 deletions google/cloud/firestore_v1/stream_generator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
2 changes: 1 addition & 1 deletion google/cloud/firestore_v1/vector_query.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion tests/unit/v1/test_base_document.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down

0 comments on commit 7439e76

Please sign in to comment.