sql-query-mcp 0.3.0__tar.gz → 0.4.0__tar.gz

{sql_query_mcp-0.3.0/sql_query_mcp.egg-info → sql_query_mcp-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sql-query-mcp
-Version: 0.3.0
+Version: 0.4.0
 Summary: Read-only SQL MCP server for PostgreSQL and MySQL.
 Author: Andy Wang
 License-Expression: MIT
@@ -62,10 +62,11 @@ without exposing raw connection strings or flattening engine-specific concepts.
 ## What AI can do with it
 
 The current tool set focuses on database discovery, controlled query workflows,
-asynchronous read-only queries, and one narrow local file import path. You can
-use it to help an AI assistant understand structure before it generates SQL,
-runs a bounded query, starts a long-running read-only query, or imports a
-prepared CSV/XLSX file into an existing table.
+asynchronous read-only queries, batched query result exports, and one narrow
+local file import path. You can use it to help an AI assistant understand
+structure before it generates SQL, runs a bounded query, starts a long-running
+read-only query, exports PostgreSQL or MySQL results to a local file, or imports
+a prepared CSV/XLSX file into an existing table.
 
 MySQL and Hive support `explain_query`. Hive uses `EXPLAIN` and
 `EXPLAIN ANALYZE` for `explain_query`.
@@ -83,19 +84,22 @@ MySQL and Hive support `explain_query`. Hive uses `EXPLAIN` and
 | `cancel_query(query_id)` | Yes | Yes | Yes | Cancel running async queries |
 | `explain_query(connection_id, sql, analyze?)` | Yes | Yes | Yes | Inspect query plans |
 | `get_table_sample(connection_id, table_name, schema?, database?, limit?)` | Yes | Yes | Yes | Fetch small table samples |
+| `export_query_file(connection_id, sql, output_path, format?, limit?, export_all?, file_name?, overwrite?)` | Yes | Yes | No | Export query results to local CSV/XLSX files |
 | `import_table_file(connection_id, table_name, file_path, schema?, database?, sheet_name?)` | Yes | Yes | Yes | Import local CSV/XLSX files |
 
 These tools are useful for tasks such as listing namespaces, inspecting table
 definitions, reviewing indexes, sampling records, running short read-only
 queries with `run_select`, running long read-only queries with `start_query`,
 `get_query`, and `cancel_query`, analyzing read-only queries with `EXPLAIN`, and
-importing prepared local files. For full request and response details, see
+exporting PostgreSQL or MySQL query results to local CSV/XLSX files. You can
+also import prepared local files. For full request and response details, see
 `docs/api-reference.md` (Chinese).
 
 ## How boundaries are constrained
 
 The product boundary is intentionally narrow today. PostgreSQL, MySQL, and Hive
-are available today. Query tools remain read-only, and the only write path is a
+are available today. Query tools remain read-only, PostgreSQL and MySQL query
+results can be exported to local files, and the only database write path is a
 controlled local CSV/XLSX import into existing tables.
 
 The service keeps those boundaries explicit in a few ways.
@@ -112,6 +116,11 @@ The service keeps those boundaries explicit in a few ways.
   queries.
 - The server accepts only `SELECT` and `WITH ... SELECT`, rejects comments and
   multi-statement input, and records audit logs for each call.
+- `export_query_file` writes files on the MCP server machine. It is synchronous
+  but reads database rows and writes CSV/XLSX files in batches. Large exports can
+  still hit your MCP client's tool timeout. For XLSX output, UUID values are
+  written as text and timezone-aware datetime values are written without the
+  timezone. Hive export is not supported yet.
 - `import_table_file` doesn't accept raw SQL. It inserts only file columns whose
   headers exactly match existing table columns.
 - Hive `import_table_file` is intended for small files only and rejects files

{sql_query_mcp-0.3.0 → sql_query_mcp-0.4.0}/README.md

@@ -30,10 +30,11 @@ without exposing raw connection strings or flattening engine-specific concepts.
 ## What AI can do with it
 
 The current tool set focuses on database discovery, controlled query workflows,
-asynchronous read-only queries, and one narrow local file import path. You can
-use it to help an AI assistant understand structure before it generates SQL,
-runs a bounded query, starts a long-running read-only query, or imports a
-prepared CSV/XLSX file into an existing table.
+asynchronous read-only queries, batched query result exports, and one narrow
+local file import path. You can use it to help an AI assistant understand
+structure before it generates SQL, runs a bounded query, starts a long-running
+read-only query, exports PostgreSQL or MySQL results to a local file, or imports
+a prepared CSV/XLSX file into an existing table.
 
 MySQL and Hive support `explain_query`. Hive uses `EXPLAIN` and
 `EXPLAIN ANALYZE` for `explain_query`.
@@ -51,19 +52,22 @@ MySQL and Hive support `explain_query`. Hive uses `EXPLAIN` and
 | `cancel_query(query_id)` | Yes | Yes | Yes | Cancel running async queries |
 | `explain_query(connection_id, sql, analyze?)` | Yes | Yes | Yes | Inspect query plans |
 | `get_table_sample(connection_id, table_name, schema?, database?, limit?)` | Yes | Yes | Yes | Fetch small table samples |
+| `export_query_file(connection_id, sql, output_path, format?, limit?, export_all?, file_name?, overwrite?)` | Yes | Yes | No | Export query results to local CSV/XLSX files |
 | `import_table_file(connection_id, table_name, file_path, schema?, database?, sheet_name?)` | Yes | Yes | Yes | Import local CSV/XLSX files |
 
 These tools are useful for tasks such as listing namespaces, inspecting table
 definitions, reviewing indexes, sampling records, running short read-only
 queries with `run_select`, running long read-only queries with `start_query`,
 `get_query`, and `cancel_query`, analyzing read-only queries with `EXPLAIN`, and
-importing prepared local files. For full request and response details, see
+exporting PostgreSQL or MySQL query results to local CSV/XLSX files. You can
+also import prepared local files. For full request and response details, see
 `docs/api-reference.md` (Chinese).
 
 ## How boundaries are constrained
 
 The product boundary is intentionally narrow today. PostgreSQL, MySQL, and Hive
-are available today. Query tools remain read-only, and the only write path is a
+are available today. Query tools remain read-only, PostgreSQL and MySQL query
+results can be exported to local files, and the only database write path is a
 controlled local CSV/XLSX import into existing tables.
 
 The service keeps those boundaries explicit in a few ways.
@@ -80,6 +84,11 @@ The service keeps those boundaries explicit in a few ways.
   queries.
 - The server accepts only `SELECT` and `WITH ... SELECT`, rejects comments and
   multi-statement input, and records audit logs for each call.
+- `export_query_file` writes files on the MCP server machine. It is synchronous
+  but reads database rows and writes CSV/XLSX files in batches. Large exports can
+  still hit your MCP client's tool timeout. For XLSX output, UUID values are
+  written as text and timezone-aware datetime values are written without the
+  timezone. Hive export is not supported yet.
 - `import_table_file` doesn't accept raw SQL. It inserts only file columns whose
   headers exactly match existing table columns.
 - Hive `import_table_file` is intended for small files only and rejects files

{sql_query_mcp-0.3.0 → sql_query_mcp-0.4.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "sql-query-mcp"
-version = "0.3.0"
+version = "0.4.0"
 description = "Read-only SQL MCP server for PostgreSQL and MySQL."
 readme = "README.md"
 requires-python = ">=3.10"

{sql_query_mcp-0.3.0 → sql_query_mcp-0.4.0}/sql_query_mcp/adapters/mysql.py

@@ -42,6 +42,11 @@ class MySQLAdapter:
         with conn.cursor() as cur:
             cur.execute("SET SESSION max_execution_time = %s", (int(timeout_ms),))
 
+    def export_cursor(self, conn: object):
+        if pymysql is None:
+            raise ConfigurationError("缺少 PyMySQL 依赖，请先安装项目依赖。")
+        return conn.cursor(pymysql.cursors.SSDictCursor)
+
     def list_databases(self, conn: object) -> List[str]:
         with conn.cursor() as cur:
             cur.execute(

{sql_query_mcp-0.3.0 → sql_query_mcp-0.4.0}/sql_query_mcp/adapters/postgres.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from contextlib import contextmanager
+from uuid import uuid4
 from typing import Iterator, List
 
 try:
@@ -37,6 +38,19 @@ class PostgresAdapter:
         with conn.cursor() as cur:
             cur.execute("SELECT set_config('statement_timeout', %s, false)", (str(timeout_ms),))
 
+    @contextmanager
+    def export_cursor(self, conn: object) -> Iterator[object]:
+        previous_autocommit = getattr(conn, "autocommit", None)
+        if previous_autocommit is True:
+            conn.autocommit = False
+        try:
+            with conn.cursor(name=f"sql_query_mcp_export_{uuid4().hex}") as cur:
+                yield cur
+        finally:
+            if previous_autocommit is True:
+                conn.rollback()
+                conn.autocommit = True
+
     def list_schemas(self, conn: object) -> List[str]:
         with conn.cursor() as cur:
             cur.execute(

{sql_query_mcp-0.3.0 → sql_query_mcp-0.4.0}/sql_query_mcp/app.py

@@ -11,6 +11,7 @@ from .audit import AuditLogger
 from .config import load_config
 from .errors import SqlQueryMCPError
 from .executor import QueryExecutor
+from .exporter import QueryExporter
 from .importer import TableFileImporter
 from .introspection import MetadataService
 from .registry import ConnectionRegistry
@@ -22,6 +23,7 @@ def create_app() -> FastMCP:
     audit_logger = AuditLogger(app_config.settings.audit_log_path)
     metadata = MetadataService(registry, app_config.settings, audit_logger)
     executor = QueryExecutor(registry, app_config.settings, audit_logger)
+    exporter = QueryExporter(registry, app_config.settings, audit_logger)
     importer = TableFileImporter(registry, app_config.settings, audit_logger)
     async_queries = AsyncQueryService(registry, app_config.settings, audit_logger)
 
@@ -90,6 +92,32 @@ def create_app() -> FastMCP:
 
         return _run_tool(lambda: executor.get_table_sample(connection_id, table_name, schema, database, limit))
 
+    @mcp.tool()
+    def export_query_file(
+        connection_id: str,
+        sql: str,
+        output_path: str,
+        format: str = "csv",
+        limit: Optional[int] = 1000,
+        export_all: bool = False,
+        file_name: Optional[str] = None,
+        overwrite: bool = False,
+    ) -> dict:
+        """Export a read-only query result to a local CSV or XLSX file."""
+
+        return _run_tool(
+            lambda: exporter.export_query_file(
+                connection_id,
+                sql,
+                output_path,
+                format,
+                limit,
+                export_all,
+                file_name,
+                overwrite,
+            )
+        )
+
     @mcp.tool()
     def start_query(connection_id: str, sql: str, limit: Optional[int] = None) -> dict:
         """Start an asynchronous read-only SELECT or CTE query."""

sql_query_mcp-0.4.0/sql_query_mcp/exporter.py

@@ -0,0 +1,264 @@
+"""Controlled query result exports."""
+
+from __future__ import annotations
+
+import csv
+import time
+import uuid
+from datetime import datetime, time as datetime_time, timezone
+from pathlib import Path
+from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, cast
+
+try:
+    from openpyxl import Workbook
+except ImportError:  # pragma: no cover - runtime dependency
+    Workbook = None
+
+from .audit import AuditLogger
+from .config import ServerSettings
+from .errors import QueryExecutionError, sanitize_error_message
+from .validator import clamp_limit, summarize_sql, validate_select_sql
+
+EXPORT_BATCH_SIZE = 1000
+SUPPORTED_EXPORT_ENGINES = {"postgres", "mysql"}
+SUPPORTED_EXPORT_FORMATS = {"csv", "xlsx"}
+
+
+class QueryExporter:
+    """Export validated read-only query results to local files."""
+
+    def __init__(self, registry: Any, settings: ServerSettings, audit_logger: AuditLogger):
+        self._registry = registry
+        self._settings = settings
+        self._audit = audit_logger
+
+    def export_query_file(
+        self,
+        connection_id: str,
+        sql_text: str,
+        output_path: str,
+        format: str = "csv",
+        limit: Optional[int] = 1000,
+        export_all: bool = False,
+        file_name: Optional[str] = None,
+        overwrite: bool = False,
+    ) -> Dict[str, object]:
+        started = time.perf_counter()
+        config = None
+        final_path: Optional[Path] = None
+        row_count = 0
+        applied_limit = None
+        sql_summary = summarize_sql(sql_text)
+        normalized_format = str(format).lower().lstrip(".")
+
+        try:
+            if normalized_format not in SUPPORTED_EXPORT_FORMATS:
+                raise QueryExecutionError("导出格式仅支持 csv 和 xlsx。")
+            config = self._registry.get_connection_config(connection_id)
+            if config.engine not in SUPPORTED_EXPORT_ENGINES:
+                raise QueryExecutionError("export_query_file 首版仅支持 PostgreSQL 和 MySQL。")
+            cleaned_sql = validate_select_sql(sql_text, config.engine)
+            sql_summary = summarize_sql(cleaned_sql)
+            final_path = _resolve_output_file(
+                output_path,
+                normalized_format,
+                file_name=file_name,
+                overwrite=overwrite,
+            )
+            query = cleaned_sql
+            if not export_all:
+                applied_limit = clamp_limit(limit, 1000, self._settings.max_limit)
+                query = _build_exact_limited_query(cleaned_sql, applied_limit)
+
+            with self._registry.connection_from_config(config) as (conn, adapter):
+                _apply_statement_timeout(adapter, conn, self._settings.statement_timeout_ms)
+                with _open_export_cursor(adapter, conn) as cur:
+                    cur.execute(query)
+                    columns = adapter.column_names(cur.description)
+                    batches = _iter_batches(cur, columns, adapter)
+                    if normalized_format == "csv":
+                        row_count = _write_csv(final_path, columns, batches)
+                    else:
+                        row_count = _write_xlsx(final_path, columns, batches)
+
+            duration_ms = _elapsed_ms(started)
+            self._audit.log(
+                tool="export_query_file",
+                connection_id=connection_id,
+                success=True,
+                duration_ms=duration_ms,
+                row_count=row_count,
+                sql_summary=sql_summary,
+                extra=_audit_extra(config, final_path, normalized_format, export_all, applied_limit),
+            )
+            return {
+                "connection_id": connection_id,
+                "engine": config.engine,
+                "file_path": str(final_path),
+                "format": normalized_format,
+                "row_count": row_count,
+                "duration_ms": duration_ms,
+                "export_all": export_all,
+                "applied_limit": applied_limit,
+            }
+        except Exception as exc:
+            duration_ms = _elapsed_ms(started)
+            sanitized = sanitize_error_message(str(exc))
+            self._audit.log(
+                tool="export_query_file",
+                connection_id=connection_id,
+                success=False,
+                duration_ms=duration_ms,
+                row_count=row_count,
+                sql_summary=sql_summary,
+                error=sanitized,
+                extra=_audit_extra(config, final_path, normalized_format, export_all, applied_limit),
+            )
+            raise QueryExecutionError(sanitized) from exc
+
+
+def _iter_batches(cur: Any, columns: Sequence[str], adapter: Any) -> Iterable[List[object]]:
+    while True:
+        rows = cur.fetchmany(EXPORT_BATCH_SIZE)
+        if not rows:
+            return
+        if hasattr(adapter, "normalize_rows"):
+            rows = adapter.normalize_rows(rows, list(columns))
+        yield cast(List[object], rows)
+
+
+def _build_exact_limited_query(sql: str, row_limit: int) -> str:
+    return f"SELECT * FROM ({sql}) AS pq_result LIMIT {int(row_limit)}"
+
+
+def _write_csv(path: Path, columns: Sequence[str], batches: Iterable[Sequence[object]]) -> int:
+    row_count = 0
+    with path.open("w", encoding="utf-8", newline="") as handle:
+        writer = csv.writer(handle)
+        writer.writerow(columns)
+        for batch in batches:
+            for row in batch:
+                writer.writerow(_row_values(row, columns))
+                row_count += 1
+    return row_count
+
+
+def _write_xlsx(path: Path, columns: Sequence[str], batches: Iterable[Sequence[object]]) -> int:
+    if Workbook is None:
+        raise QueryExecutionError("缺少 openpyxl 依赖，请先安装项目依赖。")
+    workbook = Workbook(write_only=True)
+    worksheet = workbook.create_sheet("Export")
+    worksheet.append(list(columns))
+    row_count = 0
+    for batch in batches:
+        for row in batch:
+            values = _row_values(row, columns, normalize_value=_normalize_xlsx_value)
+            try:
+                worksheet.append(values)
+            except Exception as exc:
+                raise QueryExecutionError(_format_xlsx_error(columns, values, exc)) from exc
+            row_count += 1
+    workbook.save(path)
+    return row_count
+
+
+def _row_values(
+    row: object,
+    columns: Sequence[str],
+    normalize_value: Optional[Callable[[object], object]] = None,
+) -> List[object]:
+    if isinstance(row, dict):
+        values: List[object] = [row.get(column) for column in columns]
+    elif isinstance(row, (list, tuple)):
+        values = list(row)
+    else:
+        values = [row]
+    if normalize_value is None:
+        return values
+    return [normalize_value(value) for value in values]
+
+
+def _normalize_xlsx_value(value: object) -> object:
+    if isinstance(value, uuid.UUID):
+        return str(value)
+    if isinstance(value, datetime) and value.tzinfo is not None:
+        return value.replace(tzinfo=None)
+    if isinstance(value, datetime_time) and value.tzinfo is not None:
+        return value.replace(tzinfo=None)
+    return value
+
+
+def _format_xlsx_error(columns: Sequence[str], values: Sequence[object], exc: Exception) -> str:
+    message = str(exc) or exc.__class__.__name__
+    details = ", ".join(
+        f"{column}={type(value).__name__}"
+        for column, value in zip(columns, values)
+    )
+    return f"XLSX 导出失败: {message}; columns: {details}"
+
+
+def _resolve_output_file(output_path: str, format: str, file_name: Optional[str], overwrite: bool) -> Path:
+    base = Path(output_path).expanduser()
+    suffix = f".{format}"
+    if base.exists() and base.is_dir():
+        name = file_name or f"export_{datetime.now(timezone.utc).strftime('%Y%m%d_%H%M%S')}"
+        candidate = base / _with_suffix(name, suffix)
+    else:
+        if file_name:
+            raise QueryExecutionError("output_path 为文件路径时不能同时传 file_name。")
+        candidate = Path(_with_suffix(str(base), suffix)).expanduser()
+        if not candidate.parent.exists():
+            raise QueryExecutionError("导出目录不存在。")
+    if overwrite or not candidate.exists():
+        return candidate
+    return _next_available_path(candidate)
+
+
+def _with_suffix(value: str, suffix: str) -> str:
+    return value if value.lower().endswith(suffix) else value + suffix
+
+
+def _next_available_path(path: Path) -> Path:
+    stem = path.stem
+    suffix = path.suffix
+    parent = path.parent
+    index = 1
+    while True:
+        candidate = parent / f"{stem} ({index}){suffix}"
+        if not candidate.exists():
+            return candidate
+        index += 1
+
+
+def _apply_statement_timeout(adapter: Any, conn: Any, timeout_ms: Optional[int]) -> None:
+    if timeout_ms is not None:
+        getattr(adapter, "set_statement_timeout")(conn, timeout_ms)
+
+
+def _open_export_cursor(adapter: Any, conn: Any) -> Any:
+    export_cursor = getattr(adapter, "export_cursor", None)
+    if callable(export_cursor):
+        return export_cursor(conn)
+    return conn.cursor()
+
+
+def _elapsed_ms(started: float) -> int:
+    return int((time.perf_counter() - started) * 1000)
+
+
+def _audit_extra(
+    config: Any,
+    file_path: Optional[Path],
+    format: str,
+    export_all: bool,
+    applied_limit: Optional[int],
+) -> Dict[str, object]:
+    extra: Dict[str, object] = {
+        "file_path": str(file_path) if file_path is not None else None,
+        "format": format,
+        "export_all": export_all,
+        "applied_limit": applied_limit,
+    }
+    if config is not None:
+        extra["engine"] = config.engine
+    return extra

{sql_query_mcp-0.3.0 → sql_query_mcp-0.4.0/sql_query_mcp.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sql-query-mcp
-Version: 0.3.0
+Version: 0.4.0
 Summary: Read-only SQL MCP server for PostgreSQL and MySQL.
 Author: Andy Wang
 License-Expression: MIT
@@ -62,10 +62,11 @@ without exposing raw connection strings or flattening engine-specific concepts.
 ## What AI can do with it
 
 The current tool set focuses on database discovery, controlled query workflows,
-asynchronous read-only queries, and one narrow local file import path. You can
-use it to help an AI assistant understand structure before it generates SQL,
-runs a bounded query, starts a long-running read-only query, or imports a
-prepared CSV/XLSX file into an existing table.
+asynchronous read-only queries, batched query result exports, and one narrow
+local file import path. You can use it to help an AI assistant understand
+structure before it generates SQL, runs a bounded query, starts a long-running
+read-only query, exports PostgreSQL or MySQL results to a local file, or imports
+a prepared CSV/XLSX file into an existing table.
 
 MySQL and Hive support `explain_query`. Hive uses `EXPLAIN` and
 `EXPLAIN ANALYZE` for `explain_query`.
@@ -83,19 +84,22 @@ MySQL and Hive support `explain_query`. Hive uses `EXPLAIN` and
 | `cancel_query(query_id)` | Yes | Yes | Yes | Cancel running async queries |
 | `explain_query(connection_id, sql, analyze?)` | Yes | Yes | Yes | Inspect query plans |
 | `get_table_sample(connection_id, table_name, schema?, database?, limit?)` | Yes | Yes | Yes | Fetch small table samples |
+| `export_query_file(connection_id, sql, output_path, format?, limit?, export_all?, file_name?, overwrite?)` | Yes | Yes | No | Export query results to local CSV/XLSX files |
 | `import_table_file(connection_id, table_name, file_path, schema?, database?, sheet_name?)` | Yes | Yes | Yes | Import local CSV/XLSX files |
 
 These tools are useful for tasks such as listing namespaces, inspecting table
 definitions, reviewing indexes, sampling records, running short read-only
 queries with `run_select`, running long read-only queries with `start_query`,
 `get_query`, and `cancel_query`, analyzing read-only queries with `EXPLAIN`, and
-importing prepared local files. For full request and response details, see
+exporting PostgreSQL or MySQL query results to local CSV/XLSX files. You can
+also import prepared local files. For full request and response details, see
 `docs/api-reference.md` (Chinese).
 
 ## How boundaries are constrained
 
 The product boundary is intentionally narrow today. PostgreSQL, MySQL, and Hive
-are available today. Query tools remain read-only, and the only write path is a
+are available today. Query tools remain read-only, PostgreSQL and MySQL query
+results can be exported to local files, and the only database write path is a
 controlled local CSV/XLSX import into existing tables.
 
 The service keeps those boundaries explicit in a few ways.
@@ -112,6 +116,11 @@ The service keeps those boundaries explicit in a few ways.
   queries.
 - The server accepts only `SELECT` and `WITH ... SELECT`, rejects comments and
   multi-statement input, and records audit logs for each call.
+- `export_query_file` writes files on the MCP server machine. It is synchronous
+  but reads database rows and writes CSV/XLSX files in batches. Large exports can
+  still hit your MCP client's tool timeout. For XLSX output, UUID values are
+  written as text and timezone-aware datetime values are written without the
+  timezone. Hive export is not supported yet.
 - `import_table_file` doesn't accept raw SQL. It inserts only file columns whose
   headers exactly match existing table columns.
 - Hive `import_table_file` is intended for small files only and rejects files

{sql_query_mcp-0.3.0 → sql_query_mcp-0.4.0}/sql_query_mcp.egg-info/SOURCES.txt

@@ -9,6 +9,7 @@ sql_query_mcp/audit.py
 sql_query_mcp/config.py
 sql_query_mcp/errors.py
 sql_query_mcp/executor.py
+sql_query_mcp/exporter.py
 sql_query_mcp/importer.py
 sql_query_mcp/introspection.py
 sql_query_mcp/namespace.py
@@ -30,6 +31,7 @@ tests/test_async_queries.py
 tests/test_audit.py
 tests/test_config.py
 tests/test_executor.py
+tests/test_exporter.py
 tests/test_importer.py
 tests/test_metadata.py
 tests/test_namespace.py

{sql_query_mcp-0.3.0 → sql_query_mcp-0.4.0}/tests/test_app.py

@@ -24,6 +24,13 @@ class AppTestCase(unittest.TestCase):
         self.assertIn("get_query", tool_names)
         self.assertIn("cancel_query", tool_names)
 
+    def test_create_app_registers_export_query_file_tool(self) -> None:
+        app = create_app()
+
+        tools = asyncio.run(app.list_tools())
+
+        self.assertIn("export_query_file", {tool.name for tool in tools})
+
 
 if __name__ == "__main__":
     unittest.main()

sql_query_mcp-0.4.0/tests/test_exporter.py

@@ -0,0 +1,395 @@
+from __future__ import annotations
+
+import csv
+import json
+import tempfile
+import unittest
+import uuid
+from contextlib import contextmanager
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import cast
+
+from openpyxl import load_workbook
+
+from sql_query_mcp.audit import AuditLogger
+from sql_query_mcp.config import ConnectionConfig, ServerSettings
+from sql_query_mcp.errors import QueryExecutionError
+from sql_query_mcp.exporter import QueryExporter
+
+
+class _CursorStub:
+    def __init__(self, batches, description=("id", "name")) -> None:
+        self._batches = list(batches)
+        self.description = description
+        self.executed = []
+        self.fetchmany_sizes = []
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> None:
+        return None
+
+    def execute(self, query) -> None:
+        self.executed.append(query)
+
+    def fetchmany(self, size: int):
+        self.fetchmany_sizes.append(size)
+        if self._batches:
+            return self._batches.pop(0)
+        return []
+
+
+class _ConnectionStub:
+    def __init__(self, cursor: _CursorStub) -> None:
+        self._cursor = cursor
+
+    def cursor(self):
+        return self._cursor
+
+
+class _AdapterStub:
+    def __init__(self) -> None:
+        self.set_statement_timeout_calls = []
+
+    def set_statement_timeout(self, conn: object, timeout_ms: int) -> None:
+        self.set_statement_timeout_calls.append(timeout_ms)
+
+    def column_names(self, description):
+        return list(description or [])
+
+
+class _StreamingAdapterStub(_AdapterStub):
+    def __init__(self, export_cursor) -> None:
+        super().__init__()
+        self._export_cursor = export_cursor
+        self.export_cursor_calls = 0
+
+    def export_cursor(self, conn: object):
+        self.export_cursor_calls += 1
+        return self._export_cursor
+
+
+class _RegistryStub:
+    def __init__(self, config: ConnectionConfig, adapter: object, conn: object) -> None:
+        self._config = config
+        self._adapter = adapter
+        self._conn = conn
+        self.connection_calls = 0
+
+    def get_connection_config(self, connection_id: str) -> ConnectionConfig:
+        if connection_id != self._config.connection_id:
+            raise AssertionError(connection_id)
+        return self._config
+
+    @contextmanager
+    def connection_from_config(self, config: ConnectionConfig):
+        if config != self._config:
+            raise AssertionError(config)
+        self.connection_calls += 1
+        yield self._conn, self._adapter
+
+
+def _mysql_config() -> ConnectionConfig:
+    return ConnectionConfig(
+        connection_id="crm_mysql_prod_main_ro",
+        engine="mysql",
+        label="CRM MySQL",
+        env="prod",
+        tenant="main",
+        role="ro",
+        dsn_env="MYSQL_CONN",
+        enabled=True,
+        default_database="crm",
+    )
+
+
+def _postgres_config() -> ConnectionConfig:
+    return ConnectionConfig(
+        connection_id="crm_pg_prod_main_ro",
+        engine="postgres",
+        label="CRM PostgreSQL",
+        env="prod",
+        tenant="main",
+        role="ro",
+        dsn_env="PG_CONN",
+        enabled=True,
+        default_schema="public",
+    )
+
+
+def _hive_config() -> ConnectionConfig:
+    return ConnectionConfig(
+        connection_id="warehouse_hive_prod_main_ro",
+        engine="hive",
+        label="Warehouse Hive",
+        env="prod",
+        tenant="main",
+        role="ro",
+        dsn_env="HIVE_CONN",
+        enabled=True,
+        default_database="default",
+    )
+
+
+def _build_exporter(
+    config: ConnectionConfig,
+    cursor: _CursorStub,
+    temp_dir: str,
+    adapter: _AdapterStub | None = None,
+):
+    adapter = adapter or _AdapterStub()
+    registry = _RegistryStub(config, adapter, _ConnectionStub(cursor))
+    log_path = Path(temp_dir) / "audit.jsonl"
+    exporter = QueryExporter(
+        registry=registry,
+        settings=ServerSettings(
+            default_limit=200,
+            max_limit=1000,
+            statement_timeout_ms=2500,
+            audit_log_path=log_path,
+        ),
+        audit_logger=AuditLogger(log_path),
+    )
+    return exporter, registry, adapter, log_path
+
+
+def _read_audit_records(path: Path):
+    return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines()]
+
+
+class QueryExporterTestCase(unittest.TestCase):
+    def test_export_csv_writes_multiple_fetchmany_batches(self) -> None:
+        cursor = _CursorStub(
+            batches=[
+                [(1, "Alice")],
+                [(2, "Bob")],
+                [],
+            ]
+        )
+        with tempfile.TemporaryDirectory() as temp_dir:
+            exporter, _, adapter, log_path = _build_exporter(
+                _mysql_config(), cursor, temp_dir
+            )
+
+            result = exporter.export_query_file(
+                "crm_mysql_prod_main_ro",
+                "SELECT id, name FROM users",
+                temp_dir,
+                limit=2,
+                file_name="users",
+            )
+            file_path = cast(str, result["file_path"])
+            with Path(file_path).open("r", encoding="utf-8", newline="") as handle:
+                rows = list(csv.reader(handle))
+            records = _read_audit_records(log_path)
+
+        self.assertEqual("csv", result["format"])
+        self.assertEqual(2, result["row_count"])
+        self.assertEqual(2, result["applied_limit"])
+        self.assertEqual(["id", "name"], rows[0])
+        self.assertEqual([["1", "Alice"], ["2", "Bob"]], rows[1:])
+        self.assertIn("LIMIT 2", cursor.executed[0])
+        self.assertEqual([2500], adapter.set_statement_timeout_calls)
+        self.assertGreaterEqual(len(cursor.fetchmany_sizes), 2)
+        self.assertEqual("export_query_file", records[0]["tool"])
+        self.assertEqual(2, records[0]["row_count"])
+
+    def test_export_xlsx_writes_rows(self) -> None:
+        cursor = _CursorStub(batches=[[(1, "Alice")], []])
+        with tempfile.TemporaryDirectory() as temp_dir:
+            exporter, _, _, _ = _build_exporter(_postgres_config(), cursor, temp_dir)
+
+            result = exporter.export_query_file(
+                "crm_pg_prod_main_ro",
+                "SELECT id, name FROM users",
+                temp_dir,
+                format="xlsx",
+                file_name="users",
+            )
+            workbook = load_workbook(cast(str, result["file_path"]), read_only=True)
+            worksheet = workbook.active
+            if worksheet is None:
+                raise AssertionError("workbook has no active worksheet")
+            rows = list(worksheet.iter_rows(values_only=True))
+            workbook.close()
+
+        self.assertEqual("xlsx", result["format"])
+        self.assertEqual([("id", "name"), (1, "Alice")], rows)
+
+    def test_export_all_does_not_wrap_query_with_limit(self) -> None:
+        cursor = _CursorStub(batches=[[(1, "Alice")], []])
+        with tempfile.TemporaryDirectory() as temp_dir:
+            exporter, _, _, _ = _build_exporter(_mysql_config(), cursor, temp_dir)
+            result = exporter.export_query_file(
+                "crm_mysql_prod_main_ro",
+                "SELECT id, name FROM users",
+                temp_dir,
+                export_all=True,
+                file_name="users",
+            )
+
+        self.assertEqual("SELECT id, name FROM users", cursor.executed[0])
+        self.assertIsNone(result["applied_limit"])
+
+    def test_empty_result_creates_header_only_csv(self) -> None:
+        cursor = _CursorStub(batches=[[]])
+        with tempfile.TemporaryDirectory() as temp_dir:
+            exporter, _, _, _ = _build_exporter(_mysql_config(), cursor, temp_dir)
+            result = exporter.export_query_file(
+                "crm_mysql_prod_main_ro",
+                "SELECT id, name FROM users",
+                temp_dir,
+                file_name="empty",
+            )
+            file_path = cast(str, result["file_path"])
+            with Path(file_path).open("r", encoding="utf-8", newline="") as handle:
+                rows = list(csv.reader(handle))
+
+        self.assertEqual(0, result["row_count"])
+        self.assertEqual([["id", "name"]], rows)
+
+    def test_existing_file_gets_numbered_name_when_not_overwriting(self) -> None:
+        cursor = _CursorStub(batches=[[]])
+        with tempfile.TemporaryDirectory() as temp_dir:
+            existing = Path(temp_dir) / "users.csv"
+            existing.write_text("old", encoding="utf-8")
+            exporter, _, _, _ = _build_exporter(_mysql_config(), cursor, temp_dir)
+            result = exporter.export_query_file(
+                "crm_mysql_prod_main_ro",
+                "SELECT id, name FROM users",
+                temp_dir,
+                file_name="users",
+            )
+
+        self.assertTrue(cast(str, result["file_path"]).endswith("users (1).csv"))
+
+    def test_overwrite_replaces_existing_file(self) -> None:
+        cursor = _CursorStub(batches=[[]])
+        with tempfile.TemporaryDirectory() as temp_dir:
+            existing = Path(temp_dir) / "users.csv"
+            existing.write_text("old", encoding="utf-8")
+            exporter, _, _, _ = _build_exporter(_mysql_config(), cursor, temp_dir)
+            result = exporter.export_query_file(
+                "crm_mysql_prod_main_ro",
+                "SELECT id, name FROM users",
+                str(existing),
+                overwrite=True,
+            )
+            content = existing.read_text(encoding="utf-8")
+
+        self.assertEqual(str(existing), result["file_path"])
+        self.assertNotEqual("old", content)
+
+    def test_invalid_sql_rejected_before_connecting(self) -> None:
+        cursor = _CursorStub(batches=[])
+        with tempfile.TemporaryDirectory() as temp_dir:
+            exporter, registry, _, _ = _build_exporter(_mysql_config(), cursor, temp_dir)
+            with self.assertRaises(QueryExecutionError):
+                exporter.export_query_file(
+                    "crm_mysql_prod_main_ro",
+                    "DELETE FROM users",
+                    temp_dir,
+                )
+
+        self.assertEqual(0, registry.connection_calls)
+
+    def test_hive_rejected_before_connecting(self) -> None:
+        cursor = _CursorStub(batches=[])
+        with tempfile.TemporaryDirectory() as temp_dir:
+            exporter, registry, _, _ = _build_exporter(_hive_config(), cursor, temp_dir)
+            with self.assertRaises(QueryExecutionError):
+                exporter.export_query_file(
+                    "warehouse_hive_prod_main_ro",
+                    "SELECT id, name FROM users",
+                    temp_dir,
+                )
+
+        self.assertEqual(0, registry.connection_calls)
+
+    def test_export_uses_adapter_export_cursor_when_available(self) -> None:
+        default_cursor = _CursorStub(batches=[])
+        export_cursor = _CursorStub(batches=[[(1, "Alice")], []])
+        adapter = _StreamingAdapterStub(export_cursor=export_cursor)
+        with tempfile.TemporaryDirectory() as temp_dir:
+            exporter, _, _, _ = _build_exporter(
+                _mysql_config(), default_cursor, temp_dir, adapter=adapter
+            )
+
+            result = exporter.export_query_file(
+                "crm_mysql_prod_main_ro",
+                "SELECT id, name FROM users",
+                temp_dir,
+                file_name="users",
+            )
+
+        self.assertEqual(1, adapter.export_cursor_calls)
+        self.assertEqual(1, result["row_count"])
+        self.assertEqual([], default_cursor.executed)
+        self.assertEqual(1, len(export_cursor.executed))
+
+    def test_export_xlsx_converts_uuid_values_to_text(self) -> None:
+        row_id = uuid.uuid4()
+        cursor = _CursorStub(batches=[[(row_id, "Alice")], []])
+        with tempfile.TemporaryDirectory() as temp_dir:
+            exporter, _, _, _ = _build_exporter(_postgres_config(), cursor, temp_dir)
+
+            result = exporter.export_query_file(
+                "crm_pg_prod_main_ro",
+                "SELECT id, name FROM users",
+                temp_dir,
+                format="xlsx",
+                file_name="users",
+            )
+            workbook = load_workbook(cast(str, result["file_path"]), read_only=True)
+            worksheet = workbook.active
+            if worksheet is None:
+                raise AssertionError("workbook has no active worksheet")
+            rows = list(worksheet.iter_rows(values_only=True))
+            workbook.close()
+
+        self.assertEqual(str(row_id), rows[1][0])
+
+    def test_export_xlsx_removes_datetime_timezone(self) -> None:
+        value = datetime(2026, 1, 1, 8, 30, tzinfo=timezone.utc)
+        cursor = _CursorStub(batches=[[(value, "Alice")], []], description=("update_time", "name"))
+        with tempfile.TemporaryDirectory() as temp_dir:
+            exporter, _, _, _ = _build_exporter(_postgres_config(), cursor, temp_dir)
+
+            result = exporter.export_query_file(
+                "crm_pg_prod_main_ro",
+                "SELECT update_time, name FROM users",
+                temp_dir,
+                format="xlsx",
+                file_name="users",
+            )
+            workbook = load_workbook(cast(str, result["file_path"]), read_only=True)
+            worksheet = workbook.active
+            if worksheet is None:
+                raise AssertionError("workbook has no active worksheet")
+            rows = list(worksheet.iter_rows(values_only=True))
+            workbook.close()
+
+        self.assertEqual(value.replace(tzinfo=None), rows[1][0])
+
+    def test_export_xlsx_error_includes_column_and_type(self) -> None:
+        cursor = _CursorStub(batches=[[(object(), "Alice")], []])
+        with tempfile.TemporaryDirectory() as temp_dir:
+            exporter, _, _, _ = _build_exporter(_postgres_config(), cursor, temp_dir)
+
+            with self.assertRaises(QueryExecutionError) as caught:
+                exporter.export_query_file(
+                    "crm_pg_prod_main_ro",
+                    "SELECT id, name FROM users",
+                    temp_dir,
+                    format="xlsx",
+                    file_name="users",
+                )
+
+        self.assertIn("id", str(caught.exception))
+        self.assertIn("object", str(caught.exception))
+
+
+if __name__ == "__main__":
+    unittest.main()