AirlockTools/Development/WIP/Server/datastore.py

"""
datastore.py

Utility module for saving/loading objects to/from JSON and inserting/updating/reading them in SQLite.

Supports any class with:
- a `to_dict()` method
- a constructor accepting `**kwargs`
- optionally, a `from_dict()` method

Author: Brandon Wickline, James Brotosky
License: GNU Affero General Public License v3.0
"""

import json
import sqlite3
import logging
from typing import Type, List, TypeVar

# Setup logger
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

# Generic type variable for typing
T = TypeVar('T')

# -------------------------------------------------------------------
# JSON SAVE
# -------------------------------------------------------------------
def save_to_json(obj_list: List[T], filepath: str) -> None:
    """
    Save a list of objects to a JSON file using their `to_dict()` method.

    Args:
        obj_list (List[T]): List of objects to serialize.
        filepath (str): Path to the output JSON file.

    Example:
        save_to_json(policies, "policies.json")
    """
    try:
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump([obj.to_dict() for obj in obj_list], f, ensure_ascii=False, indent=4) # pyright: ignore[reportAttributeAccessIssue]
        logger.info(f"Saved {len(obj_list)} objects to {filepath}")
    except Exception as e:
        logger.error(f"Failed to save to {filepath}: {e}")

# -------------------------------------------------------------------
# JSON LOAD
# -------------------------------------------------------------------
def load_from_json(cls: Type[T], filepath: str) -> List[T]:
    """
    Load a list of objects from a JSON file and instantiate them using the class constructor
    or a `from_dict()` method if available.

    Args:
        cls (Type[T]): Class type to instantiate.
        filepath (str): Path to the input JSON file.

    Returns:
        List[T]: List of instantiated objects.

    Example:
        loaded_agents = load_from_json(Agent, "agents.json")
    """
    try:
        with open(filepath, 'r', encoding='utf-8') as f:
            data_list = json.load(f)
        logger.info(f"Loaded {len(data_list)} records from {filepath}")
        if hasattr(cls, "from_dict"):
            return [cls.from_dict(data) for data in data_list] # pyright: ignore[reportAttributeAccessIssue]
        return [cls(**data) for data in data_list]
    except Exception as e:
        logger.error(f"Failed to load from {filepath}: {e}")
        return []

# -------------------------------------------------------------------
# SQLITE INSERT OR UPDATE
# -------------------------------------------------------------------
def insert_or_update_objects_to_sqlite(obj_list: List[T], table_name: str, db_path: str, primary_key: str) -> None:
    """
    Insert or update a list of objects into a SQLite table.
    Uses `ON CONFLICT(primary_key) DO UPDATE` for upsert behavior.

    Args:
        obj_list (List[T]): List of objects with `to_dict()` method.
        table_name (str): Name of the SQLite table.
        db_path (str): Path to the SQLite database file.
        primary_key (str): Field name to use as the primary key.

    Example:
        insert_or_update_objects_to_sqlite(hashes, "hashes", "data_store.db", primary_key="sha256")
    """
    if not obj_list:
        logger.warning("No objects to insert or update into SQLite.")
        return

    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()

        sample_dict = obj_list[0].to_dict() # pyright: ignore[reportAttributeAccessIssue]
        columns = ', '.join(sample_dict.keys())
        placeholders = ', '.join(['?'] * len(sample_dict))
        update_clause = ', '.join([f"{key}=excluded.{key}" for key in sample_dict.keys() if key != primary_key])

        # Create table if it doesn't exist
        create_stmt = f"""
        CREATE TABLE IF NOT EXISTS {table_name} (
            {', '.join([f"{key} TEXT" for key in sample_dict.keys()])},
            PRIMARY KEY ({primary_key})
        )
        """
        cursor.execute(create_stmt)

        # Insert or update each object
        for obj in obj_list:
            values = tuple(str(v) if v is not None else "" for v in obj.to_dict().values()) # pyright: ignore[reportAttributeAccessIssue]
            insert_stmt = f"""
            INSERT INTO {table_name} ({columns}) VALUES ({placeholders})
            ON CONFLICT({primary_key}) DO UPDATE SET {update_clause}
            """
            cursor.execute(insert_stmt, values)

        conn.commit()
        conn.close()
        logger.info(f"Inserted or updated {len(obj_list)} records into {table_name} table in {db_path}")
    except Exception as e:
        logger.error(f"Failed to insert or update into SQLite: {e}")

# -------------------------------------------------------------------
# SQLITE READ
# -------------------------------------------------------------------
def read_objects_from_sqlite(cls: Type[T], table_name: str, db_path: str) -> List[T]:
    """
    Generic function to read rows from a SQLite table and convert them into class instances.

    Args:
        cls (Type[T]): The class to instantiate (e.g., Agent, Policy, Hash).
        table_name (str): The name of the table to query.
        db_path (str): Path to the SQLite database file.

    Returns:
        List[T]: List of class instances.

    Example:
        agents = read_objects_from_sqlite(Agent, "agents", "data_store.db")
    """
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        cursor.execute(f"SELECT * FROM {table_name}")
        rows = cursor.fetchall()
        columns = [desc[0] for desc in cursor.description]

        objects = []
        for row in rows:
            data = dict(zip(columns, row))
            if hasattr(cls, "from_dict"):
                obj = cls.from_dict(data) # pyright: ignore[reportAttributeAccessIssue]
            else:
                obj = cls(**data)
            objects.append(obj)

        conn.close()
        logger.info(f"Read {len(objects)} records from {table_name} table in {db_path}")
        return objects
    except Exception as e:
        logger.error(f"Failed to read from {table_name} in {db_path}: {e}")
        return []