Spaces:
Running on Zero
Running on Zero
| import os | |
| import re | |
| import json | |
| import time | |
| import uuid | |
| import threading | |
| from typing import Any, Dict, List, Optional | |
| def _redact_db_url(url: Optional[str]) -> Optional[str]: | |
| """Mask the password component of a connection URL for API responses. | |
| Keeps scheme, user, host, port and database visible so the UI can still | |
| show WHERE it is connected; only the credential is masked. No-op for | |
| credential-less URLs (sqlite:///, mysql:// without auth, etc.). | |
| """ | |
| if not url or not isinstance(url, str) or "://" not in url: | |
| return url | |
| try: | |
| return re.sub(r"(://[^:/?#\s]+:)[^@/?#\s]+(@)", r"\1***\2", url) | |
| except Exception: | |
| return url | |
| 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() | |
| 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, | |
| # Password redacted: this lands on PUBLIC endpoints (/schema/current, | |
| # /schemas/add, /schemas/switch) and must never leak credentials. | |
| # The real URL stays on the profile object for reconnection. | |
| "db_url": _redact_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 | |