Agent/schema_profile.py

import os
import json
import time
import uuid
import threading
from typing import Any, Dict, List, Optional


class SchemaProfile:
    """
    Runtime schema profile representing the active database.
    Stores tables, columns, PK/FK metadata, row counts, sample rows, and
    a digest string suitable for injecting into LLM prompts.
    """

    SAMPLE_ROW_LIMIT = 3
    MAX_DIGEST_TABLES = 40
    MAX_DIGEST_COLUMNS = 25

    def __init__(self, name: str, db_url: Optional[str] = None, is_demo: bool = False,
                 notes: str = ""):
        self.name = name
        self.db_url = db_url
        self.is_demo = is_demo
        # Owner-supplied schema notes (domain hints): column semantics, business
        # rules, aliases. Injected verbatim into LLM prompts alongside the
        # introspected digest. Stays empty unless the database owner writes it.
        self.notes = (notes or "").strip()
        self.tables: List[str] = []
        self.columns_by_table: Dict[str, List[Dict[str, Any]]] = {}
        self.pk_by_table: Dict[str, List[str]] = {}
        self.fk_by_table: Dict[str, List[Dict[str, Any]]] = {}
        self.row_counts: Dict[str, int] = {}
        self.sample_rows: Dict[str, List[Dict[str, Any]]] = {}
        self.available_tables_set: set = set()

    @property
    def is_employee_db(self) -> bool:
        # Treat as employee demo if explicitly marked or name matches
        if self.is_demo:
            return True
        n = (self.name or "").lower()
        return "employee" in n or n == "demo_employee"

    # ------------------------------------------------------------------
    # Population
    # ------------------------------------------------------------------
    def populate_from_inspector(self, inspector: Any, engine: Optional[Any] = None) -> None:
        """Fill all metadata from a SQLAlchemy inspector (and optionally sample
        row counts / preview rows via a sync engine connection)."""
        try:
            tables = list(
                t for t in (inspector.get_table_names() + inspector.get_view_names())
                if not t.startswith("_")
            )
        except Exception as e:
            print(f"[SCHEMA-PROFILE] inspector table listing failed: {e}")
            tables = []

        self.tables = sorted(tables)
        self.available_tables_set = set(self.tables)
        self.columns_by_table = {}
        self.pk_by_table = {}
        self.fk_by_table = {}

        for t in self.tables:
            try:
                self.columns_by_table[t] = inspector.get_columns(t)
            except Exception:
                self.columns_by_table[t] = []
            try:
                pks = inspector.get_pk_constraint(t)
                self.pk_by_table[t] = (pks or {}).get("constrained_columns", []) or []
            except Exception:
                self.pk_by_table[t] = []
            try:
                self.fk_by_table[t] = inspector.get_foreign_keys(t) or []
            except Exception:
                self.fk_by_table[t] = []

        if engine is not None:
            self._collect_row_stats(engine)

    def _collect_row_stats(self, engine: Any) -> None:
        """Best-effort row counts + small sample rows using a sync engine.

        Dialect-aware so row previews also work on databases whose syntax for
        limiting differs (mssql uses TOP, oracle 12c+ uses FETCH FIRST, mysql
        prefers backtick-quoted identifiers, standard SQL uses LIMIT).
        """
        from sqlalchemy import text as sql_text

        # Identifier quoting + limit syntax per dialect (best-effort)
        dname = ""
        try:
            dname = (getattr(engine, "dialect", None).name or "").lower()
        except Exception:
            dname = ""

        def q(name: str) -> str:
            if dname == "mssql":
                return f"[{name}]"
            if dname in ("mysql", "mariadb"):
                return f"`{name}`"
            return f'"{name}"'

        def sample_candidates(table: str):
            limit = self.SAMPLE_ROW_LIMIT
            quoted = q(table)
            if dname == "mssql":
                return [f"SELECT TOP {limit} * FROM {quoted}"]
            if dname == "oracle":
                # 12c+ thin driver supports FETCH FIRST; legacy ROWNUM fallback
                return [
                    f"SELECT * FROM {quoted} FETCH FIRST {limit} ROWS ONLY",
                    f"SELECT * FROM {quoted} WHERE ROWNUM <= {limit}",
                ]
            # postgres / mysql / sqlite / everything else
            return [
                f"SELECT * FROM {quoted} LIMIT {limit}",
                f"SELECT * FROM {table} LIMIT {limit}",
            ]

        self.row_counts = {}
        self.sample_rows = {}
        try:
            with engine.connect() as conn:
                for t in self.tables:
                    # Row count (quoted identifier first, unquoted fallback)
                    try:
                        cnt = conn.execute(sql_text(f"SELECT COUNT(*) FROM {q(t)}")).scalar()
                        self.row_counts[t] = int(cnt) if cnt is not None else 0
                    except Exception:
                        try:
                            cnt = conn.execute(sql_text(f"SELECT COUNT(*) FROM {t}")).scalar()
                            self.row_counts[t] = int(cnt) if cnt is not None else 0
                        except Exception:
                            self.row_counts[t] = 0
                    # Sample rows (first successful candidate wins)
                    if self.row_counts.get(t, 0) > 0:
                        for sql_c in sample_candidates(t):
                            try:
                                result = conn.execute(sql_text(sql_c))
                                rows = result.mappings().all()
                                self.sample_rows[t] = [
                                    {k: _json_safe(v) for k, v in dict(r).items()} for r in rows
                                ]
                                break
                            except Exception:
                                continue
                        else:
                            self.sample_rows[t] = []
                    else:
                        self.sample_rows[t] = []
        except Exception as e:
            print(f"[SCHEMA-PROFILE] row stats collection failed: {e}")

    # ------------------------------------------------------------------
    # Digest for LLM prompts
    # ------------------------------------------------------------------
    def build_digest(self, max_tables: Optional[int] = None) -> str:
        max_tables = max_tables or self.MAX_DIGEST_TABLES
        lines: List[str] = []
        lines.append(f"Database schema name: {self.name}")
        lines.append(f"Number of tables/views: {len(self.tables)}")
        lines.append("Tables and columns:")
        for t in self.tables[:max_tables]:
            cols = self.columns_by_table.get(t, [])
            col_parts = []
            for c in cols[: self.MAX_DIGEST_COLUMNS]:
                ctype = str(c.get("type", ""))[:24]
                col_parts.append(f"{c.get('name')} {ctype}".strip())
            colstr = ", ".join(col_parts) if col_parts else "(no columns)"
            rc = self.row_counts.get(t)
            lines.append(f"  - {t} ({colstr})  [rows={rc if rc is not None else '?'}]")
            fks = self.fk_by_table.get(t, [])
            for fk in fks[:4]:
                ref = fk.get("referred_table")
                lcols = fk.get("constrained_columns") or []
                rcols = fk.get("referred_columns") or []
                if ref:
                    lines.append(f"      FK: {lcols} -> {ref}({rcols})")
            if len(cols) > self.MAX_DIGEST_COLUMNS:
                lines.append(f"      ... {len(cols) - self.MAX_DIGEST_COLUMNS} more columns")
        if len(self.tables) > max_tables:
            lines.append(f"  ... {len(self.tables) - max_tables} more tables")
        return "\n".join(lines)

    def summarize(self) -> Dict[str, Any]:
        """Compact JSON-friendly summary (for API responses / UI)."""
        return {
            "name": self.name,
            "db_url": self.db_url,
            "is_demo": self.is_demo,
            "is_employee_db": self.is_employee_db,
            "notes": self.notes,
            "tables": [
                {
                    "name": t,
                    "row_count": self.row_counts.get(t),
                    "columns": [c.get("name") for c in self.columns_by_table.get(t, [])],
                    "sample_rows": self.sample_rows.get(t, []),
                }
                for t in self.tables
            ],
        }


def _json_safe(v: Any) -> Any:
    try:
        json.dumps(v)
        return v
    except (TypeError, ValueError):
        return str(v)


class ProfileRegistry:
    """
    Registry of user-registered database profiles.

    Each registry entry stores lightweight metadata (id, name, db_url, kind).
    Full introspection happens lazily when a profile is switched to (the
    agent's initialize_database walks the schema at that moment).
    """

    def __init__(self, store_path: str):
        self.store_path = store_path
        self._lock = threading.Lock()
        self.profiles: Dict[str, Dict[str, Any]] = {}
        self.active_id: Optional[str] = None
        self.load()

    # ------------------------------------------------------------------
    def load(self) -> None:
        try:
            if os.path.exists(self.store_path):
                with open(self.store_path, "r", encoding="utf-8") as fh:
                    data = json.load(fh)
                self.profiles = data.get("profiles", {})
                self.active_id = data.get("active_id")
        except Exception as e:
            print(f"[PROFILE-REGISTRY] load failed: {e}")
            self.profiles = {}
            self.active_id = None

    def save(self) -> None:
        with self._lock:
            try:
                os.makedirs(os.path.dirname(self.store_path), exist_ok=True)
                with open(self.store_path, "w", encoding="utf-8") as fh:
                    json.dump(
                        {"profiles": self.profiles, "active_id": self.active_id},
                        fh,
                        indent=2,
                    )
            except Exception as e:
                print(f"[PROFILE-REGISTRY] save failed: {e}")

    # ------------------------------------------------------------------
    def list(self) -> List[Dict[str, Any]]:
        out = []
        for pid, p in self.profiles.items():
            out.append(
                {
                    "id": pid,
                    "name": p.get("name"),
                    "kind": p.get("kind", "user"),
                    "db_url": p.get("db_url"),
                    "created_at": p.get("created_at"),
                    "notes": p.get("notes", ""),
                    "active": pid == self.active_id,
                }
            )
        return sorted(out, key=lambda p: (not p["active"], p.get("created_at") or 0))

    def get(self, profile_id: str) -> Optional[Dict[str, Any]]:
        return self.profiles.get(profile_id)

    def active(self) -> Optional[Dict[str, Any]]:
        if self.active_id and self.active_id in self.profiles:
            return {"id": self.active_id, **self.profiles[self.active_id]}
        return None

    def add(
        self,
        name: str,
        db_url: str,
        kind: str = "user",
        activate: bool = True,
        notes: Optional[str] = None,
    ) -> Dict[str, Any]:
        pid = uuid.uuid4().hex[:12]
        entry = {
            "name": name,
            "db_url": db_url,
            "kind": kind,
            "created_at": time.time(),
            "notes": (notes or "").strip(),
        }
        with self._lock:
            self.profiles[pid] = entry
            if activate or self.active_id is None:
                self.active_id = pid
        self.save()
        return {"id": pid, **entry}

    def remove(self, profile_id: str) -> bool:
        with self._lock:
            if profile_id not in self.profiles:
                return False
            del self.profiles[profile_id]
            if self.active_id == profile_id:
                self.active_id = next(iter(self.profiles), None)
        self.save()
        return True

    def switch(self, profile_id: str) -> Optional[Dict[str, Any]]:
        with self._lock:
            if profile_id not in self.profiles:
                return None
            self.active_id = profile_id
        self.save()
        return {"id": profile_id, **self.profiles[profile_id]}

    def set_notes(self, profile_id: str, notes: str) -> bool:
        """Persist owner-supplied schema notes (domain hints) for a profile."""
        with self._lock:
            if profile_id not in self.profiles:
                return False
            self.profiles[profile_id]["notes"] = (notes or "").strip()
        self.save()
        return True