noaa/init_db.py

import argparse
import logging
import os
import sqlite3

# Setup basic configuration for logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)


def initialize_db(db_path):
    # Check if the database file already exists
    db_exists = os.path.exists(db_path)
    if db_exists:
        logging.info(f"{db_path} exists")  # Log a message if the database exists
        return
    try:
        # Using 'with' to ensure that the connection is closed automatically
        with sqlite3.connect(db_path) as conn:
              configure_database(conn)
    except sqlite3.Error as e:
        logging.error(f"Database error: {e}")  # Log any SQLite errors that occur
    except Exception as e:
        logging.error(f"Exception in initialize_db: {e}")


def configure_database(conn):
    cursor = conn.cursor()
    # Setting the journal mode to WAL for better concurrency
    cursor.execute("PRAGMA journal_mode = WAL;")
    # Ensuring foreign key constraints are enforced for data integrity
    cursor.execute("PRAGMA foreign_keys = ON;")
    # disable legacy alter table behavior as it will cause problems during
    # migrations when tables are renamed as references would otherwise be retained
    # in some locations
    # https://www.sqlite.org/pragma.html#pragma_legacy_alter_table
    cursor.execute("PRAGMA legacy_alter_table=OFF")
    # when using the WAL, we do need to sync changes on every write. sqlite
    # recommends using 'normal' mode which is much faster
    # Setting synchronous to NORMAL for a balance between speed and reliability
    cursor.execute("PRAGMA synchronous = NORMAL;")
    # Increasing the cache size to reduce the number of disk I/O operations
    cursor.execute("PRAGMA cache_size = 20000;")
    # Enabling memory-mapped I/O for potentially faster file operations
    cursor.execute("PRAGMA mmap_size = 536870912;")
    # Setting a busy timeout to prevent immediate failures when the database is locked
    # setting the value very high allows for more 'concurrency'
    # without running into errors, but may result in slow api calls
    cursor.execute("PRAGMA busy_timeout = 60000;")
    # Setting locking mode to EXCLUSIVE can enhance performance for single-user scenarios
    # cursor.execute("PRAGMA locking_mode = EXCLUSIVE;")
    conn.commit()  # Commit all PRAGMA configurations

    logging.info("Set up database with multi-user optimizations.")


def batch_transact(db_path, operations):
    try:
        with sqlite3.connect(
            db_path
        ) as conn:  # Ensure that the connection is handled properly
            cursor = conn.cursor()
            # Start a transaction for batch operations
            cursor.execute("BEGIN TRANSACTION;")  
            for operation in operations:
                # Execute each SQL operation provided in the operations list
                cursor.execute(operation)  
            cursor.execute("COMMIT;")  # Commit all operations at once
    except sqlite3.Error as e:
        logging.error(f"Database error during batch transaction: {e}")
    except Exception as e:
        logging.error(f"Exception in batch_transact: {e}")


def maintenance(db_path):
    try:
        with sqlite3.connect(db_path) as conn:
            cursor = conn.cursor()
            # Optimize the database to maintain performance
            cursor.execute("PRAGMA optimize;")
            cursor.execute("VACUUM;")  # Reclaim space and defragment the database file
    except sqlite3.Error as e:
        logging.error(f"Database error during maintenance: {e}")
    except Exception as e:
        logging.error(f"Exception in maintenance: {e}")


def parse_args():
    parser = argparse.ArgumentParser(
        description="Initialize and manage an SQLite database."
    )
    parser.add_argument("db_path", type=str, help="Path to the SQLite database file.")
    args = parser.parse_args()
    return args


if __name__ == "__main__":
    args = parse_args()  # Parse the command-line arguments for the database path
    initialize_db(args.db_path)  # Use the parsed path to initialize the database