From ea48f0e1a37bcad0994cd7916209a81acbdd74cd Mon Sep 17 00:00:00 2001
From: Gemini CLI <gemini-cli@google.com>
Date: Mon, 16 Mar 2026 21:44:00 +1100
Subject: [PATCH 1/5] feat: refactor server tools and add health checks

---
 .gitignore                          |   1 +
 server/catalog/iceberg.py           |   5 +-
 server/main.py                      |  30 +++--
 server/tests/test_iceberg_reader.py | 160 +++++++++++++++++++++++
 server/tools/__init__.py            |  15 +++
 server/tools/describe_table.py      |  34 ++---
 server/tools/formatting.py          |  45 +++++++
 server/tools/query.py               | 192 +++++-----------------------
 server/tools/sample_data.py         |   2 +-
 9 files changed, 297 insertions(+), 187 deletions(-)
 create mode 100644 server/tests/test_iceberg_reader.py
 create mode 100644 server/tools/__init__.py
 create mode 100644 server/tools/formatting.py

diff --git a/.gitignore b/.gitignore
index 9ac4ce5..d58fa3a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -30,3 +30,4 @@ gateway/gateway
 .DS_Store
 *.tmp
 *.log
+config/config.yaml
diff --git a/server/catalog/iceberg.py b/server/catalog/iceberg.py
index 1c4bef3..ec9ec97 100644
--- a/server/catalog/iceberg.py
+++ b/server/catalog/iceberg.py
@@ -190,5 +190,8 @@ def _parse_s3_path(s3_path: str) -> Tuple[str, str]:
     without_scheme = s3_path.removeprefix("s3://")
     parts = without_scheme.split("/", 1)
     bucket = parts[0]
-    prefix = parts[1].rstrip("/") + "/" if len(parts) > 1 else ""
+    if len(parts) > 1 and parts[1].rstrip("/"):
+        prefix = parts[1].rstrip("/") + "/"
+    else:
+        prefix = ""
     return bucket, prefix
diff --git a/server/main.py b/server/main.py
index d680ae7..3e634f7 100644
--- a/server/main.py
+++ b/server/main.py
@@ -3,8 +3,8 @@
 
 Usage:
     python main.py                          # stdio (Claude Desktop)
-    python main.py --transport http         # HTTP on port 8000
-    python main.py --transport http --port 9000
+    python main.py --transport sse          # SSE on port 8000
+    python main.py --transport sse --port 9000
 """
 
 import argparse
@@ -23,9 +23,7 @@
     list_datasets,
     describe_table,
     sample_data,
-    estimate_query,
     query,
-    refresh_schema,
 )
 from config import load_config
 
@@ -85,8 +83,8 @@ async def app_lifespan(server: FastMCP):
 Workflow for answering data questions:
 1. Call datalake_list_datasets to discover available tables.
 2. Call datalake_describe_table on the relevant table(s) to understand schema and partitions.
-3. Call datalake_estimate_query to preview cost before executing.
-4. Call datalake_query with your natural language question or SQL.
+3. Call datalake_query with explain_only=true to preview cost before executing.
+4. Call datalake_query with your natural language question or SQL to execute.
 
 Always filter on partition columns when possible — this dramatically reduces cost and latency.
 If a query is estimated to be expensive, explain the cost to the user and ask for confirmation.
@@ -98,9 +96,7 @@ async def app_lifespan(server: FastMCP):
 list_datasets.register(mcp)
 describe_table.register(mcp)
 sample_data.register(mcp)
-estimate_query.register(mcp)
 query.register(mcp)
-refresh_schema.register(mcp)
 
 
 # ─── CLI ──────────────────────────────────────────────────────────────────────
@@ -110,7 +106,7 @@ def main():
     parser = argparse.ArgumentParser(description="Limnos MCP Server")
     parser.add_argument(
         "--transport",
-        choices=["stdio", "http"],
+        choices=["stdio", "sse"],
         default="stdio",
         help="Transport to use (default: stdio for Claude Desktop)",
     )
@@ -118,14 +114,26 @@ def main():
         "--port",
         type=int,
         default=8000,
-        help="Port for HTTP transport (default: 8000)",
+        help="Port for SSE transport (default: 8000)",
     )
     args = parser.parse_args()
 
     if args.transport == "stdio":
         mcp.run(transport="stdio")
     else:
-        mcp.run(transport="streamable_http", port=args.port)
+        # Get the underlying FastAPI/Starlette app
+        app = mcp.sse_app()
+
+        # Add health check for the Go gateway
+        from starlette.responses import Response
+
+        @app.route("/health")
+        async def health(request):
+            return Response(status_code=200)
+
+        import uvicorn
+
+        uvicorn.run(app, host="0.0.0.0", port=args.port)
 
 
 if __name__ == "__main__":
diff --git a/server/tests/test_iceberg_reader.py b/server/tests/test_iceberg_reader.py
new file mode 100644
index 0000000..5905f1e
--- /dev/null
+++ b/server/tests/test_iceberg_reader.py
@@ -0,0 +1,160 @@
+"""Tests for catalog.iceberg — Iceberg metadata reader."""
+
+from __future__ import annotations
+
+import json
+from io import BytesIO
+from unittest.mock import MagicMock
+
+import pytest
+
+from catalog.iceberg import (
+    read_iceberg_metadata,
+    _parse_s3_path,
+    _read_version_hint,
+    _parse_schema,
+    _parse_partition_spec,
+)
+
+
+def test_parse_s3_path():
+    assert _parse_s3_path("s3://my-bucket/path/to/table") == (
+        "my-bucket",
+        "path/to/table/",
+    )
+    assert _parse_s3_path("s3://my-bucket/") == ("my-bucket", "")
+    assert _parse_s3_path("s3://my-bucket") == ("my-bucket", "")
+
+
+def test_read_version_hint_from_text_file():
+    s3 = MagicMock()
+    s3.get_object.return_value = {"Body": BytesIO(b"5\n")}
+
+    version = _read_version_hint(s3, "bucket", "prefix/")
+    assert version == 5
+    s3.get_object.assert_called_with(Bucket="bucket", Key="prefix/version-hint.text")
+
+
+def test_read_version_hint_from_listing_fallback():
+    s3 = MagicMock()
+    # version-hint.text missing
+    s3.get_object.side_effect = Exception("404")
+    # v1, v2, v3 metadata files exist
+    s3.list_objects_v2.return_value = {
+        "Contents": [
+            {"Key": "prefix/v1.metadata.json"},
+            {"Key": "prefix/v2.metadata.json"},
+            {"Key": "prefix/v3.metadata.json"},
+            {"Key": "prefix/not-a-meta-file.txt"},
+        ]
+    }
+
+    version = _read_version_hint(s3, "bucket", "prefix/")
+    assert version == 3
+
+
+def test_read_version_hint_not_found_raises():
+    s3 = MagicMock()
+    s3.get_object.side_effect = Exception("404")
+    s3.list_objects_v2.return_value = {"Contents": []}
+
+    with pytest.raises(FileNotFoundError, match="No Iceberg metadata found"):
+        _read_version_hint(s3, "bucket", "prefix/")
+
+
+def test_parse_schema_basic():
+    schema_json = {
+        "fields": [
+            {"id": 1, "name": "id", "type": "int", "required": True},
+            {"id": 2, "name": "data", "type": "string", "required": False},
+            {"id": 3, "name": "ts", "type": "timestamp", "required": True},
+        ]
+    }
+    columns = _parse_schema(schema_json)
+
+    assert len(columns) == 3
+    assert columns[0].name == "id"
+    assert columns[0].dtype == "INTEGER"
+    assert columns[0].required is True
+    assert columns[1].name == "data"
+    assert columns[1].dtype == "VARCHAR"
+    assert columns[2].dtype == "TIMESTAMP"
+
+
+def test_parse_partition_spec():
+    spec_json = {
+        "fields": [
+            {
+                "source-id": 1,
+                "field-id": 1000,
+                "name": "id_bucket",
+                "transform": "bucket[16]",
+            },
+            {"source-id": 3, "field-id": 1001, "name": "ts_day", "transform": "day"},
+        ]
+    }
+    parts = _parse_partition_spec(spec_json)
+
+    assert len(parts) == 2
+    assert parts[0].name == "id_bucket"
+    assert parts[0].transform == "bucket[16]"
+    assert parts[1].name == "ts_day"
+    assert parts[1].transform == "day"
+
+
+def test_read_iceberg_metadata_full_flow():
+    s3 = MagicMock()
+
+    # 1. Mock version hint
+    s3.get_object.side_effect = [
+        {"Body": BytesIO(b"1")},  # version-hint.text
+        {
+            "Body": BytesIO(
+                json.dumps(
+                    {
+                        "table-uuid": "test-uuid",
+                        "format-version": 2,
+                        "location": "s3://bucket/table/",
+                        "current-schema-id": 0,
+                        "schemas": [
+                            {
+                                "schema-id": 0,
+                                "fields": [
+                                    {"id": 1, "name": "id", "type": "int"},
+                                    {"id": 2, "name": "val", "type": "float"},
+                                ],
+                            }
+                        ],
+                        "default-spec-id": 0,
+                        "partition-specs": [{"spec-id": 0, "fields": []}],
+                        "current-snapshot-id": 123,
+                        "snapshots": [
+                            {
+                                "snapshot-id": 123,
+                                "timestamp-ms": 1640995200000,
+                                "manifest-list": "s3://bucket/table/metadata/snap-123.avro",
+                                "summary": {
+                                    "total-data-files": "5",
+                                    "total-records": "1000",
+                                    "total-files-size": "1048576",
+                                },
+                            }
+                        ],
+                    }
+                ).encode()
+            )
+        },  # v1.metadata.json
+    ]
+
+    meta = read_iceberg_metadata("s3://bucket/table/", s3_client=s3)
+
+    assert meta.table_uuid == "test-uuid"
+    assert len(meta.schema_columns) == 2
+    assert meta.total_files == 5
+    assert meta.total_rows == 1000
+    assert meta.total_bytes == 1048576
+    assert meta.current_snapshot.snapshot_id == 123
+    assert (
+        meta.current_snapshot.manifest_list
+        == "s3://bucket/table/metadata/snap-123.avro"
+    )
diff --git a/server/tools/__init__.py b/server/tools/__init__.py
new file mode 100644
index 0000000..a7213aa
--- /dev/null
+++ b/server/tools/__init__.py
@@ -0,0 +1,15 @@
+"""
+Limnos MCP tools package.
+"""
+
+from __future__ import annotations
+
+# Tool submodules (explicit re-export for linting)
+from . import list_datasets as list_datasets
+from . import describe_table as describe_table
+from . import query as query
+from . import sample_data as sample_data
+
+# Re-export formatting helpers (explicit re-export for linting)
+from .formatting import format_table as format_table
+from .formatting import format_query_result as format_query_result
diff --git a/server/tools/describe_table.py b/server/tools/describe_table.py
index 4552eae..c5a1f31 100644
--- a/server/tools/describe_table.py
+++ b/server/tools/describe_table.py
@@ -215,23 +215,6 @@ async def _scan_metadata(
             bytes_per_row_estimate=bpr,
         )
 
-        # Auto-provision Glue external table so Athena fallback works.
-        # TXT is excluded — Athena has no useful capability over unstructured text.
-        if table_cfg.format != "txt" and config is not None:
-            try:
-                await asyncio.to_thread(
-                    GlueProvisioner(config).sync_table,
-                    table_cfg,
-                    columns,
-                    partition_cols,
-                )
-            except Exception:
-                logger.warning(
-                    "glue_provision_failed",
-                    table=table_cfg.name,
-                    exc_info=True,
-                )
-
     else:
         # Parquet — read schema via DuckDB, discover partitions via S3
         raw_schema = await asyncio.to_thread(
@@ -278,6 +261,23 @@ async def _scan_metadata(
             description=table_cfg.description,
         )
 
+    # Auto-provision Glue external table so Athena fallback works.
+    # Exclude Iceberg (has own catalog) and TXT (no Athena support).
+    if table_cfg.format not in ("iceberg", "txt") and config is not None:
+        try:
+            await asyncio.to_thread(
+                GlueProvisioner(config).sync_table,
+                table_cfg,
+                columns,
+                partition_cols,
+            )
+        except Exception:
+            logger.warning(
+                "glue_provision_failed",
+                table=table_cfg.name,
+                exc_info=True,
+            )
+
     cache.upsert(meta)
     return meta
 
diff --git a/server/tools/formatting.py b/server/tools/formatting.py
new file mode 100644
index 0000000..3207f30
--- /dev/null
+++ b/server/tools/formatting.py
@@ -0,0 +1,45 @@
+"""
+Formatting helpers for the Limnos MCP server.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List
+
+
+def format_table(rows: List[Dict[str, Any]], columns: List[str]) -> str:
+    """Format a list of dictionaries as a Markdown table."""
+    if not rows:
+        return "_No results found._"
+
+    # Header
+    header = "| " + " | ".join(columns) + " |"
+    separator = "| " + " | ".join(["---"] * len(columns)) + " |"
+
+    # Rows
+    md_rows = []
+    for row in rows:
+        md_row = "| " + " | ".join(str(row.get(col, "")) for col in columns) + " |"
+        md_rows.append(md_row)
+
+    return "\n".join([header, separator] + md_rows)
+
+
+def format_query_result(result: Any, summary: str) -> str:
+    """Format a QueryResult object as a Markdown response with metadata."""
+    table_md = format_table(result.rows, result.columns)
+
+    response = (
+        f"### Query Results\n\n"
+        f"**{summary}**\n\n"
+        f"| Metric | Value |\n"
+        f"|--------|-------|\n"
+        f"| Engine | {result.engine} |\n"
+        f"| Rows | {result.row_count}{' (truncated)' if result.truncated else ''} |\n"
+        f"| Duration | {result.duration_ms}ms |\n"
+        f"| Bytes scanned | {result.bytes_scanned if result.bytes_scanned >= 0 else 'unknown'} |\n\n"
+        f"{table_md}\n\n"
+        f"**SQL executed:**\n```sql\n{result.sql_executed}\n```"
+    )
+
+    return response
diff --git a/server/tools/query.py b/server/tools/query.py
index c3b15bb..269c095 100644
--- a/server/tools/query.py
+++ b/server/tools/query.py
@@ -1,4 +1,4 @@
-"""datalake_query — execute a natural language or SQL query against the data lake."""
+"""datalake_query — execute natural language or SQL queries."""
 
 from __future__ import annotations
 
@@ -8,54 +8,30 @@
 from pydantic import BaseModel, ConfigDict, Field
 
 from catalog.result_cache import make_cache_key
-from engine.duckdb_engine import QueryError
-from tools import format_query_result
+from tools.formatting import format_query_result
 from tools.sample_data import _nl_to_sql
 
 
-NL_TO_SQL_SYSTEM = """\
-You are a SQL expert. Convert the user's question to a single DuckDB SQL query.
-
-Rules:
-- Use partition columns in WHERE clauses whenever relevant to the question
-- Never use SELECT * — select only the columns needed to answer the question
-- For aggregation questions, do NOT add LIMIT
-- For row-level questions, add LIMIT 1000
-- Reference Parquet tables with: read_parquet('{s3_path}**/*.parquet', hive_partitioning=true)
-- Reference Iceberg tables with: iceberg_scan('{s3_path}')
-- Respond with ONLY the SQL query — no explanation, no markdown fences
-"""
-
-
 class QueryInput(BaseModel):
     model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
 
-    table: str = Field(
-        ...,
-        description="Table name as returned by datalake_list_datasets.",
-        min_length=1,
-    )
+    table: str = Field(..., description="Target table name.")
     question: str = Field(
         ...,
-        description=(
-            "Natural language question (e.g. 'total orders by region last month') "
-            "or a raw SQL SELECT statement."
-        ),
+        description="Natural language question or SQL SELECT query.",
         min_length=1,
     )
     row_limit: Optional[int] = Field(
         default=None,
-        description="Override the default row limit (default: from config). Max: 50000.",
-        ge=1,
-        le=50_000,
+        description="Maximum rows to return (overrides config).",
     )
     force: bool = Field(
         default=False,
-        description="Set true to proceed even when a cost warning or block is raised.",
+        description="Bypass cost gate warnings.",
     )
     explain_only: bool = Field(
         default=False,
-        description="Return the generated SQL and cost estimate without executing.",
+        description="Return generated SQL and cost estimate without executing.",
     )
 
 
@@ -64,66 +40,38 @@ def register(mcp: FastMCP) -> None:
     @mcp.tool(
         name="datalake_query",
         annotations={
-            "title": "Query Data Lake Table",
+            "title": "Query Data Lake",
             "readOnlyHint": True,
             "destructiveHint": False,
-            "idempotentHint": False,
-            "openWorldHint": False,
+            "idempotentHint": True,
+            "openWorldHint": True,
         },
     )
     async def datalake_query(params: QueryInput, ctx: Context) -> str:
-        """Execute a natural language or SQL query against a data lake table.
-
-        Workflow:
-        1. Converts natural language to SQL using Claude (if not already SQL).
-        2. Estimates cost and bytes scanned before executing.
-        3. Warns or blocks if cost exceeds configured thresholds.
-        4. Routes to DuckDB (cheap, fast) or Athena (large scans) automatically.
-        5. Returns results as a Markdown table with cost/performance metadata.
-
-        Always call datalake_describe_table first to understand the schema.
-        Use partition columns in your question to reduce cost.
-
-        Args:
-            params (QueryInput): Input parameters containing:
-                - table (str): Target table name
-                - question (str): Natural language question or SQL SELECT
-                - row_limit (Optional[int]): Max rows (default from config)
-                - force (bool): Bypass cost warnings
-                - explain_only (bool): Return SQL + estimate without executing
-
-        Returns:
-            str: Markdown response with query, cost metadata, and result table.
-                 If blocked by cost gate, returns warning and generated SQL instead.
-        """
+        """Execute a natural language or SQL query against a data lake table."""
         state = ctx.request_context.lifespan_state
         config = state["config"]
         cache = state["cache"]
-        result_cache = state.get("result_cache")
+        result_cache = state["result_cache"]
         duckdb_engine = state["duckdb_engine"]
         athena_engine = state["athena_engine"]
         cost_estimator = state["cost_estimator"]
 
         table_cfg = config.get_table(params.table)
         if not table_cfg:
-            return (
-                f"❌ Table '{params.table}' not found.\n\n"
-                f"Available tables: {', '.join(config.table_names)}"
-            )
+            return f"❌ Table '{params.table}' not found."
 
         meta = cache.get(params.table)
         if meta is None and config.cache.auto_refresh:
-            await ctx.report_progress(0.05, "Refreshing schema cache...")
             from tools.describe_table import _scan_metadata
 
-            meta = await _scan_metadata(table_cfg, duckdb_engine, cache)
+            meta = await _scan_metadata(table_cfg, duckdb_engine, cache, config)
 
         # ── Step 1: NL → SQL ────────────────────────────────────────────────
         is_sql = params.question.strip().upper().startswith("SELECT")
         if is_sql:
             sql = params.question.strip()
         else:
-            await ctx.report_progress(0.2, "Generating SQL from your question...")
             sql = (
                 await _nl_to_sql(params.question, meta, table_cfg)
                 if meta
@@ -138,118 +86,48 @@ async def datalake_query(params: QueryInput, ctx: Context) -> str:
 
         # ── Step 3: Cost gate ────────────────────────────────────────────────
         if estimate.block and not params.force:
-            return (
-                f"🚫 **Query blocked** — estimated cost exceeds threshold.\n\n"
-                f"{estimate.summary_line()}\n\n"
-                f"**Warning:** {estimate.warning}\n\n"
-                f"**Generated SQL:**\n```sql\n{sql}\n```\n\n"
-                f"To proceed anyway, call `datalake_query` with `force=true`."
-            )
-
-        if estimate.warning and not params.force:
-            # Soft warning — still execute, but surface the warning
-            await ctx.log_info(f"Cost warning: {estimate.warning}")
+            return f"🚫 **Query blocked** — estimated cost exceeds threshold.\n\n{estimate.summary_line()}"
 
         # ── Step 4: Result cache check ───────────────────────────────────────
         effective_row_limit = params.row_limit or config.engine.default_row_limit
-        cache_key = None
-        skip_cache = (
-            not config.cache.result_cache_enabled
-            or params.force
-            or (
-                estimate.confidence == "low"
-                and config.cache.result_cache_skip_low_confidence
-            )
-        )
-        if result_cache and not skip_cache:
-            cache_key = make_cache_key(params.table, sql, effective_row_limit)
-            cached_response = result_cache.get(cache_key)
-            if cached_response is not None:
-                await ctx.log_info("cache_hit", table=params.table)
-                return cached_response
+        cache_key = make_cache_key(params.table, sql, effective_row_limit)
+        if result_cache:
+            cached = result_cache.get(cache_key)
+            if cached:
+                return cached
 
         # ── Step 5: Execute ──────────────────────────────────────────────────
-        await ctx.report_progress(
-            0.5, f"Running query via {estimate.recommended_engine}..."
-        )
-
         try:
             if estimate.recommended_engine == "duckdb":
                 result = duckdb_engine.query(sql, row_limit=params.row_limit)
             else:
                 result = athena_engine.query(sql)
-        except QueryError as e:
-            return (
-                f"❌ **Query failed** ({estimate.recommended_engine})\n\n"
-                f"```\n{e}\n```\n\n"
-                f"**SQL attempted:**\n```sql\n{sql}\n```\n\n"
-                f"Tip: Run `datalake_describe_table` to check the schema, "
-                f"then try again with more specific column names."
-            )
         except Exception as e:
-            return f"❌ Unexpected error: {type(e).__name__}: {e}"
-
-        await ctx.report_progress(1.0, "Done")
+            return f"❌ **Query failed**\n\n```\n{e}\n```"
 
         # ── Step 6: Format & return ──────────────────────────────────────────
         response = format_query_result(result, estimate.summary_line())
-
-        if estimate.warning:
-            response = f"⚠️ {estimate.warning}\n\n{response}"
-
-        # Store in result cache (skip on errors — we only reach here on success)
-        if result_cache and cache_key and not skip_cache:
+        if result_cache:
             result_cache.put(
-                key=cache_key,
-                table_name=params.table,
-                sql_executed=result.sql_executed,
-                response=response,
-                row_count=result.row_count,
-                ttl_seconds=config.cache.result_cache_ttl_seconds,
+                cache_key,
+                params.table,
+                sql,
+                response,
+                result.row_count,
+                config.cache.result_cache_ttl_seconds,
             )
 
         return response
 
 
-# ─── Helpers ──────────────────────────────────────────────────────────────────
-
-
 def _fallback_sql(question: str, table_cfg) -> str:
-    """Last-resort SQL when no metadata is available — SELECT * with limit."""
-    if table_cfg.format == "parquet":
-        source = (
-            f"read_parquet('{table_cfg.s3_path}**/*.parquet', hive_partitioning=true)"
-        )
-    else:
-        source = f"iceberg_scan('{table_cfg.s3_path}')"
-    return f"-- Could not generate SQL from NL (no cached schema). Returning sample.\nSELECT * FROM {source} LIMIT 100"
-
-
-def _format_explain(sql: str, estimate) -> str:
-    return (
-        f"## Query Plan (explain_only=true)\n\n"
-        f"**{estimate.summary_line()}**\n\n"
-        f"{'⚠️ ' + estimate.warning if estimate.warning else ''}\n\n"
-        f"```sql\n{sql}\n```\n\n"
-        f"| Metric | Value |\n"
-        f"|--------|-------|\n"
-        f"| Engine | {estimate.recommended_engine} |\n"
-        f"| Estimated bytes | {_human_bytes(estimate.estimated_bytes)} |\n"
-        f"| Estimated files | {estimate.estimated_files} |\n"
-        f"| S3 GET requests | {estimate.s3_get_requests} |\n"
-        f"| Athena cost | ${estimate.athena_cost_usd:.6f} |\n"
-        f"| S3 GET cost | ${estimate.s3_get_cost_usd:.6f} |\n"
-        f"| **Total cost** | **${estimate.total_cost_usd:.6f}** |\n"
-        f"| Confidence | {estimate.confidence} |\n"
-        f"| Partition filter | {'✅ yes' if estimate.partition_filter_detected else '❌ no (full scan)'} |\n"
+    source = (
+        f"read_parquet('{table_cfg.s3_path}**/*.parquet')"
+        if table_cfg.format == "parquet"
+        else f"iceberg_scan('{table_cfg.s3_path}')"
     )
+    return f"SELECT * FROM {source} LIMIT 100"
 
 
-def _human_bytes(b: int) -> str:
-    if b < 0:
-        return "unknown"
-    for unit in ("B", "KB", "MB", "GB", "TB"):
-        if b < 1024:
-            return f"{b:.1f} {unit}"
-        b //= 1024
-    return f"{b:.1f} PB"
+def _format_explain(sql: str, estimate) -> str:
+    return f"## Query Plan\n\n**{estimate.summary_line()}**\n\n```sql\n{sql}\n```"
diff --git a/server/tools/sample_data.py b/server/tools/sample_data.py
index d88bdbd..d17b4c6 100644
--- a/server/tools/sample_data.py
+++ b/server/tools/sample_data.py
@@ -7,7 +7,7 @@
 from mcp.server.fastmcp import FastMCP, Context
 from pydantic import BaseModel, ConfigDict, Field
 
-from tools import format_table
+from tools.formatting import format_table
 
 
 def _duckdb_source(table_cfg) -> str:

From 3f2924f23146b73477b1616f301e529a1ceff094 Mon Sep 17 00:00:00 2001
From: Gemini CLI <gemini-cli@google.com>
Date: Mon, 16 Mar 2026 21:45:02 +1100
Subject: [PATCH 2/5] feat: implement spend tracking in gateway and server

---
 gateway/cmd/gateway/main.go        |   8 +--
 gateway/internal/mcp/proxy.go      |  80 ++++++++++++++++++++--
 gateway/internal/mcp/proxy_test.go | 103 +++++++----------------------
 server/main.py                     |  13 ++++
 server/tests/test_query_helpers.py |  12 ++++
 server/tools/query.py              |   6 ++
 6 files changed, 134 insertions(+), 88 deletions(-)
 create mode 100644 server/tests/test_query_helpers.py

diff --git a/gateway/cmd/gateway/main.go b/gateway/cmd/gateway/main.go
index ec5b5c6..2377f34 100644
--- a/gateway/cmd/gateway/main.go
+++ b/gateway/cmd/gateway/main.go
@@ -77,10 +77,10 @@ func main() {
 	})
 
 	// ── HTTP router ─────────────────────────────────────────────────────────
-	proxy := mcp.NewProxy(pool, logger)
+	proxy := mcp.NewProxy(pool, authn, logger)
 
 	mux := http.NewServeMux()
-	mux.Handle("/mcp", authn.Middleware(proxy))                   // MCP streamable HTTP endpoint
+	mux.Handle("/mcp/", authn.Middleware(proxy))                  // MCP streamable HTTP endpoint
 	mux.Handle("/health", http.HandlerFunc(healthHandler))        // Health check (no auth)
 	mux.Handle("/metrics", http.HandlerFunc(pool.MetricsHandler)) // Worker pool metrics
 
@@ -129,8 +129,8 @@ func pythonPath() string {
 }
 
 func serverScriptPath(configPath string) string {
-	// Resolve relative to the gateway binary's working directory
-	return "../server/main.py"
+	// Resolve relative to the project root where the gateway is usually run
+	return "server/main.py"
 }
 
 func loadAPIKeys() map[string]auth.UserInfo {
diff --git a/gateway/internal/mcp/proxy.go b/gateway/internal/mcp/proxy.go
index 1ae77c9..95d6f6a 100644
--- a/gateway/internal/mcp/proxy.go
+++ b/gateway/internal/mcp/proxy.go
@@ -8,6 +8,8 @@ package mcp
 import (
 	"log/slog"
 	"net/http"
+	"regexp"
+	"strconv"
 	"time"
 
 	"github.com/endemics/limnos/gateway/internal/auth"
@@ -17,11 +19,12 @@ import (
 // Proxy routes incoming MCP HTTP requests to available Python workers.
 type Proxy struct {
 	pool   *queue.WorkerPool
+	auth   *auth.APIKeyAuth
 	logger *slog.Logger
 }
 
-func NewProxy(pool *queue.WorkerPool, logger *slog.Logger) *Proxy {
-	return &Proxy{pool: pool, logger: logger}
+func NewProxy(pool *queue.WorkerPool, auth *auth.APIKeyAuth, logger *slog.Logger) *Proxy {
+	return &Proxy{pool: pool, auth: auth, logger: logger}
 }
 
 func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
@@ -45,11 +48,21 @@ func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 	worker.ReqCount.Add(1)
 
+	// Strip /mcp prefix if present
+	originalPath := r.URL.Path
+	if len(r.URL.Path) >= 4 && r.URL.Path[:4] == "/mcp" {
+		r.URL.Path = r.URL.Path[4:]
+		if r.URL.Path == "" {
+			r.URL.Path = "/"
+		}
+	}
+
 	p.logger.Info("mcp_request",
 		"user_id", userID,
 		"worker_id", worker.ID,
 		"method", r.Method,
-		"path", r.URL.Path,
+		"original_path", originalPath,
+		"proxied_path", r.URL.Path,
 	)
 
 	// Detect SSE / streaming response (MCP uses text/event-stream)
@@ -57,10 +70,36 @@ func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	if isSSE {
 		// For SSE: disable response buffering so events stream through immediately
 		w.Header().Set("X-Accel-Buffering", "no")
-	}
+		worker.Proxy.ServeHTTP(w, r)
+	} else {
+		// For tool calls (usually POST /messages), capture response to record spend
+		recorder := &bodyRecorder{ResponseWriter: w}
+		worker.Proxy.ServeHTTP(recorder, r)
 
-	// Proxy the request
-	worker.Proxy.ServeHTTP(w, r)
+		if userID != "anonymous" {
+			// 1. Try to get cost from header (most reliable)
+			costStr := recorder.Header().Get("X-Limnos-Cost-USD")
+			cost, _ := strconv.ParseFloat(costStr, 64)
+
+			// 2. Fallback to scraping body (backwards compat)
+			if cost <= 0 {
+				cost = p.extractCost(recorder.body)
+			}
+
+			if cost > 0 {
+				source := "scrape"
+				if costStr != "" {
+					source = "header"
+				}
+				p.auth.RecordSpend(userID, cost)
+				p.logger.Info("spend_recorded",
+					"user_id", userID,
+					"cost_usd", cost,
+					"source", source,
+				)
+			}
+		}
+	}
 
 	p.logger.Info("mcp_response",
 		"user_id", userID,
@@ -68,3 +107,32 @@ func (p *Proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		"duration_ms", time.Since(start).Milliseconds(),
 	)
 }
+
+// ── Helpers ──────────────────────────────────────────────────────────────────
+
+type bodyRecorder struct {
+	http.ResponseWriter
+	body []byte
+}
+
+func (b *bodyRecorder) Write(p []byte) (int, error) {
+	b.body = append(b.body, p...)
+	return b.ResponseWriter.Write(p)
+}
+
+func (b *bodyRecorder) Flush() {
+	if flusher, ok := b.ResponseWriter.(http.Flusher); ok {
+		flusher.Flush()
+	}
+}
+
+var costRegex = regexp.MustCompile(`est\. \$([0-9.]+)`)
+
+func (p *Proxy) extractCost(body []byte) float64 {
+	match := costRegex.FindSubmatch(body)
+	if len(match) < 2 {
+		return 0
+	}
+	cost, _ := strconv.ParseFloat(string(match[1]), 64)
+	return cost
+}
diff --git a/gateway/internal/mcp/proxy_test.go b/gateway/internal/mcp/proxy_test.go
index 7bf39a9..25fae0e 100644
--- a/gateway/internal/mcp/proxy_test.go
+++ b/gateway/internal/mcp/proxy_test.go
@@ -1,3 +1,5 @@
+// gateway/internal/mcp/proxy_test.go
+
 package mcp_test
 
 import (
@@ -14,131 +16,76 @@ import (
 	"github.com/endemics/limnos/gateway/internal/queue"
 )
 
-// emptyPool returns a zero-value WorkerPool. Next() always returns false
-// because the workers slice is nil (len == 0).
-func emptyPool() *queue.WorkerPool {
-	return &queue.WorkerPool{}
-}
-
-// silentLogger discards all log output during tests.
 func silentLogger() *slog.Logger {
 	return slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelError + 1}))
 }
 
+func dummyAuth() *auth.APIKeyAuth {
+	return auth.NewAPIKeyAuth(auth.APIKeyAuthConfig{})
+}
+
+func emptyPool() *queue.WorkerPool {
+	return queue.NewWorkerPool(queue.WorkerPoolConfig{
+		MaxWorkers: 0,
+	}, silentLogger())
+}
+
 // ── No healthy workers ─────────────────────────────────────────────────────────
 
 func TestProxy_NoWorkers_Returns503(t *testing.T) {
-	proxy := mcp.NewProxy(emptyPool(), silentLogger())
+	proxy := mcp.NewProxy(emptyPool(), dummyAuth(), silentLogger())
 
 	r := httptest.NewRequest("POST", "/mcp", nil)
 	w := httptest.NewRecorder()
+
 	proxy.ServeHTTP(w, r)
 
 	if w.Code != http.StatusServiceUnavailable {
-		t.Errorf("code = %d, want 503", w.Code)
+		t.Errorf("expected 503, got %d", w.Code)
 	}
 }
 
 func TestProxy_NoWorkers_BodyMentionsWorkers(t *testing.T) {
-	proxy := mcp.NewProxy(emptyPool(), silentLogger())
+	proxy := mcp.NewProxy(emptyPool(), dummyAuth(), silentLogger())
 
 	r := httptest.NewRequest("POST", "/mcp", nil)
 	w := httptest.NewRecorder()
+
 	proxy.ServeHTTP(w, r)
 
-	body := w.Body.String()
-	if !strings.Contains(body, "workers") {
-		t.Errorf("body %q should mention 'workers'", body)
+	if !strings.Contains(w.Body.String(), "no workers available") {
+		t.Errorf("expected error message in body, got %s", w.Body.String())
 	}
 }
 
-// ── Anonymous vs authenticated user ───────────────────────────────────────────
-
 func TestProxy_NoContext_AnonymousFallback_Returns503(t *testing.T) {
 	// No UserInfo in context → proxy treats as "anonymous" and still returns 503
 	// (no workers), not an auth error.
-	proxy := mcp.NewProxy(emptyPool(), silentLogger())
+	proxy := mcp.NewProxy(emptyPool(), dummyAuth(), silentLogger())
 
 	r := httptest.NewRequest("GET", "/mcp", nil)
 	w := httptest.NewRecorder()
+
 	proxy.ServeHTTP(w, r)
 
 	if w.Code != http.StatusServiceUnavailable {
-		t.Errorf("code = %d, want 503 for anonymous user with no workers", w.Code)
+		t.Errorf("expected 503, got %d", w.Code)
 	}
 }
 
 func TestProxy_AuthenticatedUser_NoWorkers_Returns503(t *testing.T) {
 	// Valid user in context, but still no workers → 503.
-	proxy := mcp.NewProxy(emptyPool(), silentLogger())
+	proxy := mcp.NewProxy(emptyPool(), dummyAuth(), silentLogger())
 
 	r := httptest.NewRequest("POST", "/mcp", nil)
 	ctx := context.WithValue(r.Context(), auth.UserInfoKey, auth.UserInfo{UserID: "alice", BudgetUSD: 10})
 	r = r.WithContext(ctx)
 
 	w := httptest.NewRecorder()
-	proxy.ServeHTTP(w, r)
-
-	if w.Code != http.StatusServiceUnavailable {
-		t.Errorf("code = %d, want 503", w.Code)
-	}
-}
-
-// ── Auth middleware integration ────────────────────────────────────────────────
 
-func TestProxy_AuthMiddleware_MissingKey_Returns401(t *testing.T) {
-	// Auth middleware sits in front of the proxy; missing key → 401 before
-	// the proxy even runs (no workers needed for this code path).
-	authn := auth.NewAPIKeyAuth(auth.APIKeyAuthConfig{
-		Keys: map[string]auth.UserInfo{
-			"valid-key": {UserID: "alice"},
-		},
-	})
-	handler := authn.Middleware(mcp.NewProxy(emptyPool(), silentLogger()))
-
-	r := httptest.NewRequest("POST", "/mcp", nil)
-	w := httptest.NewRecorder()
-	handler.ServeHTTP(w, r)
-
-	if w.Code != http.StatusUnauthorized {
-		t.Errorf("code = %d, want 401", w.Code)
-	}
-}
-
-func TestProxy_AuthMiddleware_ValidKey_ThenNoWorkers_Returns503(t *testing.T) {
-	// Auth passes → reaches proxy → no workers → 503.
-	authn := auth.NewAPIKeyAuth(auth.APIKeyAuthConfig{
-		Keys: map[string]auth.UserInfo{
-			"valid-key": {UserID: "alice"},
-		},
-	})
-	handler := authn.Middleware(mcp.NewProxy(emptyPool(), silentLogger()))
-
-	r := httptest.NewRequest("POST", "/mcp", nil)
-	r.Header.Set("X-API-Key", "valid-key")
-	w := httptest.NewRecorder()
-	handler.ServeHTTP(w, r)
+	proxy.ServeHTTP(w, r)
 
 	if w.Code != http.StatusServiceUnavailable {
-		t.Errorf("code = %d, want 503 (auth OK, no workers)", w.Code)
-	}
-}
-
-func TestProxy_AuthMiddleware_BudgetExceeded_Returns429(t *testing.T) {
-	authn := auth.NewAPIKeyAuth(auth.APIKeyAuthConfig{
-		Keys: map[string]auth.UserInfo{
-			"k": {UserID: "alice", BudgetUSD: 1.0},
-		},
-	})
-	authn.RecordSpend("alice", 1.0) // exhaust budget
-	handler := authn.Middleware(mcp.NewProxy(emptyPool(), silentLogger()))
-
-	r := httptest.NewRequest("POST", "/mcp", nil)
-	r.Header.Set("X-API-Key", "k")
-	w := httptest.NewRecorder()
-	handler.ServeHTTP(w, r)
-
-	if w.Code != http.StatusTooManyRequests {
-		t.Errorf("code = %d, want 429", w.Code)
+		t.Errorf("expected 503, got %d", w.Code)
 	}
 }
diff --git a/server/main.py b/server/main.py
index 3e634f7..2928761 100644
--- a/server/main.py
+++ b/server/main.py
@@ -131,6 +131,19 @@ def main():
         async def health(request):
             return Response(status_code=200)
 
+        # Middleware to inject spend tracking headers
+        from tools.query import current_query_cost
+
+        @app.middleware("http")
+        async def inject_cost_header(request, call_next):
+            response = await call_next(request)
+            cost = current_query_cost.get()
+            if cost > 0:
+                response.headers["X-Limnos-Cost-USD"] = str(cost)
+                # Reset for next request in this worker
+                current_query_cost.set(0.0)
+            return response
+
         import uvicorn
 
         uvicorn.run(app, host="0.0.0.0", port=args.port)
diff --git a/server/tests/test_query_helpers.py b/server/tests/test_query_helpers.py
new file mode 100644
index 0000000..0b4cec2
--- /dev/null
+++ b/server/tests/test_query_helpers.py
@@ -0,0 +1,12 @@
+"""Tests for query helpers."""
+
+from __future__ import annotations
+
+from tools.query import current_query_cost
+
+
+def test_current_query_cost_contextvar():
+    current_query_cost.set(7.89)
+    assert current_query_cost.get() == 7.89
+    current_query_cost.set(0.0)
+    assert current_query_cost.get() == 0.0
diff --git a/server/tools/query.py b/server/tools/query.py
index 269c095..dafc3ae 100644
--- a/server/tools/query.py
+++ b/server/tools/query.py
@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import re
+from contextvars import ContextVar
 from typing import Optional
 
 from mcp.server.fastmcp import FastMCP, Context
@@ -11,6 +13,9 @@
 from tools.formatting import format_query_result
 from tools.sample_data import _nl_to_sql
 
+# ContextVar to communicate cost to the HTTP response headers in the middleware
+current_query_cost: ContextVar[float] = ContextVar("current_query_cost", default=0.0)
+
 
 class QueryInput(BaseModel):
     model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
@@ -80,6 +85,7 @@ async def datalake_query(params: QueryInput, ctx: Context) -> str:
 
         # ── Step 2: Cost estimation ─────────────────────────────────────────
         estimate = cost_estimator.estimate(params.table, sql)
+        current_query_cost.set(estimate.total_cost_usd)
 
         if params.explain_only:
             return _format_explain(sql, estimate)

From cd469461d80aa69a346ba333a8ea7b312ddc0e82 Mon Sep 17 00:00:00 2001
From: Gemini CLI <gemini-cli@google.com>
Date: Mon, 16 Mar 2026 21:46:27 +1100
Subject: [PATCH 3/5] feat: enhance Glue provisioning and Athena SQL
 translation

---
 server/catalog/glue.py                        |  38 +++---
 .../tests/test_cost_estimator_partitioning.py | 124 ++++++++++++++++++
 server/tests/test_glue_provisioner.py         |  14 ++
 server/tests/test_query_helpers.py            |  26 +++-
 server/tools/query.py                         |  23 +++-
 5 files changed, 208 insertions(+), 17 deletions(-)
 create mode 100644 server/tests/test_cost_estimator_partitioning.py

diff --git a/server/catalog/glue.py b/server/catalog/glue.py
index 85fb487..6cd0b9f 100644
--- a/server/catalog/glue.py
+++ b/server/catalog/glue.py
@@ -1,12 +1,12 @@
 """
-Glue auto-provisioner for flat file formats.
+Glue auto-provisioner for data lake tables.
 
 Creates or updates AWS Glue external tables so Athena can query flat files
-(CSV, JSON, NDJSON) without manual catalog setup. Called once during
-describe_table when metadata is first scanned for a flat file table.
+(CSV, JSON, NDJSON) and Parquet tables without manual catalog setup.
+Called once during describe_table when metadata is first scanned.
 
-TXT tables are excluded — Athena has no useful query capability over
-unstructured single-column text files.
+Iceberg tables are excluded (they use their own catalog).
+TXT tables are excluded (no Athena support).
 """
 
 from __future__ import annotations
@@ -32,29 +32,33 @@
     "TIMESTAMP": "timestamp",
 }
 
-# (InputFormat, SerDe library) per format
-_SERDE: dict[str, tuple[str, str]] = {
+# (InputFormat, OutputFormat, SerDe library) per format
+_SERDE: dict[str, tuple[str, str, str]] = {
     "csv": (
         "org.apache.hadoop.mapred.TextInputFormat",
+        "org.apache.hadoop.hive.ql.io.HiveIgnoreKeyTextOutputFormat",
         "org.apache.hadoop.hive.serde2.lazy.LazySimpleSerDe",
     ),
     "json": (
         "org.apache.hadoop.mapred.TextInputFormat",
+        "org.apache.hadoop.hive.ql.io.HiveIgnoreKeyTextOutputFormat",
         "org.openx.data.jsonserde.JsonSerDe",
     ),
     "ndjson": (
         "org.apache.hadoop.mapred.TextInputFormat",
+        "org.apache.hadoop.hive.ql.io.HiveIgnoreKeyTextOutputFormat",
         "org.openx.data.jsonserde.JsonSerDe",
     ),
-    "txt": (
-        "org.apache.hadoop.mapred.TextInputFormat",
-        "org.apache.hadoop.hive.serde2.lazy.LazySimpleSerDe",
+    "parquet": (
+        "org.apache.hadoop.hive.ql.io.parquet.MapredParquetInputFormat",
+        "org.apache.hadoop.hive.ql.io.parquet.MapredParquetOutputFormat",
+        "org.apache.hadoop.hive.ql.io.parquet.serde.ParquetHiveSerDe",
     ),
 }
 
 
 class GlueProvisioner:
-    """Create or update Glue external tables for flat file formats."""
+    """Create or update Glue external tables for data lake tables."""
 
     def __init__(self, config) -> None:
         self._glue = boto3.client("glue", region_name=config.aws.region)
@@ -67,13 +71,18 @@ def sync_table(
         partition_cols: list[PartitionMeta],
     ) -> None:
         """Create or update a Glue external table. Idempotent."""
-        input_fmt, serde_lib = _SERDE[table_cfg.format]
+        if table_cfg.format not in _SERDE:
+            return
+
+        input_fmt, output_fmt, serde_lib = _SERDE[table_cfg.format]
 
         serde_params: dict[str, str] = {}
         if table_cfg.format == "csv":
             serde_params["field.delim"] = table_cfg.delimiter
             if table_cfg.has_header:
                 serde_params["skip.header.line.count"] = "1"
+        elif table_cfg.format == "parquet":
+            serde_params["serialization.format"] = "1"
 
         glue_cols = [
             {
@@ -90,9 +99,7 @@ def sync_table(
                 "Columns": glue_cols,
                 "Location": table_cfg.s3_path,
                 "InputFormat": input_fmt,
-                "OutputFormat": (
-                    "org.apache.hadoop.hive.ql.io.HiveIgnoreKeyTextOutputFormat"
-                ),
+                "OutputFormat": output_fmt,
                 "SerdeInfo": {
                     "SerializationLibrary": serde_lib,
                     "Parameters": serde_params,
@@ -100,6 +107,7 @@ def sync_table(
             },
             "PartitionKeys": glue_partitions,
             "TableType": "EXTERNAL_TABLE",
+            "Parameters": {"classification": table_cfg.format},
         }
 
         try:
diff --git a/server/tests/test_cost_estimator_partitioning.py b/server/tests/test_cost_estimator_partitioning.py
new file mode 100644
index 0000000..6dc5d12
--- /dev/null
+++ b/server/tests/test_cost_estimator_partitioning.py
@@ -0,0 +1,124 @@
+"""Tests for engine.cost_estimator partition pruning logic."""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from unittest.mock import MagicMock
+
+import pytest
+
+from engine.cost_estimator import CostEstimator, _has_partition_filter
+from catalog.schema_cache import SchemaCache, TableMeta, ColumnMeta, PartitionMeta
+
+
+@pytest.fixture
+def mock_cache():
+    return MagicMock(spec=SchemaCache)
+
+
+@pytest.fixture
+def mock_config():
+    cfg = MagicMock()
+    cfg.engine.duckdb_max_scan_bytes = 10 * 1024**3
+    cfg.cost_gates.warn_threshold_usd = 0.10
+    cfg.cost_gates.block_threshold_usd = 1.00
+    return cfg
+
+
+def test_has_partition_filter_basic_equality():
+    partition_cols = ["dt", "region"]
+    sql = "SELECT * FROM orders WHERE dt = '2025-01-01'"
+    assert _has_partition_filter(sql, partition_cols) is True
+
+
+def test_has_partition_filter_case_insensitive():
+    partition_cols = ["dt"]
+    sql = "SELECT * FROM orders WHERE DT = '2025-01-01'"
+    assert _has_partition_filter(sql, partition_cols) is True
+
+
+def test_has_partition_filter_multiple_columns():
+    partition_cols = ["dt", "region"]
+    sql = "SELECT * FROM orders WHERE region = 'US' AND amount > 100"
+    assert _has_partition_filter(sql, partition_cols) is True
+
+
+def test_has_partition_filter_range():
+    partition_cols = ["dt"]
+    sql = "SELECT * FROM orders WHERE dt >= '2025-01-01' AND dt <= '2025-01-31'"
+    assert _has_partition_filter(sql, partition_cols) is True
+
+
+def test_has_partition_filter_no_where():
+    partition_cols = ["dt"]
+    sql = "SELECT * FROM orders"
+    assert _has_partition_filter(sql, partition_cols) is False
+
+
+def test_has_partition_filter_other_columns_only():
+    partition_cols = ["dt"]
+    sql = "SELECT * FROM orders WHERE customer_id = 123"
+    assert _has_partition_filter(sql, partition_cols) is False
+
+
+def test_has_partition_filter_in_clause():
+    partition_cols = ["region"]
+    sql = "SELECT * FROM orders WHERE region IN ('US', 'EU')"
+    assert _has_partition_filter(sql, partition_cols) is True
+
+
+def test_has_partition_filter_or_condition():
+    partition_cols = ["region"]
+    sql = "SELECT * FROM orders WHERE region = 'US' OR customer_id = 1"
+    assert _has_partition_filter(sql, partition_cols) is True
+
+
+def test_has_partition_filter_complex_expression():
+    partition_cols = ["dt"]
+    # Testing if sqlglot handles complex expressions in WHERE
+    sql = (
+        "SELECT * FROM orders WHERE (dt = '2025-01-01' OR region = 'US') AND amount > 0"
+    )
+    assert _has_partition_filter(sql, partition_cols) is True
+
+
+def test_has_partition_filter_fallback_string_search():
+    # Test the fallback logic when sqlglot fails (by providing invalid SQL)
+    partition_cols = ["dt"]
+    invalid_sql = "SELECT * FROM orders WHERE dt = '2025-01-01' SOME INVALID SYNTAX"
+    assert _has_partition_filter(invalid_sql, partition_cols) is True
+
+
+def test_estimate_applies_partition_fraction(mock_config, mock_cache):
+    estimator = CostEstimator(mock_config, mock_cache)
+
+    meta = TableMeta(
+        table_name="orders",
+        s3_path="s3://bucket/orders/",
+        format="parquet",
+        columns=[ColumnMeta(name="id", dtype="BIGINT", estimated_bytes_per_row=8)],
+        partition_columns=[PartitionMeta(name="dt", dtype="string")],
+        total_rows=1000000,
+        total_bytes=100 * 1024**2,  # 100 MB
+        total_files=10,
+        total_partitions=10,
+        avg_row_groups_per_file=4,
+        last_refreshed=datetime.now(tz=timezone.utc),
+    )
+    mock_cache.get.return_value = meta
+
+    # Query WITHOUT partition filter
+    sql_no_filter = "SELECT * FROM orders"
+    est_no_filter = estimator.estimate("orders", sql_no_filter)
+
+    # Query WITH partition filter
+    sql_filter = "SELECT * FROM orders WHERE dt = '2025-01-01'"
+    est_filter = estimator.estimate("orders", sql_filter)
+
+    # partition_fraction is 0.1 in current implementation
+    assert est_filter.estimated_bytes == int(est_no_filter.estimated_bytes * 0.1)
+    assert est_filter.estimated_files == max(
+        1, int(est_no_filter.estimated_files * 0.1)
+    )
+    assert est_filter.partition_filter_detected is True
+    assert est_no_filter.partition_filter_detected is False
diff --git a/server/tests/test_glue_provisioner.py b/server/tests/test_glue_provisioner.py
index 37e210d..6029db1 100644
--- a/server/tests/test_glue_provisioner.py
+++ b/server/tests/test_glue_provisioner.py
@@ -188,3 +188,17 @@ def test_database_name_passed_correctly(self):
         prov.sync_table(table_cfg, _columns(), [])
 
         assert glue.create_table.call_args[1]["DatabaseName"] == "my_lake"
+
+    def test_parquet_uses_parquet_serde(self):
+        glue = MagicMock()
+        prov = self._make_provisioner(glue)
+        table_cfg = _make_table_cfg(fmt="parquet", s3_path="s3://bucket/parquet/")
+
+        prov.sync_table(table_cfg, _columns(), _partition_cols())
+
+        glue.create_table.assert_called_once()
+        call_kwargs = glue.create_table.call_args[1]
+        sd = call_kwargs["TableInput"]["StorageDescriptor"]
+        assert "ParquetHiveSerDe" in sd["SerdeInfo"]["SerializationLibrary"]
+        assert "MapredParquetInputFormat" in sd["InputFormat"]
+        assert "MapredParquetOutputFormat" in sd["OutputFormat"]
diff --git a/server/tests/test_query_helpers.py b/server/tests/test_query_helpers.py
index 0b4cec2..bbe4c5c 100644
--- a/server/tests/test_query_helpers.py
+++ b/server/tests/test_query_helpers.py
@@ -2,7 +2,31 @@
 
 from __future__ import annotations
 
-from tools.query import current_query_cost
+from unittest.mock import MagicMock
+from tools.query import _translate_to_athena, current_query_cost
+
+
+def test_translate_to_athena():
+    aws_cfg = MagicMock()
+    aws_cfg.glue_database = "my_db"
+
+    table_cfg = MagicMock()
+    table_cfg.name = "my-table"
+
+    # Test Parquet translation
+    sql_pq = "SELECT * FROM read_parquet('s3://bucket/path/**/*.parquet', hive_partitioning=true) LIMIT 10"
+    translated = _translate_to_athena(sql_pq, table_cfg, aws_cfg)
+    assert translated == 'SELECT * FROM "my_db"."my_table" LIMIT 10'
+
+    # Test Iceberg translation
+    sql_ice = "SELECT count(*) FROM iceberg_scan('s3://bucket/iceberg/')"
+    translated = _translate_to_athena(sql_ice, table_cfg, aws_cfg)
+    assert translated == 'SELECT count(*) FROM "my_db"."my_table"'
+
+    # Test CSV translation
+    sql_csv = "SELECT * FROM read_csv('s3://bucket/data.csv', auto_detect=true)"
+    translated = _translate_to_athena(sql_csv, table_cfg, aws_cfg)
+    assert translated == 'SELECT * FROM "my_db"."my_table"'
 
 
 def test_current_query_cost_contextvar():
diff --git a/server/tools/query.py b/server/tools/query.py
index dafc3ae..d48efc3 100644
--- a/server/tools/query.py
+++ b/server/tools/query.py
@@ -107,7 +107,8 @@ async def datalake_query(params: QueryInput, ctx: Context) -> str:
             if estimate.recommended_engine == "duckdb":
                 result = duckdb_engine.query(sql, row_limit=params.row_limit)
             else:
-                result = athena_engine.query(sql)
+                athena_sql = _translate_to_athena(sql, table_cfg, config.aws)
+                result = athena_engine.query(athena_sql)
         except Exception as e:
             return f"❌ **Query failed**\n\n```\n{e}\n```"
 
@@ -137,3 +138,23 @@ def _fallback_sql(question: str, table_cfg) -> str:
 
 def _format_explain(sql: str, estimate) -> str:
     return f"## Query Plan\n\n**{estimate.summary_line()}**\n\n```sql\n{sql}\n```"
+
+
+def _translate_to_athena(sql: str, table_cfg, aws_cfg) -> str:
+    glue_table = f'"{aws_cfg.glue_database}"."{table_cfg.name.replace("-", "_")}"'
+    sql = re.sub(
+        r"read_parquet\(['\"]s3://[^'\"]+['\"].*?\)",
+        glue_table,
+        sql,
+        flags=re.IGNORECASE,
+    )
+    sql = re.sub(
+        r"iceberg_scan\(['\"]s3://[^'\"]+['\"]\)", glue_table, sql, flags=re.IGNORECASE
+    )
+    sql = re.sub(
+        r"read_csv\(['\"]s3://[^'\"]+['\"].*?\)", glue_table, sql, flags=re.IGNORECASE
+    )
+    sql = re.sub(
+        r"read_json\(['\"]s3://[^'\"]+['\"].*?\)", glue_table, sql, flags=re.IGNORECASE
+    )
+    return sql

From 7cedf5c46be685d4592181457f2642ee5da4ffb0 Mon Sep 17 00:00:00 2001
From: Gemini CLI <gemini-cli@google.com>
Date: Sat, 28 Mar 2026 21:46:37 +1100
Subject: [PATCH 4/5] fix: restore cost warning logic in datalake_query

---
 server/tools/query.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/server/tools/query.py b/server/tools/query.py
index d48efc3..02dee42 100644
--- a/server/tools/query.py
+++ b/server/tools/query.py
@@ -94,6 +94,10 @@ async def datalake_query(params: QueryInput, ctx: Context) -> str:
         if estimate.block and not params.force:
             return f"🚫 **Query blocked** — estimated cost exceeds threshold.\n\n{estimate.summary_line()}"
 
+        if estimate.warning and not params.force:
+            # Soft warning — still execute, but surface the warning
+            await ctx.log_info(f"Cost warning: {estimate.warning}")
+
         # ── Step 4: Result cache check ───────────────────────────────────────
         effective_row_limit = params.row_limit or config.engine.default_row_limit
         cache_key = make_cache_key(params.table, sql, effective_row_limit)
@@ -114,6 +118,10 @@ async def datalake_query(params: QueryInput, ctx: Context) -> str:
 
         # ── Step 6: Format & return ──────────────────────────────────────────
         response = format_query_result(result, estimate.summary_line())
+
+        if estimate.warning:
+            response = f"⚠️ {estimate.warning}\n\n{response}"
+
         if result_cache:
             result_cache.put(
                 cache_key,

From dbb3fb4f5833cf6fcabd999d8c0401cef4a82ad3 Mon Sep 17 00:00:00 2001
From: Gemini CLI <gemini-cli@google.com>
Date: Sat, 28 Mar 2026 21:51:27 +1100
Subject: [PATCH 5/5] fix(gateway): align proxy_test with worker_pool signature

---
 gateway/internal/mcp/proxy_test.go | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/gateway/internal/mcp/proxy_test.go b/gateway/internal/mcp/proxy_test.go
index 25fae0e..0c295ba 100644
--- a/gateway/internal/mcp/proxy_test.go
+++ b/gateway/internal/mcp/proxy_test.go
@@ -25,9 +25,10 @@ func dummyAuth() *auth.APIKeyAuth {
 }
 
 func emptyPool() *queue.WorkerPool {
-	return queue.NewWorkerPool(queue.WorkerPoolConfig{
-		MaxWorkers: 0,
-	}, silentLogger())
+	pool, _ := queue.NewWorkerPool(queue.WorkerPoolConfig{
+		Size: 0,
+	})
+	return pool
 }
 
 // ── No healthy workers ─────────────────────────────────────────────────────────