codette_memory_sql.py

import os
import logging
import pyodbc
import sqlite3
from typing import Optional, List, Tuple


def connect_to_memory_database() -> Optional[pyodbc.Connection]:
    """Connect to Azure SQL Database using HF secret environment variables."""
    try:
        server = os.getenv("AZURE_SQL_SERVER")
        database = os.getenv("AZURE_SQL_DATABASE")
        username = os.getenv("AZURE_SQL_USERNAME")
        password = os.getenv("AZURE_SQL_PASSWORD")
        driver = "{ODBC Driver 17 for SQL Server}"

        if not all([server, database, username, password]):
            logging.warning("Azure SQL credentials missing. Falling back to SQLite.")
            return None

        conn_str = (
            f"DRIVER={driver};SERVER={server};DATABASE={database};"
            f"UID={username};PWD={password}"
        )
        conn = pyodbc.connect(conn_str)
        logging.info("Connected to Azure SQL Database.")
        return conn
    except pyodbc.Error as e:
        logging.error(f"Azure SQL connection error: {e}")
        return None


def log_memory(conn, query: str, response: str) -> None:
    try:
        cursor = conn.cursor()
        cursor.execute(
            "CREATE TABLE IF NOT EXISTS codette_memory (query TEXT, response TEXT)"
        )
        cursor.execute(
            "INSERT INTO codette_memory (query, response) VALUES (?, ?)", (query, response)
        )
        conn.commit()
        logging.info("Memory logged successfully.")
    except Exception as e:
        logging.error(f"Failed to log memory: {e}")


def get_recent_entries(conn, limit: int = 5) -> List[Tuple[str, str]]:
    try:
        cursor = conn.cursor()
        cursor.execute(
            "SELECT TOP (?) query, response FROM codette_memory ORDER BY 1 DESC", (limit,)
        )
        return cursor.fetchall()
    except Exception as e:
        logging.warning(f"Failed to fetch memory: {e}")
        return []


def ping_database(conn) -> bool:
    try:
        cursor = conn.cursor()
        cursor.execute("SELECT 1")
        return True
    except Exception as e:
        logging.warning(f"Ping failed: {e}")
        return False


# SQLite fallback (in-memory or file-backed)
def get_sqlite_memory(fallback_path: Optional[str] = ":memory:") -> sqlite3.Connection:
    conn = sqlite3.connect(fallback_path)
    conn.execute("CREATE TABLE IF NOT EXISTS codette_memory (query TEXT, response TEXT)")
    return conn


# Context manager for auto-closing
class MemoryConnection:
    def __init__(self):
        self.conn = connect_to_memory_database()
        if self.conn is None:
            self.conn = get_sqlite_memory()

    def __enter__(self):
        return self.conn

    def __exit__(self, exc_type, exc_value, traceback):
        try:
            self.conn.close()
            logging.info("Memory DB connection closed.")
        except Exception as e:
            logging.warning(f"Failed to close DB: {e}")