From ad49f8c23a4a7d710754c6d63da7dd8094834254 Mon Sep 17 00:00:00 2001 From: luis Date: Wed, 22 Nov 2023 01:58:21 -0300 Subject: [PATCH 1/3] Added support for batch method --- mockfirestore/client.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/mockfirestore/client.py b/mockfirestore/client.py index 75943bd..e39d82e 100644 --- a/mockfirestore/client.py +++ b/mockfirestore/client.py @@ -59,4 +59,7 @@ def get_all(self, references: Iterable[DocumentReference], def transaction(self, **kwargs) -> Transaction: return Transaction(self, **kwargs) + def batch(self) -> Transaction: + return Transaction(self) + From e392dc2eac461517d9fafe32349a36f9db4ccfe6 Mon Sep 17 00:00:00 2001 From: luis Date: Wed, 22 Nov 2023 02:13:41 -0300 Subject: [PATCH 2/3] Fixed commit error --- mockfirestore/client.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/mockfirestore/client.py b/mockfirestore/client.py index e39d82e..4bd7c67 100644 --- a/mockfirestore/client.py +++ b/mockfirestore/client.py @@ -60,6 +60,8 @@ def transaction(self, **kwargs) -> Transaction: return Transaction(self, **kwargs) def batch(self) -> Transaction: - return Transaction(self) + batch = Transaction(self) + batch._begin() + return batch From 75ae7325bc9b7fb438fcafab1ce3c8071f6e2b26 Mon Sep 17 00:00:00 2001 From: luis Date: Wed, 22 Nov 2023 02:22:07 -0300 Subject: [PATCH 3/3] Created a Batch class beacuse batch can call commit many times --- mockfirestore/client.py | 8 +++----- mockfirestore/transaction.py | 9 ++++++++- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/mockfirestore/client.py b/mockfirestore/client.py index 4bd7c67..ebfc448 100644 --- a/mockfirestore/client.py +++ b/mockfirestore/client.py @@ -1,7 +1,7 @@ from typing import Iterable, Sequence from mockfirestore.collection import CollectionReference from mockfirestore.document import DocumentReference, DocumentSnapshot -from mockfirestore.transaction import Transaction +from mockfirestore.transaction import Transaction, Batch class MockFirestore: @@ -59,9 +59,7 @@ def get_all(self, references: Iterable[DocumentReference], def transaction(self, **kwargs) -> Transaction: return Transaction(self, **kwargs) - def batch(self) -> Transaction: - batch = Transaction(self) - batch._begin() - return batch + def batch(self) -> Batch: + return Batch(self) diff --git a/mockfirestore/transaction.py b/mockfirestore/transaction.py index 7f06d2d..4d07c93 100644 --- a/mockfirestore/transaction.py +++ b/mockfirestore/transaction.py @@ -1,6 +1,6 @@ from functools import partial -import random from typing import Iterable, Callable + from mockfirestore._helpers import generate_random_string, Timestamp from mockfirestore.document import DocumentReference, DocumentSnapshot from mockfirestore.query import Query @@ -22,6 +22,7 @@ class Transaction: This mostly follows the model from https://googleapis.dev/python/firestore/latest/transaction.html """ + def __init__(self, client, max_attempts=MAX_ATTEMPTS, read_only=False): self._client = client @@ -117,3 +118,9 @@ def __enter__(self): def __exit__(self, exc_type, exc_val, exc_tb): if exc_type is None: self.commit() + + +class Batch(Transaction): + def commit(self): + self._begin() # batch can call commit many times + super(Batch, self).commit()