rp/pr/tools/database.py

import time


def db_set(key, value, db_conn):
    if not db_conn:
        return {"status": "error", "error": "Database not initialized"}

    try:
        cursor = db_conn.cursor()
        cursor.execute(
            """INSERT OR REPLACE INTO kv_store (key, value, timestamp)
                          VALUES (?, ?, ?)""",
            (key, value, time.time()),
        )
        db_conn.commit()
        return {"status": "success", "message": f"Set {key}"}
    except Exception as e:
        return {"status": "error", "error": str(e)}


def db_get(key, db_conn):
    if not db_conn:
        return {"status": "error", "error": "Database not initialized"}

    try:
        cursor = db_conn.cursor()
        cursor.execute("SELECT value FROM kv_store WHERE key = ?", (key,))
        result = cursor.fetchone()
        if result:
            return {"status": "success", "value": result[0]}
        else:
            return {"status": "error", "error": "Key not found"}
    except Exception as e:
        return {"status": "error", "error": str(e)}


def db_query(query, db_conn):
    if not db_conn:
        return {"status": "error", "error": "Database not initialized"}

    try:
        cursor = db_conn.cursor()
        cursor.execute(query)

        if query.strip().upper().startswith("SELECT"):
            results = cursor.fetchall()
            columns = (
                [desc[0] for desc in cursor.description] if cursor.description else []
            )
            return {"status": "success", "columns": columns, "rows": results}
        else:
            db_conn.commit()
            return {"status": "success", "rows_affected": cursor.rowcount}
    except Exception as e:
        return {"status": "error", "error": str(e)}