Source code for GUIBRUSHR.GUI.DataInterface.DBSQLite3

import traceback
from pathlib import Path
import sqlite3
from typing import Optional, List, Tuple, Any, Dict
from tkinter import messagebox

from GUIBRUSHR.GUI.WIDGET.MyFigure import MyFigure
from GUIBRUSHR.General_Constants.Classes.Instrument import Instrument
from GUIBRUSHR.General_Constants.FunctionsAndConstants.Constant_Variables import ConstantVariables


[docs] class DBSQLite3: """Database management class for GUIBRUSHR application. This class handles all database operations for the application, including: - Creation and management of database tables - Instrument management (insert, delete, query) - Retrieval run management (insert, delete, query) The database consists of three main tables: 1. retrieval_run: Stores information about retrieval runs 2. instruments: Stores information about available instruments 3. observatories: Stores information about observatory locations Attributes: DB (sqlite3.Connection): Database connection object cursor (sqlite3.Cursor): Database cursor object """
[docs] def __init__(self, path) -> None: """Initialize database connection. Args: path (Path): Base path for the application """ db_path = Path(path, "GUIBRUSHR", "Files", "DB", "atmo.db") self.DB = sqlite3.connect(str(db_path)) self.cursor = self.DB.cursor()
def __enter__(self) -> 'DBSQLite3': """Context manager entry point.""" return self def __exit__(self, exc_type, exc_val, exc_tb) -> None: """Context manager exit point.""" self.close_DB()
[docs] def create_db(self) -> None: """Create database tables if they don't exist. Creates three main tables: 1. retrieval_run: For storing retrieval run information 2. instruments: For storing instrument information 3. observatories: For storing observatory information """ self._create_retrieval_run_table() self._create_instruments_table() self._create_observatories_table()
def _create_retrieval_run_table(self) -> None: """Create or update the retrieval_run table. This method: 1. Checks if the table exists 2. Creates it if it doesn't exist 3. Adds any new columns if the table exists """ # Check if table exists self.cursor.execute( "SELECT name FROM sqlite_master WHERE type='table' AND name=?;", ("retrieval_run",) ) if self.cursor.fetchone(): # Table exists, check for new columns self._update_retrieval_run_columns() else: # Create new table self._create_new_retrieval_run_table() def _update_retrieval_run_columns(self) -> None: """Update existing retrieval_run table with new columns if needed.""" table_name = "retrieval_run" for i, column_to_check in enumerate(ConstantVariables.NEW_COLUMN_RETRIEVAL_DB): self.cursor.execute(f"PRAGMA table_info(retrieval_run);") columns = self.cursor.fetchall() column_exists = any(column_to_check == col[1] for col in columns) if not column_exists: self.cursor.execute( f"ALTER TABLE {table_name} ADD COLUMN {column_to_check} {ConstantVariables.NEW_COLUMN_RETRIEVAL_DB_TYPE[i]};" ) self.cursor.execute(f"UPDATE {table_name} SET {column_to_check} = '';") self.DB.commit() def _create_new_retrieval_run_table(self) -> None: """Create a new retrieval_run table with all required columns.""" self.cursor.execute(""" CREATE TABLE IF NOT EXISTS retrieval_run ( retrieval_date character varying(20) PRIMARY KEY NOT NULL, target_id character varying(20), nights character varying(60), chemistry character varying(20), t_p_profile character varying(30), scattering boolean DEFAULT false, ecc_opi boolean DEFAULT false, order_selection character varying(30), retrieval_model character varying(30), resolution character varying(10), tell_rm_method character varying(20), rad_mode character varying(20), instruments character varying(200), user_id character varying(50), bestpars_params character varying(1000), fixed_params character varying(1000), molecules character varying(500), sigma_priors character varying(500) ); """) def _create_instruments_table(self) -> None: """Create the instruments table if it doesn't exist.""" self.cursor.execute( "SELECT name FROM sqlite_master WHERE type='table' AND name=?;", ("instruments",) ) if not self.cursor.fetchone(): self.cursor.execute(""" CREATE TABLE IF NOT EXISTS instruments ( instrument character varying(40) PRIMARY KEY NOT NULL, resolution character varying(20), hwhm_km_s double precision, wl_min_μm double precision, wl_max_μm double precision ); """) else: try: # Get information about the columns in the instruments table self.cursor.execute("PRAGMA table_info(instruments);") columns = self.cursor.fetchall() # Check if the column name is already correct column_names = [column[1] for column in columns] # Column name is the second item in each tuple if "hwhm_km_s" in column_names: pass else: # Begin transaction self.cursor.execute("BEGIN TRANSACTION;") # Step 1: Create a temporary table with the new structure self.cursor.execute(""" CREATE TABLE instruments_new ( instrument character varying(40) PRIMARY KEY NOT NULL, resolution character varying(20), hwhm_km_s double precision, wl_min_μm double precision, wl_max_μm double precision ); """) # Step 2: Copy data from old table to new table self.cursor.execute(""" INSERT INTO instruments_new(instrument, resolution, hwhm_km_s, wl_min_μm, wl_max_μm) SELECT instrument, resolution, hfwh_km_s, wl_min_μm, wl_max_μm FROM instruments; """) # Step 3: Drop the old table self.cursor.execute("DROP TABLE instruments;") # Step 4: Rename the new table to the original name self.cursor.execute("ALTER TABLE instruments_new RENAME TO instruments;") # Commit the transaction self.DB.commit() print("Column successfully renamed from 'hfwh_km_s' to 'hwhm_km_s'") except Exception as e: # Rollback in case of error self.DB.rollback() print(f"An error occurred: {e}") def _create_observatories_table(self) -> None: """Create the observatories table if it doesn't exist.""" self.cursor.execute( "SELECT name FROM sqlite_master WHERE type='table' AND name=?;", ("observatories",) ) if not self.cursor.fetchone(): self.cursor.execute(""" CREATE TABLE IF NOT EXISTS observatories ( name character varying(100) PRIMARY KEY NOT NULL, latitude double precision, longitude double precision, altitude double precision ); """) self.DB.commit()
[docs] def insert_instrument( self, instrument: str, resolution: str, hwhm_km_s, wl_min, wl_max ) -> None: """Insert a new instrument into the database. Args: instrument (str): Instrument identifier resolution (str): Instrument resolution hwhm_km_s: Half-width at half maximum in km/s wl_min : Minimum wavelength in micrometers wl_max : Maximum wavelength in micrometers """ sql = """ INSERT INTO instruments (instrument, resolution, hwhm_km_s, wl_min_μm, wl_max_μm) VALUES (?, ?, ?, ?, ?) """ try: self.cursor.execute(sql, (instrument, resolution, hwhm_km_s, wl_min, wl_max)) self.DB.commit() print(f"Instrument {instrument} inserted successfully.") messagebox.showinfo("Success", "Instrument stored in DB") except Exception as e: MyFigure("Error in insert instrument", message=str(e) + "\n" + str(traceback.format_exc())) print(str(e) + "\n" + str(traceback.format_exc()))
[docs] def delete_instrument_by_key(self, instrument: str) -> bool: """Delete an instrument from the database. Args: instrument (str): Instrument identifier to delete Returns: bool: True if deletion was successful, False otherwise """ sql = "DELETE FROM instruments WHERE instrument = ?" try: self.cursor.execute(sql, (instrument,)) self.DB.commit() if self.cursor.rowcount > 0: print(f"Instrument {instrument} deleted successfully.") return True print(f"No instrument found with the key {instrument}.") return False except sqlite3.Error as e: print(f"SQL error: {e}") self.DB.rollback() return False
[docs] def get_instruments(self): """Get all instruments from the database. Returns: Tuple containing: - List of instrument identifiers - List of resolutions - List of HWHM values - List of minimum wavelengths - List of maximum wavelengths """ try: self.cursor.execute("SELECT * FROM instruments ORDER BY instrument;") records = self.cursor.fetchall() if not records: print("No Instruments found") return ["None"], None, None, None, None return tuple(map(list, zip(*records))) except sqlite3.Error as e: print(f"SQL error: {e}") return ["None"], None, None, None, None
[docs] def get_instrument_by_key(self, instrument: str): """Get instrument details by its identifier. Args: instrument (str): Instrument identifier to look up Returns: Instrument: Instrument object with the retrieved details, or None if not found """ if "None" in instrument: return None sql = "SELECT * FROM instruments WHERE instrument = ?" try: self.cursor.execute(sql, (instrument,)) record = self.cursor.fetchone() if record: # instrument, resolution, hwhm_km_s, wl_min, wl_max instrument_obj = Instrument(record[0], None, record[1], hwhm_km_s=record[2], wl_min=record[3], wl_max=record[4], nights=None, lr_data=None) return instrument_obj print(f"No record found for instrument {instrument}.") return None except sqlite3.Error as e: print(f"SQL error: {e}") return None
[docs] def insert_retrieval_into_DB( self, retrieval_date: str, target: str, nights: str, chemistry: str, t_p_profile: str, scattering: str, ecc_opi: str, order_selection: str, retrieval_model: str, resolution: str, tell_rm_method: str, rad_mode: str, instruments: str, user_id: str, bestpars_params: str, fixed_params: str, molecules: str, sigma_priors: str ) -> bool: """Insert a new retrieval run into the database. Args: retrieval_date (str): Date of the retrieval run target (str): Target identifier nights (str): Observation nights chemistry (str): Chemistry type t_p_profile (str): Temperature-pressure profile scattering (str): Scattering flag ecc_opi (str): Eccentricity/obliquity flag order_selection (str): Order selection method retrieval_model (str): Retrieval model used resolution (str): Resolution type tell_rm_method (str): Telluric removal method rad_mode (str): Radiation mode instruments (str): Instruments used user_id (str): User identifier bestpars_params (str): Best parameters fixed_params (str): Fixed parameters molecules (str): Molecules involved sigma_priors (str): Sigma priors Returns: bool: True if insertion was successful, False otherwise """ # Prepare boolean values scattering = "TRUE" if scattering else "FALSE" ecc_opi = "TRUE" if ecc_opi else "FALSE" # Prepare resolution value resolution_H = "H" if "High" in resolution else "" resolution_L = "L" if "Low" in resolution else "" resolution = resolution_H + resolution_L sql = """ INSERT INTO retrieval_run ( retrieval_date, target_id, nights, chemistry, t_p_profile, scattering, ecc_opi, order_selection, retrieval_model, resolution, tell_rm_method, rad_mode, instruments, user_id, bestpars_params, fixed_params, molecules, sigma_priors ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """ values = ( retrieval_date, target, nights, chemistry, t_p_profile, scattering, ecc_opi, order_selection, retrieval_model, resolution, tell_rm_method, rad_mode, instruments, user_id, bestpars_params, fixed_params, molecules, sigma_priors ) try: self.cursor.execute(sql, values) self.DB.commit() print("Record inserted successfully!") return True except sqlite3.Error as e: print(f"Error occurred: {e}") self.DB.rollback() return False
[docs] def select_retrieval_by_target( self, user_idf: str, target: str, chemistry: str = "None", t_p_profile: str = "None", scattering: str = "None", ecc_opi: str = "None", resolution: str = "None", retrieval_model: str = "None", rad_mode: str = "None", instruments: str = "None" ) -> Tuple[List[Any], ...]: """Select retrieval runs by target with optional filters. Args: user_idf (str): User identifier target (str): Target identifier chemistry (str, optional): Chemistry filter t_p_profile (str, optional): Temperature-pressure profile filter scattering (str, optional): Scattering filter ecc_opi (str, optional): Eccentricity/obliquity filter resolution (str, optional): Resolution filter retrieval_model (str, optional): Retrieval model filter rad_mode (str, optional): Radiation mode filter instruments (str, optional): Instruments filter Returns: Tuple of lists containing retrieval run details """ # Prepare filter values filters = self._prepare_retrieval_filters( resolution, instruments, scattering, ecc_opi ) # Build and execute query sql, values = self._build_retrieval_query( user_idf, target, chemistry, t_p_profile, filters["scattering"], filters["ecc_opi"], filters["resolution"], retrieval_model, rad_mode, filters["instruments"] ) try: self.cursor.execute(sql, values) records = self.cursor.fetchall() if not records: print(f"No records found for user {user_idf} and target {target}") return tuple([None] * 18) return tuple(map(list, zip(*records))) except sqlite3.Error as e: print(f"Error occurred: {e}") return tuple([None] * 18)
def _prepare_retrieval_filters( self, resolution: str, instruments: str, scattering: str, ecc_opi: str ) -> Dict[str, str]: """Prepare filter values for retrieval query. Args: resolution (str): Resolution filter instruments (str): Instruments filter scattering (str): Scattering filter ecc_opi (str): Eccentricity/obliquity filter Returns: Dict containing prepared filter values """ filters = { "resolution": resolution, "instruments": instruments, "scattering": scattering, "ecc_opi": ecc_opi } if resolution != "None": resolution_H = "H" if "High" in resolution else "" resolution_L = "L" if "Low" in resolution else "" filters["resolution"] = resolution_H + resolution_L if "None" in instruments: filters["instruments"] = "None" if scattering != "None": filters["scattering"] = "TRUE" if scattering == "True" else "FALSE" if ecc_opi != "None": filters["ecc_opi"] = "TRUE" if ecc_opi == "True" else "FALSE" return filters def _build_retrieval_query( self, user_idf: str, target: str, chemistry: str, t_p_profile: str, scattering: str, ecc_opi: str, resolution: str, retrieval_model: str, rad_mode: str, instruments: str ) -> Tuple[str, List[Any]]: """Build SQL query for retrieval selection. Args: user_idf (str): User identifier target (str): Target identifier chemistry (str): Chemistry filter t_p_profile (str): Temperature-pressure profile filter scattering (str): Scattering filter ecc_opi (str): Eccentricity/obliquity filter resolution (str): Resolution filter retrieval_model (str): Retrieval model filter rad_mode (str): Radiation mode filter instruments (str): Instruments filter Returns: Tuple containing SQL query and parameter values """ sql = """ SELECT retrieval_date, target_id, nights, chemistry, t_p_profile, scattering, ecc_opi, order_selection, retrieval_model, resolution, tell_rm_method, rad_mode, instruments, user_id, bestpars_params, fixed_params, molecules, sigma_priors FROM retrieval_run WHERE user_id = ? AND target_id = ? """ values = [user_idf, target] filters = { "chemistry": chemistry, "t_p_profile": t_p_profile, "scattering": scattering, "ecc_opi": ecc_opi, "resolution": resolution, "retrieval_model": retrieval_model, "rad_mode": rad_mode, "instruments": instruments } for field, value in filters.items(): if value != "None": if field == ConstantVariables.COLUMN_WITH_MULTIPLE_VALUES: sql += f"AND (('_' || {field} || '_' LIKE '%_{value}_%')" else: sql += f" AND ({field} = ?" values.append(value) if field in ConstantVariables.NEW_COLUMN_RETRIEVAL_DB: sql += f" OR {field} = '')" else: sql += ")" sql += " ORDER BY retrieval_date DESC;" return sql, values
[docs] def delete_retrieval_by_date(self, retrieval_date: str) -> bool: """Delete a retrieval run by its date. Args: retrieval_date (str): Date of the retrieval run to delete Returns: bool: True if deletion was successful, False otherwise """ sql = "DELETE FROM retrieval_run WHERE retrieval_date = ?" try: self.cursor.execute(sql, (retrieval_date,)) self.DB.commit() rows_deleted = self.cursor.rowcount if rows_deleted > 0: print(f"Deleted {rows_deleted} record(s) for date {retrieval_date}.") return True print(f"No records found for date {retrieval_date} to delete.") return False except sqlite3.Error as e: print(f"Error occurred: {e}") self.DB.rollback() return False
[docs] def insert_observatory( self, name: str, latitude: float, longitude: float, altitude: float ) -> None: """Insert a new observatory into the database. Args: name (str): Observatory name latitude (float): Observatory latitude in degrees longitude (float): Observatory longitude in degrees altitude (float): Observatory altitude in meters """ sql = """ INSERT INTO observatories (name, latitude, longitude, altitude) VALUES (?, ?, ?, ?) """ try: self.cursor.execute(sql, (name, latitude, longitude, altitude)) self.DB.commit() print(f"Observatory {name} inserted successfully.") messagebox.showinfo("Success", "Observatory stored in DB") except Exception as e: MyFigure("Error in insert observatory", message=str(e) + "\n" + str(traceback.format_exc())) print(str(e) + "\n" + str(traceback.format_exc()))
[docs] def delete_observatory_by_key(self, name: str) -> bool: """Delete an observatory from the database. Args: name (str): Observatory name to delete Returns: bool: True if deletion was successful, False otherwise """ sql = "DELETE FROM observatories WHERE name = ?" try: self.cursor.execute(sql, (name,)) self.DB.commit() if self.cursor.rowcount > 0: print(f"Observatory {name} deleted successfully.") return True print(f"No observatory found with the name {name}.") return False except sqlite3.Error as e: print(f"SQL error: {e}") self.DB.rollback() return False
[docs] def get_observatories(self): """Get all observatories from the database. Returns: Tuple containing: - List of observatory names - List of latitudes - List of longitudes - List of altitudes """ try: self.cursor.execute("SELECT * FROM observatories ORDER BY name;") records = self.cursor.fetchall() if not records: print("No Observatories found") return ["None"], None, None, None return tuple(map(list, zip(*records))) except sqlite3.Error as e: print(f"SQL error: {e}") return ["None"], None, None, None
[docs] def get_observatory_by_key(self, name: str): """Get observatory details by its name. Args: name (str): Observatory name to look up Returns: Tuple containing observatory details: - Observatory name - Latitude - Longitude - Altitude """ sql = "SELECT * FROM observatories WHERE name = ?" try: self.cursor.execute(sql, (name,)) record = self.cursor.fetchone() if record: return record print(f"No record found for observatory {name}.") return None except sqlite3.Error as e: print(f"SQL error: {e}") return None
[docs] def close_DB(self) -> None: """Close database connection and cursor.""" self.cursor.close() self.DB.close()