sqlsaber 0.2.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

sqlsaber/agents/anthropic.py

@@ -11,13 +11,7 @@ from sqlsaber.agents.streaming import (
     build_tool_result_block,
 )
 from sqlsaber.config.settings import Config
-from sqlsaber.database.connection import (
-    BaseDatabaseConnection,
-    MySQLConnection,
-    PostgreSQLConnection,
-    SQLiteConnection,
-)
-from sqlsaber.database.schema import SchemaManager
+from sqlsaber.database.connection import BaseDatabaseConnection
 from sqlsaber.memory.manager import MemoryManager
 from sqlsaber.models.events import StreamEvent
 from sqlsaber.models.types import ToolDefinition
@@ -36,7 +30,6 @@ class AnthropicSQLAgent(BaseSQLAgent):
 
         self.client = AsyncAnthropic(api_key=config.api_key)
         self.model = config.model_name.replace("anthropic:", "")
-        self.schema_manager = SchemaManager(db_connection)
 
         self.database_name = database_name
         self.memory_manager = MemoryManager()
@@ -94,17 +87,6 @@ class AnthropicSQLAgent(BaseSQLAgent):
         # Build system prompt with memories if available
         self.system_prompt = self._build_system_prompt()
 
-    def _get_database_type_name(self) -> str:
-        """Get the human-readable database type name."""
-        if isinstance(self.db, PostgreSQLConnection):
-            return "PostgreSQL"
-        elif isinstance(self.db, MySQLConnection):
-            return "MySQL"
-        elif isinstance(self.db, SQLiteConnection):
-            return "SQLite"
-        else:
-            return "database"  # Fallback
-
     def _build_system_prompt(self) -> str:
         """Build system prompt with optional memory context."""
         db_type = self._get_database_type_name()
@@ -152,109 +134,33 @@ Guidelines:
         self.system_prompt = self._build_system_prompt()
         return memory.id
 
-    async def introspect_schema(self, table_pattern: Optional[str] = None) -> str:
-        """Introspect database schema to understand table structures."""
-        try:
-            # Pass table_pattern to get_schema_info for efficient filtering at DB level
-            schema_info = await self.schema_manager.get_schema_info(table_pattern)
-
-            # Format the schema information
-            formatted_info = {}
-            for table_name, table_info in schema_info.items():
-                formatted_info[table_name] = {
-                    "columns": {
-                        col_name: {
-                            "type": col_info["data_type"],
-                            "nullable": col_info["nullable"],
-                            "default": col_info["default"],
-                        }
-                        for col_name, col_info in table_info["columns"].items()
-                    },
-                    "primary_keys": table_info["primary_keys"],
-                    "foreign_keys": [
-                        f"{fk['column']} -> {fk['references']['table']}.{fk['references']['column']}"
-                        for fk in table_info["foreign_keys"]
-                    ],
-                }
-
-            return json.dumps(formatted_info)
-        except Exception as e:
-            return json.dumps({"error": f"Error introspecting schema: {str(e)}"})
-
-    async def list_tables(self) -> str:
-        """List all tables in the database with basic information."""
-        try:
-            tables_info = await self.schema_manager.list_tables()
-            return json.dumps(tables_info)
-        except Exception as e:
-            return json.dumps({"error": f"Error listing tables: {str(e)}"})
-
     async def execute_sql(self, query: str, limit: Optional[int] = 100) -> str:
-        """Execute a SQL query against the database."""
-        try:
-            # Security check - only allow SELECT queries unless write is enabled
-            write_error = self._validate_write_operation(query)
-            if write_error:
-                return json.dumps(
-                    {
-                        "error": write_error,
-                    }
-                )
-
-            # Add LIMIT if not present and it's a SELECT query
-            query = self._add_limit_to_query(query, limit)
-
-            # Execute the query (wrapped in a transaction for safety)
-            results = await self.db.execute_query(query)
-
-            # Format results - but also store the actual data
-            actual_limit = limit if limit is not None else len(results)
-            self._last_results = results[:actual_limit]
-            self._last_query = query
+        """Execute a SQL query against the database with streaming support."""
+        # Call parent implementation for core functionality
+        result = await super().execute_sql(query, limit)
 
-            return json.dumps(
-                {
-                    "success": True,
-                    "row_count": len(results),
-                    "results": results[:actual_limit],  # Extra safety for limit
-                    "truncated": len(results) > actual_limit,
-                }
-            )
-
-        except Exception as e:
-            error_msg = str(e)
-
-            # Provide helpful error messages
-            suggestions = []
-            if "column" in error_msg.lower() and "does not exist" in error_msg.lower():
-                suggestions.append(
-                    "Check column names using the schema introspection tool"
-                )
-            elif "table" in error_msg.lower() and "does not exist" in error_msg.lower():
-                suggestions.append(
-                    "Check table names using the schema introspection tool"
-                )
-            elif "syntax error" in error_msg.lower():
-                suggestions.append(
-                    "Review SQL syntax, especially JOIN conditions and WHERE clauses"
+        # Parse result to extract data for streaming (AnthropicSQLAgent specific)
+        try:
+            result_data = json.loads(result)
+            if result_data.get("success") and "results" in result_data:
+                # Store results for streaming
+                actual_limit = (
+                    limit if limit is not None else len(result_data["results"])
                 )
+                self._last_results = result_data["results"][:actual_limit]
+                self._last_query = query
+        except (json.JSONDecodeError, KeyError):
+            # If we can't parse the result, just continue without storing
+            pass
 
-            return json.dumps({"error": error_msg, "suggestions": suggestions})
+        return result
 
     async def process_tool_call(
         self, tool_name: str, tool_input: Dict[str, Any]
     ) -> str:
         """Process a tool call and return the result."""
-        if tool_name == "list_tables":
-            return await self.list_tables()
-        elif tool_name == "introspect_schema":
-            return await self.introspect_schema(tool_input.get("table_pattern"))
-        elif tool_name == "execute_sql":
-            return await self.execute_sql(
-                tool_input["query"], tool_input.get("limit", 100)
-            )
-        else:
-            return json.dumps({"error": f"Unknown tool: {tool_name}"})
+        # Use parent implementation for core tools
+        return await super().process_tool_call(tool_name, tool_input)
 
     async def _process_stream_events(
         self, stream, content_blocks: List[Dict], tool_use_blocks: List[Dict]

sqlsaber/agents/base.py

@@ -1,9 +1,17 @@
 """Abstract base class for SQL agents."""
 
+import json
 from abc import ABC, abstractmethod
 from typing import Any, AsyncIterator, Dict, List, Optional
 
-from sqlsaber.database.connection import BaseDatabaseConnection
+from sqlsaber.database.connection import (
+    BaseDatabaseConnection,
+    CSVConnection,
+    MySQLConnection,
+    PostgreSQLConnection,
+    SQLiteConnection,
+)
+from sqlsaber.database.schema import SchemaManager
 from sqlsaber.models.events import StreamEvent
 
 
@@ -12,6 +20,7 @@ class BaseSQLAgent(ABC):
 
     def __init__(self, db_connection: BaseDatabaseConnection):
         self.db = db_connection
+        self.schema_manager = SchemaManager(db_connection)
         self.conversation_history: List[Dict[str, Any]] = []
 
     @abstractmethod
@@ -25,12 +34,120 @@ class BaseSQLAgent(ABC):
         """Clear conversation history."""
         self.conversation_history = []
 
-    @abstractmethod
+    def _get_database_type_name(self) -> str:
+        """Get the human-readable database type name."""
+        if isinstance(self.db, PostgreSQLConnection):
+            return "PostgreSQL"
+        elif isinstance(self.db, MySQLConnection):
+            return "MySQL"
+        elif isinstance(self.db, SQLiteConnection):
+            return "SQLite"
+        elif isinstance(self.db, CSVConnection):
+            return "SQLite"  # we convert csv to in-memory sqlite
+        else:
+            return "database"  # Fallback
+
+    async def introspect_schema(self, table_pattern: Optional[str] = None) -> str:
+        """Introspect database schema to understand table structures."""
+        try:
+            # Pass table_pattern to get_schema_info for efficient filtering at DB level
+            schema_info = await self.schema_manager.get_schema_info(table_pattern)
+
+            # Format the schema information
+            formatted_info = {}
+            for table_name, table_info in schema_info.items():
+                formatted_info[table_name] = {
+                    "columns": {
+                        col_name: {
+                            "type": col_info["data_type"],
+                            "nullable": col_info["nullable"],
+                            "default": col_info["default"],
+                        }
+                        for col_name, col_info in table_info["columns"].items()
+                    },
+                    "primary_keys": table_info["primary_keys"],
+                    "foreign_keys": [
+                        f"{fk['column']} -> {fk['references']['table']}.{fk['references']['column']}"
+                        for fk in table_info["foreign_keys"]
+                    ],
+                }
+
+            return json.dumps(formatted_info)
+        except Exception as e:
+            return json.dumps({"error": f"Error introspecting schema: {str(e)}"})
+
+    async def list_tables(self) -> str:
+        """List all tables in the database with basic information."""
+        try:
+            tables_info = await self.schema_manager.list_tables()
+            return json.dumps(tables_info)
+        except Exception as e:
+            return json.dumps({"error": f"Error listing tables: {str(e)}"})
+
+    async def execute_sql(self, query: str, limit: Optional[int] = 100) -> str:
+        """Execute a SQL query against the database."""
+        try:
+            # Security check - only allow SELECT queries unless write is enabled
+            write_error = self._validate_write_operation(query)
+            if write_error:
+                return json.dumps(
+                    {
+                        "error": write_error,
+                    }
+                )
+
+            # Add LIMIT if not present and it's a SELECT query
+            query = self._add_limit_to_query(query, limit)
+
+            # Execute the query (wrapped in a transaction for safety)
+            results = await self.db.execute_query(query)
+
+            # Format results
+            actual_limit = limit if limit is not None else len(results)
+
+            return json.dumps(
+                {
+                    "success": True,
+                    "row_count": len(results),
+                    "results": results[:actual_limit],  # Extra safety for limit
+                    "truncated": len(results) > actual_limit,
+                }
+            )
+
+        except Exception as e:
+            error_msg = str(e)
+
+            # Provide helpful error messages
+            suggestions = []
+            if "column" in error_msg.lower() and "does not exist" in error_msg.lower():
+                suggestions.append(
+                    "Check column names using the schema introspection tool"
+                )
+            elif "table" in error_msg.lower() and "does not exist" in error_msg.lower():
+                suggestions.append(
+                    "Check table names using the schema introspection tool"
+                )
+            elif "syntax error" in error_msg.lower():
+                suggestions.append(
+                    "Review SQL syntax, especially JOIN conditions and WHERE clauses"
+                )
+
+            return json.dumps({"error": error_msg, "suggestions": suggestions})
+
     async def process_tool_call(
         self, tool_name: str, tool_input: Dict[str, Any]
     ) -> str:
         """Process a tool call and return the result."""
-        pass
+        if tool_name == "list_tables":
+            return await self.list_tables()
+        elif tool_name == "introspect_schema":
+            return await self.introspect_schema(tool_input.get("table_pattern"))
+        elif tool_name == "execute_sql":
+            return await self.execute_sql(
+                tool_input["query"], tool_input.get("limit", 100)
+            )
+        else:
+            return json.dumps({"error": f"Unknown tool: {tool_name}"})
 
     def _validate_write_operation(self, query: str) -> Optional[str]:
         """Validate if a write operation is allowed.

sqlsaber/agents/mcp.py

@@ -0,0 +1,21 @@
+"""Generic SQL agent implementation for MCP tools."""
+
+from typing import AsyncIterator
+from sqlsaber.agents.base import BaseSQLAgent
+from sqlsaber.database.connection import BaseDatabaseConnection
+from sqlsaber.models.events import StreamEvent
+
+
+class MCPSQLAgent(BaseSQLAgent):
+    """MCP SQL Agent for MCP tool operations without LLM-specific logic."""
+
+    def __init__(self, db_connection: BaseDatabaseConnection):
+        super().__init__(db_connection)
+
+    async def query_stream(
+        self, user_query: str, use_history: bool = True
+    ) -> AsyncIterator[StreamEvent]:
+        """Not implemented for generic agent as it's only used for tool operations."""
+        raise NotImplementedError(
+            "MCPSQLAgent does not support query streaming. Use specific agent implementations for conversation."
+        )

sqlsaber/agents/streaming.py

@@ -14,13 +14,3 @@ class StreamingResponse:
 def build_tool_result_block(tool_use_id: str, content: str) -> Dict[str, Any]:
     """Build a tool result block for the conversation."""
     return {"type": "tool_result", "tool_use_id": tool_use_id, "content": content}
-
-
-def extract_sql_from_text(text: str) -> str:
-    """Extract SQL query from markdown-formatted text."""
-    if "```sql" in text:
-        sql_start = text.find("```sql") + 6
-        sql_end = text.find("```", sql_start)
-        if sql_end > sql_start:
-            return text[sql_start:sql_end].strip()
-    return ""

sqlsaber/cli/commands.py

@@ -1,6 +1,7 @@
 """CLI command definitions and handlers."""
 
 import asyncio
+from pathlib import Path
 from typing import Optional
 
 import typer
@@ -62,15 +63,31 @@ def query(
     """Run a query against the database or start interactive mode."""
 
     async def run_session():
-        # Get database configuration
+        # Get database configuration or handle direct CSV file
         if database:
-            db_config = config_manager.get_database(database)
-            if not db_config:
-                console.print(
-                    f"[bold red]Error:[/bold red] Database connection '{database}' not found."
-                )
-                console.print("Use 'sqlsaber db list' to see available connections.")
-                raise typer.Exit(1)
+            # Check if this is a direct CSV file path
+            if database.endswith(".csv"):
+                csv_path = Path(database).expanduser().resolve()
+                if not csv_path.exists():
+                    console.print(
+                        f"[bold red]Error:[/bold red] CSV file '{database}' not found."
+                    )
+                    raise typer.Exit(1)
+                connection_string = f"csv:///{csv_path}"
+                db_name = csv_path.stem
+            else:
+                # Look up configured database connection
+                db_config = config_manager.get_database(database)
+                if not db_config:
+                    console.print(
+                        f"[bold red]Error:[/bold red] Database connection '{database}' not found."
+                    )
+                    console.print(
+                        "Use 'sqlsaber db list' to see available connections."
+                    )
+                    raise typer.Exit(1)
+                connection_string = db_config.to_connection_string()
+                db_name = db_config.name
         else:
             db_config = config_manager.get_default_database()
             if not db_config:
@@ -81,10 +98,11 @@ def query(
                     "Use 'sqlsaber db add <name>' to add a database connection."
                 )
                 raise typer.Exit(1)
+            connection_string = db_config.to_connection_string()
+            db_name = db_config.name
 
         # Create database connection
         try:
-            connection_string = db_config.to_connection_string()
             db_conn = DatabaseConnection(connection_string)
         except Exception as e:
             console.print(
@@ -93,7 +111,7 @@ def query(
             raise typer.Exit(1)
 
         # Create agent instance with database name for memory context
-        agent = AnthropicSQLAgent(db_conn, db_config.name)
+        agent = AnthropicSQLAgent(db_conn, db_name)
 
         try:
             if query_text:

sqlsaber/cli/database.py

@@ -75,7 +75,7 @@ def add_database(
         if type == "sqlite":
             # SQLite only needs database path
             database = database or questionary.path("Database file path:").ask()
-            database = str(Path(database).expanduser())
+            database = str(Path(database).expanduser().resolve())
             host = "localhost"
             port = 0
             username = "sqlite"

sqlsaber/config/database.py

@@ -4,12 +4,12 @@ import json
 import os
 import platform
 import stat
-import keyring
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Dict, List, Optional, Any
+from typing import Any, Dict, List, Optional
 from urllib.parse import quote_plus
 
+import keyring
 import platformdirs
 
 
@@ -18,7 +18,7 @@ class DatabaseConfig:
     """Database connection configuration."""
 
     name: str
-    type: str  # postgresql, mysql, sqlite
+    type: str  # postgresql, mysql, sqlite, csv
     host: Optional[str]
     port: Optional[int]
     database: str
@@ -90,6 +90,28 @@ class DatabaseConfig:
 
         elif self.type == "sqlite":
             return f"sqlite:///{self.database}"
+        elif self.type == "csv":
+            # For CSV files, database field contains the file path
+            base_url = f"csv:///{self.database}"
+
+            # Add CSV-specific parameters if they exist in schema field
+            if self.schema:
+                # Schema field can contain CSV options in JSON format
+                try:
+                    csv_options = json.loads(self.schema)
+                    params = []
+                    if "delimiter" in csv_options:
+                        params.append(f"delimiter={csv_options['delimiter']}")
+                    if "encoding" in csv_options:
+                        params.append(f"encoding={csv_options['encoding']}")
+                    if "header" in csv_options:
+                        params.append(f"header={str(csv_options['header']).lower()}")
+
+                    if params:
+                        return f"{base_url}?{'&'.join(params)}"
+                except (json.JSONDecodeError, KeyError):
+                    pass
+            return base_url
         else:
             raise ValueError(f"Unsupported database type: {self.type}")
 

sqlsaber/database/connection.py

@@ -4,10 +4,12 @@ from abc import ABC, abstractmethod
 from typing import Any, Dict, List, Optional
 from urllib.parse import urlparse, parse_qs
 import ssl
+from pathlib import Path
 
 import aiomysql
 import aiosqlite
 import asyncpg
+import pandas as pd
 
 
 class BaseDatabaseConnection(ABC):
@@ -272,6 +274,131 @@ class SQLiteConnection(BaseDatabaseConnection):
                 await conn.rollback()
 
 
+class CSVConnection(BaseDatabaseConnection):
+    """CSV file connection using in-memory SQLite database."""
+
+    def __init__(self, connection_string: str):
+        super().__init__(connection_string)
+
+        # Parse CSV file path from connection string
+        self.csv_path = connection_string.replace("csv:///", "")
+
+        # CSV parsing options
+        self.delimiter = ","
+        self.encoding = "utf-8"
+        self.has_header = True
+
+        # Parse additional options from connection string
+        parsed = urlparse(connection_string)
+        if parsed.query:
+            params = parse_qs(parsed.query)
+            self.delimiter = params.get("delimiter", [","])[0]
+            self.encoding = params.get("encoding", ["utf-8"])[0]
+            self.has_header = params.get("header", ["true"])[0].lower() == "true"
+
+        # Table name derived from filename
+        self.table_name = Path(self.csv_path).stem
+
+        # Initialize connection and flag to track if CSV is loaded
+        self._conn = None
+        self._csv_loaded = False
+
+    async def get_pool(self):
+        """Get or create the in-memory database connection."""
+        if self._conn is None:
+            self._conn = await aiosqlite.connect(":memory:")
+            self._conn.row_factory = aiosqlite.Row
+            await self._load_csv_data()
+        return self._conn
+
+    async def close(self):
+        """Close the database connection."""
+        if self._conn:
+            await self._conn.close()
+            self._conn = None
+            self._csv_loaded = False
+
+    async def _load_csv_data(self):
+        """Load CSV data into the in-memory SQLite database."""
+        if self._csv_loaded or not self._conn:
+            return
+
+        try:
+            # Read CSV file using pandas
+            df = pd.read_csv(
+                self.csv_path,
+                delimiter=self.delimiter,
+                encoding=self.encoding,
+                header=0 if self.has_header else None,
+            )
+
+            # If no header, create column names
+            if not self.has_header:
+                df.columns = [f"column_{i}" for i in range(len(df.columns))]
+
+            # Create table with proper column types
+            columns_sql = []
+            for col in df.columns:
+                # Infer SQLite type from pandas dtype
+                dtype = df[col].dtype
+                if pd.api.types.is_integer_dtype(dtype):
+                    sql_type = "INTEGER"
+                elif pd.api.types.is_float_dtype(dtype):
+                    sql_type = "REAL"
+                elif pd.api.types.is_bool_dtype(dtype):
+                    sql_type = "INTEGER"  # SQLite doesn't have BOOLEAN
+                else:
+                    sql_type = "TEXT"
+
+                columns_sql.append(f'"{col}" {sql_type}')
+
+            create_table_sql = (
+                f'CREATE TABLE "{self.table_name}" ({", ".join(columns_sql)})'
+            )
+            await self._conn.execute(create_table_sql)
+
+            # Insert data row by row
+            placeholders = ", ".join(["?" for _ in df.columns])
+            insert_sql = f'INSERT INTO "{self.table_name}" VALUES ({placeholders})'
+
+            for _, row in df.iterrows():
+                # Convert pandas values to Python native types
+                values = []
+                for val in row:
+                    if pd.isna(val):
+                        values.append(None)
+                    elif isinstance(val, (pd.Timestamp, pd.Timedelta)):
+                        values.append(str(val))
+                    else:
+                        values.append(val)
+
+                await self._conn.execute(insert_sql, values)
+
+            await self._conn.commit()
+            self._csv_loaded = True
+
+        except Exception as e:
+            raise ValueError(f"Error loading CSV file '{self.csv_path}': {str(e)}")
+
+    async def execute_query(self, query: str, *args) -> List[Dict[str, Any]]:
+        """Execute a query and return results as list of dicts.
+
+        All queries run in a transaction that is rolled back at the end,
+        ensuring no changes are persisted to the database.
+        """
+        conn = await self.get_pool()
+
+        # Start transaction
+        await conn.execute("BEGIN")
+        try:
+            cursor = await conn.execute(query, args if args else ())
+            rows = await cursor.fetchall()
+            return [dict(row) for row in rows]
+        finally:
+            # Always rollback to ensure no changes are committed
+            await conn.rollback()
+
+
 def DatabaseConnection(connection_string: str) -> BaseDatabaseConnection:
     """Factory function to create appropriate database connection based on connection string."""
     if connection_string.startswith("postgresql://"):
@@ -280,6 +407,8 @@ def DatabaseConnection(connection_string: str) -> BaseDatabaseConnection:
         return MySQLConnection(connection_string)
     elif connection_string.startswith("sqlite:///"):
         return SQLiteConnection(connection_string)
+    elif connection_string.startswith("csv:///"):
+        return CSVConnection(connection_string)
     else:
         raise ValueError(
             f"Unsupported database type in connection string: {connection_string}"

sqlsaber/database/schema.py

@@ -8,6 +8,7 @@ import aiosqlite
 
 from sqlsaber.database.connection import (
     BaseDatabaseConnection,
+    CSVConnection,
     MySQLConnection,
     PostgreSQLConnection,
     SQLiteConnection,
@@ -200,8 +201,8 @@ class PostgreSQLSchemaIntrospector(BaseSchemaIntrospector):
                     t.table_type,
                     COALESCE(ts.approximate_row_count, 0) as row_count
                 FROM information_schema.tables t
-                LEFT JOIN table_stats ts 
-                    ON t.table_schema = ts.schemaname 
+                LEFT JOIN table_stats ts
+                    ON t.table_schema = ts.schemaname
                     AND t.table_name = ts.tablename
                 WHERE t.table_schema NOT IN ('pg_catalog', 'information_schema')
                 ORDER BY t.table_schema, t.table_name;
@@ -375,15 +376,30 @@ class MySQLSchemaIntrospector(BaseSchemaIntrospector):
 class SQLiteSchemaIntrospector(BaseSchemaIntrospector):
     """SQLite-specific schema introspection."""
 
+    async def _execute_query(self, connection, query: str, params=()) -> list:
+        """Helper method to execute queries on both SQLite and CSV connections."""
+        # Handle both SQLite and CSV connections
+        if hasattr(connection, "database_path"):
+            # Regular SQLite connection
+            async with aiosqlite.connect(connection.database_path) as conn:
+                conn.row_factory = aiosqlite.Row
+                cursor = await conn.execute(query, params)
+                return await cursor.fetchall()
+        else:
+            # CSV connection - use the existing connection
+            conn = await connection.get_pool()
+            cursor = await conn.execute(query, params)
+            return await cursor.fetchall()
+
     async def get_tables_info(
         self, connection, table_pattern: Optional[str] = None
     ) -> Dict[str, Any]:
         """Get tables information for SQLite."""
-        where_clause = ""
+        where_conditions = ["type IN ('table', 'view')", "name NOT LIKE 'sqlite_%'"]
         params = ()
 
         if table_pattern:
-            where_clause = "WHERE name LIKE ?"
+            where_conditions.append("name LIKE ?")
             params = (table_pattern,)
 
         query = f"""
@@ -392,16 +408,11 @@ class SQLiteSchemaIntrospector(BaseSchemaIntrospector):
                 name as table_name,
                 type as table_type
             FROM sqlite_master
-            WHERE type IN ('table', 'view')
-            AND name NOT LIKE 'sqlite_%'
-            {where_clause}
+            WHERE {" AND ".join(where_conditions)}
             ORDER BY name;
         """
 
-        async with aiosqlite.connect(connection.database_path) as conn:
-            conn.row_factory = aiosqlite.Row
-            cursor = await conn.execute(query, params)
-            return await cursor.fetchall()
+        return await self._execute_query(connection, query, params)
 
     async def get_columns_info(self, connection, tables: list) -> list:
         """Get columns information for SQLite."""
@@ -414,26 +425,22 @@ class SQLiteSchemaIntrospector(BaseSchemaIntrospector):
 
             # Get table info using PRAGMA
             pragma_query = f"PRAGMA table_info({table_name})"
+            table_columns = await self._execute_query(connection, pragma_query)
 
-            async with aiosqlite.connect(connection.database_path) as conn:
-                conn.row_factory = aiosqlite.Row
-                cursor = await conn.execute(pragma_query)
-                table_columns = await cursor.fetchall()
-
-                for col in table_columns:
-                    columns.append(
-                        {
-                            "table_schema": "main",
-                            "table_name": table_name,
-                            "column_name": col["name"],
-                            "data_type": col["type"],
-                            "is_nullable": "YES" if not col["notnull"] else "NO",
-                            "column_default": col["dflt_value"],
-                            "character_maximum_length": None,
-                            "numeric_precision": None,
-                            "numeric_scale": None,
-                        }
-                    )
+            for col in table_columns:
+                columns.append(
+                    {
+                        "table_schema": "main",
+                        "table_name": table_name,
+                        "column_name": col["name"],
+                        "data_type": col["type"],
+                        "is_nullable": "YES" if not col["notnull"] else "NO",
+                        "column_default": col["dflt_value"],
+                        "character_maximum_length": None,
+                        "numeric_precision": None,
+                        "numeric_scale": None,
+                    }
+                )
 
         return columns
 
@@ -448,23 +455,19 @@ class SQLiteSchemaIntrospector(BaseSchemaIntrospector):
 
             # Get foreign key info using PRAGMA
             pragma_query = f"PRAGMA foreign_key_list({table_name})"
+            table_fks = await self._execute_query(connection, pragma_query)
 
-            async with aiosqlite.connect(connection.database_path) as conn:
-                conn.row_factory = aiosqlite.Row
-                cursor = await conn.execute(pragma_query)
-                table_fks = await cursor.fetchall()
-
-                for fk in table_fks:
-                    foreign_keys.append(
-                        {
-                            "table_schema": "main",
-                            "table_name": table_name,
-                            "column_name": fk["from"],
-                            "foreign_table_schema": "main",
-                            "foreign_table_name": fk["table"],
-                            "foreign_column_name": fk["to"],
-                        }
-                    )
+            for fk in table_fks:
+                foreign_keys.append(
+                    {
+                        "table_schema": "main",
+                        "table_name": table_name,
+                        "column_name": fk["from"],
+                        "foreign_table_schema": "main",
+                        "foreign_table_name": fk["table"],
+                        "foreign_column_name": fk["to"],
+                    }
+                )
 
         return foreign_keys
 
@@ -479,43 +482,64 @@ class SQLiteSchemaIntrospector(BaseSchemaIntrospector):
 
             # Get table info using PRAGMA to find primary keys
             pragma_query = f"PRAGMA table_info({table_name})"
+            table_columns = await self._execute_query(connection, pragma_query)
 
-            async with aiosqlite.connect(connection.database_path) as conn:
-                conn.row_factory = aiosqlite.Row
-                cursor = await conn.execute(pragma_query)
-                table_columns = await cursor.fetchall()
-
-                for col in table_columns:
-                    if col["pk"]:  # Primary key indicator
-                        primary_keys.append(
-                            {
-                                "table_schema": "main",
-                                "table_name": table_name,
-                                "column_name": col["name"],
-                            }
-                        )
+            for col in table_columns:
+                if col["pk"]:  # Primary key indicator
+                    primary_keys.append(
+                        {
+                            "table_schema": "main",
+                            "table_name": table_name,
+                            "column_name": col["name"],
+                        }
+                    )
 
         return primary_keys
 
     async def list_tables_info(self, connection) -> Dict[str, Any]:
         """Get list of tables with basic information for SQLite."""
-        # Get tables (SQLite doesn't have easy row count access)
+        # First get the table names
         tables_query = """
             SELECT
                 'main' as table_schema,
                 name as table_name,
-                type as table_type,
-                0 as row_count
+                type as table_type
             FROM sqlite_master
             WHERE type IN ('table', 'view')
             AND name NOT LIKE 'sqlite_%'
             ORDER BY name;
         """
 
-        async with aiosqlite.connect(connection.database_path) as conn:
-            conn.row_factory = aiosqlite.Row
-            cursor = await conn.execute(tables_query)
-            return await cursor.fetchall()
+        tables = await self._execute_query(connection, tables_query)
+
+        # Now get row counts for each table
+        result = []
+        for table in tables:
+            table_name = table["table_name"]
+            table_type = table["table_type"]
+
+            # Only count rows for tables, not views
+            if table_type.lower() == "table":
+                try:
+                    count_query = f"SELECT COUNT(*) as count FROM [{table_name}]"
+                    count_result = await self._execute_query(connection, count_query)
+                    row_count = count_result[0]["count"] if count_result else 0
+                except Exception:
+                    # If count fails (e.g., table locked), default to 0
+                    row_count = 0
+            else:
+                # For views, we don't count rows as it could be expensive
+                row_count = 0
+
+            result.append(
+                {
+                    "table_schema": table["table_schema"],
+                    "table_name": table_name,
+                    "table_type": table_type,
+                    "row_count": row_count,
+                }
+            )
+        return result
 
 
 class SchemaManager:
@@ -531,7 +555,7 @@ class SchemaManager:
             self.introspector = PostgreSQLSchemaIntrospector()
         elif isinstance(db_connection, MySQLConnection):
             self.introspector = MySQLSchemaIntrospector()
-        elif isinstance(db_connection, SQLiteConnection):
+        elif isinstance(db_connection, (SQLiteConnection, CSVConnection)):
             self.introspector = SQLiteSchemaIntrospector()
         else:
             raise ValueError(

sqlsaber/mcp/__init__.py

@@ -0,0 +1,5 @@
+"""MCP (Model Context Protocol) server implementation for SQLSaber."""
+
+from .mcp import mcp
+
+__all__ = ["mcp"]

sqlsaber/mcp/mcp.py

@@ -0,0 +1,138 @@
+"""FastMCP server implementation for SQLSaber."""
+
+import json
+from typing import Optional
+
+from fastmcp import FastMCP
+
+from sqlsaber.agents.mcp import MCPSQLAgent
+from sqlsaber.config.database import DatabaseConfigManager
+from sqlsaber.database.connection import DatabaseConnection
+
+INSTRUCTIONS = """
+This server provides helpful resources and tools that will help you address users queries on their database.
+
+- Get all databases using `get_databases()`
+- Call `list_tables()` to get a list of all tables in the database with row counts. Use this first to discover available tables.
+- Call `introspect_schema()` to introspect database schema to understand table structures.
+- Call `execute_sql()` to execute SQL queries against the database and retrieve results.
+
+Guidelines:
+- Use list_tables first, then introspect_schema for specific tables only
+- Use table patterns like 'sample%' or '%experiment%' to filter related tables
+- Use proper JOIN syntax and avoid cartesian products
+- Include appropriate WHERE clauses to limit results
+- Handle errors gracefully and suggest fixes
+"""
+
+# Create the FastMCP server instance
+mcp = FastMCP(name="SQL Assistant", instructions=INSTRUCTIONS)
+
+# Initialize the database config manager
+config_manager = DatabaseConfigManager()
+
+
+async def _create_agent_for_database(database_name: str) -> Optional[MCPSQLAgent]:
+    """Create a MCPSQLAgent for the specified database."""
+    try:
+        # Look up configured database connection
+        db_config = config_manager.get_database(database_name)
+        if not db_config:
+            return None
+        connection_string = db_config.to_connection_string()
+
+        # Create database connection
+        db_conn = DatabaseConnection(connection_string)
+
+        # Create and return the agent
+        agent = MCPSQLAgent(db_conn)
+        return agent
+
+    except Exception:
+        return None
+
+
+@mcp.tool
+def get_databases() -> dict:
+    """List all configured databases with their types."""
+    databases = []
+    for db_config in config_manager.list_databases():
+        databases.append(
+            {
+                "name": db_config.name,
+                "type": db_config.type,
+                "database": db_config.database,
+                "host": db_config.host,
+                "port": db_config.port,
+                "is_default": db_config.name == config_manager.get_default_name(),
+            }
+        )
+
+    return {"databases": databases, "count": len(databases)}
+
+
+@mcp.tool
+async def list_tables(database: str) -> str:
+    """
+    Get a list of all tables in the database with row counts. Use this first to discover available tables.
+    """
+    try:
+        agent = await _create_agent_for_database(database)
+        if not agent:
+            return json.dumps(
+                {"error": f"Database '{database}' not found or could not connect"}
+            )
+
+        result = await agent.list_tables()
+        await agent.db.close()
+        return result
+
+    except Exception as e:
+        return json.dumps({"error": f"Error listing tables: {str(e)}"})
+
+
+@mcp.tool
+async def introspect_schema(database: str, table_pattern: Optional[str] = None) -> str:
+    """
+    Introspect database schema to understand table structures. Use optional pattern to filter tables (e.g., 'public.users', 'user%', '%order%').
+    """
+    try:
+        agent = await _create_agent_for_database(database)
+        if not agent:
+            return json.dumps(
+                {"error": f"Database '{database}' not found or could not connect"}
+            )
+
+        result = await agent.introspect_schema(table_pattern)
+        await agent.db.close()
+        return result
+
+    except Exception as e:
+        return json.dumps({"error": f"Error introspecting schema: {str(e)}"})
+
+
+@mcp.tool
+async def execute_sql(database: str, query: str, limit: Optional[int] = 100) -> str:
+    """Execute a SQL query against the specified database."""
+    try:
+        agent = await _create_agent_for_database(database)
+        if not agent:
+            return json.dumps(
+                {"error": f"Database '{database}' not found or could not connect"}
+            )
+
+        result = await agent.execute_sql(query, limit)
+        await agent.db.close()
+        return result
+
+    except Exception as e:
+        return json.dumps({"error": f"Error executing SQL: {str(e)}"})
+
+
+def main():
+    """Entry point for the MCP server console script."""
+    mcp.run()
+
+
+if __name__ == "__main__":
+    main()

{sqlsaber-0.2.0.dist-info → sqlsaber-0.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlsaber
-Version: 0.2.0
+Version: 0.4.0
 Summary: SQLSaber - Agentic SQL assistant like Claude Code
 License-File: LICENSE
 Requires-Python: >=3.12
@@ -8,8 +8,10 @@ Requires-Dist: aiomysql>=0.2.0
 Requires-Dist: aiosqlite>=0.21.0
 Requires-Dist: anthropic>=0.54.0
 Requires-Dist: asyncpg>=0.30.0
+Requires-Dist: fastmcp>=2.9.0
 Requires-Dist: httpx>=0.28.1
 Requires-Dist: keyring>=25.6.0
+Requires-Dist: pandas>=2.0.0
 Requires-Dist: platformdirs>=4.0.0
 Requires-Dist: questionary>=2.1.0
 Requires-Dist: rich>=13.7.0
@@ -33,6 +35,7 @@ Ask your questions in natural language and it will gather the right context and
 - 💬 Interactive REPL mode
 - 🎨 Beautiful formatted output with syntax highlighting
 - 🗄️ Support for PostgreSQL, SQLite, and MySQL
+- 🔌 MCP (Model Context Protocol) server support
 
 ## Installation
 
@@ -139,6 +142,43 @@ saber query "show me the distribution of customer ages"
 saber query "which products had the highest sales growth last quarter?"
 ```
 
+## MCP Server Integration
+
+SQLSaber includes an MCP (Model Context Protocol) server that allows AI agents like Claude Code to directly leverage tools available in SQLSaber.
+
+### Starting the MCP Server
+
+Run the MCP server using uvx:
+
+```bash
+uvx saber-mcp
+```
+
+### Configuring MCP Clients
+
+#### Claude Code
+
+Add SQLSaber as an MCP server in Claude Code:
+
+```bash
+claude mcp add -- uvx saber-mcp
+```
+
+#### Other MCP Clients
+
+For other MCP clients, configure them to run the command: `uvx saber-mcp`
+
+### Available MCP Tools
+
+Once connected, the MCP client will have access to these tools:
+
+- `get_databases()` - Lists all configured databases
+- `list_tables(database)` - Get all tables in a database with row counts
+- `introspect_schema(database, table_pattern?)` - Get detailed schema information
+- `execute_sql(database, query, limit?)` - Execute SQL queries (read-only)
+
+The MCP server uses your existing SQLSaber database configurations, so make sure to set up your databases using `saber db add` first.
+
 ## How It Works
 
 SQLSaber uses an intelligent three-step process optimized for minimal token usage:

{sqlsaber-0.2.0.dist-info → sqlsaber-0.4.0.dist-info}/RECORD

@@ -1,12 +1,13 @@
 sqlsaber/__init__.py,sha256=QCFi8xTVMohelfi7zOV1-6oLCcGoiXoOcKQY-HNBCk8,66
 sqlsaber/__main__.py,sha256=RIHxWeWh2QvLfah-2OkhI5IJxojWfy4fXpMnVEJYvxw,78
 sqlsaber/agents/__init__.py,sha256=LWeSeEUE4BhkyAYFF3TE-fx8TtLud3oyEtyB8ojFJgo,167
-sqlsaber/agents/anthropic.py,sha256=CPNshN68QxxdIvWS5YjEuiXd_8V8mmMTL9aC2Mqkvts,17863
-sqlsaber/agents/base.py,sha256=UUSGhoJImATXrYS7yrLR2qjg1iFW4udOUdRaV3Ryk5s,2086
-sqlsaber/agents/streaming.py,sha256=0bNzd_JhLlgQB40pf9FZFMvmU9Q7W6D9BmglA1rIGqw,850
+sqlsaber/agents/anthropic.py,sha256=xAjKeQSnaut-P5VBeBISbQeqdP41epDjX6MJb2ZUXWg,14060
+sqlsaber/agents/base.py,sha256=IuVyCaA7VsA92odfQS2_lYNzwIZwPxK55mL_xRewgwQ,6943
+sqlsaber/agents/mcp.py,sha256=FKtXgDrPZ2-xqUYCw2baI5JzrWekXaC5fjkYW1_Mg50,827
+sqlsaber/agents/streaming.py,sha256=_EO390-FHUrL1fRCNfibtE9QuJz3LGQygbwG3CB2ViY,533
 sqlsaber/cli/__init__.py,sha256=qVSLVJLLJYzoC6aj6y9MFrzZvAwc4_OgxU9DlkQnZ4M,86
-sqlsaber/cli/commands.py,sha256=Adrt_0LRgykb2FZ4F0TQpuBM8Z0qgfbggn0FexcVALI,4094
-sqlsaber/cli/database.py,sha256=ZOaW94dCM8_RUT5OZzI6-lbVmSYxSjdmAdYa3Nf_d9g,12633
+sqlsaber/cli/commands.py,sha256=h418lgh_Xp7XEQ1xvjcDyplC2JON0-y98QMaDm6o29k,4919
+sqlsaber/cli/database.py,sha256=DUfyvNBDp47oFM_VAC_hXHQy_qyE7JbXtowflJpwwH8,12643
 sqlsaber/cli/display.py,sha256=5J4AgJADmMwKi9Aq5u6_MKRO1TA6unS4F4RUfml_sfU,7651
 sqlsaber/cli/interactive.py,sha256=y92rdoM49SOSwEctm9ZcrEN220fhJ_DMHPSd_7KsORg,3701
 sqlsaber/cli/memory.py,sha256=LW4ZF2V6Gw6hviUFGZ4ym9ostFCwucgBTIMZ3EANO-I,7671
@@ -14,19 +15,21 @@ sqlsaber/cli/models.py,sha256=3IcXeeU15IQvemSv-V-RQzVytJ3wuQ4YmWk89nTDcSE,7813
 sqlsaber/cli/streaming.py,sha256=5QGAYTAvg9mzQLxDEVtdDH-TIbGfYYzMOLoOYPrHPu0,3788
 sqlsaber/config/__init__.py,sha256=olwC45k8Nc61yK0WmPUk7XHdbsZH9HuUAbwnmKe3IgA,100
 sqlsaber/config/api_keys.py,sha256=kLdoExF_My9ojmdhO5Ca7-ZeowsO0v1GVa_QT5jjUPo,3658
-sqlsaber/config/database.py,sha256=hrIr5IIqhkOpjJ2A4oaeMvTqWqNDdlJiVWQTaBudVQM,10388
+sqlsaber/config/database.py,sha256=vKFOxPjVakjQhj1uoLcfzhS9ZFr6Z2F5b4MmYALQZoA,11421
 sqlsaber/config/settings.py,sha256=zjQ7nS3ybcCb88Ea0tmwJox5-q0ettChZw89ZqRVpX8,3975
 sqlsaber/database/__init__.py,sha256=a_gtKRJnZVO8-fEZI7g3Z8YnGa6Nio-5Y50PgVp07ss,176
-sqlsaber/database/connection.py,sha256=1lMQ2LhxlUIkl_msGzuNXczI_tNp0WBgfsdqDC2nxsw,10479
-sqlsaber/database/schema.py,sha256=gURfCFVE--UWIqD_0StqS2NMB9VIPpqczBEoS2GnKR4,27025
+sqlsaber/database/connection.py,sha256=s8GSFZebB8be8sVUr-N0x88-20YfkfljJFRyfoB1gH0,15154
+sqlsaber/database/schema.py,sha256=9QoH-gADzWlepq-tGz3nPU3miSUU0koWmpDaoWvz8Q0,27951
+sqlsaber/mcp/__init__.py,sha256=COdWq7wauPBp5Ew8tfZItFzbcLDSEkHBJSMhxzy8C9c,112
+sqlsaber/mcp/mcp.py,sha256=ACm1P1TnicjOptQgeLNhXg5xgZf4MYq2kqdfVdj6wh0,4477
 sqlsaber/memory/__init__.py,sha256=GiWkU6f6YYVV0EvvXDmFWe_CxarmDCql05t70MkTEWs,63
 sqlsaber/memory/manager.py,sha256=ML2NEO5Z4Aw36sEI9eOvWVnjl-qT2VOTojViJAj7Seo,2777
 sqlsaber/memory/storage.py,sha256=DvZBsSPaAfk_DqrNEn86uMD-TQsWUI6rQLfNw6PSCB8,5788
 sqlsaber/models/__init__.py,sha256=RJ7p3WtuSwwpFQ1Iw4_DHV2zzCtHqIzsjJzxv8kUjUE,287
 sqlsaber/models/events.py,sha256=55m41tDwMsFxnKKA5_VLJz8iV-V4Sq3LDfta4VoutJI,737
 sqlsaber/models/types.py,sha256=3U_30n91EB3IglBTHipwiW4MqmmaA2qfshfraMZyPps,896
-sqlsaber-0.2.0.dist-info/METADATA,sha256=EUgK7o9feWxmsx_pr9DwER6_t6db7ZaYSghwUKKYPao,3953
-sqlsaber-0.2.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-sqlsaber-0.2.0.dist-info/entry_points.txt,sha256=POwcsEskUp7xQQWabrAi6Eawz4qc5eBlB3KzAiBq-Y0,124
-sqlsaber-0.2.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-sqlsaber-0.2.0.dist-info/RECORD,,
+sqlsaber-0.4.0.dist-info/METADATA,sha256=CL1mNjOLrc6VDJqE2dSrCXO5OJz9gTMxYNoYq6jtzYE,5071
+sqlsaber-0.4.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+sqlsaber-0.4.0.dist-info/entry_points.txt,sha256=jmFo96Ylm0zIKXJBwhv_P5wQ7SXP9qdaBcnTp8iCEe8,195
+sqlsaber-0.4.0.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+sqlsaber-0.4.0.dist-info/RECORD,,

{sqlsaber-0.2.0.dist-info → sqlsaber-0.4.0.dist-info}/entry_points.txt

@@ -1,4 +1,6 @@
 [console_scripts]
 saber = sqlsaber.cli.commands:main
+saber-mcp = sqlsaber.mcp.mcp:main
 sql = sqlsaber.cli.commands:main
 sqlsaber = sqlsaber.cli.commands:main
+sqlsaber-mcp = sqlsaber.mcp.mcp:main