DatabaseManager.py

Go to the documentation of this file.

import sqlite3
import pandas as pd
import json
import re
 
class DatabaseManager:
    """
    Manager for interacting with a SQLite database. This class provides methods to add data, query data, and manage tables in the database.
 
    Parameters
    ----------
    - **db_filename** : str
        The filename of the SQLite database.
 
    Methods
    -------
    - **add_conditions** : Adds a DataFrame of conditions to the 'condition' table in the database.
    - **add_parameters** : Adds a list of parameters to the 'parameters' table in the database.
    - **add_data** : Adds a DataFrame of data to a specified table in the database.
    - **add_error** : Adds a DataFrame of error data to the 'error' table in the database.
    - **get** : Retrieves records from a table based on filter conditions.
    - **getByID** : Retrieves records from a table based on ID_PARAMETER and/or ID_CONDITION.
    - **getPareto** : Retrieves the Pareto front for given errors.
    - **getBestScore** : Retrieves the best scores for a given error.
    - **getBest** : Retrieves the best records from a table based on a given error.
    - **summarize** : Returns a summary of the database.
 
    Example
    -------
    ```python
    conditions = pd.DataFrame([
        [0.5, 1],
        [0.2, 2],
        [0.2, 5]
    ], columns=['info1', 'info2'])
    parameters_list = LHSIterator(param,10)
 
    with DatabaseManager("example.db") as db:
        db.add_parameters(parameters_list)
        db.add_conditions(conditions)
        
        for states,other in MultiSimu(run_simu,parameters_list,conditions=conditions.values.tolist(),replicat=5,withTqdm=True,batch_size_level='param'):
            db.add_data('states', states.reset_index())
            db.add_data('otherData', other.reset_index())
            error = your_error_function(states)
            db.add_error(error)
    ```
    """
 
    def __init__(self, db_filename: str):
        self.db_filename = db_filename
        self._table_columns_cache = {}
 
    def open(self):
        """
        Opens a connection to the database.
        """
        self.conn = sqlite3.connect(self.db_filename)
        cursor = self.conn.cursor()
        cursor.execute("PRAGMA journal_mode=WAL;")
        cursor.execute("PRAGMA synchronous = NORMAL;")
        cursor.execute("PRAGMA cache_size = -50000;")
        self._check_tables()
        return self
 
    def close(self):
        """
        Closes the connection to the database.
        """
        self.conn.close()
 
    def __enter__(self):
        self.open()
        return self
 
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
 
    def add_conditions(self, conditions_df: pd.DataFrame, id_column: str = None):
        """
        Adds a DataFrame of conditions to the 'condition' table in the database.
 
        Parameters
        ----------
        - **conditions_df** : pd.DataFrame
            The DataFrame containing the conditions to add.
        - **id_column** : str, optional
            The column to use as the ID_CONDITION. If not provided, the index will be used.
        """
        for i in range(conditions_df.shape[1]):
            if hasattr(conditions_df.iloc[0, i], '__iter__'):
                conditions_df.iloc[:, i] = conditions_df.iloc[:, i].apply(lambda x: json.dumps(x))
        if id_column is None:
            conditions_df = conditions_df.reset_index().rename(columns={'index': 'ID_CONDITION'})
        elif id_column != 'ID_CONDITION':
            conditions_df = conditions_df.rename(columns={id_column: 'ID_CONDITION'})
        
        with self.conn:
            conditions_df.to_sql('condition', self.conn, if_exists='fail', index=False)
            self._check_tables()
            self._create_index('condition', 'ID_CONDITION')
 
    def add_error(self, data_df: pd.DataFrame):
        """
        Adds a DataFrame of error data to the 'error' table in the database.
 
        Parameters
        ----------
        - **data_df** : pd.DataFrame
            The DataFrame containing the error data to add.
        """
        already_exists = self.available_tables and 'error' in self.available_tables
        with self.conn:
            data_df.to_sql('error', self.conn, if_exists='append', index=False)
            if not already_exists:
                self._check_tables()
                if 'ID_PARAMETER' in data_df.columns:
                    self._create_index('error', 'ID_PARAMETER')
                if 'ID_CONDITION' in data_df.columns:
                    self._create_index('error', 'ID_CONDITION')
 
    def add_parameters(self, parameters_list: list):
        """
        Adds a list of parameters to the 'parameters' table in the database.
 
        Parameters
        ----------
        - **parameters_list** : list
            The list of parameters to add.
        """
        flattened_parameters = [self._flatten_dict(p) for p in parameters_list]
        parameters_df = pd.DataFrame(flattened_parameters)
        parameters_df = parameters_df.reset_index().rename(columns={'index': 'ID_PARAMETER'})
        with self.conn:
            parameters_df.to_sql('parameters', self.conn, if_exists='fail', index=False)
            self._check_tables()
            self._create_index('parameters', 'ID_PARAMETER')
 
    def add_data(self, table_name: str, data_df: pd.DataFrame):
        """
        Adds a DataFrame of data to a specified table in the database.
 
        Parameters
        ----------
        - **table_name** : str
            The name of the table to add data to.
        - **data_df** : pd.DataFrame
            The DataFrame containing the data to add.
        """
        already_exists = self.available_tables and table_name in self.available_tables
        with self.conn:
            data_df.to_sql(table_name, self.conn, if_exists='append', index=False)
            if not already_exists:
                self._check_tables()
 
                if 'ID_PARAMETER' in data_df.columns:
                    self._create_index(table_name, 'ID_PARAMETER')
                if 'ID_CONDITION' in data_df.columns:
                    self._create_index(table_name, 'ID_CONDITION')
 
 
    def _parse_json_conditions(self, filter_conditions):
        json_path_regex = re.compile(r"(\w+)\.(\w+(\.\w+)*)")
        def replace_json_path(match):
            table_alias = match.group(1)
            json_path = match.group(2).replace('.', '$')
            return f"{table_alias}.{json_path}"
        converted_conditions = json_path_regex.sub(replace_json_path, filter_conditions)
        return converted_conditions
 
    def get_parameters(self, filter_conditions: str = None, raw: bool = False) -> pd.DataFrame:
        """
        Retrieves the parameters based on filter conditions.
 
        Parameters
        ----------
        - **filter_conditions** : str, optional
            The SQL filter conditions to apply.
        - **raw** : bool, optional
            If True, returns the raw DataFrame without JSON parsing.
 
        Returns
        -------
        - **pd.DataFrame** : The DataFrame containing the retrieved records.
        """
        return self.get('parameters', filter_conditions)
 
    def get_conditions(self, filter_conditions: str = None) -> pd.DataFrame:
        """
        Retrieves records from the 'condition' table based on filter conditions.
 
        Parameters
        ----------
        - **filter_conditions** : str, optional
            The SQL filter conditions to apply.
 
        Returns
        -------
        - **pd.DataFrame** : The DataFrame containing the retrieved records.
        """
        return self.get('condition', filter_conditions, rawParameters=rawParameters)
 
 
    def _build_clauses(self, table: str, filter_conditions: str = None, alias_prefix: str = 'main_table'):
        """
        Build JOIN and WHERE clauses based on filter_conditions, similar to get().
        Returns: join_sql, where_sql, alias_dict
        """
        alias_dict = {table: alias_prefix}
        alias_count = 0
        join_clauses = []
        referenced = set()
 
        if filter_conditions:
            cond = self._parse_json_conditions(filter_conditions)
            for tbl in re.findall(r'(\w+)\.', cond):
                referenced.add(tbl)
            filter_conditions = cond
 
        main_cols = self._get_table_columns(table)
        # build joins for each referenced table
        for ref in referenced:
            if ref in self.available_tables and ref != table:
                alias = f"t_{ref}_{alias_count}"
                alias_count += 1
                alias_dict[ref] = alias
                ref_cols = self._get_table_columns(ref)
                if 'ID_PARAMETER' in main_cols and 'ID_PARAMETER' in ref_cols:
                    join_clauses.append(f"LEFT JOIN {ref} AS {alias} ON main_table.ID_PARAMETER = {alias}.ID_PARAMETER")
                elif 'ID_CONDITION' in main_cols and 'ID_CONDITION' in ref_cols:
                    join_clauses.append(f"LEFT JOIN {ref} AS {alias} ON main_table.ID_CONDITION = {alias}.ID_CONDITION")
                elif 'ID_PARAMETER' in ref_cols:
                    join_clauses.append(f"LEFT JOIN {ref} AS {alias} ON {alias}.ID_PARAMETER = main_table.ID_PARAMETER")
                elif 'ID_CONDITION' in ref_cols:
                    join_clauses.append(f"LEFT JOIN {ref} AS {alias} ON {alias}.ID_CONDITION = main_table.ID_CONDITION")
 
        # apply aliases in filter
        where_sql = ''
        if filter_conditions:
            for tbl, al in alias_dict.items():
                filter_conditions = re.sub(rf"\b{tbl}\.", f"{al}.", filter_conditions)
            where_sql = f"WHERE {filter_conditions}"
 
        join_sql = ' '.join(join_clauses)
        return join_sql, where_sql, alias_dict
 
    def get(self, table: str, filter_conditions: str = None, rawParameters: bool = False) -> pd.DataFrame:
        """
        Retrieves records from a table based on filter conditions.
 
        Parameters
        ----------
        - **table** : str
            The name of the table to retrieve records from.
        - **filter_conditions** : str, optional
            The SQL filter conditions to apply.
        - **rawParameters** : bool, optional
            If True, returns the raw DataFrame without JSON parsing.
 
        Returns
        -------
        - **pd.DataFrame** : The DataFrame containing the retrieved records.
        """
        if table not in self.available_tables:
            raise ValueError(f"Table '{table}' does not exist.")
        join_sql, where_sql, _ = self._build_clauses(table, filter_conditions)
        query = f"SELECT DISTINCT main_table.* FROM {table} AS main_table {join_sql} {where_sql}".strip()
        df = pd.read_sql(query, self.conn)
 
        if table == 'condition':
            for col in df.columns:
                if df[col].dtype == object and df[col].str.startswith(('[','{')).any():
                    df[col] = df[col].apply(lambda x: json.loads(x) if isinstance(x, str) else x)
        if table == 'parameters' and not rawParameters :
            return pd.DataFrame([self._unflat(row.to_dict()) for _, row in df.iterrows()])
        elif table == 'parameters' and rawParameters:
            return df.rename(columns=lambda x: x.replace('$', '.'))
        return df
    
    def iterOn(self, table: str, group, filter_conditions: str = None):
        # normalize group
        group_cols = (group,) if isinstance(group, str) else tuple(group)
        join_sql, where_sql, _ = self._build_clauses(table, filter_conditions)
        # distinct keys
        cols = ', '.join(f"main_table.\"{c}\"" for c in group_cols)
        distinct_sql = f"SELECT DISTINCT {cols} FROM {table} AS main_table {join_sql} {where_sql}".strip()
 
        all_cols = list(self._get_table_columns(table))
        cols_sql = ', '.join([f'"{c}"' for c in all_cols if c not in group_cols])
        where_clauses = ' AND '.join(f'{table}."{c}" = ?' for c in group_cols)
        iter_query = f"SELECT {cols_sql} FROM {table} WHERE {where_clauses}"
 
        for key in self.conn.execute(distinct_sql):
            df_part = pd.read_sql(iter_query, self.conn, params=list(key))
            yield key, df_part
 
    def getBestScore(self, errorName: str, n: int = 1, forEachCondition: bool = False) -> pd.DataFrame:
        """
        Retrieves the best scores for a given error.
 
        Parameters
        ----------
        - **errorName** : str
            The name of the error column to base the scores on.
        - **n** : int, optional
            The number of top scores to retrieve.
        - **forEachCondition** : bool, optional
            Whether to retrieve the best scores for each condition separately.
 
        Returns
        -------
        - **pd.DataFrame** : The DataFrame containing the best scores.
        """
        if forEachCondition:
            query = f'''WITH ranked_errors AS (
                            SELECT
                                ID_PARAMETER,
                                ID_CONDITION,
                                "{errorName}",
                                ROW_NUMBER() OVER (PARTITION BY ID_CONDITION ORDER BY "{errorName}" ASC, ID_PARAMETER ASC) AS rn
                            FROM error)
                        SELECT
                            ID_PARAMETER,
                            ID_CONDITION,
                            "{errorName}"
                        FROM
                            ranked_errors
                        WHERE
                            rn <= {n};'''
        else:
            query = f'''SELECT ID_PARAMETER, AVG("{errorName}") AS avg_error
                        FROM error
                        GROUP BY ID_PARAMETER
                        ORDER BY avg_error ASC
                        LIMIT {n};'''
        return pd.read_sql(query, self.conn)
 
    def getBest(self, table: str, errorName: str, n: int = 1, forEachCondition: bool = False) -> pd.DataFrame:
        """
        Retrieves the best records from a table based on a given error.
 
        Parameters
        ----------
        - **table** : str
            The name of the table to retrieve records from.
        - **errorName** : str
            The name of the error column to base the records on.
        - **n** : int, optional
            The number of top records to retrieve.
        - **forEachCondition** : bool, optional
            Whether to retrieve the best records for each condition separately.
 
        Returns
        -------
        - **pd.DataFrame** : The DataFrame containing the best records.
        """
        if forEachCondition:
            table_columns = self._get_table_columns(table)
            query = f'''WITH ranked_errors AS (
                            SELECT
                                ID_PARAMETER,
                                ID_CONDITION,
                                "{errorName}",
                                ROW_NUMBER() OVER (PARTITION BY ID_CONDITION ORDER BY "{errorName}" ASC) AS rn
                            FROM
                                error
                        )
                        SELECT *
                        FROM {table} t
                        WHERE {'(t.ID_PARAMETER, t.ID_CONDITION)' if 'ID_CONDITION' in table_columns else 't.ID_PARAMETER' } IN (
                            SELECT {'ID_PARAMETER, ID_CONDITION' if 'ID_CONDITION' in table_columns else 'ID_PARAMETER' }
                            FROM ranked_errors
                            WHERE rn <= {n}
                        );'''
        else:
            query = f'''SELECT *
                        FROM {table}
                        WHERE ID_PARAMETER IN (
                            SELECT ID_PARAMETER
                            FROM error
                            GROUP BY ID_PARAMETER
                            ORDER BY SUM("{errorName}") ASC
                            LIMIT {n}
                        );'''
        return pd.read_sql(query, self.conn)
 
    def getByID(self, table: str, ID_PARAMETER: int = None, ID_CONDITION: int = None) -> pd.DataFrame:
        """
        Retrieves records from a table based on ID_PARAMETER and/or ID_CONDITION.
 
        Parameters
        ----------
        - **table** : str
            The name of the table to retrieve records from.
        - **ID_PARAMETER** : int or list of int, optional
            The ID_PARAMETER to filter by.
        - **ID_CONDITION** : int or list of int, optional
            The ID_CONDITION to filter by.
 
        Returns
        -------
        - **pd.DataFrame** : The DataFrame containing the retrieved records.
        """
        whereP = ''
        whereC = ''
        if isinstance(ID_PARAMETER, list):
            whereP = f'ID_PARAMETER IN ({",".join(map(str, ID_PARAMETER))})'
        elif isinstance(ID_PARAMETER, int):
            whereP = f'ID_PARAMETER = {ID_PARAMETER}'
        elif ID_PARAMETER is not None:
            raise ValueError("ID_PARAMETER must be an integer or a list of integers.")
        if isinstance(ID_CONDITION, list):
            whereC = f'ID_CONDITION IN ({",".join(map(str, ID_CONDITION))})'
        elif isinstance(ID_CONDITION, int):
            whereC = f'ID_CONDITION = {ID_CONDITION}'
        elif ID_CONDITION is not None:
            raise ValueError("ID_CONDITION must be an integer or a list of integers.")
 
        where_clauses = [clause for clause in [whereP, whereC] if clause]
        where = ' AND '.join(where_clauses)
 
        query = f'''SELECT *
                    FROM {table}
                    WHERE {where};'''
        return pd.read_sql(query, self.conn)
 
    def getPareto(self, errorNameList: list, forEachCondition: bool = False) -> pd.DataFrame:
        """
        Retrieves the Pareto front for given errors.
 
        Parameters
        ----------
        - **errorNameList** : list
            The list of error column names to consider.
        - **forEachCondition** : bool, optional
            Whether to retrieve the Pareto front for each condition separately.
 
        Returns
        -------
        - **pd.DataFrame** : The DataFrame containing the Pareto front.
        """
        if forEachCondition:
            query = f'''WITH avg_errors AS (
                            SELECT
                                e.ID_PARAMETER,
                                e.ID_CONDITION,
                                {', '.join(['AVG(e."'+en+'") AS '+en for en in errorNameList])}
                            FROM
                                error e
                            GROUP BY
                                e.ID_PARAMETER, e.ID_CONDITION
                        )
                        SELECT t1.*
                        FROM avg_errors t1
                        LEFT JOIN avg_errors t2
                        ON t1.ID_CONDITION = t2.ID_CONDITION
                        AND t1.ID_PARAMETER != t2.ID_PARAMETER
                        AND {' AND '.join(['t1."'+en+'" >= t2."'+en+'"' for en in errorNameList])}
                        AND ({' OR '.join(['t1."'+en+'" > t2."'+en+'"' for en in errorNameList])})
                        WHERE t2.ID_PARAMETER IS NULL;'''
        else:
            query = f'''WITH avg_errors AS (
                            SELECT
                                e.ID_PARAMETER,
                                {', '.join(['AVG(e."'+en+'") AS '+en for en in errorNameList])}
                            FROM
                                error e
                            GROUP BY
                                e.ID_PARAMETER
                        )
                        SELECT t1.*
                        FROM avg_errors t1
                        LEFT JOIN avg_errors t2
                        ON t1.ID_PARAMETER != t2.ID_PARAMETER
                        AND {' AND '.join(['t1."'+en+'" >= t2."'+en+'"' for en in errorNameList])}
                        AND ({' OR '.join(['t1."'+en+'" > t2."'+en+'"' for en in errorNameList])})
                        WHERE t2.ID_PARAMETER IS NULL;'''
        
        return pd.read_sql(query, self.conn)
 
    def query(self, query: str) -> pd.DataFrame:
        """
        Executes a raw SQL query and returns the result as a DataFrame.
 
        Parameters
        ----------
        - **query** : str
            The SQL query to execute.
 
        Returns
        -------
        - **pd.DataFrame** : The DataFrame containing the query results.
        """
        df = pd.read_sql(query, self.conn)
        return df
 
    def __getitem__(self, table: str) -> pd.DataFrame:
        df = pd.read_sql(f"SELECT * FROM {table}", self.conn)
        df.columns = [c.replace('$', '.') for c in df.columns.values]
        return df
 
    def summarize(self) -> str:
        """
        Returns a summary of the database.
 
        Returns
        -------
        - **str** : The summary of the database.
        """
        summary = []
        summary.append("Database Summary:")
        for table in self.available_tables:
            df = pd.read_sql(f"SELECT * FROM {table} LIMIT 5", self.conn)
            count_query = f"SELECT COUNT(*) FROM {table}"
            count = self.conn.execute(count_query).fetchone()[0]
            summary.append(f"\nTable: {table}")
            summary.append(f"Number of rows: {count}")
            summary.append("Preview:")
            summary.append(df.head().to_string())
        return "\n".join(summary)
        
    def __str__(self) -> str:
        return self.summarize()
 
 
    def _check_tables(self):
        query = "SELECT name FROM sqlite_master WHERE type='table';"
        tables = self.conn.execute(query).fetchall()
        self.available_tables = {table[0] for table in tables}
 
    def _get_table_columns(self, table_name):
        if table_name in self._table_columns_cache:
            return self._table_columns_cache[table_name]
 
        # Sinon, on interroge SQLite et on stocke
        query = f"PRAGMA table_info({table_name})"
        cols = {col[1] for col in self.conn.execute(query).fetchall()}
        self._table_columns_cache[table_name] = cols
        return cols
 
    def _create_index(self, table_name, column_name):
        try:
            query = f"CREATE INDEX IF NOT EXISTS idx_{table_name}_{column_name} ON {table_name} ({column_name});"
            self.conn.execute(query)
        except sqlite3.Error as e:
            print(f"Failed to create index on {table_name}({column_name}): {e}")
 
    def _flatten_dict(self, d, parent_key='', sep='$'):
        items = []
        for k, v in d.items():
            new_key = f"{parent_key}{sep}{k}" if parent_key else k
            if isinstance(v, dict):
                items.extend(self._flatten_dict(v, new_key, sep=sep).items())
            elif isinstance(v, list):
                items.append((new_key, json.dumps(v)))  # Convert lists to JSON strings
            else:
                items.append((new_key, v))
        return dict(items)
 
    def _unflat(self, d, sep='$'):
        result_dict = {}
        for key, value in d.items():
            parts = key.split(sep)
            d = result_dict
            for part in parts[:-1]:
                if part not in d:
                    d[part] = {}
                d = d[part]
            # Convert back JSON strings to lists if needed
            try:
                if isinstance(value, str) and (value.startswith('[') or value.startswith('{')):
                    value = json.loads(value)
            except (TypeError, json.JSONDecodeError):
                pass
            d[parts[-1]] = value
        return result_dict