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/gateway/internal/mcp/proxy_test.go b/gateway/internal/mcp/proxy_test.go index 7bf39a9..247ba92 100644 --- a/gateway/internal/mcp/proxy_test.go +++ b/gateway/internal/mcp/proxy_test.go @@ -14,17 +14,22 @@ 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 { + pool, _ := queue.NewWorkerPool(queue.WorkerPoolConfig{ + Size: 0, + }) + return pool +} + // ── No healthy workers ───────────────────────────────────────────────────────── func TestProxy_NoWorkers_Returns503(t *testing.T) { 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..837d7b4 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,13 +86,7 @@ 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`." - ) + 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 @@ -152,44 +94,20 @@ async def datalake_query(params: QueryInput, ctx: Context) -> str: # ── 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()) @@ -197,59 +115,27 @@ async def datalake_query(params: QueryInput, ctx: Context) -> str: 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: