Skip to content

Commit

Permalink
feat: method get_document_data has an optional parameter to get sub-c…
Browse files Browse the repository at this point in the history
…ollections too
  • Loading branch information
AntonioVentilii committed Apr 7, 2024
1 parent d3dfce4 commit cd7c99f
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 7 deletions.
25 changes: 19 additions & 6 deletions firestore_wrapper/document_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,12 @@

from google.cloud.firestore_v1 import DocumentReference, DocumentSnapshot

from .firestore_base import FirestoreBase
from .collection_manager import CollectionManager

Validator = Callable[[dict], None]


class DocumentManager(FirestoreBase):
class DocumentManager(CollectionManager):

def __init__(self, credentials_path: str, database: str = None):
"""
Expand Down Expand Up @@ -74,6 +74,7 @@ def add_documents_batch(self, collection_name: str, data_list: list[dict], docum
:param validator: Optional callable to validate each document's data.
: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.
:param parent_ref: Optional parent document reference.
"""
parent_ref = parent_ref or self.db
existing_doc_names = self.get_document_names(collection_name, parent_ref=parent_ref)
Expand Down Expand Up @@ -168,20 +169,32 @@ def get_document_names(self, collection_name: str, parent_ref: DocumentReference
return [doc.id for doc in parent_ref.collection(collection_name).stream()]

def get_document_data(self, collection_name: str, document_name: str, with_id: bool = False,
parent_ref: DocumentReference = None) -> dict:
parent_ref: DocumentReference = None, with_subcollections: bool = False) -> dict:
"""
Retrieves the data of a specified document in a collection.
:param collection_name: The name of the collection.
:param document_name: The name of the document.
:param with_id: If True, includes the document's ID in the returned dictionary.
:param parent_ref: Optional parent document reference.
:param with_subcollections: If True, includes data for all subcollections.
:return: A dictionary containing the document's data.
"""
parent_ref = parent_ref or self.db
document = parent_ref.collection(collection_name).document(document_name).get()
data = document.to_dict()
if with_id:
return {'id': document.id, **document.to_dict()}
else:
return document.to_dict()
data['id'] = document.id
if with_subcollections:
doc_ref = self.create_doc_ref(collection_name, document_name, parent_ref=parent_ref)
subcollections = doc_ref.collections()
for subcollection in subcollections:
sub_document_names = self.get_document_names(subcollection.id, parent_ref=doc_ref)
sub_data = {}
for sub_document_name in sub_document_names:
sub_data[sub_document_name] = self.get_document_data(subcollection.id, sub_document_name,
with_id=with_id, parent_ref=doc_ref,
with_subcollections=with_subcollections)
data[subcollection.id] = sub_data
return data
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

setup(
name='firestore_wrapper',
version='0.3.1',
version='0.3.2',
packages=find_packages(),
description='A custom wrapper for Google Firestore.',
long_description=open('README.md').read(),
Expand Down

0 comments on commit cd7c99f

Please sign in to comment.