|
import argparse |
|
import csv |
|
from datetime import datetime |
|
from os import PathLike |
|
from pathlib import Path |
|
from typing import Generator, Tuple, Union |
|
|
|
import firebase_admin |
|
from firebase_admin import firestore |
|
|
|
|
|
def _generate_filename(collection_name) -> str: |
|
return f"{collection_name}-{datetime.now().strftime('%Y-%m-%d-%H-%M-%S')}.csv" |
|
|
|
def export( |
|
file: Union[PathLike, str], |
|
collection: firestore.firestore.CollectionReference |
|
) -> Generator[Tuple[str, int], None, None]: |
|
""" |
|
Export a collection to given file. |
|
This is a generator function which yields a tuple each time a row is written to the csv file. |
|
The first element of the tuple is the id of last document written and the second is rows |
|
written to the CSV file so far. |
|
|
|
Note: To ensure the CSV file gets written to completely, the generator must be exhausted |
|
""" |
|
db = firestore.client() |
|
documents = collection.stream() |
|
|
|
with open(file, "w") as fp: |
|
try: |
|
first = next(documents) |
|
except StopIteration: |
|
return # no rows |
|
|
|
first_obj = first.to_dict() |
|
|
|
writer = csv.DictWriter(fp, first_obj.keys()) |
|
writer.writeheader() |
|
writer.writerow(first_obj) |
|
yield first.id, 1 |
|
|
|
for count, document in enumerate(documents, 2): |
|
writer.writerow(document.to_dict()) |
|
yield document.id, count |
|
|
|
def _interactive_export( |
|
file: Union[PathLike, str], |
|
collection: firestore.firestore.CollectionReference |
|
): |
|
|
|
print(f'Exporting collection "{collection.id}" to file "{file}"') |
|
|
|
for doc_name, doc_count in export(file, collection): |
|
print(f'Retrieved document "{doc_name}" ({doc_count})') |
|
|
|
if __name__ == "__main__": |
|
firebase_admin.initialize_app() |
|
|
|
argument_parser = argparse.ArgumentParser("exportCSV") |
|
argument_parser.add_argument("collection", help="The collection to export") |
|
argument_parser.add_argument( |
|
"--output", |
|
"-o", |
|
required=False, |
|
help=( |
|
"Output file to export to. " |
|
"'[collection]-export-[time].csv' is used if this is left out" |
|
) |
|
) |
|
|
|
args = argument_parser.parse_args() |
|
|
|
file = Path(args.output or _generate_filename(args.collection)) |
|
|
|
db = firestore.client() |
|
collection = db.collection(args.collection) |
|
|
|
_interactive_export(file, collection) |
