1# Copyright 2021 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 
    17from google.api_core import gapic_v1 
    18from google.api_core import retry as retries 
    19 
    20from google.cloud.firestore_v1 import _helpers 
    21from google.cloud.firestore_v1.base_batch import BaseBatch 
    22from google.cloud.firestore_v1.types.firestore import BatchWriteResponse 
    23 
    24 
    25class BulkWriteBatch(BaseBatch): 
    26    """Accumulate write operations to be sent in a batch. Use this over 
    27    `WriteBatch` for higher volumes (e.g., via `BulkWriter`) and when the order 
    28    of operations within a given batch is unimportant. 
    29 
    30    Because the order in which individual write operations are applied to the database 
    31    is not guaranteed, `batch_write` RPCs can never contain multiple operations 
    32    to the same document. If calling code detects a second write operation to a 
    33    known document reference, it should first cut off the previous batch and 
    34    send it, then create a new batch starting with the latest write operation. 
    35    In practice, the [Async]BulkWriter classes handle this. 
    36 
    37    This has the same set of methods for write operations that 
    38    :class:`~google.cloud.firestore_v1.document.DocumentReference` does, 
    39    e.g. :meth:`~google.cloud.firestore_v1.document.DocumentReference.create`. 
    40 
    41    Args: 
    42        client (:class:`~google.cloud.firestore_v1.client.Client`): 
    43            The client that created this batch. 
    44    """ 
    45 
    46    def __init__(self, client) -> None: 
    47        super(BulkWriteBatch, self).__init__(client=client) 
    48 
    49    def commit( 
    50        self, 
    51        retry: retries.Retry | object | None = gapic_v1.method.DEFAULT, 
    52        timeout: float | None = None, 
    53    ) -> BatchWriteResponse: 
    54        """Writes the changes accumulated in this batch. 
    55 
    56        Write operations are not guaranteed to be applied in order and must not 
    57        contain multiple writes to any given document. Preferred over `commit` 
    58        for performance reasons if these conditions are acceptable. 
    59 
    60        Args: 
    61            retry (google.api_core.retry.Retry): Designation of what errors, if any, 
    62                should be retried.  Defaults to a system-specified policy. 
    63            timeout (float): The timeout for this request.  Defaults to a 
    64                system-specified value. 
    65 
    66        Returns: 
    67            :class:`google.cloud.firestore_v1.write.BatchWriteResponse`: 
    68            Container holding the write results corresponding to the changes 
    69            committed, returned in the same order as the changes were applied to 
    70            this batch. An individual write result contains an ``update_time`` 
    71            field. 
    72        """ 
    73        request, kwargs = self._prep_commit(retry, timeout) 
    74 
    75        _api = self._client._firestore_api 
    76        save_response: BatchWriteResponse = _api.batch_write( 
    77            request=request, 
    78            metadata=self._client._rpc_metadata, 
    79            **kwargs, 
    80        ) 
    81 
    82        self._write_pbs = [] 
    83        self.write_results = list(save_response.write_results) 
    84 
    85        return save_response 
    86 
    87    def _prep_commit(self, retry: retries.Retry | object | None, timeout: float | None): 
    88        request = { 
    89            "database": self._client._database_string, 
    90            "writes": self._write_pbs, 
    91            "labels": None, 
    92        } 
    93        kwargs = _helpers.make_retry_timeout_kwargs(retry, timeout) 
    94        return request, kwargs