Read database info from configuration file (#330)

Author: EZoni

dashboard/app.py

@@ -33,8 +33,6 @@
 opt_manager = None
 cal_manager = None
 
-# load database
-db = load_database()
 # list of available experiments
 experiments = load_experiments()
 
@@ -60,12 +58,13 @@ def update(
     global par_manager
     global opt_manager
     global cal_manager
-    # load data
-    exp_data, sim_data = load_data(db)
     # load input and output variables
     input_variables, output_variables, simulation_calibration = load_variables(
         state.experiment
     )
+    # load data
+    db = load_database(state.experiment)
+    exp_data, sim_data = load_data(db)
     # reset output
     if reset_output:
         out_manager = OutputManager(output_variables)
@@ -232,7 +231,7 @@ def find_simulation(event, db):
         print(msg)
 
 
-def open_simulation_dialog(event):
+def open_simulation_dialog(event, db):
     try:
         data_directory, file_path = find_simulation(event, db)
         state.simulation_video = file_path.endswith(".mp4")
@@ -314,7 +313,10 @@ def home_route():
                             figure = plotly.Figure(
                                 display_mode_bar="true",
                                 config={"responsive": True},
-                                click=(open_simulation_dialog, "[utils.safe($event)]"),
+                                click=(
+                                    open_simulation_dialog,
+                                    "[utils.safe($event), db]",
+                                ),
                             )
                             ctrl.figure_update = figure.update
 

dashboard/utils.py

@@ -106,40 +106,29 @@ def verify_input_variables(model_file, experiment):
 
 
 @timer
-def load_database():
+def load_database(experiment):
     print("Loading database...")
-    # load database
-    db_defaults = {
-        "host": "mongodb05.nersc.gov",
-        "port": 27017,
-        "name": "bella_sf",
-        "auth": "bella_sf",
-        "user": "bella_sf_ro",
-    }
-    # read database information from environment variables (if unset, use defaults)
-    db_host = os.getenv("SF_DB_HOST", db_defaults["host"])
-    db_port = int(os.getenv("SF_DB_PORT", db_defaults["port"]))
-    db_name = os.getenv("SF_DB_NAME", db_defaults["name"])
-    db_auth = os.getenv("SF_DB_AUTH_SOURCE", db_defaults["auth"])
-    db_user = os.getenv("SF_DB_USER", db_defaults["user"])
-    # read database password from environment variable (no default provided)
-    db_password = os.getenv("SF_DB_READONLY_PASSWORD")
+    # load configuration dictionary
+    config_dict = load_config_dict(experiment)
+    # read database information from configuration dictionary
+    db_host = config_dict["database"]["host"]
+    db_port = config_dict["database"]["port"]
+    db_name = config_dict["database"]["name"]
+    db_auth = config_dict["database"]["auth"]
+    db_username = config_dict["database"]["username_ro"]
+    db_password_env = config_dict["database"]["password_ro_env"]
+    db_password = os.getenv(db_password_env)
     if db_password is None:
-        raise RuntimeError("Environment variable SF_DB_READONLY_PASSWORD must be set!")
-    # SSH forward?
-    if db_host == "localhost" or db_host == "127.0.0.1":
-        direct_connection = True
-    else:
-        direct_connection = False
+        raise RuntimeError(f"Environment variable {db_password_env} must be set!")
     # get database instance
     print(f"Connecting to database {db_name}@{db_host}:{db_port}...")
     db = pymongo.MongoClient(
         host=db_host,
         port=db_port,
-        username=db_user,
-        password=db_password,
         authSource=db_auth,
-        directConnection=direct_connection,
+        username=db_username,
+        password=db_password,
+        directConnection=(db_host in ["localhost", "127.0.0.1"]),  # SSH forwarding
     )[db_name]
     return db
 

ml/train_model.py

@@ -63,20 +63,6 @@
 if model_type not in ["NN", "ensemble_NN", "GP"]:
     raise ValueError(f"Invalid model type: {model_type}")
 
-###############################################
-# Open credential file for database
-###############################################
-with open(os.path.join(os.getenv("HOME"), "db.profile")) as f:
-    db_profile = f.read()
-
-# Connect to the MongoDB database with read-only access
-db = pymongo.MongoClient(
-    host="mongodb05.nersc.gov",
-    username="bella_sf_admin",
-    password=re.findall("SF_DB_ADMIN_PASSWORD='(.+)'", db_profile)[0],
-    authSource="bella_sf",
-)["bella_sf"]
-
 # Extract configurations of experiments & models
 yaml_dict = None
 current_file_directory = os.path.dirname(os.path.abspath(__file__))
@@ -94,6 +80,26 @@
 if yaml_dict is None:
     raise RuntimeError("File config.yaml not found.")
 
+# Connect to the MongoDB database with read-write access
+db_host = yaml_dict["database"]["host"]
+db_name = yaml_dict["database"]["name"]
+db_auth = yaml_dict["database"]["auth"]
+db_username = yaml_dict["database"]["username_rw"]
+db_password_env = yaml_dict["database"]["password_rw_env"]
+# Look for the password in the profile file
+with open(os.path.join(os.getenv("HOME"), "db.profile")) as f:
+    db_profile = f.read()
+match = re.search(f"{db_password_env}='([^']*)'", db_profile)
+if not match:
+    raise RuntimeError(f"Environment variable {db_password_env} must be set")
+db_password = match.group(1)
+db = pymongo.MongoClient(
+    host=db_host,
+    authSource=db_auth,
+    username=db_username,
+    password=db_password,
+)[db_name]
+
 input_variables = yaml_dict["inputs"]
 input_names = [v["name"] for v in input_variables.values()]
 output_variables = yaml_dict["outputs"]