dbagent-cli 0.1.0__py3-none-any.whl

dbagent/connectors/relational.py

@@ -0,0 +1,472 @@
+"""
+Universal Relational Database Connector using SQLAlchemy.
+Supports PostgreSQL, MySQL, MariaDB, SQLite, SQL Server (MSSQL), Oracle, DuckDB, Snowflake.
+Enhanced with table name caching, fuzzy matching, and speed optimizations.
+"""
+
+from typing import Dict, Any, List, Optional, Tuple, Set
+import datetime
+import decimal
+import re
+import uuid
+import sqlalchemy as sa
+from sqlalchemy import inspect, text
+from sqlalchemy.engine import Engine
+
+from dbagent.connectors.base import BaseConnector
+from dbagent.schema.models import (
+    DatabaseSchema,
+    TableModel,
+    ColumnModel,
+    ForeignKeyModel,
+    IndexModel,
+)
+
+
+def _serialize_sample_val(val: Any) -> Any:
+    """Convert database values into clean displayable Python types."""
+    if val is None:
+        return None
+    if isinstance(val, (datetime.date, datetime.datetime, datetime.time)):
+        return val.isoformat()
+    if isinstance(val, decimal.Decimal):
+        return float(val)
+    if isinstance(val, uuid.UUID):
+        return str(val)
+    if isinstance(val, bytes):
+        return f"<bytes len={len(val)}>"
+    return val
+
+
+def _fuzzy_match_tables(table_names: List[str], user_prompt: str) -> Tuple[List[str], List[str]]:
+    """
+    Match table names from user prompt with fuzzy/partial/singular-plural normalization.
+    Returns (exact_matches, fuzzy_matches).
+    """
+    prompt_lower = user_prompt.lower()
+    exact: List[str] = []
+    fuzzy: List[str] = []
+
+    # Extract meaningful words from the prompt (3+ chars, skip SQL keywords)
+    sql_keywords = {
+        "select", "from", "where", "join", "left", "right", "inner", "outer",
+        "order", "group", "having", "limit", "offset", "insert", "update",
+        "delete", "into", "values", "set", "and", "not", "the", "all",
+        "show", "list", "find", "get", "recent", "last", "first", "count",
+        "with", "like", "between", "table", "column", "create", "alter",
+        "drop", "index", "primary", "foreign", "key", "null", "default",
+        "hour", "minute", "day", "week", "month", "year", "ago", "today",
+        "give", "tell", "want", "need", "please", "display", "fetch",
+    }
+    prompt_words = set()
+    for word in re.findall(r"\b[a-z_][a-z0-9_]*\b", prompt_lower):
+        if len(word) >= 3 and word not in sql_keywords:
+            prompt_words.add(word)
+
+    for t_name in table_names:
+        t_lower = t_name.lower()
+
+        # 1. Exact word boundary match
+        if re.search(r"\b" + re.escape(t_lower) + r"\b", prompt_lower):
+            exact.append(t_name)
+            continue
+
+        # 2. Substring match (table name mentioned in prompt)
+        if t_lower in prompt_lower:
+            exact.append(t_name)
+            continue
+
+        # 3. Singular/plural normalization
+        # "users" matches "user", "categories" matches "category"
+        t_singular = t_lower.rstrip("s")
+        t_no_ies = re.sub(r"ies$", "y", t_lower)
+        for word in prompt_words:
+            w_singular = word.rstrip("s")
+            w_no_ies = re.sub(r"ies$", "y", word)
+            # Check various forms
+            if (word == t_lower or
+                word == t_singular or
+                w_singular == t_lower or
+                w_singular == t_singular or
+                word == t_no_ies or
+                w_no_ies == t_lower):
+                fuzzy.append(t_name)
+                break
+        else:
+            # 4. Partial prefix/suffix match (e.g. "user" matches "user_logins", "app_users")
+            for word in prompt_words:
+                if len(word) >= 4:
+                    if t_lower.startswith(word) or t_lower.endswith(word):
+                        fuzzy.append(t_name)
+                        break
+                    # Check if word is a component of underscore-separated name
+                    t_parts = t_lower.split("_")
+                    if word in t_parts or word.rstrip("s") in t_parts:
+                        fuzzy.append(t_name)
+                        break
+
+    return exact, fuzzy
+
+
+class RelationalConnector(BaseConnector):
+    """Universal connector for all SQL databases supported by SQLAlchemy."""
+
+    def __init__(self, connection_url: str):
+        self.connection_url = connection_url
+        self._engine: Optional[Engine] = None
+        self._cached_table_names: Optional[List[str]] = None
+        self._cached_view_names: Optional[List[str]] = None
+        self._init_engine()
+
+    def _init_engine(self) -> None:
+        url = self.connection_url
+        if url.startswith("sqlite:///") or url.startswith("sqlite://"):
+            self._engine = sa.create_engine(url)
+        elif url.startswith("duckdb:///"):
+            self._engine = sa.create_engine(url)
+        else:
+            self._engine = sa.create_engine(
+                url,
+                pool_pre_ping=True,
+                pool_size=2,
+                max_overflow=3,
+                connect_args={"connect_timeout": 10} if "mysql" in url or "postgres" in url else {},
+            )
+
+    @property
+    def engine(self) -> Engine:
+        if self._engine is None:
+            self._init_engine()
+        return self._engine
+
+    def test_connection(self) -> Tuple[bool, str]:
+        """Test database connection with a lightweight SELECT 1 query."""
+        try:
+            with self.engine.connect() as conn:
+                conn.execute(text("SELECT 1"))
+            return True, f"Successfully connected to {self.engine.dialect.name} database."
+        except Exception as e:
+            return False, f"Connection failed: {str(e)}"
+
+    def get_table_names(self) -> List[str]:
+        """Fast lookup of all table names with caching."""
+        if self._cached_table_names is not None:
+            return self._cached_table_names
+        insp = inspect(self.engine)
+        try:
+            self._cached_table_names = insp.get_table_names()
+        except Exception:
+            self._cached_table_names = []
+        return self._cached_table_names
+
+    def get_view_names(self) -> List[str]:
+        """Fast lookup of all view names with caching."""
+        if self._cached_view_names is not None:
+            return self._cached_view_names
+        insp = inspect(self.engine)
+        try:
+            self._cached_view_names = insp.get_view_names()
+        except Exception:
+            self._cached_view_names = []
+        return self._cached_view_names
+
+    def refresh_cache(self) -> None:
+        """Invalidate cached table/view names."""
+        self._cached_table_names = None
+        self._cached_view_names = None
+
+    def inspect_schema(
+        self,
+        table_names: Optional[List[str]] = None,
+        include_samples: bool = True,
+        max_samples: int = 2,
+        include_views: bool = True,
+        include_row_counts: bool = False,
+    ) -> DatabaseSchema:
+        """Introspect schema metadata for all or specified tables."""
+        insp = inspect(self.engine)
+        dialect_name = self.engine.dialect.name
+        db_name = self.engine.url.database or "default"
+
+        server_version = None
+        try:
+            server_version = ".".join(str(v) for v in self.engine.dialect.server_version_info)
+        except Exception:
+            pass
+
+        tables: List[TableModel] = []
+
+        # 1. Inspect regular tables
+        all_table_names = self.get_table_names()
+
+        target_tables = all_table_names
+        if table_names is not None:
+            lower_targets = {t.lower() for t in table_names}
+            target_tables = [t for t in all_table_names if t.lower() in lower_targets]
+
+        for t_name in target_tables:
+            table_model = self._inspect_single_table(
+                insp=insp,
+                table_name=t_name,
+                is_view=False,
+                include_samples=include_samples,
+                max_samples=max_samples,
+                include_row_counts=include_row_counts,
+            )
+            tables.append(table_model)
+
+        # 2. Inspect views
+        if include_views:
+            all_view_names = self.get_view_names()
+
+            target_views = all_view_names
+            if table_names is not None:
+                lower_targets = {t.lower() for t in table_names}
+                target_views = [v for v in all_view_names if v.lower() in lower_targets]
+
+            for v_name in target_views:
+                view_model = self._inspect_single_table(
+                    insp=insp,
+                    table_name=v_name,
+                    is_view=True,
+                    include_samples=include_samples,
+                    max_samples=max_samples,
+                    include_row_counts=include_row_counts,
+                )
+                try:
+                    view_model.view_definition = insp.get_view_definition(v_name)
+                except Exception:
+                    pass
+                tables.append(view_model)
+
+        return DatabaseSchema(
+            dialect_name=dialect_name,
+            database_name=db_name,
+            server_version=server_version,
+            tables=tables,
+        )
+
+    def inspect_targeted(
+        self,
+        user_prompt: str,
+        max_tables: int = 8,
+        include_samples: bool = True,
+    ) -> DatabaseSchema:
+        """
+        On-demand targeted inspection with fuzzy matching.
+        Returns (schema, exact_matches, fuzzy_matches) for conflict resolution.
+        """
+        all_tables = self.get_table_names()
+        exact, fuzzy = _fuzzy_match_tables(all_tables, user_prompt)
+
+        # Combine exact + fuzzy, but prefer exact
+        matched = exact[:]
+        for f in fuzzy:
+            if f not in matched:
+                matched.append(f)
+
+        if matched:
+            # Also pull in FK-related tables for completeness
+            fk_related = self._get_fk_related_tables(matched[:max_tables])
+            all_to_inspect = matched[:max_tables]
+            for r in fk_related:
+                if r not in all_to_inspect and len(all_to_inspect) < max_tables:
+                    all_to_inspect.append(r)
+
+            return self.inspect_schema(
+                table_names=all_to_inspect,
+                include_samples=include_samples,
+                max_samples=2,
+                include_views=False,
+                include_row_counts=False,
+            )
+
+        # Fallback: inspect top max_tables
+        return self.inspect_schema(
+            table_names=all_tables[:max_tables],
+            include_samples=include_samples,
+            max_samples=2,
+            include_views=True,
+            include_row_counts=False,
+        )
+
+    def resolve_tables(self, user_prompt: str) -> Tuple[List[str], List[str]]:
+        """
+        Resolve which tables the user is referring to.
+        Returns (exact_matches, fuzzy_matches) for the pipeline to handle ambiguity.
+        """
+        all_tables = self.get_table_names()
+        return _fuzzy_match_tables(all_tables, user_prompt)
+
+    def _get_fk_related_tables(self, table_names: List[str]) -> List[str]:
+        """Find tables related by foreign keys to the given tables."""
+        related: Set[str] = set()
+        insp = inspect(self.engine)
+        for t_name in table_names:
+            try:
+                fks = insp.get_foreign_keys(t_name)
+                for fk in fks:
+                    ref = fk.get("referred_table")
+                    if ref and ref not in table_names:
+                        related.add(ref)
+            except Exception:
+                pass
+        return list(related)
+
+    def _inspect_single_table(
+        self,
+        insp: Any,
+        table_name: str,
+        is_view: bool,
+        include_samples: bool,
+        max_samples: int,
+        include_row_counts: bool = False,
+    ) -> TableModel:
+        """Inspect a single table/view."""
+        # Columns
+        raw_columns = []
+        try:
+            raw_columns = insp.get_columns(table_name)
+        except Exception:
+            pass
+
+        # Primary keys
+        pk_cols = []
+        try:
+            pk_constraint = insp.get_pk_constraint(table_name)
+            pk_cols = pk_constraint.get("constrained_columns", []) or []
+        except Exception:
+            pass
+
+        columns: List[ColumnModel] = []
+        for col in raw_columns:
+            c_name = col.get("name", "")
+            c_type = str(col.get("type", "UNKNOWN"))
+            is_pk = c_name in pk_cols or col.get("primary_key", False)
+            col_model = ColumnModel(
+                name=c_name,
+                data_type=c_type,
+                is_nullable=col.get("nullable", True),
+                is_primary_key=is_pk,
+                default_value=str(col.get("default")) if col.get("default") is not None else None,
+                comment=col.get("comment"),
+                is_autoincrement=bool(col.get("autoincrement", False)),
+            )
+            columns.append(col_model)
+
+        # Foreign Keys
+        foreign_keys: List[ForeignKeyModel] = []
+        if not is_view:
+            try:
+                raw_fks = insp.get_foreign_keys(table_name)
+                for fk in raw_fks:
+                    foreign_keys.append(
+                        ForeignKeyModel(
+                            name=fk.get("name"),
+                            constrained_columns=fk.get("constrained_columns", []),
+                            referred_table=fk.get("referred_table", ""),
+                            referred_columns=fk.get("referred_columns", []),
+                        )
+                    )
+            except Exception:
+                pass
+
+        # Indexes
+        indexes: List[IndexModel] = []
+        if not is_view:
+            try:
+                raw_indexes = insp.get_indexes(table_name)
+                for idx in raw_indexes:
+                    indexes.append(
+                        IndexModel(
+                            name=idx.get("name") or "unnamed_idx",
+                            columns=idx.get("column_names", []),
+                            is_unique=idx.get("unique", False),
+                        )
+                    )
+            except Exception:
+                pass
+
+        # Table Comment
+        comment = None
+        try:
+            tbl_comment = insp.get_table_comment(table_name)
+            if tbl_comment and isinstance(tbl_comment, dict):
+                comment = tbl_comment.get("text")
+        except Exception:
+            pass
+
+        # Row count & Sample rows
+        row_count = None
+        sample_rows = []
+        try:
+            with self.engine.connect() as conn:
+                if include_row_counts:
+                    try:
+                        count_res = conn.execute(text(f'SELECT COUNT(*) FROM "{table_name}"')).scalar()
+                        row_count = int(count_res) if count_res is not None else None
+                    except Exception:
+                        try:
+                            count_res = conn.execute(text(f"SELECT COUNT(*) FROM `{table_name}`")).scalar()
+                            row_count = int(count_res) if count_res is not None else None
+                        except Exception:
+                            pass
+
+                if include_samples:
+                    try:
+                        q = text(f'SELECT * FROM "{table_name}" LIMIT {max_samples}')
+                        rows = conn.execute(q).mappings().all()
+                        for r in rows:
+                            sample_rows.append({k: _serialize_sample_val(v) for k, v in dict(r).items()})
+                    except Exception:
+                        try:
+                            q = text(f"SELECT * FROM `{table_name}` LIMIT {max_samples}")
+                            rows = conn.execute(q).mappings().all()
+                            for r in rows:
+                                sample_rows.append({k: _serialize_sample_val(v) for k, v in dict(r).items()})
+                        except Exception:
+                            pass
+        except Exception:
+            pass
+
+        return TableModel(
+            name=table_name,
+            is_view=is_view,
+            columns=columns,
+            primary_key=pk_cols,
+            foreign_keys=foreign_keys,
+            indexes=indexes,
+            row_count=row_count,
+            sample_rows=sample_rows,
+            comment=comment,
+        )
+
+    def execute_query(
+        self,
+        query: str,
+        limit: int = 100,
+    ) -> Tuple[List[str], List[Dict[str, Any]], Optional[str]]:
+        """Safely execute a query with row limit."""
+        try:
+            with self.engine.connect() as conn:
+                result = conn.execute(text(query))
+                if not result.returns_rows:
+                    # DML: commit and report affected rows
+                    try:
+                        conn.commit()
+                    except Exception:
+                        pass
+                    return ["status"], [{"status": "Query executed successfully."}], None
+                columns = list(result.keys())
+                raw_rows = result.fetchmany(limit)
+                rows = [
+                    {k: _serialize_sample_val(v) for k, v in zip(columns, row)}
+                    for row in raw_rows
+                ]
+                return columns, rows, None
+        except Exception as e:
+            return [], [], str(e)
+
+    def close(self) -> None:
+        if self._engine:
+            self._engine.dispose()