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"""Helpers for batch requests to the Google Cloud Firestore API.""" 
    16from __future__ import annotations 
    17 
    18from google.api_core import gapic_v1 
    19from google.api_core import retry_async as retries 
    20 
    21from google.cloud.firestore_v1.base_batch import BaseWriteBatch 
    22 
    23 
    24class AsyncWriteBatch(BaseWriteBatch): 
    25    """Accumulate write operations to be sent in a batch. 
    26 
    27    This has the same set of methods for write operations that 
    28    :class:`~google.cloud.firestore_v1.async_document.AsyncDocumentReference` does, 
    29    e.g. :meth:`~google.cloud.firestore_v1.async_document.AsyncDocumentReference.create`. 
    30 
    31    Args: 
    32        client (:class:`~google.cloud.firestore_v1.async_client.AsyncClient`): 
    33            The client that created this batch. 
    34    """ 
    35 
    36    def __init__(self, client) -> None: 
    37        super(AsyncWriteBatch, self).__init__(client=client) 
    38 
    39    async def commit( 
    40        self, 
    41        retry: retries.AsyncRetry | object | None = gapic_v1.method.DEFAULT, 
    42        timeout: float | None = None, 
    43    ) -> list: 
    44        """Commit the changes accumulated in this batch. 
    45 
    46        Args: 
    47            retry (google.api_core.retry.Retry): Designation of what errors, if any, 
    48                should be retried.  Defaults to a system-specified policy. 
    49            timeout (float): The timeout for this request.  Defaults to a 
    50                system-specified value. 
    51 
    52        Returns: 
    53            List[:class:`google.cloud.firestore_v1.write.WriteResult`, ...]: 
    54            The write results corresponding to the changes committed, returned 
    55            in the same order as the changes were applied to this batch. A 
    56            write result contains an ``update_time`` field. 
    57        """ 
    58        request, kwargs = self._prep_commit(retry, timeout) 
    59 
    60        commit_response = await self._client._firestore_api.commit( 
    61            request=request, 
    62            metadata=self._client._rpc_metadata, 
    63            **kwargs, 
    64        ) 
    65 
    66        self._write_pbs = [] 
    67        self.write_results = results = list(commit_response.write_results) 
    68        self.commit_time = commit_response.commit_time 
    69 
    70        return results 
    71 
    72    async def __aenter__(self): 
    73        return self 
    74 
    75    async def __aexit__(self, exc_type, exc_value, traceback): 
    76        if exc_type is None: 
    77            await self.commit()