spice-mcp 0.1.0__py3-none-any.whl

spice_mcp/__init__.py

@@ -0,0 +1,2 @@
+def hello() -> str:
+    return "Hello from spice-mcp!"

spice_mcp/adapters/dune/__init__.py

@@ -0,0 +1,10 @@
+"""Dune adapter built from the vendored spice client.
+
+This module provides a thin façade used by the new service layer while
+keeping the battle-tested logic that the original spice client offered.
+"""
+
+from . import urls  # re-export for callers needing low-level helpers
+from .extract import async_query, query  # noqa: F401
+
+__all__ = ["query", "async_query", "urls"]

spice_mcp/adapters/dune/admin.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+from ..http_client import HttpClient, HttpClientConfig
+from . import urls
+
+
+class DuneAdminAdapter:
+    """Lightweight client for Dune saved query management."""
+
+    def __init__(
+        self,
+        api_key: str,
+        *,
+        http_client: HttpClient | None = None,
+        http_config: HttpClientConfig | None = None,
+    ):
+        self.api_key = api_key
+        config = http_config or HttpClientConfig()
+        self._http = http_client or HttpClient(config)
+
+    def _headers(self) -> Mapping[str, str]:
+        return {
+            "X-Dune-API-Key": self.api_key,
+            "User-Agent": "spice-mcp-admin/1",
+            "Content-Type": "application/json",
+        }
+
+    def get(self, query_id: int) -> dict[str, Any]:
+        url = urls.url_templates["query"].format(query_id=query_id)
+        resp = self._http.request("GET", url, headers=self._headers())
+        return resp.json()
+
+    def create(
+        self,
+        *,
+        name: str,
+        query_sql: str,
+        description: str | None = None,
+        tags: list[str] | None = None,
+        parameters: list[dict] | None = None,
+        is_private: bool = False,
+    ) -> dict[str, Any]:
+        body: dict[str, Any] = {"name": name, "query_sql": query_sql, "is_private": is_private}
+        if description:
+            body["description"] = description
+        if tags is not None:
+            body["tags"] = tags
+        if parameters is not None:
+            body["parameters"] = parameters
+        resp = self._http.request(
+            "POST",
+            urls.url_templates["query_create"],
+            headers=self._headers(),
+            json=body,
+            timeout=20,
+        )
+        return resp.json()
+
+    def update(self, query_id: int, *, name: str | None = None, query_sql: str | None = None, description: str | None = None, tags: list[str] | None = None, parameters: list[dict] | None = None) -> dict[str, Any]:
+        body: dict[str, Any] = {}
+        if name is not None:
+            body["name"] = name
+        if query_sql is not None:
+            body["query_sql"] = query_sql
+        if description is not None:
+            body["description"] = description
+        if tags is not None:
+            body["tags"] = tags
+        if parameters is not None:
+            body["parameters"] = parameters
+        url = urls.url_templates["query"].format(query_id=query_id)
+        resp = self._http.request(
+            "PATCH",
+            url,
+            headers=self._headers(),
+            json=body,
+            timeout=20,
+        )
+        return resp.json()
+
+    def fork(self, source_query_id: int, *, name: str | None = None) -> dict[str, Any]:
+        url = urls.url_templates.get("query_fork", f"https://api.dune.com/api/v1/query/{source_query_id}/fork").format(query_id=source_query_id)
+        body: dict[str, Any] = {"name": name} if name else {}
+        resp = self._http.request(
+            "POST",
+            url,
+            headers=self._headers(),
+            json=body,
+            timeout=20,
+        )
+        return resp.json()

spice_mcp/adapters/dune/cache.py

@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING
+
+from . import extract as _extract
+
+if TYPE_CHECKING:
+    import polars as pl
+
+    from .types import Execution
+
+
+cache_template = (
+    '{query_id}__{execution_id}__{parameter_hash}__{timestamp}.parquet'
+)
+
+
+def load_from_cache(
+    execute_kwargs: _extract.ExecuteKwargs,
+    result_kwargs: _extract.RetrievalKwargs,
+    output_kwargs: _extract.OutputKwargs,
+) -> tuple[
+    pl.DataFrame | tuple[pl.DataFrame, Execution] | None, Execution | None
+]:
+    # get latest execution id
+    execution = _extract.get_latest_execution(execute_kwargs)
+    if execution is None:
+        return None, None
+
+    # build cache path
+    cache_path = _build_cache_path(
+        execution=execution,
+        execute_kwargs=execute_kwargs,
+        result_kwargs=result_kwargs,
+        cache_dir=output_kwargs['cache_dir'],
+    )
+
+    # load cache result
+    if os.path.exists(cache_path):
+        import polars as pl
+
+        if result_kwargs['verbose']:
+            print('loading dune query result from cache')
+        df = pl.read_parquet(cache_path)
+        if output_kwargs['include_execution']:
+            return (df, execution), execution
+        else:
+            return df, execution
+    else:
+        return None, execution
+
+
+async def async_load_from_cache(
+    execute_kwargs: _extract.ExecuteKwargs,
+    result_kwargs: _extract.RetrievalKwargs,
+    output_kwargs: _extract.OutputKwargs,
+) -> tuple[
+    pl.DataFrame | tuple[pl.DataFrame, Execution] | None, Execution | None
+]:
+    # get latest execution
+    execution = await _extract.async_get_latest_execution(execute_kwargs)
+    if execution is None:
+        return None, None
+
+    # build cache path
+    cache_path = _build_cache_path(
+        execution=execution,
+        execute_kwargs=execute_kwargs,
+        result_kwargs=result_kwargs,
+        cache_dir=output_kwargs['cache_dir'],
+    )
+
+    # load cache result
+    if os.path.exists(cache_path):
+        import polars as pl
+
+        if result_kwargs['verbose']:
+            print('loading dune query result from cache')
+        df = await pl.scan_parquet(cache_path).collect_async()
+        if output_kwargs['include_execution']:
+            return (df, execution), execution
+        else:
+            return df, execution
+    else:
+        return None, execution
+
+
+def save_to_cache(
+    df: pl.DataFrame,
+    execution: Execution,
+    execute_kwargs: _extract.ExecuteKwargs,
+    result_kwargs: _extract.RetrievalKwargs,
+    cache_dir: str | None,
+) -> None:
+    import secrets
+    import shutil
+
+    if result_kwargs['verbose']:
+        print('saving result to cache')
+
+    # build cache path
+    cache_path = _build_cache_path(
+        execution=execution,
+        execute_kwargs=execute_kwargs,
+        result_kwargs=result_kwargs,
+        cache_dir=cache_dir,
+    )
+
+    # create dir
+    os.makedirs(os.path.dirname(cache_path), exist_ok=True)
+
+    # save to cache
+    tmp_path = (
+        cache_path + '_tmp_' + secrets.token_hex(8)
+    )  # add for if running in parallel
+    df.write_parquet(tmp_path)
+    shutil.move(tmp_path, cache_path)
+
+
+def _preserialize_types(
+    result_kwargs: _extract.RetrievalKwargs,
+    key: str,
+) -> list[str | list[str]] | None:
+    raw = result_kwargs[key]  # type: ignore
+    if raw is None:
+        types: list[str | list[str]] | None = None
+    else:
+        types = []
+        if isinstance(raw, list):
+            for type in raw:
+                types.append(str(type))
+        elif isinstance(raw, dict):
+            for name, type in raw.items():
+                types.append([name, str(type)])
+        else:
+            raise Exception('invalid format for ' + key)
+    return types
+
+
+def _build_cache_path(
+    execution: Execution,
+    execute_kwargs: _extract.ExecuteKwargs,
+    result_kwargs: _extract.RetrievalKwargs,
+    cache_dir: str | None,
+) -> str:
+    import hashlib
+    import json
+
+    # get parameter hash
+    types = _preserialize_types(result_kwargs, 'types')
+    all_types = _preserialize_types(result_kwargs, 'all_types')
+    hash_params = {
+        'spice_version': 'spice-mcp',
+        'execution_id': execution['execution_id'],
+        'query_id': execute_kwargs['query_id'],
+        'parameters': execute_kwargs['parameters'],
+        'limit': result_kwargs['limit'],
+        'offset': result_kwargs['offset'],
+        'sample_count': result_kwargs['sample_count'],
+        'sort_by': result_kwargs['sort_by'],
+        'columns': result_kwargs['columns'],
+        'extras': result_kwargs['extras'],
+        'types': types,
+        'all_types': all_types,
+    }
+    md5_hash = hashlib.md5()
+    md5_hash.update(json.dumps(hash_params, sort_keys=True).encode('utf-8'))
+    hash_str = md5_hash.hexdigest()[:16]
+
+    # build from template
+    timestamp = execution['timestamp']
+    if timestamp is None:
+        raise Exception('need completion timestamp on execution')
+    cache_filename = cache_template.format(
+        query_id=execute_kwargs['query_id'],
+        execution_id=execution['execution_id'],
+        parameter_hash=hash_str,
+        timestamp=int(timestamp),
+    )
+
+    if cache_dir is None:
+        cache_dir = '/tmp/dune_spice'
+
+    return os.path.join(cache_dir, cache_filename)

spice_mcp/adapters/dune/client.py

@@ -0,0 +1,255 @@
+from __future__ import annotations
+
+import os
+import re
+import time
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+import polars as pl
+
+from ...config import Config
+from ...core.models import (
+    QueryRequest,
+    QueryResult,
+    ResultMetadata,
+    ResultPreview,
+    SchemaMatch,
+    TableColumn,
+    TableDescription,
+    TableSummary,
+)
+from ...core.ports import CatalogExplorer, QueryExecutor
+from ...polars_utils import collect_preview
+from ..http_client import HttpClient, HttpClientConfig
+from . import extract, urls
+
+
+class DuneAdapter(QueryExecutor, CatalogExplorer):
+    """Thin façade around the vendored extract module."""
+
+    def __init__(
+        self,
+        config: Config,
+        *,
+        http_client: HttpClient | None = None,
+    ):
+        self.config = config
+        http_config: HttpClientConfig = config.http
+        self._http = http_client or HttpClient(http_config)
+
+    # QueryExecutor -----------------------------------------------------------------
+    def execute(self, request: QueryRequest) -> QueryResult:
+        self._ensure_api_key()
+        start = time.time()
+        q = request.query
+        if isinstance(q, str):
+            q_rewritten = _maybe_rewrite_show_sql(q)
+            if q_rewritten is not None:
+                q = q_rewritten
+        result = extract.query(
+            query_or_execution=q,
+            verbose=False,
+            refresh=request.refresh,
+            max_age=request.max_age,
+            parameters=request.parameters,
+            api_key=self._api_key(),
+            performance=request.performance or "medium",
+            poll=request.poll,
+            timeout_seconds=request.timeout_seconds,
+            limit=request.limit,
+            offset=request.offset,
+            sample_count=request.sample_count,
+            sort_by=request.sort_by,
+            columns=request.columns,
+            cache_dir=self.config.cache.cache_dir,
+            include_execution=request.include_execution,
+            http_client=self._http,
+        )
+        if request.include_execution:
+            df, execution = result
+        else:
+            df = result
+            execution = {}
+        duration_ms = int((time.time() - start) * 1000)
+        columns = list(df.columns)
+        rowcount = int(len(df))
+        lazyframe = df.lazy()
+        preview = _build_preview(lazyframe, columns, rowcount)
+        del df
+        # Use rewritten SQL for metadata determination too
+        req_for_meta = request
+        try:
+            from dataclasses import replace
+
+            if isinstance(q, str) and q != request.query:
+                req_for_meta = replace(request, query=q)  # type: ignore[arg-type]
+        except Exception:
+            pass
+        meta = self.fetch_metadata(req_for_meta, execution=execution)
+        meta.duration_ms = duration_ms
+        return QueryResult(preview=preview, info=meta, lazyframe=lazyframe)
+
+    def fetch_metadata(
+        self, request: QueryRequest, *, execution: Mapping[str, Any] | None = None
+    ) -> ResultMetadata:
+        self._ensure_api_key()
+        query_id, _, effective_params = extract._determine_input_type(  # type: ignore[attr-defined]
+            request.query,
+            request.parameters,
+        )
+
+        payload: dict[str, Any] = {}
+        next_uri: str | None = None
+        next_offset: int | None = None
+
+        if query_id is not None:
+            params: dict[str, Any] = {}
+            if effective_params is not None:
+                params["query_parameters"] = effective_params
+            params.update(
+                {
+                    "limit": request.limit,
+                    "offset": request.offset,
+                    "sample_count": request.sample_count,
+                    "sort_by": request.sort_by,
+                    "columns": list(request.columns) if request.columns else None,
+                }
+            )
+            if request.extras:
+                try:
+                    params.update(request.extras)
+                except Exception:
+                    pass
+
+            url = urls.get_query_results_url(query_id, parameters=params, csv=False)
+            headers = {
+                "X-Dune-API-Key": self._api_key(),
+                "User-Agent": extract.get_user_agent(),
+            }
+            try:
+                resp = self._http.request("GET", url, headers=headers)
+                data = resp.json()
+                if isinstance(data, dict):
+                    payload = data.get("result", {}).get("metadata") or {}
+                    next_uri = data.get("next_uri")
+                    next_offset = data.get("next_offset")
+                    if "error" in data:
+                        payload = {**payload, "error": data["error"]}
+                    if "state" in data:
+                        payload = {**payload, "state": data["state"]}
+            except Exception:
+                payload = {}
+
+        execution_meta: dict[str, Any] = {}
+        if execution:
+            execution_meta = dict(execution)
+
+        return ResultMetadata(
+            execution=execution_meta,
+            duration_ms=0,
+            metadata=payload or None,
+            next_offset=next_offset,
+            next_uri=next_uri,
+        )
+
+    # CatalogExplorer ---------------------------------------------------------------
+    def find_schemas(self, keyword: str) -> Sequence[SchemaMatch]:
+        sql = f"SHOW SCHEMAS LIKE '%{keyword}%'"
+        df = self._run_sql(sql)
+        return [SchemaMatch(schema=row.get("Schema", "")) for row in df.to_dicts()]
+
+    def list_tables(self, schema: str, limit: int | None = None) -> Sequence[TableSummary]:
+        sql = f"SHOW TABLES FROM {schema}"
+        df = self._run_sql(sql, limit=limit)
+        return [
+            TableSummary(schema=schema, table=row.get("Table", row.get("name", "")))
+            for row in df.to_dicts()
+        ]
+
+    def describe_table(self, schema: str, table: str) -> TableDescription:
+        fq = f"{schema}.{table}"
+
+        try:
+            df = self._run_sql(f"SHOW COLUMNS FROM {fq}")
+            rows = df.to_dicts()
+            columns = [
+                TableColumn(
+                    name=row.get("Column") or row.get("column_name") or "",
+                    dune_type=row.get("Type") or row.get("data_type"),
+                    polars_dtype=row.get("Type") or None,
+                    comment=row.get("Comment"),
+                    extra=row.get("Extra"),
+                )
+                for row in rows
+            ]
+            return TableDescription(fully_qualified_name=fq, columns=columns)
+        except Exception:
+            df = self._run_sql(f"SELECT * FROM {fq} LIMIT 1")
+            columns = [
+                TableColumn(name=name, polars_dtype=str(dtype))
+                for name, dtype in zip(df.columns, df.dtypes)
+            ]
+            return TableDescription(fully_qualified_name=fq, columns=columns)
+
+    # Internal helpers --------------------------------------------------------------
+    def _run_sql(self, sql: str, *, limit: int | None = None) -> pl.DataFrame:
+        self._ensure_api_key()
+        sql_eff = _maybe_rewrite_show_sql(sql) or sql
+        df = extract.query(
+            sql_eff,
+            verbose=False,
+            performance="medium",
+            timeout_seconds=self.config.default_timeout_seconds,
+            limit=limit,
+            cache_dir=self.config.cache.cache_dir,
+            http_client=self._http,
+        )
+        if limit is not None and len(df) > limit:
+            return df.head(limit)
+        return df
+
+    def _ensure_api_key(self) -> None:
+        if not os.getenv("DUNE_API_KEY"):
+            os.environ["DUNE_API_KEY"] = self._api_key()
+
+    def _api_key(self) -> str:
+        return self.config.dune.api_key
+
+
+def _build_preview(lf: pl.LazyFrame, columns: list[str], rowcount: int) -> ResultPreview:
+    preview_rows = min(10, rowcount)
+    data_preview = collect_preview(lf, preview_rows)
+    return ResultPreview(
+        rowcount=rowcount,
+        columns=columns,
+        data_preview=data_preview,
+    )
+
+
+def _maybe_rewrite_show_sql(sql: str) -> str | None:
+    """Rewrite certain SHOW statements to information_schema SELECTs for portability.
+
+    This allows running discovery-style commands through the parameterized raw SQL
+    template which expects SELECT statements.
+    """
+    s = sql.strip()
+    m = re.match(r"^SHOW\s+SCHEMAS\s+LIKE\s+'([^']+)'\s*;?$", s, flags=re.IGNORECASE)
+    if m:
+        pat = m.group(1)
+        return (
+            "SELECT schema_name AS Schema FROM information_schema.schemata "
+            f"WHERE schema_name LIKE '{pat}'"
+        )
+    if re.match(r"^SHOW\s+SCHEMAS\s*;?$", s, flags=re.IGNORECASE):
+        return "SELECT schema_name AS Schema FROM information_schema.schemata"
+
+    m = re.match(r"^SHOW\s+TABLES\s+FROM\s+([A-Za-z0-9_\.]+)\s*;?$", s, flags=re.IGNORECASE)
+    if m:
+        schema = m.group(1)
+        return (
+            "SELECT table_name AS Table FROM information_schema.tables "
+            f"WHERE table_schema = '{schema}'"
+        )
+
+    return None