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()