OriginDex-DataManager/database/db_controller.py

import sqlite3
import threading
import json

class DBController:
    def __init__(self, db_path=':memory:', max_connections=10):
        self.db_path = db_path
        self.lock = threading.Lock()
        self.conn = sqlite3.connect(db_path, check_same_thread=False)
        self.conn.row_factory = sqlite3.Row
        self.cursor = self.conn.cursor()
        self.init_database()

    def init_database(self):
        disk_conn = sqlite3.connect('pokemon_forms.db')
        disk_cursor = disk_conn.cursor()

        # Create tables in the file-based database
        self.create_pokemon_forms_table(disk_cursor)

        # Commit changes to the file-based database
        disk_conn.commit()

        # Copy the file-based database to the in-memory database
        disk_conn.backup(self.conn)

        # Close the file-based database connection
        disk_conn.close()

    def save_changes(self):
        with self.lock:
            # Count the number of records before backup for verification
            self.cursor.execute('SELECT COUNT(*) FROM pokemon_forms')
            count = self.cursor.fetchone()[0]
            print(f"Records in memory before backup: {count}")

            # Back up the master connection to disk
            disk_conn = sqlite3.connect('pokemon_forms.db')
            with disk_conn:
                self.conn.backup(disk_conn)
            disk_conn.close()

    def close(self):
        self.save_changes()
        self.conn.close()

    def create_pokemon_forms_table(self, cursor):
        cursor.execute('''
        CREATE TABLE IF NOT EXISTS pokemon_forms (
            PFIC TEXT PRIMARY KEY,
            data JSON NOT NULL
        )
        ''')

    def add_pokemon_form(self, pfic, name, form_name, national_dex, generation, sprite_url, gender_relevant):
        data = {
            "name": name,
            "form_name": form_name,
            "national_dex": national_dex,
            "generation": generation,
            "sprite_url": sprite_url,
            "is_baby_form": False,
            "storable_in_home": False,
            "gender_relevant": gender_relevant
        }

        with self.lock:
            self.cursor.execute('''
                INSERT OR REPLACE INTO pokemon_forms (PFIC, data) VALUES (?, ?)
            ''', (pfic, json.dumps(data)))
            self.conn.commit()
            print(f"Added: {pfic}, {name}")

    def craft_pokemon_json_query(self, fields_to_include, pfic = None):
        query = f"SELECT "
        extracts = []
        for field in fields_to_include:
            if field == "pfic":
                extracts.append("PFIC as pfic")
            else:
                extracts.append(f"JSON_EXTRACT(data, '$.{field}') AS {field}")
        query = query + ", ".join(extracts)
        query = query + " FROM pokemon_forms"

        if pfic is not None:
            query = query + f" WHERE PFIC = '{pfic}'"

        return query

    def get_pokemon_details(self, pfic):
        fields = [
            "name",
            "form_name",
            "national_dex",
            "generation",
            "is_baby_form",
            "storable_in_home",
        ]
        query = self.craft_pokemon_json_query(fields, pfic)
        self.cursor.execute(query)
        results = self.cursor.fetchone()
        return dict(results)
    
    def get_list_of_pokemon_forms(self):
        fields = [
            "pfic",
            "name",
            "form_name",
            "national_dex",
            "generation",
            "is_baby_form",
            "storable_in_home",
            "gender_relevant"
        ]

        query = self.craft_pokemon_json_query(fields)
        self.cursor.execute(query)
        results = self.cursor.fetchall()

        return [dict(row) for row in results]
    
    def update_home_status(self, pfic, status):
        self.update_pokemon_field(pfic, "storable_in_home", status)
        pass

    def update_pokemon_field(self, pfic, field_name, new_value):
        # Fetch the existing record
        self.cursor.execute('SELECT data FROM pokemon_forms WHERE PFIC = ?', (pfic,))
        result = self.cursor.fetchone()
        
        if result:
            # Load the JSON data and update the field
            data = json.loads(result[0])
            data[field_name] = new_value

            # Update the record with the modified JSON
            updated_data_str = json.dumps(data)
            self.cursor.execute('''
                UPDATE pokemon_forms
                SET data = ?
                WHERE PFIC = ?
            ''', (updated_data_str, pfic))
            self.conn.commit()