/
get_cdp_infrastructure_stack.py
109 lines (86 loc) · 3.09 KB
/
get_cdp_infrastructure_stack.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
#!/usr/bin/env python
# -*- coding: utf-8 -*-
import argparse
import json
import logging
import shutil
import sys
import traceback
from pathlib import Path
from cdp_backend.database import DATABASE_MODELS
from cdp_backend.infrastructure import INFRA_DIR
###############################################################################
logging.basicConfig(
level=logging.INFO,
format="[%(levelname)4s: %(module)s:%(lineno)4s %(asctime)s] %(message)s",
)
log = logging.getLogger(__name__)
###############################################################################
class Args(argparse.Namespace):
def __init__(self) -> None:
self.__parse()
def __parse(self) -> None:
p = argparse.ArgumentParser(
prog="get_cdp_infrastructure_stack",
description=(
"Generate or copy all the files needed for a new CDP infrastructure."
),
)
p.add_argument(
"output_dir",
type=Path,
help=(
"Path to where the infrastructure files should be copied "
"or generated."
),
)
p.parse_args(namespace=self)
###############################################################################
def _generate_indexes_json(outfile: Path) -> None:
# All indexes
indexes = []
for model_cls in DATABASE_MODELS:
for idx_field_set in model_cls._INDEXES:
indexes.append(
{
"collectionGroup": model_cls.collection_name,
"queryScope": "COLLECTION",
"fields": idx_field_set.to_dict()["fields"],
}
)
# Add indexes to the normal JSON format
indexes_full_json = {
"indexes": indexes,
"fieldOverrides": [],
}
# Write out the file
outfile = outfile.resolve()
with open(outfile, "w") as open_f:
json.dump(indexes_full_json, open_f, indent=2)
log.info(f"Wrote out CDP firestore.indexes.json to: '{outfile}'")
def _copy_infra_files(output_dir: Path) -> None:
# Copy each file in the infra dir to the output dir
output_dir.mkdir(parents=True, exist_ok=True)
for f in INFRA_DIR.iterdir():
if f.name not in [
"__pycache__",
"__init__.py",
]:
shutil.copy(f, output_dir / f.name)
def main() -> None:
try:
args = Args()
output_dir = args.output_dir.expanduser().resolve()
_copy_infra_files(output_dir=output_dir)
_generate_indexes_json(outfile=output_dir / "firestore.indexes.json")
except Exception as e:
log.error("=============================================")
log.error("\n\n" + traceback.format_exc())
log.error("=============================================")
log.error("\n\n" + str(e) + "\n")
log.error("=============================================")
sys.exit(1)
###############################################################################
# Allow caller to directly run this module (usually in development scenarios)
if __name__ == "__main__":
main()