base_aggregation.py

# Copyright 2023 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.

"""Classes for representing aggregation queries for the Google Cloud Firestore API.

A :class:`~google.cloud.firestore_v1.aggregation.AggregationQuery` can be created directly from
a :class:`~google.cloud.firestore_v1.collection.Collection` and that can be
a more common way to create an aggregation query than direct usage of the constructor.
"""
from __future__ import annotations

import abc
import itertools

from abc import ABC
from typing import TYPE_CHECKING, Any, Coroutine, List, Optional, Tuple, Union, Iterable

from google.api_core import gapic_v1
from google.api_core import retry as retries

from google.cloud.firestore_v1 import _helpers
from google.cloud.firestore_v1.field_path import FieldPath
from google.cloud.firestore_v1.types import (
    StructuredAggregationQuery,
)
from google.cloud.firestore_v1.pipeline_expressions import AggregateFunction
from google.cloud.firestore_v1.pipeline_expressions import Count
from google.cloud.firestore_v1.pipeline_expressions import AliasedExpression
from google.cloud.firestore_v1.pipeline_expressions import Field

# Types needed only for Type Hints
if TYPE_CHECKING:  # pragma: NO COVER
    from google.cloud.firestore_v1 import transaction
    from google.cloud.firestore_v1.async_stream_generator import AsyncStreamGenerator
    from google.cloud.firestore_v1.query_profile import ExplainOptions
    from google.cloud.firestore_v1.query_results import QueryResultsList
    from google.cloud.firestore_v1.stream_generator import (
        StreamGenerator,
    )
    from google.cloud.firestore_v1.pipeline_source import PipelineSource

    import datetime


class AggregationResult(object):
    """
    A class representing result from Aggregation Query
    :type alias: str
    :param alias: The alias for the aggregation.
    :type value: int
    :param value: The resulting value from the aggregation.
    :type read_time:
    :param value: The resulting read_time
    """

    def __init__(self, alias: str, value: float, read_time=None):
        self.alias = alias
        self.value = value
        self.read_time = read_time

    def __repr__(self):
        return f"<Aggregation alias={self.alias}, value={self.value}, readtime={self.read_time}>"

    def _to_dict(self):
        return {self.alias: self.value}


class BaseAggregation(ABC):
    def __init__(self, alias: str | None = None):
        self.alias = alias

    @abc.abstractmethod
    def _to_protobuf(self):
        """Convert this instance to the protobuf representation"""

    @abc.abstractmethod
    def _to_pipeline_expr(
        self, autoindexer: Iterable[int]
    ) -> AliasedExpression[AggregateFunction]:
        """
        Convert this instance to a pipeline expression for use with pipeline.aggregate()

        Args:
          autoindexer: If an alias isn't supplied, one should be created with the format "field_n"
            The autoindexer is an iterable that provides the `n` value to use for each expression
        """

    def _pipeline_alias(self, autoindexer):
        """
        Helper to build the alias for the pipeline expression
        """
        if self.alias is not None:
            return self.alias
        else:
            return f"field_{next(autoindexer)}"


class CountAggregation(BaseAggregation):
    def __init__(self, alias: str | None = None):
        super(CountAggregation, self).__init__(alias=alias)

    def _to_protobuf(self):
        """Convert this instance to the protobuf representation"""
        aggregation_pb = StructuredAggregationQuery.Aggregation()
        if self.alias:
            aggregation_pb.alias = self.alias
        aggregation_pb.count = StructuredAggregationQuery.Aggregation.Count()
        return aggregation_pb

    def _to_pipeline_expr(self, autoindexer: Iterable[int]):
        return Count().as_(self._pipeline_alias(autoindexer))


class SumAggregation(BaseAggregation):
    def __init__(self, field_ref: str | FieldPath, alias: str | None = None):
        # convert field path to string if needed
        field_str = (
            field_ref.to_api_repr() if isinstance(field_ref, FieldPath) else field_ref
        )
        self.field_ref: str = field_str
        super(SumAggregation, self).__init__(alias=alias)

    def _to_protobuf(self):
        """Convert this instance to the protobuf representation"""
        aggregation_pb = StructuredAggregationQuery.Aggregation()
        if self.alias:
            aggregation_pb.alias = self.alias
        aggregation_pb.sum = StructuredAggregationQuery.Aggregation.Sum()
        aggregation_pb.sum.field.field_path = self.field_ref
        return aggregation_pb

    def _to_pipeline_expr(self, autoindexer: Iterable[int]):
        return Field.of(self.field_ref).sum().as_(self._pipeline_alias(autoindexer))


class AvgAggregation(BaseAggregation):
    def __init__(self, field_ref: str | FieldPath, alias: str | None = None):
        # convert field path to string if needed
        field_str = (
            field_ref.to_api_repr() if isinstance(field_ref, FieldPath) else field_ref
        )
        self.field_ref: str = field_str
        super(AvgAggregation, self).__init__(alias=alias)

    def _to_protobuf(self):
        """Convert this instance to the protobuf representation"""
        aggregation_pb = StructuredAggregationQuery.Aggregation()
        if self.alias:
            aggregation_pb.alias = self.alias
        aggregation_pb.avg = StructuredAggregationQuery.Aggregation.Avg()
        aggregation_pb.avg.field.field_path = self.field_ref
        return aggregation_pb

    def _to_pipeline_expr(self, autoindexer: Iterable[int]):
        return Field.of(self.field_ref).average().as_(self._pipeline_alias(autoindexer))


def _query_response_to_result(
    response_pb,
) -> List[AggregationResult]:
    results = [
        AggregationResult(
            alias=key,
            value=response_pb.result.aggregate_fields[key].integer_value
            or response_pb.result.aggregate_fields[key].double_value,
            read_time=response_pb.read_time,
        )
        for key in response_pb.result.aggregate_fields.pb.keys()
    ]

    return results


class BaseAggregationQuery(ABC):
    """Represents an aggregation query to the Firestore API."""

    def __init__(self, nested_query, alias: str | None = None) -> None:
        self._nested_query = nested_query
        self._alias = alias
        self._collection_ref = nested_query._parent
        self._aggregations: List[BaseAggregation] = []

    @property
    def _client(self):
        return self._collection_ref._client

    def count(self, alias: str | None = None):
        """
        Adds a count over the nested query
        """
        count_aggregation = CountAggregation(alias=alias)
        self._aggregations.append(count_aggregation)
        return self

    def sum(self, field_ref: str | FieldPath, alias: str | None = None):
        """
        Adds a sum over the nested query
        """
        sum_aggregation = SumAggregation(field_ref, alias=alias)
        self._aggregations.append(sum_aggregation)
        return self

    def avg(self, field_ref: str | FieldPath, alias: str | None = None):
        """
        Adds an avg over the nested query
        """
        avg_aggregation = AvgAggregation(field_ref, alias=alias)
        self._aggregations.append(avg_aggregation)
        return self

    def add_aggregation(self, aggregation: BaseAggregation) -> None:
        """
        Adds an aggregation operation to the nested query

        :type aggregation: :class:`google.cloud.firestore_v1.aggregation.BaseAggregation`
        :param aggregation: An aggregation operation, e.g. a CountAggregation
        """
        self._aggregations.append(aggregation)

    def add_aggregations(self, aggregations: List[BaseAggregation]) -> None:
        """
        Adds a list of aggregations to the nested query

        :type aggregations: list
        :param aggregations: a list of aggregation operations
        """
        self._aggregations.extend(aggregations)

    def _to_protobuf(self) -> StructuredAggregationQuery:
        pb = StructuredAggregationQuery()
        pb.structured_query = self._nested_query._to_protobuf()

        for aggregation in self._aggregations:
            aggregation_pb = aggregation._to_protobuf()
            pb.aggregations.append(aggregation_pb)
        return pb

    def _prep_stream(
        self,
        transaction=None,
        retry: Union[retries.Retry, retries.AsyncRetry, None, object] = None,
        timeout: float | None = None,
        explain_options: Optional[ExplainOptions] = None,
        read_time: Optional[datetime.datetime] = None,
    ) -> Tuple[dict, dict]:
        parent_path, expected_prefix = self._collection_ref._parent_info()
        request = {
            "parent": parent_path,
            "structured_aggregation_query": self._to_protobuf(),
            "transaction": _helpers.get_transaction_id(transaction),
        }
        if explain_options:
            request["explain_options"] = explain_options._to_dict()
        if read_time is not None:
            request["read_time"] = read_time
        kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout)

        return request, kwargs

    @abc.abstractmethod
    def get(
        self,
        transaction=None,
        retry: Union[
            retries.Retry, retries.AsyncRetry, None, object
        ] = gapic_v1.method.DEFAULT,
        timeout: float | None = None,
        *,
        explain_options: Optional[ExplainOptions] = None,
        read_time: Optional[datetime.datetime] = None,
    ) -> (
        QueryResultsList[AggregationResult]
        | Coroutine[Any, Any, List[List[AggregationResult]]]
    ):
        """Runs the aggregation query.

        This sends a ``RunAggregationQuery`` RPC and returns a list of
        aggregation results in the stream of ``RunAggregationQueryResponse``
        messages.

        Args:
            transaction
                (Optional[:class:`~google.cloud.firestore_v1.transaction.Transaction`]):
                An existing transaction that this query will run in.
                If a ``transaction`` is used and it already has write operations
                added, this method cannot be used (i.e. read-after-write is not
                allowed).
            retry (google.api_core.retry.Retry): Designation of what errors, if any,
                should be retried.  Defaults to a system-specified policy.
            timeout (float): The timeout for this request.  Defaults to a
                system-specified value.
            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.
            read_time (Optional[datetime.datetime]): If set, reads documents as they were at the given
                time. This must be a timestamp within the past one hour, or if Point-in-Time Recovery
                is enabled, can additionally be a whole minute timestamp within the past 7 days. If no
                timezone is specified in the :class:`datetime.datetime` object, it is assumed to be UTC.

        Returns:
            (QueryResultsList[List[AggregationResult]] | Coroutine[Any, Any, List[List[AggregationResult]]]):
            The aggregation query results.
        """

    @abc.abstractmethod
    def stream(
        self,
        transaction: Optional[transaction.Transaction] = None,
        retry: retries.Retry
        | retries.AsyncRetry
        | object
        | None = gapic_v1.method.DEFAULT,
        timeout: Optional[float] = None,
        *,
        explain_options: Optional[ExplainOptions] = None,
        read_time: Optional[datetime.datetime] = None,
    ) -> (
        StreamGenerator[List[AggregationResult]]
        | AsyncStreamGenerator[List[AggregationResult]]
    ):
        """Runs the aggregation query.

        This sends a``RunAggregationQuery`` RPC and returns a generator in the stream of ``RunAggregationQueryResponse`` messages.

        Args:
            transaction
                (Optional[:class:`~google.cloud.firestore_v1.transaction.Transaction`]):
                An existing transaction that this query will run in.
            retry (Optional[google.api_core.retry.Retry]): Designation of what
                errors, if any, should be retried.  Defaults to a
                system-specified policy.
            timeout (Optinal[float]): The timeout for this request.  Defaults
                to a system-specified value.
            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.
            read_time (Optional[datetime.datetime]): If set, reads documents as they were at the given
                time. This must be a timestamp within the past one hour, or if Point-in-Time Recovery
                is enabled, can additionally be a whole minute timestamp within the past 7 days. If no
                timezone is specified in the :class:`datetime.datetime` object, it is assumed to be UTC.

        Returns:
            StreamGenerator[List[AggregationResult]] | AsyncStreamGenerator[List[AggregationResult]]:
            A generator of the query results.
        """

    def _build_pipeline(self, source: "PipelineSource"):
        """
        Convert this query into a Pipeline

        Queries containing a `cursor` or `limit_to_last` are not currently supported

        Args:
            source: the PipelineSource to build the pipeline off of
        Raises:
            - NotImplementedError: raised if the query contains a `cursor` or `limit_to_last`
        Returns:
            a Pipeline representing the query
        """
        # use autoindexer to keep track of which field number to use for un-aliased fields
        autoindexer = itertools.count(start=1)
        exprs = [a._to_pipeline_expr(autoindexer) for a in self._aggregations]
        return self._nested_query._build_pipeline(source).aggregate(*exprs)