Get data from FxA firestore

get_fs_data.py

#! /usr/bin/env python3
"""
Get data from FxA firestore.

Runs from within the Google Application environment using baked in credentials.
"""
from __future__ import annotations

import argparse
import json
import logging
import os
import sys
from typing import Dict

logger = logging.getLogger(__name__)

# Environment variable for customer collection name
ENV_CUSTOMER_COLLECTION_NAME = "FS_CUSTOMER_COLLECTION"


def get_data(collection, doc_id: str) -> Dict[str, Dict]:
    """
    Recursively gather data from Firestore.

    Arguments:
    - collection: A Firestore collection
    - doc_id: A document ID in the collection

    Returns:
    A dictionary with keys {collection.id:doc.id} and document data as values
    """
    ret_data = {}

    doc_ref = collection.document(doc_id)
    key = f"{collection.id}:{doc_id}"
    logger.info(f"Getting {key}...")
    doc = doc_ref.get()
    ret_data[key] = doc.to_dict()
    for sub_collection in doc_ref.collections():
        for sub_doc in sub_collection.stream():
            sub_ret = get_data(sub_collection, sub_doc.id)
            for sub_key, sub_data in sub_ret.items():
                if sub_key not in ret_data:
                    ret_data[sub_key] = sub_data
    return ret_data


def get_all_data(collection, limit=10) -> Dict[str, Dict]:
    """
    Get all records in a collection

    Arguments:
    - collection: A Firestore collection
    - limit: A limit of top-level items to get
    """
    ret_data = {}
    count = 0
    for doc in collection.stream():
        ret_data.update(get_data(collection, doc.id))
        count += 1
        if count >= limit:
            break
    return ret_data


def db_client_from_app_default_credentials():
    """
    Create a Firestore DB client from Google Application Default Credentials.

    Documentation:
    https://firebase.google.com/docs/admin/setup#initialize-without-parameters

    To use from Cloud Shell:
    * Load Google Cloud Console, https://console.cloud.google.com
    * Select the desired Organization and Product
    * Navigate to Firestore product, interactive view
    * Click "Activate Cloud Shell" icon from top menu bar
    * In Cloud Shell terminal, run "pip3 install --upgrade firebase-admin"
    * Upload this file,Run it!
    * In the Cloud Editor (click "Open Editor" in Cloud Shell), you may need to
      run "gcloud auth login" in the Editor shell.
    """
    from firebase_admin import firestore, initialize_app

    initialize_app()
    return firestore.client()


def get_parser():

    # Get defaults from environment
    default_customer_collection = os.environ.get(ENV_CUSTOMER_COLLECTION_NAME)

    parser = argparse.ArgumentParser(description="Get data from FxA Firestore.")
    parser.add_argument("fxa_id", nargs="*", help="FxA ID of the customer")
    if default_customer_collection:
        parser.add_argument(
            "-c",
            "--customer-collection",
            help=(
                "Name of Firestore collection for Stripe Customer data"
                f" (default '{default_customer_collection}' from env variable {ENV_CUSTOMER_COLLECTION_NAME})"
            ),
            default=default_customer_collection,
        )
    else:
        parser.add_argument(
            "-c",
            "--customer-collection",
            help=(
                "Name of Firestore collection for Stripe Customer data"
                f" (default can be set as env variable {ENV_CUSTOMER_COLLECTION_NAME})"
            ),
            required=True,
        )
    parser.add_argument(
        "--all",
        help="Get all items in the collection, up to limit",
        action="store_true",
    )
    parser.add_argument(
        "--limit",
        help="With --all, set the number of items to get",
        type=int,
        default=10,
    )
    parser.add_argument(
        "-v",
        "--verbose",
        default=1,
        action="count",
        help="Print debugging information, repeat for more detail",
    )
    return parser


if __name__ == "__main__":
    parser = get_parser()
    args = parser.parse_args()

    customer_collection = args.customer_collection

    if args.verbose and args.verbose >= 2:
        level = "DEBUG"
    elif args.verbose:
        level = "INFO"
    else:
        level = "WARNING"
    logging.basicConfig(level=level)

    db = db_client_from_app_default_credentials()
    customers = db.collection(customer_collection)

    if args.all:
        data = get_all_data(customers, args.limit)
    elif args.fxa_id:
        data = {}
        for fxa_id in args.fxa_id:
            data.update(get_data(customers, fxa_id))
    else:
        print("Must specify FxA IDs, or --all")
        parser.print_help()
        sys.exit(1)

    print(json.dumps(data, indent=4))