# SPDX-FileCopyrightText: 2023-2025 Helmholtz-Zentrum Dresden-Rossendorf e.V (HZDR)
# SPDX-License-Identifier: Apache-2.0
"""
This module provides utilities for managing a MongoDB database connection and
initializing the database with data from JSON files.
The module contains functions to:
- Establish a connection to the MongoDB database.
- Close the database connection.
- Initialize the database with data from JSON files.
- Add a command to the Flask CLI for initializing the database.
Imports:
- json: Used for loading data from JSON files.
- quote from urllib.parse: Used for URL-encoding MongoDB credentials.
- click: Used for creating and managing Flask CLI commands.
- current_app, g from flask: Used for accessing the current Flask app context and storing data.
- MongoClient from pymongo: Used for connecting to the MongoDB database.
Functions:
- get_db: Establishes and returns a connection to the MongoDB database.
- close_db: Closes the MongoDB database connection.
- init_db: Initializes the database with data from JSON files.
- init_db_command: Flask CLI command to call the init_db function.
- init_app: Registers the database functions with the Flask app context.
"""
import concurrent.futures
import logging
import os
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from urllib.parse import quote
import click
from flask import current_app, g, has_request_context
from pymongo import MongoClient
from labfrog.common import resolve_config_name
RESET_CHOICES = ("none", "all", "keep-shots", "metadata-only")
REBUILD_CHOICES = ("none", "params", "diags", "both")
SOURCE_CHOICES = ("auto", "default", "full-custom")
[docs]
@dataclass(frozen=True)
class InitPlan:
reset: str
rebuild: str
source: str
SHOT_NAV_SORT = [("shot_number", -1), ("date_time", -1), ("_id", -1)]
SET_NAV_SORT = [("set_number", -1), ("_id", -1)]
# How long a Mongo operation may take before the app gives up and treats the
# database as unreachable. Short on purpose: an operator waiting on a frozen
# page has no idea whether their shot was recorded.
DEFAULT_MONGO_TIMEOUT_MS = 4000
MONGOMOCK_CLIENT_EXTENSION_KEY = "mongomock_client"
MONGO_CLIENT_EXTENSION_KEY = "labfrog_mongo_client"
READ_ONLY_DB_ENV = "LABFROG_READ_ONLY_DB"
NAV_STATE_PROJECTION = {
"_id": 1,
"date_time": 1,
"shot_number": 1,
"set_number": 1,
"shot_number_list": 1,
}
READ_ONLY_BLOCKED_COLLECTION_METHODS = {
"insert_one",
"insert_many",
"replace_one",
"update_one",
"update_many",
"delete_one",
"delete_many",
"find_one_and_update",
"find_one_and_replace",
"find_one_and_delete",
"bulk_write",
"create_index",
"create_indexes",
"drop_index",
"drop_indexes",
"drop",
"rename",
}
READ_ONLY_BLOCKED_DATABASE_METHODS = {
"create_collection",
"drop_collection",
}
def _raise_read_only_write(method_name: str) -> None:
raise RuntimeError(
f"MongoDB write operation '{method_name}' is disabled while {READ_ONLY_DB_ENV}=1."
)
class _ReadOnlyCollectionProxy:
"""Collection proxy that blocks write-style methods in read-only mode."""
def __init__(self, collection):
self._collection = collection
def __getitem__(self, name):
# PyMongo collections are subscriptable for dotted sub-collections
# (for example db["diagnostics"]["SampleDiag"]). Preserve that
# behavior while keeping the read-only write-method guard active.
return _ReadOnlyCollectionProxy(self._collection[name])
def __getattr__(self, name):
attr = getattr(self._collection, name)
if callable(attr) and name in READ_ONLY_BLOCKED_COLLECTION_METHODS:
return lambda *_args, **_kwargs: _raise_read_only_write(name)
return attr
class _ReadOnlyDatabaseProxy:
"""Database proxy that returns read-only collection proxies."""
def __init__(self, database):
self._database = database
def __getitem__(self, name):
return _ReadOnlyCollectionProxy(self._database[name])
def get_collection(self, *args, **kwargs):
return _ReadOnlyCollectionProxy(self._database.get_collection(*args, **kwargs))
def __getattr__(self, name):
attr = getattr(self._database, name)
if callable(attr) and name in READ_ONLY_BLOCKED_DATABASE_METHODS:
return lambda *_args, **_kwargs: _raise_read_only_write(name)
return attr
[docs]
def navigation_sort_for_mode(mode: str | None) -> list[tuple[str, int]]:
"""Return the stable sort used by campaign navigation."""
if mode == "set":
return SET_NAV_SORT
return SHOT_NAV_SORT
def _app_cache_get(name: str, key, loader, *, ttl_seconds: float):
app_ext = getattr(current_app, "extensions", None)
if app_ext is None:
app_ext = {}
current_app.extensions = app_ext
bucket = app_ext.setdefault(name, {})
now = time.monotonic()
cached = bucket.get(key)
if cached is not None:
expires_at, value = cached
if expires_at >= now:
return value
value = loader()
bucket[key] = (now + ttl_seconds, value)
return value
[docs]
def clear_navigation_caches() -> None:
"""Clear cached navigation query results after shot/set writes."""
app_ext = getattr(current_app, "extensions", None)
if not app_ext:
return
app_ext.pop("labfrog_navigation_docs_cache", None)
app_ext.pop("labfrog_navigation_count_cache", None)
def _navigation_cache_ttl_seconds() -> float:
value = current_app.config.get("NAVIGATION_CACHE_SECONDS", 10.0)
try:
return max(float(value), 0.0)
except (TypeError, ValueError):
return 10.0
def _freeze_cache_value(value):
if isinstance(value, dict):
return tuple(
sorted(
(
str(key),
_freeze_cache_value(inner_value),
)
for key, inner_value in value.items()
)
)
if isinstance(value, (list, tuple)):
return tuple(_freeze_cache_value(item) for item in value)
if isinstance(value, set):
return tuple(sorted(_freeze_cache_value(item) for item in value))
if isinstance(value, datetime):
normalized = _coerce_nav_datetime(value)
return (
"datetime",
normalized.isoformat() if normalized is not None else str(value),
)
try:
hash(value)
return value
except Exception:
return repr(value)
def _navigation_docs_cache_key(
shots_collection,
query,
*,
mode: str,
projection: dict | None,
):
return (
current_app.config.get("MONGODB_DATABASE"),
getattr(shots_collection, "name", "shots"),
mode,
_freeze_cache_value(query),
_freeze_cache_value(projection),
)
def _navigation_count_cache_key(shots_collection, query):
return (
current_app.config.get("MONGODB_DATABASE"),
getattr(shots_collection, "name", "shots"),
_freeze_cache_value(query),
)
def _cached_navigation_docs(
shots_collection,
query,
*,
mode: str,
projection: dict | None,
) -> list[dict]:
ttl_seconds = _navigation_cache_ttl_seconds()
return _app_cache_get(
"labfrog_navigation_docs_cache",
_navigation_docs_cache_key(
shots_collection,
query,
mode=mode,
projection=projection,
),
lambda: fetch_sorted_navigation_docs(
shots_collection,
query,
mode=mode,
projection=projection,
),
ttl_seconds=ttl_seconds,
)
def _cached_navigation_count(shots_collection, query) -> int:
ttl_seconds = _navigation_cache_ttl_seconds()
return _app_cache_get(
"labfrog_navigation_count_cache",
_navigation_count_cache_key(shots_collection, query),
lambda: shots_collection.count_documents(query),
ttl_seconds=ttl_seconds,
)
def _coerce_nav_datetime(value) -> datetime | None:
if isinstance(value, str):
try:
value = datetime.fromisoformat(value)
except ValueError:
return None
if not isinstance(value, datetime):
return None
if value.tzinfo is not None:
return value.astimezone(timezone.utc).replace(tzinfo=None)
return value
def _coerce_nav_number(value) -> int:
try:
return int(value)
except (TypeError, ValueError):
return -1
def _datetime_key(value: datetime | None) -> tuple[int, int, int, int, int, int, int]:
if value is None:
return (-1, -1, -1, -1, -1, -1, -1)
return (
value.year,
value.month,
value.day,
value.hour,
value.minute,
value.second,
value.microsecond,
)
def _shot_nav_sort_docs(docs: list[dict]) -> list[dict]:
"""Split shot navigation only when numbering restarts at 1."""
return [doc for block in _shot_nav_blocks(docs) for doc in block]
def _shot_nav_blocks(docs: list[dict]) -> list[list[dict]]:
"""Return shot navigation blocks split only at chronological restarts to 1."""
docs_with_dt = []
docs_without_dt = []
for doc in docs:
normalized = dict(doc)
dt_value = _coerce_nav_datetime(normalized.get("date_time"))
if dt_value is None:
docs_without_dt.append(normalized)
else:
normalized["_nav_date_time"] = dt_value
docs_with_dt.append(normalized)
docs_with_dt.sort(
key=lambda doc: (
_datetime_key(doc.get("_nav_date_time")),
str(doc.get("_id") or ""),
),
)
number_blocks: list[list[dict]] = []
current_block: list[dict] = []
previous_number: int | None = None
for doc in docs_with_dt:
current_number = _coerce_nav_number(doc.get("shot_number"))
if current_block and current_number == 1 and (previous_number or 0) > 1:
number_blocks.append(current_block)
current_block = []
current_block.append(doc)
previous_number = current_number
if current_block:
number_blocks.append(current_block)
sorted_blocks: list[list[dict]] = []
for block in reversed(number_blocks):
block.sort(
key=lambda doc: (
_coerce_nav_number(doc.get("shot_number")),
_datetime_key(doc.get("_nav_date_time")),
str(doc.get("_id") or ""),
),
reverse=True,
)
for doc in block:
doc.pop("_nav_date_time", None)
sorted_blocks.append(block)
docs_without_dt.sort(
key=lambda doc: (
_coerce_nav_number(doc.get("shot_number")),
str(doc.get("_id") or ""),
),
reverse=True,
)
if docs_without_dt:
sorted_blocks.append(docs_without_dt)
return sorted_blocks
[docs]
def fetch_shot_navigation_blocks(
shots_collection,
query,
projection: dict | None = None,
) -> list[list[dict]]:
"""Return shot-mode navigation blocks in the same order shown in the UI."""
docs = list(shots_collection.find(query, projection))
return _shot_nav_blocks(docs)
[docs]
def fetch_sorted_navigation_docs(
shots_collection,
query,
*,
mode: str,
projection: dict | None = None,
) -> list[dict]:
"""Return navigation docs in the order shown to the user."""
if mode == "shot":
return [
doc
for block in fetch_shot_navigation_blocks(
shots_collection,
query,
projection=projection,
)
for doc in block
]
return list(
shots_collection.find(query, projection).sort(navigation_sort_for_mode(mode))
)
def _coerce_list(value) -> list:
if isinstance(value, list):
return value
if isinstance(value, (tuple, set)):
return list(value)
return []
def _field_sections_dict_is_valid(value) -> bool:
if not isinstance(value, dict):
return False
for section, raw in value.items():
if not isinstance(section, str) or not section:
return False
if isinstance(raw, dict):
if not isinstance(raw.get("fields", []), list):
return False
continue
if isinstance(raw, list):
continue
return False
return True
def _normalize_field_sections_dict(value) -> dict:
if not isinstance(value, dict):
return {}
normalized = {}
for idx, (section, raw) in enumerate(value.items()):
if not isinstance(section, str) or not section:
continue
if isinstance(raw, dict):
fields = [f for f in raw.get("fields", []) if isinstance(f, str) and f]
normalized[section] = {
"fields": fields,
"order": raw.get("order", idx),
"color": raw.get("color"),
}
continue
if isinstance(raw, list):
fields = [f for f in raw if isinstance(f, str) and f]
normalized[section] = {"fields": fields, "order": idx, "color": None}
return normalized
def _flatten_field_sections_dict(sections: dict) -> list[str]:
ordered = []
for _, value in sorted(
sections.items(), key=lambda entry: entry[1].get("order", 0)
):
ordered.extend(value.get("fields", []))
return ordered
def _minimal_default_layout_doc(mode: str) -> dict:
if mode == "set":
fields = ["set_number", "set_length", "date_time", "Campaign"]
else:
fields = ["shot_number", "date_time", "Campaign"]
return {
"layout_name": "DEFAULT",
"mode": mode,
"date_time": datetime.now(timezone.utc),
"responsible_person": "system",
"description": "Auto-generated fallback default layout",
"always_include": ["Campaign"],
"diagnostics_list": [],
"field_sections_dict": {
"Shot Details": {"fields": fields, "order": 0, "color": None}
},
"selected_fields": fields,
}
[docs]
def get_db():
"""
Establish and return a connection to the MongoDB database.
If a connection already exists in the Flask app context, it is reused.
Otherwise, a new connection is established using the app's configuration.
Returns:
pymongo.database.Database: The MongoDB database instance.
"""
if "db" not in g:
# Prefer an in-memory mongomock client when running tests to avoid
# accidental network calls to a real MongoDB server.
use_mock = current_app.config.get("USE_MONGOMOCK")
if use_mock is None:
# Default to True when TESTING is enabled unless explicitly disabled.
use_mock = bool(current_app.config.get("TESTING"))
app_ext = getattr(current_app, "extensions", None)
if app_ext is None:
app_ext = {}
current_app.extensions = app_ext
if use_mock:
try:
import mongomock
except ImportError: # pragma: no cover - safety fallback
mongomock = None
if mongomock is None:
logging.warning("USE_MONGOMOCK enabled but mongomock is not installed")
else:
# Reuse a single in-memory client across requests to preserve data
# seeded during app/test setup.
client = app_ext.setdefault(
MONGOMOCK_CLIENT_EXTENSION_KEY, mongomock.MongoClient()
)
g.db = client
if "db" not in g:
host = current_app.config["MONGODB_HOST"]
port = current_app.config["MONGODB_PORT"]
auth_source = current_app.config["MONGODB_AUTH_SOURCE"]
# URL-encoded ("quoted") to be safe to use in the connection string
username = quote(current_app.config["MONGODB_USERNAME"])
password = quote(current_app.config["MONGODB_PASSWORD"])
connection_string = (
f"mongodb://{username}:{password}@{host}:{port}/"
f"?authMechanism=DEFAULT&authSource={auth_source}"
)
# Mongo client timeouts (ms). These have defaults rather than being
# opt-in: PyMongo otherwise waits 30s to select a server and will
# wait forever on a socket read, so a dropped network connection
# froze the page instead of reporting a problem. Failing in a few
# seconds lets the route spool the entry locally and say so.
client_kwargs = {}
timeout_ms = current_app.config.get(
"MONGODB_SERVER_SELECTION_TIMEOUT_MS", DEFAULT_MONGO_TIMEOUT_MS
)
if timeout_ms is not None:
try:
# serverSelectionTimeoutMS controls how long PyMongo waits
# to select an available server. Use the same value for
# connectTimeoutMS to ensure initial TCP connect attempts
# fail fast on flaky networks.
timeout_value = int(timeout_ms)
client_kwargs["serverSelectionTimeoutMS"] = timeout_value
client_kwargs["connectTimeoutMS"] = timeout_value
# socketTimeoutMS bounds network reads after a connection
# is established. Default to the selection timeout unless
# callers provide a specific socket timeout override.
socket_timeout_ms = current_app.config.get(
"MONGODB_SOCKET_TIMEOUT_MS", timeout_value
)
client_kwargs["socketTimeoutMS"] = int(socket_timeout_ms)
except Exception:
# ignore misconfigured timeout values and continue
pass
client = app_ext.get(MONGO_CLIENT_EXTENSION_KEY)
if client is None:
client = MongoClient(connection_string, **client_kwargs)
app_ext[MONGO_CLIENT_EXTENSION_KEY] = client
g.db = client
database = current_app.config["MONGODB_DATABASE"]
resolved_db = g.db[database]
if os.environ.get(READ_ONLY_DB_ENV):
return _ReadOnlyDatabaseProxy(resolved_db)
return resolved_db
[docs]
def ensure_runtime_indexes() -> None:
"""Create the hot-path indexes used by request-time navigation queries."""
if os.environ.get(READ_ONLY_DB_ENV):
return
app_ext = getattr(current_app, "extensions", None)
if app_ext is None:
app_ext = {}
current_app.extensions = app_ext
indexed_dbs = app_ext.setdefault("labfrog_runtime_indexes", set())
database_name = current_app.config["MONGODB_DATABASE"]
if database_name in indexed_dbs:
return
try:
shots = get_db()["shots"]
shots.create_index(
[("Campaign", 1), ("mode", 1), ("date_time", -1), ("_id", -1)],
name="lf_campaign_mode_date_desc",
)
shots.create_index(
[("Campaign", 1), ("mode", 1), ("shot_number", -1)],
name="lf_campaign_mode_shot_number_desc",
)
shots.create_index(
[("Campaign", 1), ("mode", 1), ("set_number", -1)],
name="lf_campaign_mode_set_number_desc",
)
except Exception as exc: # pragma: no cover - defensive logging
logging.info("Unable to ensure runtime indexes: %s", exc)
return
indexed_dbs.add(database_name)
[docs]
def fetch_shot_navigation_state(
shots_collection,
query,
*,
current_index: int,
mode: str,
) -> dict:
"""Fetch the documents needed for navigation on tab1/tab2."""
sorted_docs = _cached_navigation_docs(
shots_collection,
query,
mode=mode,
projection=NAV_STATE_PROJECTION,
)
total_entries = len(sorted_docs)
if total_entries == 0:
return {
"total_entries": 0,
"current_index": 0,
"prev_index": 0,
"next_index": 0,
"first_entry": 0,
"last_entry": 0,
"latest_doc": None,
"current_doc": None,
}
bounded_index = max(0, min(current_index, total_entries - 1))
latest_doc = dict(sorted_docs[0])
current_nav_doc = dict(sorted_docs[bounded_index])
current_doc = shots_collection.find_one({"_id": current_nav_doc["_id"]})
if current_doc is None:
current_doc = current_nav_doc
return {
"total_entries": _cached_navigation_count(shots_collection, query),
"current_index": bounded_index,
"prev_index": min(bounded_index + 1, total_entries - 1),
"next_index": max(bounded_index - 1, 0),
"first_entry": total_entries - 1,
"last_entry": 0,
"latest_doc": latest_doc,
"current_doc": current_doc,
}
[docs]
def close_db(e=None):
"""
Close the MongoDB database connection.
If a connection exists in the Flask app context, it is closed.
Args:
e (Exception, optional): An exception that triggered the function call. Defaults to None.
"""
db = g.pop("db", None)
if db is not None:
app_ext = getattr(current_app, "extensions", None) or {}
if db is app_ext.get(MONGOMOCK_CLIENT_EXTENSION_KEY):
return
if db is app_ext.get(MONGO_CLIENT_EXTENSION_KEY):
return
# Closing a PyMongo MongoClient may perform network I/O (ending
# sessions) which can block if the server is unreachable. Perform
# the close in a short-lived background thread and give it a small
# timeout so teardown handlers do not hang the request/CLI run.
try:
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as ex:
fut = ex.submit(db.close)
try:
fut.result(timeout=1.0)
except KeyboardInterrupt:
logging.info(
"Interrupted while closing MongoClient; skipping clean shutdown."
)
except concurrent.futures.TimeoutError:
logging.info("Timed out while closing MongoClient; skipping.")
except Exception as exc: # pragma: no cover - best-effort close
logging.info("Error while closing MongoClient: %s", exc)
except Exception as exc:
logging.info("Failed to shutdown MongoClient cleanly: %s", exc)
[docs]
def close_app_mongo_client(app=None):
"""Close the shared app-level real Mongo client, if one exists."""
app = app or current_app
app_ext = getattr(app, "extensions", None) or {}
client = app_ext.pop(MONGO_CLIENT_EXTENSION_KEY, None)
if client is None:
return
try:
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as ex:
fut = ex.submit(client.close)
try:
fut.result(timeout=1.0)
except KeyboardInterrupt:
logging.info(
"Interrupted while closing shared MongoClient; skipping clean shutdown."
)
except concurrent.futures.TimeoutError:
logging.info("Timed out while closing shared MongoClient; skipping.")
except Exception as exc: # pragma: no cover - best-effort close
logging.info("Error while closing shared MongoClient: %s", exc)
except Exception as exc:
logging.info("Failed to shutdown shared MongoClient cleanly: %s", exc)
def _coerce_choice(
value: str | None, choices: tuple[str, ...], default: str, label: str
) -> str:
if not value:
return default
lowered = str(value).strip().lower()
if lowered in choices:
return lowered
logging.warning("Unknown %s '%s'; using default '%s'", label, value, default)
return default
def _resolve_plan(reset=None, rebuild=None, source=None) -> InitPlan:
config = current_app.config
resolved_reset = _coerce_choice(
reset or config.get("INIT_RESET"),
RESET_CHOICES,
"none",
"reset plan",
)
resolved_rebuild = _coerce_choice(
rebuild or config.get("INIT_REBUILD"),
REBUILD_CHOICES,
"none",
"rebuild plan",
)
resolved_source = _coerce_choice(
source or config.get("INIT_SOURCE"),
SOURCE_CHOICES,
"auto",
"init source",
)
return InitPlan(resolved_reset, resolved_rebuild, resolved_source)
def _drop_collections(db, plan: InitPlan) -> None:
try:
all_collections = db.list_collection_names()
except Exception as exc: # pragma: no cover - defensive logging
logging.info("Unable to list collections before drop: %s", exc)
return
if plan.reset == "none":
return
skip = set()
targets: set[str] | None = None
if plan.reset == "keep-shots":
skip = {"shots"}
elif plan.reset == "metadata-only":
# Only drop metadata collections (diagnostics + settings)
targets = {
"diagnostics",
current_app.config.get("MONGODB_COLLECTION_FOR_SETTINGS", ""),
}
for collection in all_collections:
if collection in skip:
continue
if targets is not None:
if collection in targets or collection.startswith("diagnostics."):
pass
else:
continue
try:
db[collection].drop()
logging.info("Dropped collection: %s", collection)
except Exception as exc: # pragma: no cover - defensive logging
logging.info("Error dropping collection %s: %s", collection, exc)
def _get_source_flags(plan: InitPlan) -> bool:
if plan.source == "auto":
return current_app.config.get("USE_FULL_CUSTOM", False)
return plan.source == "full-custom"
def _load_custom_documents(config_name: str):
from labfrog.customize.diagnostics_maker import (
create_new_diagnostics,
get_data_and_docs,
make_shot_and_set_defaults,
)
data, shot_document, set_document, aliases = get_data_and_docs(config_name)
return (
data,
shot_document,
set_document,
aliases,
create_new_diagnostics,
make_shot_and_set_defaults,
)
[docs]
def init_db(reset=None, rebuild=None, source=None):
"""
Initialize the MongoDB database with a simple plan-based API.
Args:
reset (str | None): One of RESET_CHOICES to control collection drops.
rebuild (str | None): One of REBUILD_CHOICES to control metadata rebuild.
source (str | None): One of SOURCE_CHOICES to choose metadata source.
"""
if os.environ.get(READ_ONLY_DB_ENV):
raise RuntimeError(f"init_db is disabled while {READ_ONLY_DB_ENV} is enabled.")
db = get_db()
config_name = resolve_config_name(current_app.config)
plan = _resolve_plan(reset=reset, rebuild=rebuild, source=source)
logging.info(
"Init DB plan: reset=%s, rebuild=%s, source=%s",
plan.reset,
plan.rebuild,
plan.source,
)
_drop_collections(db, plan)
if plan.reset != "none":
clear_runtime_metadata_caches()
if plan.rebuild == "none":
return
use_full_custom = _get_source_flags(plan)
if not use_full_custom:
current_app.config["MONGODB_COLLECTION_FOR_SETTINGS"] = current_app.config[
"DEFAULT_COLLECTION_FOR_SETTINGS"
]
collection_names = set(db.list_collection_names())
diagnostics_initialized = "diagnostics" in collection_names
should_rebuild_diags = (
plan.rebuild in ("diags", "both") or not diagnostics_initialized
)
should_rebuild_params = plan.rebuild in ("params", "both")
if not use_full_custom:
if should_rebuild_params or should_rebuild_diags:
logging.info(
"Skipping rebuild because source=%s and USE_FULL_CUSTOM is disabled",
plan.source,
)
return
(
data,
shot_document,
set_document,
aliases,
create_new_diagnostics,
make_shot_and_set_defaults,
) = _load_custom_documents(config_name)
parameters_collection = db[current_app.config["MONGODB_COLLECTION_FOR_SETTINGS"]]
diagnostics_collection = db["diagnostics"]
if should_rebuild_params:
make_shot_and_set_defaults(parameters_collection, shot_document, set_document)
logging.info("Rebuilt parameters from %s", config_name)
if should_rebuild_diags:
create_new_diagnostics(db, diagnostics_collection, data)
logging.info("Rebuilt diagnostics from %s", config_name)
if aliases and (should_rebuild_params or should_rebuild_diags):
alias_collection = db["field_aliases"]
alias_collection.delete_many({})
alias_collection.insert_many(aliases)
logging.info("Seeded %d field aliases from %s", len(aliases), config_name)
if should_rebuild_params or should_rebuild_diags:
clear_runtime_metadata_caches()
@click.command(
"init-db",
help="Initialize the MongoDB database.",
)
@click.option(
"--reset",
type=click.Choice(RESET_CHOICES),
help="Which collections to drop before initialization.",
)
@click.option(
"--rebuild",
type=click.Choice(REBUILD_CHOICES),
help="Which metadata to rebuild after dropping collections.",
)
@click.option(
"--source",
type=click.Choice(SOURCE_CHOICES),
help="Where to source parameters/diagnostics from.",
)
@click.option(
"--yes",
is_flag=True,
help="Proceed without confirmation.",
)
def init_db_command(
reset: str | None, rebuild: str | None, source: str | None, yes: bool
):
"""
Flask CLI command to initialize the MongoDB database.
When executed, it calls the init_db function and prints a confirmation message.
"""
plan = _resolve_plan(reset=reset, rebuild=rebuild, source=source)
click.echo(
f"Init DB plan -> reset: {plan.reset}, rebuild: {plan.rebuild}, source: {plan.source}"
)
if not yes:
if not click.confirm("Proceed with this initialization plan?", default=False):
click.echo("Aborted.")
return
init_db(reset=reset, rebuild=rebuild, source=source)
click.echo("Database initialized.")
[docs]
def init_app(app):
"""
Register the database functions with the Flask app context.
Args:
app (Flask): The Flask app instance.
"""
app.teardown_appcontext(close_db)
app.cli.add_command(init_db_command)