From f3305d8bbe3a7f7d6a84b0822a685b5951348647 Mon Sep 17 00:00:00 2001 From: Chris Cason Date: Wed, 15 Jul 2026 07:54:51 -0600 Subject: [PATCH] Add Redash HTTP API as a DB_MODE backend. Introduce RedashRunner for ad-hoc SQL via the Redash REST API, wire it through settings/config/doctor, and improve web startup UX for slow remote schema introspection during auto-load. --- AGENTS.md | 2 +- docs/reference/configuration.md | 17 +- frontend/src/App.svelte | 58 +++- frontend/src/lib/api/projects.ts | 16 + .../src/lib/components/PlotlyChart.svelte | 52 ++- src/datasight/agent.py | 5 + src/datasight/cli_commands/doctor.py | 39 ++- src/datasight/config.py | 34 +- src/datasight/explore.py | 4 +- src/datasight/export.py | 8 +- src/datasight/prompts.py | 25 +- src/datasight/redash_runner.py | 295 ++++++++++++++++++ src/datasight/runner.py | 4 +- src/datasight/settings.py | 42 ++- src/datasight/sql_validation.py | 4 + src/datasight/templates/env.template | 13 +- .../html/export_session_python.mustache | 12 + src/datasight/web/app.py | 45 ++- tests/test_config_extra.py | 40 +++ tests/test_export.py | 9 + tests/test_redash_runner.py | 213 +++++++++++++ tests/test_settings.py | 31 +- 22 files changed, 923 insertions(+), 45 deletions(-) create mode 100644 src/datasight/redash_runner.py create mode 100644 tests/test_redash_runner.py diff --git a/AGENTS.md b/AGENTS.md index 49d28bc2..455c6f16 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -10,7 +10,7 @@ datasight is an AI-powered data exploration tool. Users ask questions in natural - **FastAPI + uvicorn** for the web server with SSE streaming - **Svelte 5 + TypeScript + Tailwind CSS** frontend — built with Vite, served by FastAPI - **LLM backends**: Anthropic (default), GitHub Models, Ollama — all via a common `LLMClient` abstraction in `datasight.llm` -- **Database backends**: DuckDB (default), SQLite, PostgreSQL, Flight SQL — all via `SqlRunner` implementations in `datasight.runner` +- **Database backends**: DuckDB (default), SQLite, PostgreSQL, Flight SQL, Spark Connect, Redash HTTP API — via `SqlRunner` / `RedashRunner` in `datasight.runner` / `datasight.redash_runner` - **Click CLI** with commands: `run`, `ask`, `init`, `demo`, `generate`, `verify`, `profile`, `quality`, `doctor`, `export`, `log` ## Repository layout diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index 737bbcf9..36223089 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -81,7 +81,7 @@ For help picking a provider, see [Choosing an LLM](../use/concepts/choosing-an-l | Variable | Default | Description | |----------|---------|-------------| -| `DB_MODE` | `duckdb` | Database type: `duckdb`, `sqlite`, `postgres`, `flightsql`, or `spark` | +| `DB_MODE` | `duckdb` | Database type: `duckdb`, `sqlite`, `postgres`, `flightsql`, `spark`, or `redash` | | `DB_PATH` | `./database.duckdb` | Path to DuckDB or SQLite file (used when `DB_MODE=duckdb` or `sqlite`) | #### PostgreSQL settings (when `DB_MODE=postgres`) @@ -118,6 +118,21 @@ Requires the `spark` extra: `pip install 'datasight[spark]'`. | `SPARK_TOKEN` | — | Optional bearer token for Spark Connect auth | | `SPARK_MAX_RESULT_BYTES` | `104857600` | Client-side cap on Arrow result size (100 MiB default). Results above this are truncated to protect the web server on multi-TB backends. | +#### Redash API settings (when `DB_MODE=redash`) + +Execute SQL through [Redash](https://redash.io/)’s HTTP API instead of connecting to a warehouse directly. Use a **user API key** with permission to run queries on the chosen data source. + +| Variable | Default | Description | +|----------|---------|-------------| +| `REDASH_BASE_URL` | *(required)* | Redash server origin, e.g. `https://redash.example.com` (no trailing slash) | +| `REDASH_API_KEY` | *(required)* | User API key (`Authorization: Key …`) | +| `REDASH_DATA_SOURCE_ID` | *(required)* | Numeric Redash data source ID SQL should run against | +| `REDASH_SQL_DIALECT` | `postgres` | Warehouse dialect for prompts and sqlglot validation (`postgres`, `mysql`, `spark`, `bigquery`, `snowflake`, etc.). Must match the engine behind the Redash data source. | +| `REDASH_QUERY_TIMEOUT` | `120` | Seconds to wait for Redash job completion (same scale as other backends’ query timeouts). | +| `REDASH_POLL_INTERVAL` | `0.5` | Seconds between `GET /api/jobs/…` polls while a query is running. | + +Schema introspection issues one Redash job per probe query; on large warehouses prefer narrowing with `schema.yaml`. Exported Python sessions cannot embed API secrets — replays need a direct warehouse connection or your own Redash client code. + ### Other settings | Variable | Default | Description | diff --git a/frontend/src/App.svelte b/frontend/src/App.svelte index dd3b3fb6..b008d5df 100644 --- a/frontend/src/App.svelte +++ b/frontend/src/App.svelte @@ -27,7 +27,11 @@ import { paletteStore } from "$lib/stores/palette.svelte"; import { tidyStore } from "$lib/stores/tidy.svelte"; import { groundingStore } from "$lib/stores/grounding.svelte"; - import { exitExploreSession, getProjectStatus } from "$lib/api/projects"; + import { + exitExploreSession, + getProjectStatus, + waitForProjectAutoLoad, + } from "$lib/api/projects"; import { loadSettings, loadLlmConfig } from "$lib/api/settings"; import { loadSchema, loadQueries, loadRecipes } from "$lib/api/schema"; import { @@ -70,6 +74,10 @@ let exportMode = $state(false); let exportExcludeIndices = $state(new Set()); let booting = $state(true); + /** After minimal bootstrap; server-side `--project-dir` load still running (`loading: true`). */ + let awaitingAutoLoad = $state(false); + /** Shown under “Loading…” during slow startup (e.g. Redash schema introspection). */ + let bootHint = $state(""); let fromLanding = $state(false); function toggleTheme() { @@ -267,26 +275,40 @@ loadSettings(), loadLlmConfig(), ]); + booting = false; - if (status.loaded) { + let resolved = status; + if (status.loading) { + awaitingAutoLoad = true; + bootHint = + "Discovering database schema — remote backends can take several minutes."; + resolved = await waitForProjectAutoLoad(status); + bootHint = ""; + awaitingAutoLoad = false; + } + + if (resolved.loaded) { sessionStore.projectLoaded = true; - sessionStore.currentProjectPath = status.path; - sessionStore.isEphemeralSession = status.is_ephemeral; - sessionStore.hasTimeSeries = Boolean(status.has_time_series); - if (status.sql_dialect) { - sessionStore.sqlDialect = status.sql_dialect; + sessionStore.currentProjectPath = resolved.path; + sessionStore.isEphemeralSession = resolved.is_ephemeral; + sessionStore.hasTimeSeries = Boolean(resolved.has_time_series); + if (resolved.sql_dialect) { + sessionStore.sqlDialect = resolved.sql_dialect; } - if (status.tables) { - sessionStore.ephemeralTablesInfo = status.tables; + if (resolved.tables) { + sessionStore.ephemeralTablesInfo = resolved.tables; } - await onProjectReady(status.path ?? undefined); + await onProjectReady(resolved.path ?? undefined); // Restore previous conversation on reload await maybeRestoreSession(); } } catch { // Server not running — show landing page + awaitingAutoLoad = false; } finally { + bootHint = ""; booting = false; + awaitingAutoLoad = false; } }); @@ -546,8 +568,22 @@
{#if booting} -
+
Loading...
+ {#if bootHint} +
{bootHint}
+ {/if} +
+ {:else if awaitingAutoLoad} +
+
Connecting to project…
+ {#if bootHint} +
{bootHint}
+ {/if}
{:else if sessionStore.projectLoaded} (measureEditorOpen = true)} /> diff --git a/frontend/src/lib/api/projects.ts b/frontend/src/lib/api/projects.ts index 1ca3b6c2..e8ee11c0 100644 --- a/frontend/src/lib/api/projects.ts +++ b/frontend/src/lib/api/projects.ts @@ -8,6 +8,8 @@ import type { TableInfo } from "$lib/stores/schema.svelte"; export interface ProjectStatus { loaded: boolean; + /** True while `--project-dir` auto-load is still introspecting in the background. */ + loading?: boolean; path: string | null; name: string | null; is_ephemeral: boolean; @@ -40,6 +42,20 @@ export async function getProjectStatus(): Promise { return fetchJson("/api/project"); } +/** Poll ``/api/project`` until auto-load introspection finishes or ``maxWaitMs`` elapses. */ +export async function waitForProjectAutoLoad( + initial: ProjectStatus, + maxWaitMs: number = 3_600_000, +): Promise { + let s = initial; + const deadline = Date.now() + maxWaitMs; + while (s.loading && Date.now() < deadline) { + await new Promise((r) => setTimeout(r, 500)); + s = await getProjectStatus(); + } + return s; +} + export async function loadRecentProjects(): Promise { const data = await fetchJson<{ projects: RecentProject[] }>( "/api/projects/recent", diff --git a/frontend/src/lib/components/PlotlyChart.svelte b/frontend/src/lib/components/PlotlyChart.svelte index 7f512d87..f0c4796e 100644 --- a/frontend/src/lib/components/PlotlyChart.svelte +++ b/frontend/src/lib/components/PlotlyChart.svelte @@ -76,6 +76,32 @@ }; } + const STALE_CHUNK_RELOAD_KEY = "datasight-stale-chunk-reload-once"; + + function isStaleDynamicImportError(error: unknown): boolean { + const msg = error instanceof Error ? error.message : String(error); + return ( + msg.includes("Failed to fetch dynamically imported module") || + msg.includes("error loading dynamically imported module") + ); + } + + /** After a frontend rebuild, Vite chunk hashes change; a long-lived tab may 404 the old plotly chunk. */ + async function importPlotly() { + try { + const mod = await import("plotly.js-dist-min"); + sessionStorage.removeItem(STALE_CHUNK_RELOAD_KEY); + return mod.default; + } catch (error) { + if (isStaleDynamicImportError(error) && !sessionStorage.getItem(STALE_CHUNK_RELOAD_KEY)) { + sessionStorage.setItem(STALE_CHUNK_RELOAD_KEY, "1"); + location.reload(); + return new Promise(() => {}); + } + throw error; + } + } + function animationFrame(): Promise { return new Promise((resolve) => requestAnimationFrame(() => resolve())); } @@ -136,7 +162,7 @@ if (cancelled) return; isRendering = true; try { - const Plotly = (await import("plotly.js-dist-min")).default; + const Plotly = await importPlotly(); if (cancelled) return; const cloned = cloneSpec(renderSpec); const layout = themedLayout(cloned.layout || {}, currentTheme()); @@ -166,9 +192,11 @@ render(); resizeObserver = new ResizeObserver(() => { if (!hasRendered) return; - void import("plotly.js-dist-min").then(({ default: Plotly }) => { - if (!cancelled) Plotly.Plots?.resize(plotEl); - }); + void importPlotly() + .then((Plotly) => { + if (!cancelled) Plotly.Plots?.resize(plotEl); + }) + .catch(() => {}); }); resizeObserver.observe(chart); themeObserver = new MutationObserver(() => void render()); @@ -182,13 +210,15 @@ resizeObserver?.disconnect(); themeObserver?.disconnect(); plotEl.removeAllListeners?.("plotly_click"); - void import("plotly.js-dist-min").then(({ default: Plotly }) => { - try { - Plotly.purge(plotEl); - } catch { - // ignore cleanup failures - } - }); + void importPlotly() + .then((Plotly) => { + try { + Plotly.purge(plotEl); + } catch { + // ignore cleanup failures + } + }) + .catch(() => {}); }; }); diff --git a/src/datasight/agent.py b/src/datasight/agent.py index 76bf5702..fa53d2dd 100644 --- a/src/datasight/agent.py +++ b/src/datasight/agent.py @@ -277,6 +277,11 @@ def sql_error_hint(error_msg: str, dialect: str = "duckdb") -> str: hints.append( "HINT: PostgreSQL syntax — use DATE_TRUNC, EXTRACT, TO_CHAR, col::type." ) + case "mysql": + hints.append( + "HINT: MySQL syntax — use DATE_FORMAT, YEAR/MONTH/DAY, STR_TO_DATE; " + "backticks for reserved names." + ) case "sqlite": hints.append( "HINT: SQLite syntax — use strftime(), date(), datetime(). " diff --git a/src/datasight/cli_commands/doctor.py b/src/datasight/cli_commands/doctor.py index b14d9738..cfe7d9f7 100644 --- a/src/datasight/cli_commands/doctor.py +++ b/src/datasight/cli_commands/doctor.py @@ -14,6 +14,33 @@ from datasight.cli_helpers import format_epilog +async def _async_probe_select_one_and_close(sql_runner: object) -> None: + """Run ``SELECT 1`` then release async HTTP clients on the same event loop. + + ``RedashRunner`` uses ``httpx.AsyncClient``. If we ``asyncio.run(run_sql)`` + and then call sync ``close()`` on the runner, ``RedashRunner.close()`` may + call ``asyncio.run(aclose())`` on a client still bound to the *previous* + (now closed) loop, surfacing ``RuntimeError: Event loop is closed``. + Performing ``aclose()`` here avoids that. + """ + try: + run_sql = getattr(sql_runner, "run_sql") + await run_sql("SELECT 1 AS ok") + finally: + inner: object | None = sql_runner + seen: set[int] = set() + while inner is not None and id(inner) not in seen: + seen.add(id(inner)) + aclose_fn = getattr(inner, "aclose", None) + if callable(aclose_fn): + await aclose_fn() + break + inner = getattr(inner, "_inner", None) + close_fn = getattr(sql_runner, "close", None) + if callable(close_fn): + close_fn() + + @click.command( epilog=format_epilog( """ @@ -98,6 +125,16 @@ def add_check(name: str, ok: bool, detail: str) -> None: elif settings.database.mode == "spark": db_ok = bool(settings.database.spark_remote) db_detail = settings.database.spark_remote + elif settings.database.mode == "redash": + db_ok = bool( + settings.database.redash_base_url.strip() + and settings.database.redash_api_key.strip() + and settings.database.redash_data_source_id > 0 + ) + db_detail = ( + f"{settings.database.redash_base_url.rstrip('/')}" + f" (data_source_id={settings.database.redash_data_source_id})" + ) add_check("Database config", db_ok, db_detail or settings.database.mode) for name in ("schema_description.md", "queries.yaml"): @@ -116,7 +153,7 @@ def add_check(name: str, ok: bool, detail: str) -> None: try: sql_runner = create_sql_runner_from_settings(settings.database, str(project_path)) - asyncio.run(sql_runner.run_sql("SELECT 1 AS ok")) + asyncio.run(_async_probe_select_one_and_close(sql_runner)) add_check("Database connectivity", True, "SELECT 1") except Exception as exc: add_check("Database connectivity", False, str(exc)) diff --git a/src/datasight/config.py b/src/datasight/config.py index a636b991..8806fcfc 100644 --- a/src/datasight/config.py +++ b/src/datasight/config.py @@ -15,7 +15,9 @@ from loguru import logger from datasight.exceptions import ConfigurationError +from datasight.redash_runner import RedashRunner from datasight.runner import ( + DEFAULT_QUERY_TIMEOUT, DEFAULT_SPARK_MAX_RESULT_BYTES, CachingSqlRunner, DuckDBRunner, @@ -88,13 +90,18 @@ def create_sql_runner( spark_remote: str = "sc://localhost:15002", spark_token: str | None = None, spark_max_result_bytes: int = DEFAULT_SPARK_MAX_RESULT_BYTES, + redash_base_url: str = "", + redash_api_key: str = "", + redash_data_source_id: int = 0, + redash_query_timeout: float = DEFAULT_QUERY_TIMEOUT, + redash_poll_interval: float = 0.5, ) -> SqlRunner: """Create the appropriate SqlRunner based on db_mode. Parameters ---------- db_mode: - Database mode: "duckdb", "sqlite", "postgres", "flightsql", or "spark". + Database mode: "duckdb", "sqlite", "postgres", "flightsql", "spark", or "redash". db_path: Path to database file (for duckdb and sqlite modes). flight_*: @@ -104,6 +111,8 @@ def create_sql_runner( spark_*: Spark Connect parameters. ``spark_max_result_bytes`` caps how much Arrow data is materialized client-side before truncation. + redash_*: + Redash REST API parameters (``DB_MODE=redash``). Returns ------- @@ -155,6 +164,24 @@ def create_sql_runner( token=spark_token, max_result_bytes=spark_max_result_bytes, ) + case "redash": + base = (redash_base_url or "").strip() + if not base: + raise ConfigurationError("REDASH_BASE_URL is required for Redash mode") + if not (redash_api_key or "").strip(): + raise ConfigurationError("REDASH_API_KEY is required for Redash mode") + if redash_data_source_id <= 0: + raise ConfigurationError( + "REDASH_DATA_SOURCE_ID must be a positive integer for Redash mode" + ) + logger.info(f"Connecting to Redash: {base} (data_source_id={redash_data_source_id})") + return RedashRunner( + base_url=base, + api_key=redash_api_key.strip(), + data_source_id=redash_data_source_id, + query_timeout=redash_query_timeout, + poll_interval=redash_poll_interval, + ) case _: msg = f"Invalid database mode: {db_mode}" raise ConfigurationError(msg) @@ -203,6 +230,11 @@ def create_sql_runner_from_settings( spark_remote=settings.spark_remote, spark_token=settings.spark_token, spark_max_result_bytes=settings.spark_max_result_bytes, + redash_base_url=settings.redash_base_url, + redash_api_key=settings.redash_api_key, + redash_data_source_id=settings.redash_data_source_id, + redash_query_timeout=settings.redash_query_timeout, + redash_poll_interval=settings.redash_poll_interval, ) if sql_cache_max_bytes > 0: runner = CachingSqlRunner(runner, max_bytes=sql_cache_max_bytes) diff --git a/src/datasight/explore.py b/src/datasight/explore.py index 17e51318..5d57b629 100644 --- a/src/datasight/explore.py +++ b/src/datasight/explore.py @@ -831,8 +831,8 @@ def create_files_session_for_settings( files already live on a shared filesystem. - ``duckdb`` / ``sqlite`` / no settings: local ephemeral DuckDB session (the long-standing default). - - ``postgres`` / ``flightsql``: fall back to DuckDB with a warning — - those backends can't read arbitrary local files. + - ``postgres`` / ``flightsql`` / ``redash``: fall back to DuckDB with a + warning — those backends can't read arbitrary local files. """ if import_mode not in _VALID_IMPORT_MODES: msg = f"Unsupported import mode: {import_mode}" diff --git a/src/datasight/export.py b/src/datasight/export.py index bc08d390..38f6cddf 100644 --- a/src/datasight/export.py +++ b/src/datasight/export.py @@ -538,8 +538,9 @@ def export_session_python( Default value baked into the script as ``DEFAULT_DB_PATH``. The user can override at runtime with ``--db``. db_mode: - ``"duckdb"`` (default), ``"sqlite"``, or another value (renders a - scaffold whose ``run_sql`` raises NotImplementedError). + ``"duckdb"`` (default), ``"sqlite"``, ``"redash"``, or another value + (non-file backends render a scaffold whose ``run_sql`` raises + ``NotImplementedError`` unless documented otherwise). exclude_indices: Optional set of 0-based turn indices to skip — same semantics as ``export_session_html``. @@ -559,7 +560,8 @@ def export_session_python( "db_mode": db_mode, "is_duckdb": db_mode == "duckdb", "is_sqlite": db_mode == "sqlite", - "is_unknown_mode": db_mode not in ("duckdb", "sqlite"), + "is_redash": db_mode == "redash", + "is_unknown_mode": db_mode not in ("duckdb", "sqlite", "redash"), "turns": [c for c in contexts if c is not None], }, ) diff --git a/src/datasight/prompts.py b/src/datasight/prompts.py index 147551ff..b09d240b 100644 --- a/src/datasight/prompts.py +++ b/src/datasight/prompts.py @@ -99,6 +99,13 @@ "Postgres regression: regr_slope(y, x), regr_intercept(y, x), " "regr_r2(y, x), corr(y, x).\n" ), + "mysql": ( + "Use MySQL SQL syntax.\n" + "MySQL dates: DATE_FORMAT(), YEAR(), MONTH(), STR_TO_DATE(), CURDATE(). " + "Use backticks for identifiers when needed. LIKE (not ILIKE) for patterns.\n" + "Aggregates and GROUP BY behave as in standard MySQL — " + "no FILTER clause; use conditional SUM/CASE patterns.\n" + ), "sqlite": ( "Use SQLite SQL syntax.\n" "SQLite dates: strftime(), date(), datetime(). " @@ -128,11 +135,27 @@ "will see a 'truncated' warning — the cure is more aggregation, not " "a larger LIMIT.\n" ), + "bigquery": ( + "Use BigQuery Standard SQL syntax.\n" + "Dates: DATE_TRUNC, EXTRACT, FORMAT_TIMESTAMP, TIMESTAMP/DATE literals.\n" + "Prefer SAFE_ prefixed casts/parses when handling untyped strings.\n" + ), + "snowflake": ( + "Use Snowflake SQL syntax.\n" + "Dates: DATE_TRUNC, TO_TIMESTAMP, TO_DATE; casts with ``::TYPE`` optional.\n" + ), } def dialect_hint(dialect: str) -> str: - return _DIALECT_HINTS.get(dialect, _DIALECT_HINTS["duckdb"]) + hinted = _DIALECT_HINTS.get(dialect) + if hinted: + return hinted + return ( + f"Use `{dialect}` SQL syntax appropriate for the warehouse behind this " + "connection (including Redash-backed engines). Stay within this dialect — " + "do not mix functions from other engines.\n" + ) _BASE_VERIFY_PROMPT = ( diff --git a/src/datasight/redash_runner.py b/src/datasight/redash_runner.py new file mode 100644 index 00000000..8992dfca --- /dev/null +++ b/src/datasight/redash_runner.py @@ -0,0 +1,295 @@ +""" +Execute SQL via the Redash REST API (ad-hoc POST /api/query_results). + +Polls ``/api/jobs/`` until completion, then loads rows from +``GET /api/query_results/``. +""" + +from __future__ import annotations + +import asyncio +import time +from typing import Any + +import httpx +import pandas as pd +from loguru import logger + +from datasight.exceptions import ConnectionError, QueryError, QueryTimeoutError +from datasight.runner import DEFAULT_QUERY_TIMEOUT, _sql_preview + +# Redash API job.status values (legacy numeric mapping used by the REST API). +_JOB_PENDING = 1 +_JOB_STARTED = 2 +_JOB_SUCCESS = 3 +_JOB_FAILURE = 4 +_JOB_CANCELLED = 5 +_JOB_DEFERRED = 6 +_JOB_SCHEDULED = 7 + +_JOB_WAIT_STATUSES = {_JOB_PENDING, _JOB_STARTED, _JOB_DEFERRED, _JOB_SCHEDULED} + +_CLOSED_MSG = "RedashRunner is closed" + + +def _normalize_base_url(url: str) -> str: + return url.rstrip("/") + + +def _query_result_to_dataframe(query_result: dict[str, Any]) -> pd.DataFrame: + data = query_result.get("data") or {} + cols_meta = data.get("columns") or [] + col_names = [ + str(c.get("name") or c.get("friendly_name") or f"col_{i}") + for i, c in enumerate(cols_meta) + ] + rows = data.get("rows") or [] + if not rows: + return pd.DataFrame(columns=col_names) + return pd.DataFrame(rows) + + +def _job_query_result_id(job: dict[str, Any]) -> int | None: + qrid = job.get("query_result_id") + if qrid is None: + qrid = job.get("result") + if qrid is None: + return None + return int(qrid) + + +class RedashRunner: + """Run SQL against a fixed Redash data source using the HTTP API.""" + + def __init__( + self, + base_url: str, + api_key: str, + data_source_id: int, + *, + query_timeout: float = DEFAULT_QUERY_TIMEOUT, + poll_interval: float = 0.5, + http_timeout: float = 60.0, + verify_ssl: bool = True, + client: httpx.AsyncClient | None = None, + ): + self._owns_client = client is None + self._base_url = _normalize_base_url(base_url) + self._api_key = api_key + self._data_source_id = int(data_source_id) + self._query_timeout = float(query_timeout) + self._poll_interval = float(poll_interval) + self._headers = {"Authorization": f"Key {api_key}", "Content-Type": "application/json"} + self._client: httpx.AsyncClient | None = None + + if client is None: + self._client = httpx.AsyncClient( + base_url=self._base_url, + headers=self._headers, + timeout=http_timeout, + verify=verify_ssl, + ) + else: + self._client = client + + logger.info( + f"Redash runner: base={self._base_url}, data_source_id={self._data_source_id}" + ) + + def close(self) -> None: + """Release the HTTP client when no asyncio loop is running. + + When called under a running event loop, prefer ``await aclose()`` or + ``async with RedashRunner(...)``. + """ + if self._client is None or not self._owns_client: + return + client = self._client + try: + asyncio.get_running_loop() + except RuntimeError: + asyncio.run(client.aclose()) + self._client = None + else: + logger.warning( + "RedashRunner.close() skipped closing the HTTP client while an event loop " + "is running — use `async with` or `await runner.aclose()`" + ) + + async def aclose(self) -> None: + if self._client is not None and self._owns_client: + await self._client.aclose() + self._client = None + + def __enter__(self) -> RedashRunner: + return self + + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + self.close() + + async def __aenter__(self) -> RedashRunner: + return self + + async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + await self.aclose() + + async def get_row_count(self, table: str) -> int | None: # noqa: ARG002 + """Skip COUNT(*) via Redash during schema introspection.""" + return None + + def _require_client(self) -> httpx.AsyncClient: + if self._client is None: + raise ConnectionError(_CLOSED_MSG) + return self._client + + async def _post_query_results(self, sql: str) -> dict[str, Any]: + client = self._require_client() + payload = {"query": sql, "data_source_id": self._data_source_id, "max_age": 0} + logger.debug(f"Redash ad-hoc query: {_sql_preview(sql)}") + try: + resp = await client.post("/api/query_results", json=payload) + except httpx.RequestError as e: + msg = f"Redash request failed: {e}" + raise ConnectionError(msg) from e + + try: + body = resp.json() + except ValueError as e: + msg = f"Redash returned non-JSON ({resp.status_code}): {resp.text[:500]}" + raise QueryError(msg) from e + + if resp.status_code >= 400: + err = "" + if isinstance(body, dict): + job = body.get("job") if isinstance(body.get("job"), dict) else {} + err = str(job.get("error") or body.get("message") or resp.text)[:2000] + msg = f"Redash HTTP {resp.status_code}: {err}" + raise QueryError(msg) + + if not isinstance(body, dict): + msg = "Redash returned unexpected JSON (expected object)" + raise QueryError(msg) + return body + + async def _fetch_job(self, job_id: str) -> dict[str, Any]: + client = self._require_client() + try: + resp = await client.get(f"/api/jobs/{job_id}") + except httpx.RequestError as e: + msg = f"Redash job poll failed: {e}" + raise ConnectionError(msg) from e + + try: + body = resp.json() + except ValueError as e: + msg = f"Redash job poll non-JSON ({resp.status_code})" + raise QueryError(msg) from e + + if resp.status_code >= 400: + msg = f"Redash job HTTP {resp.status_code}: {body}" + raise QueryError(msg) + + job = body.get("job") if isinstance(body, dict) else None + if not isinstance(job, dict): + msg = "Redash job response missing 'job' object" + raise QueryError(msg) + return job + + async def _poll_job(self, job_id: str) -> int: + deadline = time.monotonic() + self._query_timeout + while time.monotonic() < deadline: + job = await self._fetch_job(job_id) + status = int(job.get("status") or 0) + if status in (_JOB_FAILURE, _JOB_CANCELLED): + msg = str(job.get("error") or "Query execution failed") + raise QueryError(msg) + if status == _JOB_SUCCESS: + qrid = _job_query_result_id(job) + if qrid is None: + msg = "Redash job succeeded but returned no query_result_id" + raise QueryError(msg) + return qrid + + if status not in _JOB_WAIT_STATUSES: + msg = f"Unknown Redash job status: {status}" + raise QueryError(msg) + + await asyncio.sleep(self._poll_interval) + + msg = ( + f"Redash query timed out after {self._query_timeout:.0f}s waiting for job {job_id}. " + "Try a simpler query or increase REDASH_QUERY_TIMEOUT." + ) + raise QueryTimeoutError(msg) + + async def _fetch_query_result(self, query_result_id: int) -> pd.DataFrame: + client = self._require_client() + try: + resp = await client.get(f"/api/query_results/{query_result_id}") + except httpx.RequestError as e: + msg = f"Redash fetch result failed: {e}" + raise ConnectionError(msg) from e + + try: + body = resp.json() + except ValueError as e: + msg = f"Redash result non-JSON ({resp.status_code})" + raise QueryError(msg) from e + + if resp.status_code >= 400: + msg = f"Redash result HTTP {resp.status_code}: {body}" + raise QueryError(msg) + + qr = body.get("query_result") if isinstance(body, dict) else None + if not isinstance(qr, dict): + msg = "Redash result JSON missing 'query_result'" + raise QueryError(msg) + return _query_result_to_dataframe(qr) + + async def run_sql(self, sql: str) -> pd.DataFrame: + try: + return await asyncio.wait_for( + self._execute_once(sql), + timeout=self._query_timeout, + ) + except TimeoutError: + msg = ( + f"Redash query timed out after {self._query_timeout:.0f}s. " + "Try increasing REDASH_QUERY_TIMEOUT or simplifying the SQL." + ) + raise QueryTimeoutError(msg) from None + + async def _execute_once(self, sql: str) -> pd.DataFrame: + body = await self._post_query_results(sql) + + if "query_result" in body and isinstance(body["query_result"], dict): + df = _query_result_to_dataframe(body["query_result"]) + logger.debug(f"Redash cache hit: {len(df)} rows") + return df + + job_wrapped = body.get("job") + if not isinstance(job_wrapped, dict): + msg = "Redash response had neither query_result nor job" + raise QueryError(msg) + + job_id = job_wrapped.get("id") + if job_id is None: + msg = "Redash job response missing id" + raise QueryError(msg) + + status = int(job_wrapped.get("status") or 0) + if status == _JOB_SUCCESS: + qrid = _job_query_result_id(job_wrapped) + if qrid is None: + msg = "Redash job already succeeded but missing query_result_id" + raise QueryError(msg) + return await self._fetch_query_result(qrid) + + if status in (_JOB_FAILURE, _JOB_CANCELLED): + msg = str(job_wrapped.get("error") or "Redash query failed") + raise QueryError(msg) + + query_result_id = await self._poll_job(str(job_id)) + df = await self._fetch_query_result(query_result_id) + logger.debug(f"Redash returned {len(df)} rows, {len(df.columns)} cols") + return df diff --git a/src/datasight/runner.py b/src/datasight/runner.py index ac660d68..1ed0f463 100644 --- a/src/datasight/runner.py +++ b/src/datasight/runner.py @@ -2,8 +2,8 @@ SQL runners for datasight. Provides a common async interface for executing SQL queries against local -DuckDB files, remote Flight SQL servers, PostgreSQL, SQLite databases, or -Apache Spark via Spark Connect. +DuckDB files, remote Flight SQL servers, PostgreSQL, SQLite databases, +Apache Spark via Spark Connect, or Redash's HTTP API. """ from __future__ import annotations diff --git a/src/datasight/settings.py b/src/datasight/settings.py index 32c670ae..4e1b2fbd 100644 --- a/src/datasight/settings.py +++ b/src/datasight/settings.py @@ -15,7 +15,7 @@ from dotenv import load_dotenv from datasight.exceptions import ConfigurationError -from datasight.runner import DEFAULT_SPARK_MAX_RESULT_BYTES +from datasight.runner import DEFAULT_QUERY_TIMEOUT, DEFAULT_SPARK_MAX_RESULT_BYTES def _safe_int(value: str, default: int) -> int: @@ -74,6 +74,12 @@ def _safe_float(value: str, default: float) -> float: "SPARK_REMOTE", "SPARK_TOKEN", "SPARK_MAX_RESULT_BYTES", + "REDASH_BASE_URL", + "REDASH_API_KEY", + "REDASH_DATA_SOURCE_ID", + "REDASH_SQL_DIALECT", + "REDASH_QUERY_TIMEOUT", + "REDASH_POLL_INTERVAL", # LLM settings "LLM_PROVIDER", "LLM_TIMEOUT", @@ -169,7 +175,7 @@ def restore_original_env() -> None: # exists only for static typing and must be updated by hand when a provider # is added or removed (Python can't derive Literals from a runtime set). LLMProvider = Literal["anthropic", "ollama", "github", "openai"] -DBMode = Literal["duckdb", "sqlite", "postgres", "flightsql", "spark"] +DBMode = Literal["duckdb", "sqlite", "postgres", "flightsql", "spark", "redash"] # Mapping from database mode to SQL dialect for query generation DB_MODE_TO_DIALECT: dict[str, str] = { @@ -178,6 +184,7 @@ def restore_original_env() -> None: "postgres": "postgres", "flightsql": "duckdb", # Flight SQL uses DuckDB dialect "spark": "spark", + "redash": "postgres", # Default warehouse dialect; overridden by REDASH_SQL_DIALECT } @@ -276,9 +283,19 @@ class DatabaseSettings: spark_token: str | None = None spark_max_result_bytes: int = DEFAULT_SPARK_MAX_RESULT_BYTES + # Redash REST API (DB_MODE=redash) + redash_base_url: str = "" + redash_api_key: str = "" + redash_data_source_id: int = 0 + redash_sql_dialect: str = "postgres" + redash_query_timeout: float = DEFAULT_QUERY_TIMEOUT + redash_poll_interval: float = 0.5 + @property def sql_dialect(self) -> str: """Get the SQL dialect for the current mode.""" + if self.mode == "redash": + return self.redash_sql_dialect.strip() or "postgres" return DB_MODE_TO_DIALECT.get(self.mode, "duckdb") @@ -354,8 +371,10 @@ def from_env(cls, env_path: str | Path | None = None, *, override: bool = False) db_mode = "flightsql" case "spark": db_mode = "spark" + case "redash": + db_mode = "redash" case _: - valid_modes = "duckdb, sqlite, postgres, flightsql, spark" + valid_modes = "duckdb, sqlite, postgres, flightsql, spark, redash" msg = f"Invalid DB_MODE: {db_mode_raw!r}. Valid modes: {valid_modes}" raise ConfigurationError(msg) @@ -397,6 +416,15 @@ def from_env(cls, env_path: str | Path | None = None, *, override: bool = False) os.environ.get("SPARK_MAX_RESULT_BYTES", ""), DEFAULT_SPARK_MAX_RESULT_BYTES, ), + redash_base_url=os.environ.get("REDASH_BASE_URL", ""), + redash_api_key=os.environ.get("REDASH_API_KEY", ""), + redash_data_source_id=_safe_int(os.environ.get("REDASH_DATA_SOURCE_ID", ""), 0), + redash_sql_dialect=os.environ.get("REDASH_SQL_DIALECT", "postgres"), + redash_query_timeout=_safe_float( + os.environ.get("REDASH_QUERY_TIMEOUT", ""), + DEFAULT_QUERY_TIMEOUT, + ), + redash_poll_interval=_safe_float(os.environ.get("REDASH_POLL_INTERVAL", ""), 0.5), ), app=AppSettings( port=_safe_int(os.environ.get("PORT", ""), 8084), @@ -430,4 +458,12 @@ def validate(self) -> list[str]: if not self.database.path: errors.append("DB_PATH is not set") + if self.database.mode == "redash": + if not self.database.redash_base_url.strip(): + errors.append("REDASH_BASE_URL is not set") + if not self.database.redash_api_key.strip(): + errors.append("REDASH_API_KEY is not set") + if self.database.redash_data_source_id <= 0: + errors.append("REDASH_DATA_SOURCE_ID must be a positive integer") + return errors diff --git a/src/datasight/sql_validation.py b/src/datasight/sql_validation.py index 847e50f8..00d5d98b 100644 --- a/src/datasight/sql_validation.py +++ b/src/datasight/sql_validation.py @@ -166,6 +166,10 @@ def validate_sql( # noqa: C901 "duckdb": "duckdb", "postgres": "postgres", "sqlite": "sqlite", + "mysql": "mysql", + "spark": "spark", + "bigquery": "bigquery", + "snowflake": "snowflake", } sqlglot_dialect = _SQLGLOT_DIALECTS.get(dialect, "duckdb") diff --git a/src/datasight/templates/env.template b/src/datasight/templates/env.template index a3f09b15..bed25ad9 100644 --- a/src/datasight/templates/env.template +++ b/src/datasight/templates/env.template @@ -39,7 +39,7 @@ # GITHUB_MODELS_MODEL=gpt-4o # GITHUB_MODELS_BASE_URL=https://models.github.ai/inference -# Database mode: "duckdb" (default), "sqlite", "postgres", or "flightsql" +# Database mode: "duckdb" (default), "sqlite", "postgres", "flightsql", "spark", or "redash" DB_MODE=duckdb # DuckDB settings (DB_MODE=duckdb) @@ -65,6 +65,17 @@ DB_PATH=./my_database.duckdb # FLIGHT_SQL_USERNAME= # FLIGHT_SQL_PASSWORD= +# Redash REST API (DB_MODE=redash) — uses your Redash user API key to run ad-hoc SQL +# against one configured data source. Set REDASH_SQL_DIALECT to match that warehouse +# (e.g. postgres, bigquery, snowflake) so generated SQL validates correctly. +# REDASH_BASE_URL=https://redash.example.com +# REDASH_API_KEY=your-user-api-key +# REDASH_DATA_SOURCE_ID=1 +# REDASH_SQL_DIALECT=postgres +# or e.g. mysql when the Redash data source is MySQL +# REDASH_QUERY_TIMEOUT=120 +# REDASH_POLL_INTERVAL=0.5 + # Optional: paths to context files for the AI agent. # These are always local paths on the machine running datasight, even when # connecting to a remote database via Flight SQL. diff --git a/src/datasight/templates/html/export_session_python.mustache b/src/datasight/templates/html/export_session_python.mustache index e7373e7f..81a3d021 100644 --- a/src/datasight/templates/html/export_session_python.mustache +++ b/src/datasight/templates/html/export_session_python.mustache @@ -61,6 +61,18 @@ conn = sqlite3.connect(DB_PATH) def run_sql(sql: str) -> pd.DataFrame: return pd.read_sql_query(sql, conn) {{/is_sqlite}} +{{#is_redash}} +# Database mode: redash — queries executed via the Redash HTTP API during the session. +# This export cannot embed REDASH_API_KEY or REDASH_DATA_SOURCE_ID. To replay, either +# implement POST /api/query_results + job polling with httpx, or set DB_MODE to your +# warehouse (e.g. postgres) and connect directly. + + +def run_sql(sql: str) -> pd.DataFrame: + raise NotImplementedError( + "Wire Redash (or the backing warehouse) — exported scripts do not include API secrets." + ) +{{/is_redash}} {{#is_unknown_mode}} # Database mode: {{{db_mode}}} # This session ran against a non-file backend. Replace the run_sql body diff --git a/src/datasight/web/app.py b/src/datasight/web/app.py index 62f80ae0..27ccfb98 100644 --- a/src/datasight/web/app.py +++ b/src/datasight/web/app.py @@ -610,6 +610,10 @@ def __init__(self) -> None: # tidy_apply. Consumed by /api/tidy/grounding/repair so the LLM # repair prompt can show both shapes; cleared on project change. self.pre_tidy_schema: dict[str, set[str]] | None = None + # Set only during background auto-load (see _startup). Not cleared by + # clear_project — load_project calls clear_project first, then needs + # this flag until introspection finishes. + self.auto_loading_path: str | None = None def clear_project(self) -> None: """Clear project-specific state.""" @@ -1177,11 +1181,25 @@ async def _startup() -> None: # Check for auto-load project auto_load_project = os.environ.get("DATASIGHT_AUTO_LOAD_PROJECT") if auto_load_project: - try: - await load_project(auto_load_project, _state) - logger.info(f"Auto-loaded project: {auto_load_project}") - except ProjectError as e: - logger.error(f"Failed to auto-load project {auto_load_project}: {e}") + auto_path = auto_load_project.strip() + + async def _auto_load_project_task(project_path: str) -> None: + """Load project without blocking ASGI startup (slow remote introspection).""" + _state.auto_loading_path = project_path + err_msg: str | None = None + try: + await load_project(project_path, _state) + logger.info(f"Auto-loaded project: {project_path}") + except ProjectError as e: + err_msg = str(e) + logger.error(f"Failed to auto-load project {project_path}: {e}") + except Exception as e: + err_msg = str(e) + logger.exception(f"Unexpected auto-load failure for {project_path}") + finally: + _state.auto_loading_path = None + + asyncio.create_task(_auto_load_project_task(auto_path)) socket_path = os.environ.get("DATASIGHT_UNIX_SOCKET", "").strip() port = os.environ.get("PORT", "8084") @@ -1875,7 +1893,13 @@ async def _handle_sql_confirmation( @app.get("/", response_class=HTMLResponse) async def index(): - return FileResponse(_INDEX_HTML, media_type="text/html") + # Avoid caching the SPA shell: hashed /static assets change on each build; a stale + # index.html or long-lived tab can otherwise request removed chunk files (e.g. plotly.min-*.js). + return FileResponse( + _INDEX_HTML, + media_type="text/html", + headers={"Cache-Control": "no-cache, must-revalidate"}, + ) @app.get("/api/schema") @@ -2139,6 +2163,15 @@ async def get_current_project(state: AppState = Depends(get_state)): "tables": state.ephemeral_tables_info, "sql_dialect": state.sql_dialect, } + if state.auto_loading_path: + return { + "loaded": False, + "loading": True, + "path": state.auto_loading_path, + "name": None, + "is_ephemeral": False, + "sql_dialect": state.sql_dialect, + } if not state.project_loaded or state.project_dir is None: return { "loaded": False, diff --git a/tests/test_config_extra.py b/tests/test_config_extra.py index 07da11d0..da13b602 100644 --- a/tests/test_config_extra.py +++ b/tests/test_config_extra.py @@ -16,6 +16,7 @@ normalize_db_mode, ) from datasight.exceptions import ConfigurationError, ConnectionError as DsConnectionError +from datasight.redash_runner import RedashRunner from datasight.runner import DuckDBRunner, SQLiteRunner from datasight.settings import DatabaseSettings @@ -74,6 +75,45 @@ def test_create_sql_runner_local_alias_still_requires_path(): create_sql_runner("local", db_path="") +def test_create_sql_runner_redash_branch(): + runner = create_sql_runner( + "redash", + redash_base_url="https://redash.example.com", + redash_api_key="secret", + redash_data_source_id=3, + ) + assert isinstance(runner, RedashRunner) + runner.close() + + +def test_create_sql_runner_redash_missing_base_url(): + with pytest.raises(ConfigurationError, match="REDASH_BASE_URL"): + create_sql_runner( + "redash", + redash_api_key="x", + redash_data_source_id=1, + ) + + +def test_create_sql_runner_redash_missing_api_key(): + with pytest.raises(ConfigurationError, match="REDASH_API_KEY"): + create_sql_runner( + "redash", + redash_base_url="https://redash.example.com", + redash_data_source_id=1, + ) + + +def test_create_sql_runner_redash_invalid_data_source_id(): + with pytest.raises(ConfigurationError, match="REDASH_DATA_SOURCE_ID"): + create_sql_runner( + "redash", + redash_base_url="https://redash.example.com", + redash_api_key="k", + redash_data_source_id=0, + ) + + def test_create_sql_runner_duckdb_with_path(test_duckdb_path): runner = create_sql_runner("duckdb", db_path=test_duckdb_path) assert isinstance(runner, DuckDBRunner) diff --git a/tests/test_export.py b/tests/test_export.py index 701f6eb6..b6f7d49b 100644 --- a/tests/test_export.py +++ b/tests/test_export.py @@ -497,6 +497,15 @@ def test_export_python_unknown_db_mode_emits_todo_scaffold(): assert "raise NotImplementedError" in script +def test_export_python_redash_mode_documents_api_stub(): + events, _ = _python_export_events() + script = export_session_python(events, title="t", db_path="", db_mode="redash") + ast.parse(script) + assert "Database mode: redash" in script + assert "Redash" in script + assert "raise NotImplementedError" in script + + def test_export_python_respects_exclude_indices(): events, _ = _python_export_events() script = export_session_python( diff --git a/tests/test_redash_runner.py b/tests/test_redash_runner.py new file mode 100644 index 00000000..7e7f6a9f --- /dev/null +++ b/tests/test_redash_runner.py @@ -0,0 +1,213 @@ +"""Tests for Redash REST SQL runner (mocked HTTP).""" + +from __future__ import annotations + +import json + +import httpx +import pandas as pd +import pytest + +from datasight.exceptions import QueryError, QueryTimeoutError +from datasight.redash_runner import RedashRunner + + +def _json_response(obj: dict, status: int = 200) -> httpx.Response: + return httpx.Response(status, json=obj) + + +@pytest.mark.asyncio +async def test_redash_immediate_query_result(): + """POST returns cached query_result without job polling.""" + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST" and request.url.path == "/api/query_results": + body = { + "query_result": { + "data": { + "columns": [{"name": "ok"}], + "rows": [{"ok": 1}], + } + } + } + return _json_response(body) + return httpx.Response(404) + + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport, base_url="http://redash.test") as http_client: + runner = RedashRunner( + base_url="http://redash.test", + api_key="k", + data_source_id=9, + query_timeout=5.0, + poll_interval=0.01, + client=http_client, + ) + try: + df = await runner.run_sql("SELECT 1 AS ok") + finally: + await runner.aclose() + + assert list(df.columns) == ["ok"] + assert df.iloc[0]["ok"] == 1 + + +@pytest.mark.asyncio +async def test_redash_job_poll_then_fetch(): + calls = {"n": 0} + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST" and request.url.path == "/api/query_results": + return _json_response({"job": {"id": "job-1", "status": 1}}) + if request.method == "GET" and request.url.path == "/api/jobs/job-1": + calls["n"] += 1 + if calls["n"] < 2: + return _json_response({"job": {"id": "job-1", "status": 2}}) + return _json_response( + {"job": {"id": "job-1", "status": 3, "query_result_id": 42, "result": 42}} + ) + if request.method == "GET" and request.url.path == "/api/query_results/42": + return _json_response( + { + "query_result": { + "data": { + "columns": [{"name": "x"}], + "rows": [{"x": "a"}, {"x": "b"}], + } + } + } + ) + return httpx.Response(404) + + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport, base_url="http://redash.test") as http_client: + runner = RedashRunner( + base_url="http://redash.test", + api_key="k", + data_source_id=1, + query_timeout=10.0, + poll_interval=0.01, + client=http_client, + ) + try: + df = await runner.run_sql("SELECT x FROM t") + finally: + await runner.aclose() + + assert isinstance(df, pd.DataFrame) + assert len(df) == 2 + + +@pytest.mark.asyncio +async def test_redash_job_failure(): + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST" and request.url.path == "/api/query_results": + return _json_response({"job": {"id": "j", "status": 4, "error": "boom"}}) + return httpx.Response(404) + + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport, base_url="http://r.test") as http_client: + runner = RedashRunner( + base_url="http://r.test", + api_key="k", + data_source_id=1, + query_timeout=5.0, + client=http_client, + ) + try: + with pytest.raises(QueryError, match="boom"): + await runner.run_sql("SELECT bad") + finally: + await runner.aclose() + + +@pytest.mark.asyncio +async def test_redash_poll_timeout(): + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST" and request.url.path == "/api/query_results": + return _json_response({"job": {"id": "slow", "status": 1}}) + if request.method == "GET" and request.url.path == "/api/jobs/slow": + return _json_response({"job": {"id": "slow", "status": 1}}) + return httpx.Response(404) + + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport, base_url="http://r.test") as http_client: + runner = RedashRunner( + base_url="http://r.test", + api_key="k", + data_source_id=1, + query_timeout=0.15, + poll_interval=0.05, + client=http_client, + ) + try: + with pytest.raises(QueryTimeoutError): + await runner.run_sql("SELECT 1") + finally: + await runner.aclose() + + +@pytest.mark.asyncio +async def test_redash_post_http_error(): + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(403, json={"message": "denied"}) + + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport, base_url="http://r.test") as http_client: + runner = RedashRunner( + base_url="http://r.test", + api_key="k", + data_source_id=1, + client=http_client, + ) + try: + with pytest.raises(QueryError, match="403"): + await runner.run_sql("SELECT 1") + finally: + await runner.aclose() + + +@pytest.mark.asyncio +async def test_redash_post_payload_includes_data_source(): + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + if request.method == "POST" and request.url.path == "/api/query_results": + captured["body"] = json.loads(request.content.decode()) + return _json_response( + {"query_result": {"data": {"columns": [{"name": "n"}], "rows": [{"n": 0}]}}} + ) + return httpx.Response(404) + + transport = httpx.MockTransport(handler) + async with httpx.AsyncClient(transport=transport, base_url="http://r.test") as http_client: + runner = RedashRunner( + base_url="http://r.test", + api_key="secret", + data_source_id=7, + client=http_client, + ) + try: + await runner.run_sql("SELECT 0 AS n") + finally: + await runner.aclose() + + assert captured["body"]["data_source_id"] == 7 + assert captured["body"]["max_age"] == 0 + assert "SELECT 0" in captured["body"]["query"] + + +@pytest.mark.asyncio +async def test_redash_get_row_count_returns_none(): + transport = httpx.MockTransport(lambda r: httpx.Response(404)) + async with httpx.AsyncClient(transport=transport, base_url="http://r.test") as http_client: + runner = RedashRunner( + base_url="http://r.test", + api_key="k", + data_source_id=1, + client=http_client, + ) + try: + assert await runner.get_row_count("any") is None + finally: + await runner.aclose() diff --git a/tests/test_settings.py b/tests/test_settings.py index 5188dab2..158ac009 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -165,7 +165,7 @@ def test_invalid_db_mode_raises_error(self, monkeypatch): assert "Invalid DB_MODE" in str(exc_info.value) assert "postgress" in str(exc_info.value) - assert "duckdb, sqlite, postgres, flightsql" in str(exc_info.value) + assert "duckdb, sqlite, postgres, flightsql, spark, redash" in str(exc_info.value) def test_local_alias_for_duckdb(self, monkeypatch): """'local' should be accepted as alias for 'duckdb'.""" @@ -176,6 +176,35 @@ def test_local_alias_for_duckdb(self, monkeypatch): assert settings.database.mode == "duckdb" + def test_db_mode_redash_from_env(self, monkeypatch): + monkeypatch.setenv("DB_MODE", "redash") + monkeypatch.setenv("REDASH_BASE_URL", "https://r.example.com/") + monkeypatch.setenv("REDASH_API_KEY", "rk") + monkeypatch.setenv("REDASH_DATA_SOURCE_ID", "12") + monkeypatch.setenv("REDASH_SQL_DIALECT", "snowflake") + + settings = Settings.from_env() + + assert settings.database.mode == "redash" + assert settings.database.redash_data_source_id == 12 + assert settings.database.sql_dialect == "snowflake" + + def test_redash_validate_requires_credentials(self, monkeypatch): + monkeypatch.setenv("DB_MODE", "redash") + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-test") + + errs = Settings.from_env().validate() + assert any("REDASH_BASE_URL" in e for e in errs) + + def test_redash_validate_ok_when_configured(self, monkeypatch): + monkeypatch.setenv("DB_MODE", "redash") + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-test") + monkeypatch.setenv("REDASH_BASE_URL", "https://r.example.com") + monkeypatch.setenv("REDASH_API_KEY", "rk") + monkeypatch.setenv("REDASH_DATA_SOURCE_ID", "3") + + assert Settings.from_env().validate() == [] + def test_safe_int_for_ports(self, monkeypatch): """Invalid port values should use defaults instead of crashing.""" monkeypatch.setenv("PORT", "not-a-port")