db-connect-mcp 0.1.0__py3-none-any.whl

db_connect_mcp/adapters/mysql.py

@@ -0,0 +1,288 @@
+"""MySQL adapter with good feature support."""
+
+import json
+from typing import Any, Optional
+
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncConnection
+
+from db_connect_mcp.adapters.base import BaseAdapter
+from db_connect_mcp.models.capabilities import DatabaseCapabilities
+from db_connect_mcp.models.database import SchemaInfo
+from db_connect_mcp.models.statistics import ColumnStats, Distribution
+from db_connect_mcp.models.table import TableInfo
+
+
+class MySQLAdapter(BaseAdapter):
+    """MySQL adapter with good feature support."""
+
+    @property
+    def capabilities(self) -> DatabaseCapabilities:
+        """MySQL has good but not comprehensive support."""
+        return DatabaseCapabilities(
+            foreign_keys=True,
+            indexes=True,
+            views=True,
+            materialized_views=False,  # MySQL doesn't have native materialized views
+            partitions=True,
+            advanced_stats=False,  # No percentile functions in MySQL
+            explain_plans=True,
+            profiling=False,  # Basic profiling only
+            comments=True,
+            schemas=True,  # MySQL calls them databases
+            transactions=True,
+            stored_procedures=True,
+            triggers=True,
+        )
+
+    async def enrich_schema_info(
+        self, conn: AsyncConnection, schema_info: SchemaInfo
+    ) -> SchemaInfo:
+        """Add MySQL-specific schema metadata."""
+        query = text("""
+            SELECT
+                SUM(data_length + index_length) as size_bytes
+            FROM information_schema.TABLES
+            WHERE table_schema = :schema_name
+        """)
+
+        result = await conn.execute(query, {"schema_name": schema_info.name})
+        row = result.fetchone()
+
+        if row and row[0]:
+            schema_info.size_bytes = int(row[0])
+
+        return schema_info
+
+    async def enrich_table_info(
+        self, conn: AsyncConnection, table_info: TableInfo
+    ) -> TableInfo:
+        """Add MySQL-specific table metadata."""
+        query = text("""
+            SELECT
+                engine,
+                table_rows,
+                data_length,
+                index_length,
+                table_comment,
+                create_time,
+                update_time
+            FROM information_schema.TABLES
+            WHERE table_schema = DATABASE()
+              AND table_name = :table_name
+        """)
+
+        result = await conn.execute(query, {"table_name": table_info.name})
+        row = result.fetchone()
+
+        if row:
+            table_info.row_count = int(row[1]) if row[1] else None
+            table_info.size_bytes = int(row[2]) if row[2] else None
+            table_info.index_size_bytes = int(row[3]) if row[3] else None
+            table_info.comment = row[4] if row[4] else None
+            table_info.created_at = str(row[5]) if row[5] else None
+            table_info.updated_at = str(row[6]) if row[6] else None
+
+            # MySQL-specific: storage engine
+            table_info.extra_info["engine"] = row[0]
+
+        return table_info
+
+    async def get_column_statistics(
+        self,
+        conn: AsyncConnection,
+        table_name: str,
+        column_name: str,
+        schema: Optional[str],
+    ) -> ColumnStats:
+        """Get MySQL column statistics (basic stats only)."""
+        table_ref = self._build_table_reference(table_name, schema)
+
+        # MySQL doesn't support percentile functions, so we get basic stats only
+        query = text(f"""
+            SELECT
+                COUNT(*) as total_rows,
+                COUNT(*) - COUNT(`{column_name}`) as null_count,
+                COUNT(DISTINCT `{column_name}`) as distinct_count,
+                MIN(`{column_name}`) as min_val,
+                MAX(`{column_name}`) as max_val,
+                AVG(`{column_name}`) as avg_val,
+                STD(`{column_name}`) as stddev_val
+            FROM {table_ref}
+        """)
+
+        try:
+            result = await conn.execute(query)
+            row = result.fetchone()
+
+            if not row:
+                return ColumnStats(
+                    column=column_name,
+                    data_type="unknown",
+                    total_rows=0,
+                    null_count=0,
+                    sample_size=0,
+                    warning="No data found",
+                )
+
+            # Get most common values
+            mcv_query = text(f"""
+                SELECT `{column_name}` as value, COUNT(*) as count
+                FROM {table_ref}
+                WHERE `{column_name}` IS NOT NULL
+                GROUP BY `{column_name}`
+                ORDER BY count DESC
+                LIMIT 10
+            """)
+
+            mcv_result = await conn.execute(mcv_query)
+            mcv_rows = mcv_result.fetchall()
+            most_common = [{"value": str(r[0]), "count": int(r[1])} for r in mcv_rows]
+
+            # Get data type from information_schema
+            type_query = text("""
+                SELECT data_type
+                FROM information_schema.COLUMNS
+                WHERE table_schema = DATABASE()
+                  AND table_name = :table_name
+                  AND column_name = :column_name
+            """)
+
+            type_result = await conn.execute(
+                type_query, {"table_name": table_name, "column_name": column_name}
+            )
+            type_row = type_result.fetchone()
+            data_type = type_row[0] if type_row else "unknown"
+
+            return ColumnStats(
+                column=column_name,
+                data_type=data_type,
+                total_rows=int(row[0]),
+                null_count=int(row[1]),
+                distinct_count=int(row[2]) if row[2] else None,
+                min_value=row[3],
+                max_value=row[4],
+                avg_value=float(row[5]) if row[5] is not None else None,
+                stddev_value=float(row[6]) if row[6] is not None else None,
+                most_common_values=most_common,
+                sample_size=int(row[0]),
+                warning="Advanced statistics (percentiles) not available in MySQL",
+            )
+
+        except Exception as e:
+            return ColumnStats(
+                column=column_name,
+                data_type="unknown",
+                total_rows=0,
+                null_count=0,
+                sample_size=0,
+                warning=f"Statistics unavailable: {str(e)}",
+            )
+
+    async def get_value_distribution(
+        self,
+        conn: AsyncConnection,
+        table_name: str,
+        column_name: str,
+        schema: Optional[str],
+        limit: int,
+    ) -> Distribution:
+        """Get value distribution for MySQL."""
+        table_ref = self._build_table_reference(table_name, schema)
+
+        stats_query = text(f"""
+            SELECT
+                COUNT(*) as total_rows,
+                COUNT(DISTINCT `{column_name}`) as unique_values,
+                COUNT(*) - COUNT(`{column_name}`) as null_count
+            FROM {table_ref}
+        """)
+
+        stats_result = await conn.execute(stats_query)
+        stats_row = stats_result.fetchone()
+
+        top_query = text(f"""
+            SELECT `{column_name}` as value, COUNT(*) as count
+            FROM {table_ref}
+            WHERE `{column_name}` IS NOT NULL
+            GROUP BY `{column_name}`
+            ORDER BY count DESC
+            LIMIT :limit
+        """)
+
+        top_result = await conn.execute(top_query, {"limit": limit})
+        top_rows = top_result.fetchall()
+
+        top_values = [{"value": str(r[0]), "count": int(r[1])} for r in top_rows]
+
+        if not stats_row:
+            return Distribution(
+                column=column_name,
+                total_rows=0,
+                unique_values=0,
+                null_count=0,
+                top_values=[],
+                sample_size=0,
+            )
+
+        return Distribution(
+            column=column_name,
+            total_rows=int(stats_row[0]),
+            unique_values=int(stats_row[1]),
+            null_count=int(stats_row[2]),
+            top_values=top_values,
+            sample_size=int(stats_row[0]),
+        )
+
+    async def get_sample_query(
+        self, table_name: str, schema: Optional[str], limit: int
+    ) -> str:
+        """Generate MySQL sampling query."""
+        table_ref = self._build_table_reference(table_name, schema)
+        return f"SELECT * FROM {table_ref} LIMIT {limit}"
+
+    async def get_explain_query(self, query: str, analyze: bool) -> str:
+        """Generate MySQL EXPLAIN query."""
+        if analyze:
+            return f"EXPLAIN ANALYZE {query}"
+        return f"EXPLAIN FORMAT=JSON {query}"
+
+    async def parse_explain_plan(
+        self, plan_text: str, analyzed: bool
+    ) -> dict[str, Any]:
+        """Parse MySQL EXPLAIN output."""
+        try:
+            plan_data = json.loads(plan_text)
+
+            result: dict[str, Any] = {
+                "json": plan_data,
+                "warnings": [],
+                "recommendations": [],
+            }
+
+            # MySQL EXPLAIN has different structure
+            if "query_block" in plan_data:
+                query_block = plan_data["query_block"]
+
+                # Extract cost if available
+                if "cost_info" in query_block:
+                    cost_info = query_block["cost_info"]
+                    result["estimated_cost"] = float(cost_info.get("query_cost", 0))
+
+                # Check for table scans
+                if "table" in query_block:
+                    table = query_block["table"]
+                    if table.get("access_type") == "ALL":
+                        result["warnings"].append("Full table scan detected")
+                        result["recommendations"].append("Consider adding indexes")
+
+            return result
+
+        except (json.JSONDecodeError, KeyError):
+            pass
+
+        return {
+            "json": None,
+            "warnings": [],
+            "recommendations": [],
+        }

db_connect_mcp/adapters/postgresql.py

@@ -0,0 +1,351 @@
+"""PostgreSQL adapter with full feature support."""
+
+import json
+from typing import Any, Optional
+
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncConnection
+
+from db_connect_mcp.adapters.base import BaseAdapter
+from db_connect_mcp.models.capabilities import DatabaseCapabilities
+from db_connect_mcp.models.database import SchemaInfo
+from db_connect_mcp.models.statistics import ColumnStats, Distribution
+from db_connect_mcp.models.table import TableInfo
+
+
+class PostgresAdapter(BaseAdapter):
+    """PostgreSQL adapter with comprehensive feature support."""
+
+    @property
+    def capabilities(self) -> DatabaseCapabilities:
+        """PostgreSQL supports all features."""
+        return DatabaseCapabilities(
+            foreign_keys=True,
+            indexes=True,
+            views=True,
+            materialized_views=True,
+            partitions=True,
+            advanced_stats=True,
+            explain_plans=True,
+            profiling=True,
+            comments=True,
+            schemas=True,
+            transactions=True,
+            stored_procedures=True,
+            triggers=True,
+        )
+
+    async def enrich_schema_info(
+        self, conn: AsyncConnection, schema_info: SchemaInfo
+    ) -> SchemaInfo:
+        """Add PostgreSQL-specific schema metadata."""
+        query = text("""
+            SELECT
+                pg_catalog.pg_get_userbyid(n.nspowner) as owner,
+                pg_catalog.obj_description(n.oid, 'pg_namespace') as comment
+            FROM pg_catalog.pg_namespace n
+            WHERE n.nspname = :schema_name
+        """)
+
+        result = await conn.execute(query, {"schema_name": schema_info.name})
+        row = result.fetchone()
+
+        if row:
+            schema_info.owner = row[0]
+            schema_info.comment = row[1]
+
+        # Get schema size
+        size_query = text("""
+            SELECT SUM(pg_total_relation_size(quote_ident(schemaname) || '.' || quote_ident(tablename)))::bigint
+            FROM pg_tables
+            WHERE schemaname = :schema_name
+        """)
+
+        result = await conn.execute(size_query, {"schema_name": schema_info.name})
+        row = result.fetchone()
+        if row and row[0]:
+            schema_info.size_bytes = int(row[0])
+
+        return schema_info
+
+    async def enrich_table_info(
+        self, conn: AsyncConnection, table_info: TableInfo
+    ) -> TableInfo:
+        """Add PostgreSQL-specific table metadata."""
+        table_ref = self._build_table_reference(table_info.name, table_info.schema)
+
+        query = text("""
+            SELECT
+                pg_total_relation_size(:table_ref::regclass)::bigint as total_size,
+                pg_relation_size(:table_ref::regclass)::bigint as table_size,
+                pg_indexes_size(:table_ref::regclass)::bigint as indexes_size,
+                (SELECT reltuples::bigint FROM pg_class WHERE oid = :table_ref::regclass::oid) as row_count,
+                obj_description(:table_ref::regclass, 'pg_class') as comment
+        """)
+
+        try:
+            result = await conn.execute(query, {"table_ref": table_ref})
+            row = result.fetchone()
+
+            if row:
+                table_info.size_bytes = int(row[1]) if row[1] else None
+                table_info.index_size_bytes = int(row[2]) if row[2] else None
+                table_info.row_count = int(row[3]) if row[3] else None
+                table_info.comment = row[4]
+
+            # Add PostgreSQL-specific extras
+            extras_query = text("""
+                SELECT
+                    c.relkind as table_kind,
+                    c.relpersistence as persistence,
+                    c.relispartition as is_partition
+                FROM pg_class c
+                JOIN pg_namespace n ON n.oid = c.relnamespace
+                WHERE c.relname = :table_name
+                  AND n.nspname = COALESCE(:schema_name, 'public')
+            """)
+
+            result = await conn.execute(
+                extras_query,
+                {"table_name": table_info.name, "schema_name": table_info.schema},
+            )
+            row = result.fetchone()
+
+            if row:
+                table_info.extra_info["relkind"] = row[0]
+                table_info.extra_info["persistence"] = row[1]
+                table_info.extra_info["is_partition"] = row[2]
+
+        except Exception:
+            # If enrichment fails, return basic info
+            pass
+
+        return table_info
+
+    async def get_column_statistics(
+        self,
+        conn: AsyncConnection,
+        table_name: str,
+        column_name: str,
+        schema: Optional[str],
+    ) -> ColumnStats:
+        """Get comprehensive PostgreSQL column statistics."""
+        table_ref = self._build_table_reference(table_name, schema)
+
+        # Basic stats query with PostgreSQL-specific functions
+        query = text(f"""
+            WITH stats AS (
+                SELECT
+                    COUNT(*) as total_rows,
+                    COUNT("{column_name}") as non_null_count,
+                    COUNT(*) - COUNT("{column_name}") as null_count,
+                    COUNT(DISTINCT "{column_name}") as distinct_count,
+                    MIN("{column_name}") as min_val,
+                    MAX("{column_name}") as max_val,
+                    pg_typeof("{column_name}")::text as data_type
+                FROM {table_ref}
+            ),
+            numeric_stats AS (
+                SELECT
+                    AVG("{column_name}")::float as avg_val,
+                    STDDEV("{column_name}")::float as stddev_val,
+                    PERCENTILE_CONT(0.25) WITHIN GROUP (ORDER BY "{column_name}") as p25,
+                    PERCENTILE_CONT(0.50) WITHIN GROUP (ORDER BY "{column_name}") as p50,
+                    PERCENTILE_CONT(0.75) WITHIN GROUP (ORDER BY "{column_name}") as p75,
+                    PERCENTILE_CONT(0.95) WITHIN GROUP (ORDER BY "{column_name}") as p95,
+                    PERCENTILE_CONT(0.99) WITHIN GROUP (ORDER BY "{column_name}") as p99
+                FROM {table_ref}
+                WHERE "{column_name}" IS NOT NULL
+            )
+            SELECT
+                s.total_rows,
+                s.null_count,
+                s.distinct_count,
+                s.min_val,
+                s.max_val,
+                s.data_type,
+                n.avg_val,
+                n.stddev_val,
+                n.p25,
+                n.p50,
+                n.p75,
+                n.p95,
+                n.p99
+            FROM stats s
+            LEFT JOIN numeric_stats n ON true
+        """)
+
+        try:
+            result = await conn.execute(query)
+            row = result.fetchone()
+
+            if not row:
+                return ColumnStats(
+                    column=column_name,
+                    data_type="unknown",
+                    total_rows=0,
+                    null_count=0,
+                    sample_size=0,
+                    warning="No data found",
+                )
+
+            # Get most common values
+            mcv_query = text(f"""
+                SELECT "{column_name}" as value, COUNT(*) as count
+                FROM {table_ref}
+                WHERE "{column_name}" IS NOT NULL
+                GROUP BY "{column_name}"
+                ORDER BY count DESC
+                LIMIT 10
+            """)
+
+            mcv_result = await conn.execute(mcv_query)
+            mcv_rows = mcv_result.fetchall()
+            most_common = [{"value": str(r[0]), "count": int(r[1])} for r in mcv_rows]
+
+            return ColumnStats(
+                column=column_name,
+                data_type=str(row[5]),
+                total_rows=int(row[0]),
+                null_count=int(row[1]),
+                distinct_count=int(row[2]) if row[2] else None,
+                min_value=row[3],
+                max_value=row[4],
+                avg_value=float(row[6]) if row[6] is not None else None,
+                stddev_value=float(row[7]) if row[7] is not None else None,
+                percentile_25=row[8],
+                median_value=row[9],
+                percentile_75=row[10],
+                percentile_95=row[11],
+                percentile_99=row[12],
+                most_common_values=most_common,
+                sample_size=int(row[0]),
+            )
+
+        except Exception as e:
+            return ColumnStats(
+                column=column_name,
+                data_type="unknown",
+                total_rows=0,
+                null_count=0,
+                sample_size=0,
+                warning=f"Statistics unavailable: {str(e)}",
+            )
+
+    async def get_value_distribution(
+        self,
+        conn: AsyncConnection,
+        table_name: str,
+        column_name: str,
+        schema: Optional[str],
+        limit: int,
+    ) -> Distribution:
+        """Get value distribution for PostgreSQL."""
+        table_ref = self._build_table_reference(table_name, schema)
+
+        query = text(f"""
+            WITH stats AS (
+                SELECT
+                    COUNT(*) as total_rows,
+                    COUNT(DISTINCT "{column_name}") as unique_values,
+                    COUNT(*) - COUNT("{column_name}") as null_count
+                FROM {table_ref}
+            ),
+            top_values AS (
+                SELECT "{column_name}" as value, COUNT(*) as count
+                FROM {table_ref}
+                WHERE "{column_name}" IS NOT NULL
+                GROUP BY "{column_name}"
+                ORDER BY count DESC
+                LIMIT :limit
+            )
+            SELECT
+                s.total_rows,
+                s.unique_values,
+                s.null_count,
+                json_agg(json_build_object('value', t.value::text, 'count', t.count)) as top_values
+            FROM stats s
+            LEFT JOIN top_values t ON true
+            GROUP BY s.total_rows, s.unique_values, s.null_count
+        """)
+
+        result = await conn.execute(query, {"limit": limit})
+        row = result.fetchone()
+
+        if not row:
+            return Distribution(
+                column=column_name,
+                total_rows=0,
+                unique_values=0,
+                null_count=0,
+                top_values=[],
+                sample_size=0,
+            )
+
+        top_values_data = json.loads(row[3]) if row[3] else []
+
+        return Distribution(
+            column=column_name,
+            total_rows=int(row[0]),
+            unique_values=int(row[1]),
+            null_count=int(row[2]),
+            top_values=top_values_data,
+            sample_size=int(row[0]),
+        )
+
+    async def get_sample_query(
+        self, table_name: str, schema: Optional[str], limit: int
+    ) -> str:
+        """Generate PostgreSQL sampling query with TABLESAMPLE."""
+        table_ref = self._build_table_reference(table_name, schema)
+        # Use simple LIMIT for smaller limits, TABLESAMPLE for larger datasets
+        return f"SELECT * FROM {table_ref} LIMIT {limit}"
+
+    async def get_explain_query(self, query: str, analyze: bool) -> str:
+        """Generate PostgreSQL EXPLAIN query."""
+        if analyze:
+            return f"EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) {query}"
+        return f"EXPLAIN (FORMAT JSON) {query}"
+
+    async def parse_explain_plan(
+        self, plan_text: str, analyzed: bool
+    ) -> dict[str, Any]:
+        """Parse PostgreSQL EXPLAIN JSON output."""
+        try:
+            plan_data = json.loads(plan_text)
+
+            if isinstance(plan_data, list) and len(plan_data) > 0:
+                plan = plan_data[0].get("Plan", {})
+
+                result: dict[str, Any] = {
+                    "json": plan,
+                    "estimated_cost": plan.get("Total Cost"),
+                    "estimated_rows": plan.get("Plan Rows"),
+                    "warnings": [],
+                    "recommendations": [],
+                }
+
+                if analyzed:
+                    result["actual_time_ms"] = plan.get("Actual Total Time")
+                    result["actual_rows"] = plan.get("Actual Rows")
+
+                # Add warnings based on plan analysis
+                if "Seq Scan" in str(plan):
+                    result["warnings"].append(
+                        "Sequential scan detected - may be slow on large tables"
+                    )
+                    result["recommendations"].append(
+                        "Consider adding appropriate indexes"
+                    )
+
+                return result
+
+        except (json.JSONDecodeError, KeyError):
+            pass
+
+        # Fallback for non-JSON format
+        return {
+            "json": None,
+            "warnings": [],
+            "recommendations": [],
+        }

db_connect_mcp/core/__init__.py

@@ -0,0 +1,13 @@
+"""Core database operations layer."""
+
+from .analyzer import StatisticsAnalyzer
+from .connection import DatabaseConnection
+from .executor import QueryExecutor
+from .inspector import MetadataInspector
+
+__all__ = [
+    "DatabaseConnection",
+    "MetadataInspector",
+    "QueryExecutor",
+    "StatisticsAnalyzer",
+]