sqlsaber 0.25.0__py3-none-any.whl → 0.26.0__py3-none-any.whl

sqlsaber/database/duckdb.py

@@ -0,0 +1,313 @@
+"""DuckDB database connection and schema introspection."""
+
+import asyncio
+from typing import Any
+
+import duckdb
+
+from .base import (
+    DEFAULT_QUERY_TIMEOUT,
+    BaseDatabaseConnection,
+    BaseSchemaIntrospector,
+    QueryTimeoutError,
+)
+
+
+def _execute_duckdb_transaction(
+    conn: duckdb.DuckDBPyConnection, query: str, args: tuple[Any, ...]
+) -> list[dict[str, Any]]:
+    """Run a DuckDB query inside a transaction and return list of dicts."""
+    conn.execute("BEGIN TRANSACTION")
+    try:
+        if args:
+            conn.execute(query, args)
+        else:
+            conn.execute(query)
+
+        if conn.description is None:
+            rows: list[dict[str, Any]] = []
+        else:
+            columns = [col[0] for col in conn.description]
+            data = conn.fetchall()
+            rows = [dict(zip(columns, row)) for row in data]
+
+        conn.execute("ROLLBACK")
+        return rows
+    except Exception:
+        conn.execute("ROLLBACK")
+        raise
+
+
+class DuckDBConnection(BaseDatabaseConnection):
+    """DuckDB database connection using duckdb Python API."""
+
+    def __init__(self, connection_string: str):
+        super().__init__(connection_string)
+        if connection_string.startswith("duckdb:///"):
+            db_path = connection_string.replace("duckdb:///", "", 1)
+        elif connection_string.startswith("duckdb://"):
+            db_path = connection_string.replace("duckdb://", "", 1)
+        else:
+            db_path = connection_string
+
+        self.database_path = db_path or ":memory:"
+
+    async def get_pool(self):
+        """DuckDB creates connections per query, return database path."""
+        return self.database_path
+
+    async def close(self):
+        """DuckDB connections are created per query, no persistent pool to close."""
+        pass
+
+    async def execute_query(
+        self, query: str, *args, timeout: float | None = None
+    ) -> 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.
+        """
+        effective_timeout = timeout or DEFAULT_QUERY_TIMEOUT
+
+        args_tuple = tuple(args) if args else tuple()
+
+        def _run_query() -> list[dict[str, Any]]:
+            conn = duckdb.connect(self.database_path)
+            try:
+                return _execute_duckdb_transaction(conn, query, args_tuple)
+            finally:
+                conn.close()
+
+        try:
+            return await asyncio.wait_for(
+                asyncio.to_thread(_run_query), timeout=effective_timeout
+            )
+        except asyncio.TimeoutError as exc:
+            raise QueryTimeoutError(effective_timeout or 0) from exc
+
+
+class DuckDBSchemaIntrospector(BaseSchemaIntrospector):
+    """DuckDB-specific schema introspection."""
+
+    async def _execute_query(
+        self,
+        connection,
+        query: str,
+        params: tuple[Any, ...] = (),
+    ) -> list[dict[str, Any]]:
+        """Run a DuckDB query on a thread and return list of dictionaries."""
+
+        params_tuple = tuple(params)
+
+        def fetch_rows(conn: duckdb.DuckDBPyConnection) -> list[dict[str, Any]]:
+            cursor = conn.execute(query, params_tuple)
+            if cursor.description is None:
+                return []
+
+            columns = [col[0] for col in cursor.description]
+            rows = conn.fetchall()
+            return [dict(zip(columns, row)) for row in rows]
+
+        # Handle CSV connections differently
+        if hasattr(connection, "execute_query") and hasattr(connection, "csv_path"):
+            return await connection.execute_query(query, *params_tuple)
+
+        def run_query() -> list[dict[str, Any]]:
+            conn = duckdb.connect(connection.database_path)
+            try:
+                return fetch_rows(conn)
+            finally:
+                conn.close()
+
+        return await asyncio.to_thread(run_query)
+
+    async def get_tables_info(
+        self, connection, table_pattern: str | None = None
+    ) -> list[dict[str, Any]]:
+        """Get tables information for DuckDB."""
+        where_conditions = [
+            "table_schema NOT IN ('information_schema', 'pg_catalog', 'duckdb_catalog')"
+        ]
+        params: list[Any] = []
+
+        if table_pattern:
+            if "." in table_pattern:
+                schema_pattern, table_name_pattern = table_pattern.split(".", 1)
+                where_conditions.append("(table_schema LIKE ? AND table_name LIKE ?)")
+                params.extend([schema_pattern, table_name_pattern])
+            else:
+                where_conditions.append(
+                    "(table_name LIKE ? OR table_schema || '.' || table_name LIKE ?)"
+                )
+                params.extend([table_pattern, table_pattern])
+
+        query = f"""
+            SELECT
+                table_schema,
+                table_name,
+                table_type
+            FROM information_schema.tables
+            WHERE {" AND ".join(where_conditions)}
+            ORDER BY table_schema, table_name;
+        """
+
+        return await self._execute_query(connection, query, tuple(params))
+
+    async def get_columns_info(self, connection, tables: list) -> list[dict[str, Any]]:
+        """Get columns information for DuckDB."""
+        if not tables:
+            return []
+
+        table_filters = []
+        for table in tables:
+            table_filters.append("(table_schema = ? AND table_name = ?)")
+
+        params: list[Any] = []
+        for table in tables:
+            params.extend([table["table_schema"], table["table_name"]])
+
+        query = f"""
+            SELECT
+                table_schema,
+                table_name,
+                column_name,
+                data_type,
+                is_nullable,
+                column_default,
+                character_maximum_length,
+                numeric_precision,
+                numeric_scale
+            FROM information_schema.columns
+            WHERE {" OR ".join(table_filters)}
+            ORDER BY table_schema, table_name, ordinal_position;
+        """
+
+        return await self._execute_query(connection, query, tuple(params))
+
+    async def get_foreign_keys_info(
+        self, connection, tables: list
+    ) -> list[dict[str, Any]]:
+        """Get foreign keys information for DuckDB."""
+        if not tables:
+            return []
+
+        table_filters = []
+        params: list[Any] = []
+        for table in tables:
+            table_filters.append("(kcu.table_schema = ? AND kcu.table_name = ?)")
+            params.extend([table["table_schema"], table["table_name"]])
+
+        query = f"""
+            SELECT
+                kcu.table_schema,
+                kcu.table_name,
+                kcu.column_name,
+                ccu.table_schema AS foreign_table_schema,
+                ccu.table_name AS foreign_table_name,
+                ccu.column_name AS foreign_column_name
+            FROM information_schema.referential_constraints AS rc
+            JOIN information_schema.key_column_usage AS kcu
+                ON rc.constraint_schema = kcu.constraint_schema
+                AND rc.constraint_name = kcu.constraint_name
+            JOIN information_schema.key_column_usage AS ccu
+                ON rc.unique_constraint_schema = ccu.constraint_schema
+                AND rc.unique_constraint_name = ccu.constraint_name
+                AND ccu.ordinal_position = kcu.position_in_unique_constraint
+            WHERE {" OR ".join(table_filters)}
+            ORDER BY kcu.table_schema, kcu.table_name, kcu.ordinal_position;
+        """
+
+        return await self._execute_query(connection, query, tuple(params))
+
+    async def get_primary_keys_info(
+        self, connection, tables: list
+    ) -> list[dict[str, Any]]:
+        """Get primary keys information for DuckDB."""
+        if not tables:
+            return []
+
+        table_filters = []
+        params: list[Any] = []
+        for table in tables:
+            table_filters.append("(tc.table_schema = ? AND tc.table_name = ?)")
+            params.extend([table["table_schema"], table["table_name"]])
+
+        query = f"""
+            SELECT
+                tc.table_schema,
+                tc.table_name,
+                kcu.column_name
+            FROM information_schema.table_constraints AS tc
+            JOIN information_schema.key_column_usage AS kcu
+                ON tc.constraint_name = kcu.constraint_name
+                AND tc.constraint_schema = kcu.constraint_schema
+            WHERE tc.constraint_type = 'PRIMARY KEY'
+                AND ({" OR ".join(table_filters)})
+            ORDER BY tc.table_schema, tc.table_name, kcu.ordinal_position;
+        """
+
+        return await self._execute_query(connection, query, tuple(params))
+
+    async def get_indexes_info(self, connection, tables: list) -> list[dict[str, Any]]:
+        """Get indexes information for DuckDB."""
+        if not tables:
+            return []
+
+        indexes: list[dict[str, Any]] = []
+        for table in tables:
+            schema = table["table_schema"]
+            table_name = table["table_name"]
+            query = """
+                SELECT
+                    schema_name,
+                    table_name,
+                    index_name,
+                    sql
+                FROM duckdb_indexes()
+                WHERE schema_name = ? AND table_name = ?;
+            """
+            rows = await self._execute_query(connection, query, (schema, table_name))
+
+            for row in rows:
+                sql_text = (row.get("sql") or "").strip()
+                upper_sql = sql_text.upper()
+                unique = "UNIQUE" in upper_sql.split("(")[0]
+
+                columns: list[str] = []
+                if "(" in sql_text and ")" in sql_text:
+                    column_section = sql_text[
+                        sql_text.find("(") + 1 : sql_text.rfind(")")
+                    ]
+                    columns = [
+                        col.strip().strip('"')
+                        for col in column_section.split(",")
+                        if col.strip()
+                    ]
+
+                indexes.append(
+                    {
+                        "table_schema": row.get("schema_name") or schema or "main",
+                        "table_name": row.get("table_name") or table_name,
+                        "index_name": row.get("index_name"),
+                        "is_unique": unique,
+                        "index_type": None,
+                        "column_names": columns,
+                    }
+                )
+
+        return indexes
+
+    async def list_tables_info(self, connection) -> list[dict[str, Any]]:
+        """Get list of tables with basic information for DuckDB."""
+        query = """
+            SELECT
+                table_schema,
+                table_name,
+                table_type
+            FROM information_schema.tables
+            WHERE table_schema NOT IN ('information_schema', 'pg_catalog', 'duckdb_catalog')
+            ORDER BY table_schema, table_name;
+        """
+
+        return await self._execute_query(connection, query)

sqlsaber/database/mysql.py

@@ -0,0 +1,345 @@
+"""MySQL database connection and schema introspection."""
+
+import asyncio
+import ssl
+from typing import Any
+from urllib.parse import parse_qs, urlparse
+
+import aiomysql
+
+from .base import (
+    DEFAULT_QUERY_TIMEOUT,
+    BaseDatabaseConnection,
+    BaseSchemaIntrospector,
+    QueryTimeoutError,
+)
+
+
+class MySQLConnection(BaseDatabaseConnection):
+    """MySQL database connection using aiomysql."""
+
+    def __init__(self, connection_string: str):
+        super().__init__(connection_string)
+        self._pool: aiomysql.Pool | None = None
+        self._parse_connection_string()
+
+    def _parse_connection_string(self):
+        """Parse MySQL connection string into components."""
+        parsed = urlparse(self.connection_string)
+        self.host = parsed.hostname or "localhost"
+        self.port = parsed.port or 3306
+        self.database = parsed.path.lstrip("/") if parsed.path else ""
+        self.user = parsed.username or ""
+        self.password = parsed.password or ""
+
+        # Parse SSL parameters
+        self.ssl_params = {}
+        if parsed.query:
+            params = parse_qs(parsed.query)
+
+            ssl_mode = params.get("ssl_mode", [None])[0]
+            if ssl_mode:
+                # Map SSL modes to aiomysql SSL parameters
+                if ssl_mode.upper() == "DISABLED":
+                    self.ssl_params["ssl"] = None
+                elif ssl_mode.upper() in [
+                    "PREFERRED",
+                    "REQUIRED",
+                    "VERIFY_CA",
+                    "VERIFY_IDENTITY",
+                ]:
+                    ssl_context = ssl.create_default_context()
+
+                    if ssl_mode.upper() == "REQUIRED":
+                        ssl_context.check_hostname = False
+                        ssl_context.verify_mode = ssl.CERT_NONE
+                    elif ssl_mode.upper() == "VERIFY_CA":
+                        ssl_context.check_hostname = False
+                        ssl_context.verify_mode = ssl.CERT_REQUIRED
+                    elif ssl_mode.upper() == "VERIFY_IDENTITY":
+                        ssl_context.check_hostname = True
+                        ssl_context.verify_mode = ssl.CERT_REQUIRED
+
+                    # Load certificates if provided
+                    ssl_ca = params.get("ssl_ca", [None])[0]
+                    ssl_cert = params.get("ssl_cert", [None])[0]
+                    ssl_key = params.get("ssl_key", [None])[0]
+
+                    if ssl_ca:
+                        ssl_context.load_verify_locations(ssl_ca)
+
+                    if ssl_cert and ssl_key:
+                        ssl_context.load_cert_chain(ssl_cert, ssl_key)
+
+                    self.ssl_params["ssl"] = ssl_context
+
+    async def get_pool(self) -> aiomysql.Pool:
+        """Get or create connection pool."""
+        if self._pool is None:
+            pool_kwargs = {
+                "host": self.host,
+                "port": self.port,
+                "user": self.user,
+                "password": self.password,
+                "db": self.database,
+                "minsize": 1,
+                "maxsize": 10,
+                "autocommit": False,
+            }
+
+            # Add SSL parameters if configured
+            pool_kwargs.update(self.ssl_params)
+
+            self._pool = await aiomysql.create_pool(**pool_kwargs)
+        return self._pool
+
+    async def close(self):
+        """Close the connection pool."""
+        if self._pool:
+            self._pool.close()
+            await self._pool.wait_closed()
+            self._pool = None
+
+    async def execute_query(
+        self, query: str, *args, timeout: float | None = None
+    ) -> 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.
+        """
+        effective_timeout = timeout or DEFAULT_QUERY_TIMEOUT
+        pool = await self.get_pool()
+
+        async with pool.acquire() as conn:
+            async with conn.cursor(aiomysql.DictCursor) as cursor:
+                # Start transaction
+                await conn.begin()
+                try:
+                    # Set server-side timeout if specified
+                    if effective_timeout:
+                        # Clamp timeout to sane range (10ms to 5 minutes) and validate
+                        timeout_ms = max(10, min(int(effective_timeout * 1000), 300000))
+                        await cursor.execute(
+                            f"SET SESSION MAX_EXECUTION_TIME = {timeout_ms}"
+                        )
+
+                    # Execute query with client-side timeout
+                    if effective_timeout:
+                        await asyncio.wait_for(
+                            cursor.execute(query, args if args else None),
+                            timeout=effective_timeout,
+                        )
+                        rows = await asyncio.wait_for(
+                            cursor.fetchall(), timeout=effective_timeout
+                        )
+                    else:
+                        await cursor.execute(query, args if args else None)
+                        rows = await cursor.fetchall()
+
+                    return [dict(row) for row in rows]
+                except asyncio.TimeoutError as exc:
+                    raise QueryTimeoutError(effective_timeout or 0) from exc
+                finally:
+                    # Always rollback to ensure no changes are committed
+                    await conn.rollback()
+
+
+class MySQLSchemaIntrospector(BaseSchemaIntrospector):
+    """MySQL-specific schema introspection."""
+
+    def _build_table_filter_clause(self, tables: list) -> tuple[str, list]:
+        """Build row constructor with bind parameters for table filtering.
+
+        Args:
+            tables: List of table dictionaries with table_schema and table_name keys
+
+        Returns:
+            Tuple of (placeholders, params) for use in SQL queries
+        """
+        if not tables:
+            return "", []
+
+        table_pairs = [(table["table_schema"], table["table_name"]) for table in tables]
+        placeholders = ", ".join(["(%s, %s)"] * len(table_pairs))
+        params = [value for pair in table_pairs for value in pair]
+        return placeholders, params
+
+    async def get_tables_info(
+        self, connection, table_pattern: str | None = None
+    ) -> dict[str, Any]:
+        """Get tables information for MySQL."""
+        pool = await connection.get_pool()
+        async with pool.acquire() as conn:
+            async with conn.cursor() as cursor:
+                # Build WHERE clause for filtering
+                where_conditions = [
+                    "table_schema NOT IN ('information_schema', 'performance_schema', 'mysql', 'sys')"
+                ]
+                params = []
+
+                if table_pattern:
+                    # Support patterns like 'schema.table' or just 'table'
+                    if "." in table_pattern:
+                        schema_pattern, table_name_pattern = table_pattern.split(".", 1)
+                        where_conditions.append(
+                            "(table_schema LIKE %s AND table_name LIKE %s)"
+                        )
+                        params.extend([schema_pattern, table_name_pattern])
+                    else:
+                        where_conditions.append(
+                            "(table_name LIKE %s OR CONCAT(table_schema, '.', table_name) LIKE %s)"
+                        )
+                        params.extend([table_pattern, table_pattern])
+
+                # Get tables
+                tables_query = f"""
+                    SELECT
+                        table_schema,
+                        table_name,
+                        table_type
+                    FROM information_schema.tables
+                    WHERE {" AND ".join(where_conditions)}
+                    ORDER BY table_schema, table_name;
+                """
+                await cursor.execute(tables_query, params)
+                return await cursor.fetchall()
+
+    async def get_columns_info(self, connection, tables: list) -> list:
+        """Get columns information for MySQL."""
+        if not tables:
+            return []
+
+        pool = await connection.get_pool()
+        async with pool.acquire() as conn:
+            async with conn.cursor() as cursor:
+                placeholders, params = self._build_table_filter_clause(tables)
+
+                columns_query = f"""
+                    SELECT
+                        c.table_schema,
+                        c.table_name,
+                        c.column_name,
+                        c.data_type,
+                        c.is_nullable,
+                        c.column_default,
+                        c.character_maximum_length,
+                        c.numeric_precision,
+                        c.numeric_scale
+                    FROM information_schema.columns c
+                    WHERE (c.table_schema, c.table_name) IN ({placeholders})
+                    ORDER BY c.table_schema, c.table_name, c.ordinal_position;
+                """
+                await cursor.execute(columns_query, params)
+                return await cursor.fetchall()
+
+    async def get_foreign_keys_info(self, connection, tables: list) -> list:
+        """Get foreign keys information for MySQL."""
+        if not tables:
+            return []
+
+        pool = await connection.get_pool()
+        async with pool.acquire() as conn:
+            async with conn.cursor() as cursor:
+                placeholders, params = self._build_table_filter_clause(tables)
+
+                fk_query = f"""
+                    SELECT
+                        tc.table_schema,
+                        tc.table_name,
+                        kcu.column_name,
+                        rc.unique_constraint_schema AS foreign_table_schema,
+                        rc.referenced_table_name AS foreign_table_name,
+                        kcu.referenced_column_name AS foreign_column_name
+                    FROM information_schema.table_constraints AS tc
+                    JOIN information_schema.key_column_usage AS kcu
+                        ON tc.constraint_name = kcu.constraint_name
+                        AND tc.table_schema = kcu.table_schema
+                    JOIN information_schema.referential_constraints AS rc
+                        ON tc.constraint_name = rc.constraint_name
+                        AND tc.table_schema = rc.constraint_schema
+                    WHERE tc.constraint_type = 'FOREIGN KEY'
+                        AND (tc.table_schema, tc.table_name) IN ({placeholders});
+                """
+                await cursor.execute(fk_query, params)
+                return await cursor.fetchall()
+
+    async def get_primary_keys_info(self, connection, tables: list) -> list:
+        """Get primary keys information for MySQL."""
+        if not tables:
+            return []
+
+        pool = await connection.get_pool()
+        async with pool.acquire() as conn:
+            async with conn.cursor() as cursor:
+                placeholders, params = self._build_table_filter_clause(tables)
+
+                pk_query = f"""
+                    SELECT
+                        tc.table_schema,
+                        tc.table_name,
+                        kcu.column_name
+                    FROM information_schema.table_constraints AS tc
+                    JOIN information_schema.key_column_usage AS kcu
+                        ON tc.constraint_name = kcu.constraint_name
+                        AND tc.table_schema = kcu.table_schema
+                    WHERE tc.constraint_type = 'PRIMARY KEY'
+                        AND (tc.table_schema, tc.table_name) IN ({placeholders})
+                    ORDER BY tc.table_schema, tc.table_name, kcu.ordinal_position;
+                """
+                await cursor.execute(pk_query, params)
+                return await cursor.fetchall()
+
+    async def get_indexes_info(self, connection, tables: list) -> list:
+        """Get indexes information for MySQL."""
+        if not tables:
+            return []
+
+        pool = await connection.get_pool()
+        async with pool.acquire() as conn:
+            async with conn.cursor() as cursor:
+                placeholders, params = self._build_table_filter_clause(tables)
+
+                idx_query = f"""
+                    SELECT
+                        TABLE_SCHEMA   AS table_schema,
+                        TABLE_NAME     AS table_name,
+                        INDEX_NAME     AS index_name,
+                        (NON_UNIQUE = 0) AS is_unique,
+                        INDEX_TYPE     AS index_type,
+                        GROUP_CONCAT(COLUMN_NAME ORDER BY SEQ_IN_INDEX) AS column_names
+                    FROM INFORMATION_SCHEMA.STATISTICS
+                    WHERE (TABLE_SCHEMA, TABLE_NAME) IN ({placeholders})
+                    GROUP BY table_schema, table_name, index_name, is_unique, index_type
+                    ORDER BY table_schema, table_name, index_name;
+                """
+                await cursor.execute(idx_query, params)
+                return await cursor.fetchall()
+
+    async def list_tables_info(self, connection) -> list[dict[str, Any]]:
+        """Get list of tables with basic information for MySQL."""
+        pool = await connection.get_pool()
+        async with pool.acquire() as conn:
+            async with conn.cursor() as cursor:
+                # Get tables without row counts for better performance
+                tables_query = """
+                    SELECT
+                        t.table_schema,
+                        t.table_name,
+                        t.table_type
+                    FROM information_schema.tables t
+                    WHERE t.table_schema NOT IN ('information_schema', 'performance_schema', 'mysql', 'sys')
+                    ORDER BY t.table_schema, t.table_name;
+                """
+                await cursor.execute(tables_query)
+                rows = await cursor.fetchall()
+
+                # Convert rows to dictionaries
+                return [
+                    {
+                        "table_schema": row["table_schema"],
+                        "table_name": row["table_name"],
+                        "table_type": row["table_type"],
+                    }
+                    for row in rows
+                ]