gammasimtools 0.18.0__py3-none-any.whl → 0.20.0__py3-none-any.whl

simtools/db/db_handler.py

@@ -1,5 +1,6 @@
 """Module to handle interaction with DB."""
 
+import io
 import logging
 import re
 from collections import defaultdict
@@ -8,14 +9,15 @@ from threading import Lock
 
 import gridfs
 import jsonschema
+from astropy.table import Table
 from bson.objectid import ObjectId
 from packaging.version import Version
 from pymongo import MongoClient
 
 from simtools.data_model import validate_data
-from simtools.io_operations import io_handler
+from simtools.io import ascii_handler, io_handler
 from simtools.simtel import simtel_table_reader
-from simtools.utils import general, names, value_conversion
+from simtools.utils import names, value_conversion
 
 __all__ = ["DatabaseHandler"]
 
@@ -50,8 +52,19 @@ jsonschema_db_dict = {
             "type": "string",
             "description": "Name of simulation model database",
         },
+        "db_simulation_model_version": {
+            "type": "string",
+            "description": "Version of simulation model database",
+        },
     },
-    "required": ["db_server", "db_api_port", "db_api_user", "db_api_pw", "db_simulation_model"],
+    "required": [
+        "db_server",
+        "db_api_port",
+        "db_api_user",
+        "db_api_pw",
+        "db_simulation_model",
+        "db_simulation_model_version",
+    ],
 }
 
 
@@ -81,6 +94,14 @@ class DatabaseHandler:
 
         self._set_up_connection()
         self._find_latest_simulation_model_db()
+        self.db_name = (
+            self.get_db_name(
+                model_version=self.mongo_db_config.get("db_simulation_model_version"),
+                model_name=self.mongo_db_config.get("db_simulation_model"),
+            )
+            if self.mongo_db_config
+            else None
+        )
 
     def _set_up_connection(self):
         """Open the connection to MongoDB."""
@@ -89,6 +110,16 @@ class DatabaseHandler:
             with lock:
                 DatabaseHandler.db_client = self._open_mongo_db()
 
+    def get_db_name(self, db_name=None, model_version=None, model_name=None):
+        """Build DB name from configuration."""
+        if db_name:
+            return db_name
+        if model_version and model_name:
+            return f"{model_name}-{model_version.replace('.', '-')}"
+        if model_version or model_name:
+            return None
+        return None if (model_version or model_name) else self.db_name
+
     def _validate_mongo_db_config(self, mongo_db_config):
         """Validate the MongoDB configuration."""
         if mongo_db_config is None or all(value is None for value in mongo_db_config.values()):
@@ -137,8 +168,7 @@ class DatabaseHandler:
         """
         Find the latest released version of the simulation model and update the DB config.
 
-        This is indicated by adding "LATEST" to the name of the simulation model database
-        (field "db_simulation_model" in the database configuration dictionary).
+        This is indicated by "LATEST" to the simulation model data base version.
         Only released versions are considered, pre-releases are ignored.
 
         Raises
@@ -148,27 +178,27 @@ class DatabaseHandler:
 
         """
         try:
+            db_simulation_model_version = self.mongo_db_config["db_simulation_model_version"]
             db_simulation_model = self.mongo_db_config["db_simulation_model"]
-            if not db_simulation_model.endswith("LATEST"):
+            if db_simulation_model_version != "LATEST":
                 return
-        except TypeError:  # db_simulation_model is None
+        except TypeError:  # db_simulation_model_version is None
             return
 
-        prefix = db_simulation_model.replace("LATEST", "")
         list_of_db_names = self.db_client.list_database_names()
-        filtered_list_of_db_names = [s for s in list_of_db_names if s.startswith(prefix)]
-        versioned_strings = []
-        version_pattern = re.compile(
-            rf"{re.escape(prefix)}v?(\d+)-(\d+)-(\d+)(?:-([a-zA-Z0-9_.]+))?"
-        )
+        filtered_list_of_db_names = [
+            s for s in list_of_db_names if s.startswith(db_simulation_model)
+        ]
+        pattern = re.compile(rf"{re.escape(db_simulation_model)}-v(\d+)-(\d+)-(\d+)(?:-(.+))?$")
 
+        versioned_strings = []
         for s in filtered_list_of_db_names:
-            match = version_pattern.search(s)
-            # A version is considered a pre-release if it contains a '-' character (re group 4)
-            if match and match.group(4) is None:
-                version_str = match.group(1) + "." + match.group(2) + "." + match.group(3)
-                version = Version(version_str)
-                versioned_strings.append((s, version))
+            m = pattern.match(s)
+            if m:
+                # skip pre-releases (have suffix)
+                if m.group(4) is None:
+                    version_str = f"{m.group(1)}.{m.group(2)}.{m.group(3)}"
+                    versioned_strings.append((s, Version(version_str)))
 
         if versioned_strings:
             latest_string, _ = max(versioned_strings, key=lambda x: x[1])
@@ -177,7 +207,29 @@ class DatabaseHandler:
                 f"Updated the DB simulation model to the latest version {latest_string}"
             )
         else:
-            raise ValueError("Found LATEST in the DB name but no matching versions found in DB.")
+            raise ValueError("LATEST requested but no released versions found in DB.")
+
+    def generate_compound_indexes(self, db_name=None):
+        """
+        Generate compound indexes for the MongoDB collections.
+
+        Indexes based on the typical query patterns.
+        """
+        db_name = db_name or self.db_name
+        collection_names = [
+            "telescopes",
+            "sites",
+            "configuration_sim_telarray",
+            "configuration_corsika",
+            "calibration_devices",
+        ]
+        for collection_name in collection_names:
+            db_collection = self.get_collection(collection_name, db_name=db_name)
+            db_collection.create_index(
+                [("instrument", 1), ("site", 1), ("parameter", 1), ("parameter_version", 1)]
+            )
+        db_collection = self.get_collection("production_tables", db_name=db_name)
+        db_collection.create_index([("collection", 1), ("model_version", 1)])
 
     def get_model_parameter(
         self,
@@ -333,24 +385,23 @@ class DatabaseHandler:
         )
         return DatabaseHandler.model_parameters_cached[cache_key]
 
-    def get_collection(self, db_name, collection_name):
+    def get_collection(self, collection_name, db_name=None):
         """
         Get a collection from the DB.
 
         Parameters
         ----------
-        db_name: str
-            Name of the DB.
         collection_name: str
             Name of the collection.
+        db_name: str
+            Name of the DB.
 
         Returns
         -------
         pymongo.collection.Collection
             The collection from the DB.
-
         """
-        db_name = self._get_db_name(db_name)
+        db_name = db_name or self.db_name
         return DatabaseHandler.db_client[db_name][collection_name]
 
     def get_collections(self, db_name=None, model_collections_only=False):
@@ -370,7 +421,7 @@ class DatabaseHandler:
             List of collection names
 
         """
-        db_name = db_name or self._get_db_name()
+        db_name = db_name or self.db_name
         if db_name not in self.list_of_collections:
             self.list_of_collections[db_name] = DatabaseHandler.db_client[
                 db_name
@@ -451,7 +502,7 @@ class DatabaseHandler:
         file_id: dict of GridOut._id
             Dict of database IDs of files.
         """
-        db_name = self._get_db_name(db_name)
+        db_name = db_name or self.db_name
 
         if file_names:
             file_names = [file_names] if not isinstance(file_names, list) else file_names
@@ -467,8 +518,8 @@ class DatabaseHandler:
             if Path(dest).joinpath(file_name).exists():
                 instance_ids[file_name] = "file exists"
             else:
-                file_path_instance = self._get_file_mongo_db(self._get_db_name(), file_name)
-                self._write_file_from_mongo_to_disk(self._get_db_name(), dest, file_path_instance)
+                file_path_instance = self._get_file_mongo_db(db_name, file_name)
+                self._write_file_from_mongo_to_disk(db_name, dest, file_path_instance)
                 instance_ids[file_name] = file_path_instance._id  # pylint: disable=protected-access
         return instance_ids
 
@@ -509,8 +560,7 @@ class DatabaseHandler:
         ValueError
             if query returned no results.
         """
-        db_name = self._get_db_name()
-        collection = self.get_collection(db_name, collection_name)
+        collection = self.get_collection(collection_name, db_name=self.db_name)
         posts = list(collection.find(query))
         if not posts:
             raise ValueError(
@@ -547,7 +597,7 @@ class DatabaseHandler:
             pass
 
         query = {"model_version": model_version, "collection": collection_name}
-        collection = self.get_collection(self._get_db_name(), "production_tables")
+        collection = self.get_collection("production_tables", db_name=self.db_name)
         post = collection.find_one(query)
         if not post:
             raise ValueError(f"The following query returned zero results: {query}")
@@ -574,7 +624,7 @@ class DatabaseHandler:
         list
             List of model versions
         """
-        collection = self.get_collection(self._get_db_name(), "production_tables")
+        collection = self.get_collection("production_tables", db_name=self.db_name)
         return sorted(
             {post["model_version"] for post in collection.find({"collection": collection_name})}
         )
@@ -751,27 +801,56 @@ class DatabaseHandler:
         with open(Path(path).joinpath(file.filename), "wb") as output_file:
             fs_output.download_to_stream_by_name(file.filename, output_file)
 
-    def add_production_table(self, db_name, production_table):
+    def get_ecsv_file_as_astropy_table(self, file_name, db_name=None):
         """
-        Add a production table to the DB.
+        Read contents of an ECSV file from the database and return it as an Astropy Table.
+
+        Files are not written to disk.
 
         Parameters
         ----------
+        file_name: str
+            The name of the ECSV file.
         db_name: str
-            the name of the DB.
+            The name of the database.
+
+        Returns
+        -------
+        astropy.table.Table
+            The contents of the ECSV file as an Astropy Table.
+        """
+        db = DatabaseHandler.db_client[db_name or self.db_name]
+        fs = gridfs.GridFSBucket(db)
+
+        buf = io.BytesIO()
+        try:
+            fs.download_to_stream_by_name(file_name, buf)
+        except gridfs.errors.NoFile as exc:
+            raise FileNotFoundError(f"ECSV file '{file_name}' not found in DB.") from exc
+        buf.seek(0)
+        return Table.read(buf.getvalue().decode("utf-8"), format="ascii.ecsv")
+
+    def add_production_table(self, production_table, db_name=None):
+        """
+        Add a production table to the DB.
+
+        Parameters
+        ----------
         production_table: dict
             The production table to add to the DB.
+        db_name: str
+            the name of the DB.
         """
-        db_name = self._get_db_name(db_name)
-        collection = self.get_collection(db_name, "production_tables")
+        db_name = db_name or self.db_name
+        collection = self.get_collection("production_tables", db_name=db_name or self.db_name)
         self._logger.debug(f"Adding production for {production_table.get('collection')} to to DB")
         collection.insert_one(production_table)
         DatabaseHandler.production_table_cached.clear()
 
     def add_new_parameter(
         self,
-        db_name,
         par_dict,
+        db_name=None,
         collection_name="telescopes",
         file_prefix=None,
     ):
@@ -783,10 +862,10 @@ class DatabaseHandler:
 
         Parameters
         ----------
-        db_name: str
-            the name of the DB
         par_dict: dict
             dictionary with parameter data
+        db_name: str
+            the name of the DB
         collection_name: str
             The name of the collection to add a parameter to.
         file_prefix: str or Path
@@ -794,8 +873,8 @@ class DatabaseHandler:
         """
         par_dict = validate_data.DataValidator.validate_model_parameter(par_dict)
 
-        db_name = self._get_db_name(db_name)
-        collection = self.get_collection(db_name, collection_name)
+        db_name = db_name or self.db_name
+        collection = self.get_collection(collection_name, db_name=db_name)
 
         par_dict["value"], _base_unit, _ = value_conversion.get_value_unit_type(
             value=par_dict["value"], unit_str=par_dict.get("unit", None)
@@ -810,7 +889,7 @@ class DatabaseHandler:
                     f"corresponding to the {par_dict['parameter']} parameter, must be provided."
                 )
             file_path = Path(file_prefix).joinpath(par_dict["value"])
-            if not general.is_utf8_file(file_path):
+            if not ascii_handler.is_utf8_file(file_path):
                 raise ValueError(f"File is not UTF-8 encoded: {file_path}")
             files_to_add_to_db.add(f"{file_path}")
 
@@ -825,22 +904,6 @@ class DatabaseHandler:
 
         self._reset_parameter_cache()
 
-    def _get_db_name(self, db_name=None):
-        """
-        Return database name. If not provided, return the default database name.
-
-        Parameters
-        ----------
-        db_name: str
-            Database name
-
-        Returns
-        -------
-        str
-            Database name
-        """
-        return self.mongo_db_config["db_simulation_model"] if db_name is None else db_name
-
     def insert_file_to_db(self, file_name, db_name=None, **kwargs):
         """
         Insert a file to the DB.
@@ -853,7 +916,7 @@ class DatabaseHandler:
             the name of the DB
         **kwargs (optional): keyword arguments for file creation.
             The full list of arguments can be found in, \
-            https://docs.mongodb.com/manual/core/gridfs/#the-files-collection
+            https://www.mongodb.com/docs/manual/core/gridfs/
             mostly these are unnecessary though.
 
         Returns
@@ -863,7 +926,7 @@ class DatabaseHandler:
             "newly created DB GridOut._id.
 
         """
-        db_name = self._get_db_name(db_name)
+        db_name = db_name or self.db_name
         db = DatabaseHandler.db_client[db_name]
         file_system = gridfs.GridFS(db)
 

simtools/db/db_model_upload.py

@@ -3,56 +3,53 @@
 import logging
 from pathlib import Path
 
-import simtools.utils.general as gen
+from simtools.io import ascii_handler
 from simtools.utils import names
 
 logger = logging.getLogger(__name__)
 
 
-def add_values_from_json_to_db(file, collection, db, db_name, file_prefix):
+def add_values_from_json_to_db(file, collection, db, file_prefix):
     """
     Upload new model parameter from json files to db.
 
     Parameters
     ----------
     file : list
-        Json file to be uploaded to the DB.
+        JSON file to be uploaded to the DB.
     collection : str
         The DB collection to which to add the file.
     db : DatabaseHandler
         Database handler object.
-    db_name : str
-        Name of the database to be created.
     file_prefix : str
         Path to location of all additional files to be uploaded.
     """
-    par_dict = gen.collect_data_from_file(file_name=file)
+    par_dict = ascii_handler.collect_data_from_file(file_name=file)
     logger.debug(
         f"Adding the following parameter to the DB: {par_dict['parameter']} "
         f"version {par_dict['parameter_version']} "
-        f"(collection {collection} in database {db_name})"
+        f"(collection {collection} in database {db.get_db_name()})"
     )
 
     db.add_new_parameter(
-        db_name=db_name,
         par_dict=par_dict,
         collection_name=collection,
         file_prefix=file_prefix,
     )
 
 
-def add_model_parameters_to_db(args_dict, db):
+def add_model_parameters_to_db(input_path, db):
     """
     Read model parameters from a directory and upload them to the database.
 
     Parameters
     ----------
-    args_dict : dict
-        Command line arguments.
+    input_path : Path, str
+        Path to the directory containing the model parameters.
     db : DatabaseHandler
         Database handler object.
     """
-    input_path = Path(args_dict["input_path"])
+    input_path = Path(input_path)
     logger.info(f"Reading model parameters from repository path {input_path}")
     array_elements = [d for d in input_path.iterdir() if d.is_dir()]
     for element in array_elements:
@@ -67,12 +64,11 @@ def add_model_parameters_to_db(args_dict, db):
                 file=file,
                 collection=collection,
                 db=db,
-                db_name=args_dict["db_name"],
                 file_prefix=input_path / "Files",
             )
 
 
-def add_production_tables_to_db(args_dict, db):
+def add_production_tables_to_db(input_path, db):
     """
     Read production tables from a directory and upload them to the database.
 
@@ -81,12 +77,12 @@ def add_production_tables_to_db(args_dict, db):
 
     Parameters
     ----------
-    args_dict : dict
-        Command line arguments.
+    input_path : Path, str
+        Path to the directory containing the production tables.
     db : DatabaseHandler
         Database handler object.
     """
-    input_path = Path(args_dict["input_path"])
+    input_path = Path(input_path)
     logger.info(f"Reading production tables from repository path {input_path}")
 
     for model in filter(Path.is_dir, input_path.iterdir()):
@@ -101,7 +97,6 @@ def add_production_tables_to_db(args_dict, db):
                 continue
             logger.info(f"Adding production table for {collection} to the database")
             db.add_production_table(
-                db_name=args_dict["db_name"],
                 production_table=data,
             )
 
@@ -119,7 +114,7 @@ def _read_production_table(model_dict, file, model_name):
             "design_model": {},
         },
     )
-    parameter_dict = gen.collect_data_from_file(file_name=file)
+    parameter_dict = ascii_handler.collect_data_from_file(file_name=file)
     logger.debug(f"Reading production table for {array_element} (collection {collection})")
     try:
         if array_element in ("configuration_corsika", "configuration_sim_telarray"):