From 25a72b30025863c1f9ce61055a971c066d3a7334 Mon Sep 17 00:00:00 2001 From: VD15 Date: Mon, 26 May 2025 20:49:21 +0100 Subject: [PATCH] Add database migration system --- db.py | 64 --------------------- migrations/0000_setup.sql | 28 +++++++++ migrations/0001_fix_notif_id.sql | 1 + setup_db.py | 99 ++++++++++++++++++++++++++++++++ 4 files changed, 128 insertions(+), 64 deletions(-) delete mode 100644 db.py create mode 100644 migrations/0000_setup.sql create mode 100644 migrations/0001_fix_notif_id.sql create mode 100644 setup_db.py diff --git a/db.py b/db.py deleted file mode 100644 index ffcfe2e..0000000 --- a/db.py +++ /dev/null @@ -1,64 +0,0 @@ -import sqlite3 - -# Connect to SQLite database (or create it if it doesn't exist) -conn = sqlite3.connect('gacha_game.db') -cursor = conn.cursor() - -# Create tables -cursor.execute(''' -CREATE TABLE IF NOT EXISTS users ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - username TEXT UNIQUE NOT NULL, - has_rolled BOOLEAN NOT NULL DEFAULT 0 -) -''') - -cursor.execute(''' -CREATE TABLE IF NOT EXISTS characters ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - name TEXT NOT NULL, - rarity INTEGER NOT NULL, - weight REAL NOT NULL, - file_id TEXT NOT NULL -) -''') - -cursor.execute(''' -CREATE TABLE IF NOT EXISTS pulls ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - user_id INTEGER, - character_id INTEGER, - timestamp DATETIME DEFAULT CURRENT_TIMESTAMP, - FOREIGN KEY (user_id) REFERENCES users(id), - FOREIGN KEY (character_id) REFERENCES characters(id) -) -''') - -cursor.execute(""" - CREATE TABLE IF NOT EXISTS config ( - key TEXT PRIMARY KEY, - value TEXT - ) - """) - -# Initialize essential config key -cursor.execute('INSERT INTO config VALUES ("last_seen_notif_id", 0)') - -""" # Insert example characters into the database if they don't already exist -characters = [ - ('Murakami-san', 1, 0.35), - ('Mastodon-kun', 2, 0.25), - ('Pleroma-tan', 3, 0.2), - ('Misskey-tan', 4, 0.15), - ('Syuilo-mama', 5, 0.05) -] - - -cursor.executemany(''' -INSERT OR IGNORE INTO characters (name, rarity, weight) VALUES (?, ?, ?) -''', characters) -""" - -# Commit changes and close -conn.commit() -conn.close() diff --git a/migrations/0000_setup.sql b/migrations/0000_setup.sql new file mode 100644 index 0000000..4dcbc64 --- /dev/null +++ b/migrations/0000_setup.sql @@ -0,0 +1,28 @@ +CREATE TABLE IF NOT EXISTS users ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + username TEXT UNIQUE NOT NULL, + has_rolled BOOLEAN NOT NULL DEFAULT 0 +); + +CREATE TABLE IF NOT EXISTS characters ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + name TEXT NOT NULL, + rarity INTEGER NOT NULL, + weight REAL NOT NULL, + file_id TEXT NOT NULL +); + +CREATE TABLE IF NOT EXISTS pulls ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + user_id INTEGER, + character_id INTEGER, + timestamp DATETIME DEFAULT CURRENT_TIMESTAMP, + FOREIGN KEY (user_id) REFERENCES users(id), + FOREIGN KEY (character_id) REFERENCES characters(id) +); + +CREATE TABLE IF NOT EXISTS config ( + key TEXT PRIMARY KEY, + value TEXT +); +INSERT OR IGNORE INTO config VALUES ("schema_version", 0); diff --git a/migrations/0001_fix_notif_id.sql b/migrations/0001_fix_notif_id.sql new file mode 100644 index 0000000..3494d6e --- /dev/null +++ b/migrations/0001_fix_notif_id.sql @@ -0,0 +1 @@ +INSERT OR IGNORE INTO config VALUES ("last_seen_notif_id", 0); diff --git a/setup_db.py b/setup_db.py new file mode 100644 index 0000000..e247386 --- /dev/null +++ b/setup_db.py @@ -0,0 +1,99 @@ +import sqlite3 +import os +import argparse +from configparser import ConfigParser +from typing import List, Tuple + +class DBNotFoundError(Exception): + pass + +class InvalidMigrationError(Exception): + pass + +def get_migrations() -> List[Tuple[int, str]] | InvalidMigrationError: + '''Returns a list of migration files in numeric order.''' + # Store transaction id and filename separately + sql_files: List[Tuple[int, str]] = [] + migrations_dir = 'migrations' + + for filename in os.listdir(migrations_dir): + joined_path = os.path.join(migrations_dir, filename) + + # Ignore anything that isn't a .sql file + if not (os.path.isfile(joined_path) and filename.endswith('.sql')): + print(f'{filename} is not a .sql file, ignoring...') + continue + + parts = filename.split('_', 1) + + # Invalid filename format + if len(parts) < 2 or not parts[0].isdigit(): + raise InvalidMigrationError(f'Invalid migration file: {filename}') + + sql_files.append((int(parts[0]), joined_path)) + + # Get sorted list of files by migration number + sql_files.sort(key=lambda x: x[0]) + return sql_files + +def perform_migration(cursor: sqlite3.Cursor, migration: tuple[int, str]) -> None: + '''Performs a migration on the DB''' + print(f'Performing migration {migration[1]}...') + + # Open and execute the sql script + with open(migration[1], encoding='utf-8') as file: + script = file.read() + cursor.executescript(script) + # Update the schema version + cursor.execute('UPDATE config SET value = ? WHERE key = "schema_version"', (migration[0],)) + +def get_db_path() -> str | DBNotFoundError: + '''Gets the DB path from config.ini''' + config = ConfigParser() + config.read('config.ini') + db_path = config['application']['DatabaseLocation'] + if not db_path: + raise DBNotFoundError + return db_path + +def get_current_migration(cursor: sqlite3.Cursor) -> int: + '''Gets the current schema version of the database''' + try: + cursor.execute('SELECT value FROM config WHERE key = ?', ('schema_version',)) + version = cursor.fetchone() + return -1 if not version else int(version[0]) + except sqlite3.Error: + print('Error getting schema version') + # Database has not been initialized yet + return -1 + +def main(): + '''Does the thing''' + # Connect to the DB + db_path = get_db_path() + conn = sqlite3.connect(db_path, autocommit=False) + conn.row_factory = sqlite3.Row + cursor = conn.cursor() + + # Obtain list of migrations to run + migrations = get_migrations() + # Determine schema version + current_migration = get_current_migration(cursor) + print(f'Current schema version: {current_migration}') + + # Run any migrations newer than current schema + for migration in migrations: + if migration[0] <= current_migration: + print(f'Migration already up: {migration[1]}') + continue + try: + perform_migration(cursor, migration) + conn.commit() + except Exception as ex: + print(f'An error occurred while applying migration: {ex}, aborting...') + conn.rollback() + break + conn.close() + +if __name__ == '__main__': + main()