1# Copyright 2020 Google LLC All rights reserved. 
    2# 
    3# Licensed under the Apache License, Version 2.0 (the "License"); 
    4# you may not use this file except in compliance with the License. 
    5# You may obtain a copy of the License at 
    6# 
    7#     http://www.apache.org/licenses/LICENSE-2.0 
    8# 
    9# Unless required by applicable law or agreed to in writing, software 
    10# distributed under the License is distributed on an "AS IS" BASIS, 
    11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 
    12# See the License for the specific language governing permissions and 
    13# limitations under the License. 
    14 
    15"""Classes for representing collections for the Google Cloud Firestore API.""" 
    16from __future__ import annotations 
    17 
    18from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional, Tuple 
    19 
    20from google.api_core import gapic_v1 
    21from google.api_core import retry_async as retries 
    22 
    23from google.cloud.firestore_v1 import ( 
    24    async_aggregation, 
    25    async_query, 
    26    async_vector_query, 
    27    transaction, 
    28) 
    29from google.cloud.firestore_v1.base_collection import ( 
    30    BaseCollectionReference, 
    31    _item_to_document_ref, 
    32) 
    33 
    34if TYPE_CHECKING:  # pragma: NO COVER 
    35    import datetime 
    36    from google.cloud.firestore_v1.async_document import AsyncDocumentReference 
    37    from google.cloud.firestore_v1.async_stream_generator import AsyncStreamGenerator 
    38    from google.cloud.firestore_v1.base_document import DocumentSnapshot 
    39    from google.cloud.firestore_v1.query_profile import ExplainOptions 
    40    from google.cloud.firestore_v1.query_results import QueryResultsList 
    41 
    42 
    43class AsyncCollectionReference(BaseCollectionReference[async_query.AsyncQuery]): 
    44    """A reference to a collection in a Firestore database. 
    45 
    46    The collection may already exist or this class can facilitate creation 
    47    of documents within the collection. 
    48 
    49    Args: 
    50        path (Tuple[str, ...]): The components in the collection path. 
    51            This is a series of strings representing each collection and 
    52            sub-collection ID, as well as the document IDs for any documents 
    53            that contain a sub-collection. 
    54        kwargs (dict): The keyword arguments for the constructor. The only 
    55            supported keyword is ``client`` and it must be a 
    56            :class:`~google.cloud.firestore_v1.client.Client` if provided. It 
    57            represents the client that created this collection reference. 
    58 
    59    Raises: 
    60        ValueError: if 
    61 
    62            * the ``path`` is empty 
    63            * there are an even number of elements 
    64            * a collection ID in ``path`` is not a string 
    65            * a document ID in ``path`` is not a string 
    66        TypeError: If a keyword other than ``client`` is used. 
    67    """ 
    68 
    69    def __init__(self, *path, **kwargs) -> None: 
    70        super(AsyncCollectionReference, self).__init__(*path, **kwargs) 
    71 
    72    def _query(self) -> async_query.AsyncQuery: 
    73        """Query factory. 
    74 
    75        Returns: 
    76            :class:`~google.cloud.firestore_v1.query.Query` 
    77        """ 
    78        return async_query.AsyncQuery(self) 
    79 
    80    def _aggregation_query(self) -> async_aggregation.AsyncAggregationQuery: 
    81        """AsyncAggregationQuery factory. 
    82 
    83        Returns: 
    84            :class:`~google.cloud.firestore_v1.async_aggregation.AsyncAggregationQuery 
    85        """ 
    86        return async_aggregation.AsyncAggregationQuery(self._query()) 
    87 
    88    def _vector_query(self) -> async_vector_query.AsyncVectorQuery: 
    89        """AsyncVectorQuery factory. 
    90 
    91        Returns: 
    92            :class:`~google.cloud.firestore_v1.async_vector_query.AsyncVectorQuery` 
    93        """ 
    94        return async_vector_query.AsyncVectorQuery(self._query()) 
    95 
    96    async def _chunkify(self, chunk_size: int): 
    97        async for page in self._query()._chunkify(chunk_size): 
    98            yield page 
    99 
    100    async def add( 
    101        self, 
    102        document_data: dict, 
    103        document_id: str | None = None, 
    104        retry: retries.AsyncRetry | object | None = gapic_v1.method.DEFAULT, 
    105        timeout: float | None = None, 
    106    ) -> Tuple[Any, Any]: 
    107        """Create a document in the Firestore database with the provided data. 
    108 
    109        Args: 
    110            document_data (dict): Property names and values to use for 
    111                creating the document. 
    112            document_id (Optional[str]): The document identifier within the 
    113                current collection. If not provided, an ID will be 
    114                automatically assigned by the server (the assigned ID will be 
    115                a random 20 character string composed of digits, 
    116                uppercase and lowercase letters). 
    117            retry (google.api_core.retry.Retry): Designation of what errors, if any, 
    118                should be retried.  Defaults to a system-specified policy. 
    119            timeout (float): The timeout for this request.  Defaults to a 
    120                system-specified value. 
    121 
    122        Returns: 
    123            Tuple[:class:`google.protobuf.timestamp_pb2.Timestamp`, \ 
    124                :class:`~google.cloud.firestore_v1.async_document.AsyncDocumentReference`]: 
    125                Pair of 
    126 
    127                * The ``update_time`` when the document was created/overwritten. 
    128                * A document reference for the created document. 
    129 
    130        Raises: 
    131            :class:`google.cloud.exceptions.Conflict`: 
    132                If ``document_id`` is provided and the document already exists. 
    133        """ 
    134        document_ref, kwargs = self._prep_add( 
    135            document_data, 
    136            document_id, 
    137            retry, 
    138            timeout, 
    139        ) 
    140        write_result = await document_ref.create(document_data, **kwargs) 
    141        return write_result.update_time, document_ref 
    142 
    143    def document(self, document_id: str | None = None) -> AsyncDocumentReference: 
    144        """Create a sub-document underneath the current collection. 
    145 
    146        Args: 
    147            document_id (Optional[str]): The document identifier 
    148                within the current collection. If not provided, will default 
    149                to a random 20 character string composed of digits, 
    150                uppercase and lowercase and letters. 
    151 
    152        Returns: 
    153            :class:`~google.cloud.firestore_v1.document.async_document.AsyncDocumentReference`: 
    154            The child document. 
    155        """ 
    156        return super(AsyncCollectionReference, self).document(document_id) 
    157 
    158    async def list_documents( 
    159        self, 
    160        page_size: int | None = None, 
    161        retry: retries.AsyncRetry | object | None = gapic_v1.method.DEFAULT, 
    162        timeout: float | None = None, 
    163        *, 
    164        read_time: datetime.datetime | None = None, 
    165    ) -> AsyncGenerator[AsyncDocumentReference, None]: 
    166        """List all subdocuments of the current collection. 
    167 
    168        Args: 
    169            page_size (Optional[int]]): The maximum number of documents 
    170                in each page of results from this request. Non-positive values 
    171                are ignored. Defaults to a sensible value set by the API. 
    172            retry (google.api_core.retry.Retry): Designation of what errors, if any, 
    173                should be retried.  Defaults to a system-specified policy. 
    174            timeout (float): The timeout for this request.  Defaults to a 
    175                system-specified value. 
    176            read_time (Optional[datetime.datetime]): If set, reads documents as they were at the given 
    177                time. This must be a timestamp within the past one hour, or if Point-in-Time Recovery 
    178                is enabled, can additionally be a whole minute timestamp within the past 7 days. If no 
    179                timezone is specified in the :class:`datetime.datetime` object, it is assumed to be UTC. 
    180 
    181        Returns: 
    182            Sequence[:class:`~google.cloud.firestore_v1.collection.DocumentReference`]: 
    183                iterator of subdocuments of the current collection. If the 
    184                collection does not exist at the time of `snapshot`, the 
    185                iterator will be empty 
    186        """ 
    187        request, kwargs = self._prep_list_documents( 
    188            page_size, retry, timeout, read_time 
    189        ) 
    190 
    191        iterator = await self._client._firestore_api.list_documents( 
    192            request=request, 
    193            metadata=self._client._rpc_metadata, 
    194            **kwargs, 
    195        ) 
    196        async for i in iterator: 
    197            yield _item_to_document_ref(self, i) 
    198 
    199    async def get( 
    200        self, 
    201        transaction: Optional[transaction.Transaction] = None, 
    202        retry: retries.AsyncRetry | object | None = gapic_v1.method.DEFAULT, 
    203        timeout: Optional[float] = None, 
    204        *, 
    205        explain_options: Optional[ExplainOptions] = None, 
    206        read_time: Optional[datetime.datetime] = None, 
    207    ) -> QueryResultsList[DocumentSnapshot]: 
    208        """Read the documents in this collection. 
    209 
    210        This sends a ``RunQuery`` RPC and returns a list of documents 
    211        returned in the stream of ``RunQueryResponse`` messages. 
    212 
    213        Args: 
    214            transaction 
    215                (Optional[:class:`~google.cloud.firestore_v1.transaction.Transaction`]): 
    216                An existing transaction that this query will run in. 
    217            retry (Optional[google.api_core.retry.Retry]): Designation of what 
    218                errors, if any, should be retried.  Defaults to a 
    219                system-specified policy. 
    220            timeout (Otional[float]): The timeout for this request.  Defaults 
    221                to a system-specified value. 
    222            explain_options 
    223                (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): 
    224                Options to enable query profiling for this query. When set, 
    225                explain_metrics will be available on the returned generator. 
    226            read_time (Optional[datetime.datetime]): If set, reads documents as they were at the given 
    227                time. This must be a timestamp within the past one hour, or if Point-in-Time Recovery 
    228                is enabled, can additionally be a whole minute timestamp within the past 7 days. If no 
    229                timezone is specified in the :class:`datetime.datetime` object, it is assumed to be UTC. 
    230 
    231        If a ``transaction`` is used and it already has write operations added, 
    232        this method cannot be used (i.e. read-after-write is not allowed). 
    233 
    234        Returns: 
    235            QueryResultsList[DocumentSnapshot]: 
    236            The documents in this collection that match the query. 
    237        """ 
    238        query, kwargs = self._prep_get_or_stream(retry, timeout) 
    239        if explain_options is not None: 
    240            kwargs["explain_options"] = explain_options 
    241        if read_time is not None: 
    242            kwargs["read_time"] = read_time 
    243 
    244        return await query.get(transaction=transaction, **kwargs) 
    245 
    246    def stream( 
    247        self, 
    248        transaction: Optional[transaction.Transaction] = None, 
    249        retry: retries.AsyncRetry | object | None = gapic_v1.method.DEFAULT, 
    250        timeout: Optional[float] = None, 
    251        *, 
    252        explain_options: Optional[ExplainOptions] = None, 
    253        read_time: Optional[datetime.datetime] = None, 
    254    ) -> AsyncStreamGenerator[DocumentSnapshot]: 
    255        """Read the documents in this collection. 
    256 
    257        This sends a ``RunQuery`` RPC and then returns a generator which 
    258        consumes each document returned in the stream of ``RunQueryResponse`` 
    259        messages. 
    260 
    261        .. note:: 
    262 
    263           The underlying stream of responses will time out after 
    264           the ``max_rpc_timeout_millis`` value set in the GAPIC 
    265           client configuration for the ``RunQuery`` API.  Snapshots 
    266           not consumed from the iterator before that point will be lost. 
    267 
    268        If a ``transaction`` is used and it already has write operations 
    269        added, this method cannot be used (i.e. read-after-write is not 
    270        allowed). 
    271 
    272        Args: 
    273            transaction (Optional[:class:`~google.cloud.firestore_v1.transaction.\ 
    274                Transaction`]): 
    275                An existing transaction that the query will run in. 
    276            retry (Optional[google.api_core.retry.Retry]): Designation of what 
    277                errors, if any, should be retried.  Defaults to a 
    278                system-specified policy. 
    279            timeout (Optional[float]): The timeout for this request. Defaults 
    280                to a system-specified value. 
    281            explain_options 
    282                (Optional[:class:`~google.cloud.firestore_v1.query_profile.ExplainOptions`]): 
    283                Options to enable query profiling for this query. When set, 
    284                explain_metrics will be available on the returned generator. 
    285            read_time (Optional[datetime.datetime]): If set, reads documents as they were at the given 
    286                time. This must be a timestamp within the past one hour, or if Point-in-Time Recovery 
    287                is enabled, can additionally be a whole minute timestamp within the past 7 days. If no 
    288                timezone is specified in the :class:`datetime.datetime` object, it is assumed to be UTC. 
    289 
    290        Returns: 
    291            `AsyncStreamGenerator[DocumentSnapshot]`: A generator of the query 
    292            results. 
    293        """ 
    294        query, kwargs = self._prep_get_or_stream(retry, timeout) 
    295        if explain_options: 
    296            kwargs["explain_options"] = explain_options 
    297        if read_time is not None: 
    298            kwargs["read_time"] = read_time 
    299 
    300        return query.stream(transaction=transaction, **kwargs)