feat: included parent_ref parameter in add_batch method too

AntonioVentilii · Apr 7, 2024 · d3dfce4 · d3dfce4
1 parent 06f048e
commit d3dfce4
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/firestore_wrapper/document_manager.py b/firestore_wrapper/document_manager.py
@@ -63,7 +63,8 @@ def add_document(self, collection_name: str, data: dict, document_name: str = No
         return document_name
 
     def add_documents_batch(self, collection_name: str, data_list: list[dict], document_names: list[str] = None,
-                            validator: Validator = None, overwrite: bool = False, verbose: bool = False):
+                            validator: Validator = None, overwrite: bool = False, verbose: bool = False,
+                            parent_ref: DocumentReference = None):
         """
         Adds or updates multiple documents in a specified collection using batch operations.
 
@@ -74,15 +75,16 @@ def add_documents_batch(self, collection_name: str, data_list: list[dict], docum
         :param overwrite: If True, existing documents will be overwritten. If False, they will be ignored.
         :param verbose: If True, print additional information during the operation.
         """
-        existing_doc_names = self.get_document_names(collection_name)
+        parent_ref = parent_ref or self.db
+        existing_doc_names = self.get_document_names(collection_name, parent_ref=parent_ref)
         max_batch_size = 500  # Firestore's limit
 
         chunks = [data_list[i:i + max_batch_size] for i in range(0, len(data_list), max_batch_size)]
         name_chunks = [document_names[i:i + max_batch_size] if document_names else None for i in
                        range(0, len(document_names or []), max_batch_size)]
 
         for chunk_index, chunk in enumerate(chunks):
-            batch = self.db.batch()
+            batch = parent_ref.batch()
             names_chunk = name_chunks[chunk_index] if name_chunks else [None] * len(chunk)
 
             for index, data in enumerate(chunk):
@@ -91,7 +93,7 @@ def add_documents_batch(self, collection_name: str, data_list: list[dict], docum
 
                 document_name = names_chunk[index] if names_chunk else None
                 if document_name is None or document_name not in existing_doc_names or overwrite:
-                    doc_ref = self.db.collection(collection_name).document(document_name)
+                    doc_ref = parent_ref.collection(collection_name).document(document_name)
                     batch.set(doc_ref, data)
                 elif verbose:
                     print(f"Document with name '{document_name}' already exists. Skipping...")

diff --git a/setup.py b/setup.py
@@ -2,7 +2,7 @@
 
 setup(
     name='firestore_wrapper',
-    version='0.3.0',
+    version='0.3.1',
     packages=find_packages(),
     description='A custom wrapper for Google Firestore.',
     long_description=open('README.md').read(),