SQLite: use primary keys and "INSERT OR REPLACE"

Also, split code into functions

Author: claudiodsf

CHANGELOG.md

@@ -7,6 +7,8 @@ Copyright (c) 2011-2023 Claudio Satriano <[email protected]>
 ## unreleased
 
 This release requires at least Python 3.7.
+Warning: the SQLite database used by this version is not compatible with
+previous versions.
 
 ### Input/output
 

sourcespec/ssp_sqlite_output.py

@@ -15,62 +15,96 @@
 from sourcespec.ssp_setup import ssp_exit
 from sourcespec._version import get_versions
 logger = logging.getLogger(__name__.split('.')[-1])
+# Current supported DB version
+DB_VERSION = 2
 
 
-def _log_db_write_error(db_err, db_file):
-    logger.error(f'Unable to insert values: {db_err}')
-    logger.info('Maybe your sqlite database has an old format.')
-    logger.info('Try to remove or rename your database file.')
-    logger.info(f'(Current database file: {db_file})')
-    ssp_exit(1)
+def _db_file_exists(db_file):
+    """
+    Check if SQLite database file exists.
 
+    :param db_file: SQLite database file
+    :type db_file: str
+    :return: True if file exists, False otherwise
+    :rtype: bool
+    """
+    return os.path.isfile(db_file)
 
-def write_sqlite(config, sspec_output):
-    try:
-        database_file = config.database_file
-    except KeyError:
-        database_file = None
-    if not database_file:
-        return
 
-    event = config.event
-    evid = event.event_id
-    runid = config.options.run_id
-
-    # Current supported DB version
-    DB_VERSION = 2
+def _open_sqlite_db(db_file):
+    """
+    Open SQLite database.
 
-    # only check version if database_file exists
-    check_version = os.path.isfile(database_file)
-
-    # Open SQLite database
+    :param db_file: SQLite database file
+    :type db_file: str
+    :return: SQLite connection and cursor
+    :rtype: tuple
+    """
     try:
-        conn = sqlite3.connect(database_file, timeout=60)
+        conn = sqlite3.connect(db_file, timeout=60)
     except Exception as msg:
         logger.error(msg)
         logger.info(
-            f'Please check whether "{database_file}" is a valid SQLite file.')
+            f'Please check whether "{db_file}" is a valid SQLite file.')
         ssp_exit(1)
-    c = conn.cursor()
-
-    if check_version:
-        # Get current DB version
-        db_version = c.execute('PRAGMA user_version').fetchone()[0]
-        if db_version < DB_VERSION:
-            logger.error(
-                f'"{database_file}" has an old database version: '
-                f'"{db_version}" Current supported version is "{DB_VERSION}".'
-            )
-            logger.info(
-                'Remove or rename your old database file, '
-                'so that a new one can be created.'
-            )
-            ssp_exit(1)
-    else:
-        # Set the DB version
-        c.execute('PRAGMA user_version = {v:d}'.format(v=DB_VERSION))
+    return conn, conn.cursor()
+
+
+def _check_db_version(cursor, db_file):
+    """
+    Check database version.
+
+    :param cursor: SQLite cursor
+    :type cursor: sqlite3.Cursor
+    :param db_file: SQLite database file
+    :type db_file: str
+    """
+    db_version = cursor.execute('PRAGMA user_version').fetchone()[0]
+    if db_version < DB_VERSION:
+        logger.error(
+            f'"{db_file}" has an old database version: '
+            f'"{db_version}" Current supported version is "{DB_VERSION}".'
+        )
+        logger.info(
+            'Remove or rename your old database file, '
+            'so that a new one can be created.'
+        )
+        ssp_exit(1)
+
+
+def _set_db_version(cursor):
+    """
+    Set database version.
+
+    :param cursor: SQLite cursor
+    :type cursor: sqlite3.Cursor
+    """
+    cursor.execute('PRAGMA user_version = {v:d}'.format(v=DB_VERSION))
+
+
+def _log_db_write_error(db_err, db_file):
+    """
+    Log database write error.
 
-    # Create Station table
+    :param db_err: database error
+    :type db_err: Exception
+    :param db_file: SQLite database file
+    :type db_file: str
+    """
+    logger.error(f'Unable to insert values: {db_err}')
+    logger.info('Maybe your sqlite database has an old format.')
+    logger.info('Try to remove or rename your database file.')
+    logger.info(f'(Current database file: {db_file})')
+    ssp_exit(1)
+
+
+def _create_stations_table(cursor, db_file):
+    """
+    Create Stations table.
+
+    :param cursor: SQLite cursor
+    :type cursor: sqlite3.Cursor
+    """
     sql_create_stations_table = """CREATE TABLE IF NOT EXISTS Stations (
         stid TEXT,
         evid TEXT,
@@ -106,24 +140,36 @@ def write_sqlite(config, sspec_output):
         Er REAL,
         Er_is_outlier INT,
         dist REAL,
-        azimuth REAL
+        azimuth REAL,
+        PRIMARY KEY (stid, evid, runid)
     );"""
-    c.execute(sql_create_stations_table)
-    # Write station source parameters to database
-    nobs = 0
+    try:
+        cursor.execute(sql_create_stations_table)
+    except Exception as db_err:
+        _log_db_write_error(db_err, db_file)
+
+
+def _write_stations_table(cursor, db_file, sspec_output, config):
+    """
+    Write station source parameters to database.
+
+    :param cursor: SQLite cursor
+    :type cursor: sqlite3.Cursor
+    :param db_file: SQLite database file
+    :type db_file: str
+    :param sspec_output: sspec output object
+    :type sspec_output: ssp_data_types.SourceSpecOutput
+    :param config: sspec configuration object
+    :type config: config.Config
+    """
+    event = config.event
+    evid = event.event_id
+    runid = config.options.run_id
     stationpar = sspec_output.station_parameters
-    sql_delete_from_stations =\
-        'DELETE FROM Stations WHERE stid=? AND evid=? AND runid=?;'
+    nobs = 0
     for statId in sorted(stationpar.keys()):
         nobs += 1
         par = stationpar[statId]
-        # Remove existing line, if present
-        t = (statId, evid, runid)
-        try:
-            c.execute(sql_delete_from_stations, t)
-        except Exception as msg:
-            _log_db_write_error(msg, database_file)
-            ssp_exit(1)
         # Insert new line
         t = (
             statId, evid, runid,
@@ -148,16 +194,23 @@ def write_sqlite(config, sspec_output):
         )
         # Create a string like ?,?,?,?
         values = ','.join('?' * len(t))
-        sql_insert_into_stations = f'INSERT INTO Stations VALUES({values});'
+        sql_insert_into_stations =\
+            f'INSERT OR REPLACE INTO Stations VALUES({values});'
         try:
-            c.execute(sql_insert_into_stations, t)
+            cursor.execute(sql_insert_into_stations, t)
         except Exception as msg:
-            _log_db_write_error(msg, database_file)
+            _log_db_write_error(msg, db_file)
             ssp_exit(1)
-    # Commit changes
-    conn.commit()
+    return nobs
+
 
-    # Create Event table
+def _create_events_table(cursor, db_file):
+    """
+    Create Events table.
+
+    :param cursor: SQLite cursor
+    :type cursor: sqlite3.Cursor
+    """
     sql_create_events_table = """CREATE TABLE IF NOT EXISTS Events (
         /* Event info */
             evid TEXT,
@@ -263,9 +316,33 @@ def write_sqlite(config, sspec_output):
             author_email TEXT,
             agency_full_name TEXT,
             agency_short_name TEXT,
-            agency_url TEXT
+            agency_url TEXT,
+            PRIMARY KEY (evid, runid)
     );"""
-    c.execute(sql_create_events_table)
+    try:
+        cursor.execute(sql_create_events_table)
+    except Exception as db_err:
+        _log_db_write_error(db_err, db_file)
+
+
+def _write_events_table(cursor, db_file, sspec_output, config, nobs):
+    """
+    Write Events table.
+
+    :param cursor: SQLite cursor
+    :type cursor: sqlite3.Cursor
+    :param db_file: SQLite database file
+    :type db_file: str
+    :param sspec_output: SSP output object
+    :type sspec_output: ssp_data_types.SourceSpecOutput
+    :param config: SSP configuration object
+    :type config: config.Config
+    :param nobs: Number of observations
+    :type nobs: int
+    """
+    event = config.event
+    evid = event.event_id
+    runid = config.options.run_id
     means = sspec_output.mean_values()
     mean_errors = sspec_output.mean_uncertainties()
     wmeans = sspec_output.weighted_mean_values()
@@ -274,14 +351,6 @@ def write_sqlite(config, sspec_output):
     percentile_errors = sspec_output.percentiles_uncertainties()
     run_completed = f'{config.end_of_run} {config.end_of_run_tz}'
     ssp_version = get_versions()['version']
-    # Remove event from Event table, if present
-    t = (evid, runid)
-    sql_delete_from_events = 'DELETE FROM Events WHERE evid=? AND runid=?;'
-    try:
-        c.execute(sql_delete_from_events, t)
-    except Exception as msg:
-        _log_db_write_error(msg, database_file)
-        ssp_exit(1)
     ev_lon = event.hypocenter.longitude.value_in_deg
     ev_lat = event.hypocenter.latitude.value_in_deg
     ev_depth = event.hypocenter.depth.value_in_km
@@ -370,13 +439,45 @@ def write_sqlite(config, sspec_output):
     )
     # Create a string like ?,?,?,?
     values = ','.join('?' * len(t))
-    sql_insert_into_events = f'INSERT INTO Events VALUES({values});'
+    sql_insert_into_events = f'INSERT OR REPLACE INTO Events VALUES({values});'
     try:
-        c.execute(sql_insert_into_events, t)
+        cursor.execute(sql_insert_into_events, t)
     except Exception as msg:
-        _log_db_write_error(msg, database_file)
+        _log_db_write_error(msg, db_file)
         ssp_exit(1)
+
+
+def write_sqlite(config, sspec_output):
+    """
+    Write SSP output to SQLite database.
+
+    :param config: SSP configuration object
+    :type config: config.Config
+    :param sspec_output: SSP output object
+    :type sspec_output: ssp_data_types.SourceSpecOutput
+    """
+    db_file = config.get('database_file', None)
+    if not db_file:
+        return
+
+    db_file_exists = _db_file_exists(db_file)
+    conn, cursor = _open_sqlite_db(db_file)
+    if db_file_exists:
+        _check_db_version(cursor, db_file)
+    else:
+        _set_db_version(cursor)
+
+    # Create Stations table
+    _create_stations_table(cursor, db_file)
+    # Write station source parameters to database
+    nobs = _write_stations_table(cursor, db_file, sspec_output, config)
+    # Commit changes
+    conn.commit()
+    # Create Events table
+    _create_events_table(cursor, db_file)
+    # Write event source parameters to database
+    _write_events_table(cursor, db_file, sspec_output, config, nobs)
     # Commit changes and close database
     conn.commit()
     conn.close()
-    logger.info(f'Output written to SQLite database: {database_file}')
+    logger.info(f'Output written to SQLite database: {db_file}')