From 19baddbcab6c70bf10e0557081df5982195366a4 Mon Sep 17 00:00:00 2001 From: Eddie A Tejeda <669988+eddietejeda@users.noreply.github.com> Date: Sat, 24 Oct 2026 15:31:42 -0700 Subject: [PATCH] Use latest Hotdata API: inline Arrow query results, pagination, error_message - Request Arrow from POST /v1/query so fast queries take one round trip; fall back to the Results API on 202 (async or too large). - Read error_message from failed query runs and results; show code/message from Hotdata error bodies. - Follow next_cursor pagination on /v1/databases and /v1/information_schema. - Handle Results API 202 (Retry-After) and 409 responses. - Health check pings GET /v1/databases?limit=1. --- CHANGELOG.md | 8 ++ SPEC.md | 44 +++--- pkg/cmd/mockapi/main.go | 23 +++- pkg/hotdata/client.go | 249 ++++++++++++++++++++++++---------- pkg/plugin/arrow.go | 7 +- pkg/plugin/datasource.go | 13 +- pkg/plugin/datasource_test.go | 118 ++++++++++++++-- 7 files changed, 336 insertions(+), 126 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b00db5d..12e087f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,13 @@ # Changelog +## Unreleased + +- Read query results as Arrow directly from `POST /v1/query`, so a fast query takes one request instead of two. Async and too-large results still come from the Results API. +- Show the real error for failed queries: read `error_message` from query runs and failed results, and show the `message` and `code` from Hotdata error bodies instead of raw JSON. +- Follow pagination on `/v1/databases` and `/v1/information_schema`, so the database picker and SQL autocomplete list every database and table, not just the first page. +- Handle the Results API's `202` (still persisting, honoring `Retry-After`) and `409` (failed) responses. +- The health check validates credentials with one `GET /v1/databases?limit=1` instead of listing every database. + ## 1.0.4 - Clear the high-severity osv-scanner findings that failed Grafana plugin validation for 1.0.3 (all in dev/build tooling): force basic-ftp to 6.2.3 and source-map-js to 1.2.2, and replace micromatch with picomatch under eslint-webpack-plugin, since braces has no fixed release. diff --git a/SPEC.md b/SPEC.md index 817bca0..8a674fd 100644 --- a/SPEC.md +++ b/SPEC.md @@ -17,7 +17,7 @@ Hotdata's engine is Arrow-native — `information_schema` reports Arrow type nam So we **skip** the `database/sql` + `sqlds` route. The inline JSON rows from `POST /v1/query` carry no type metadata (timestamps arrive as bare strings without timezone), which would force lossy type inference. Instead: - Implement `backend.QueryDataHandler` directly. -- Always fetch results from `GET /v1/results/{id}?format=arrow` and convert record batches to frames (~200 lines; InfluxData's deprecated Flight SQL plugin is a reference implementation). +- Always read results as Arrow — inline from `POST /v1/query` (`Accept: application/vnd.apache.arrow.stream`), or from `GET /v1/results/{id}?format=arrow` when the query goes async — and convert record batches to frames (~200 lines; InfluxData's deprecated Flight SQL plugin is a reference implementation). - Reuse `grafana-plugin-sdk-go/data/sqlutil` for macro interpolation only, and `data.LongToWide` for time-series shaping. Typed, complete, no JSON round-trip. @@ -28,22 +28,14 @@ All requests: `Authorization: Bearer `, `x-workspace-id: work…`. Quer | Endpoint | Purpose | |---|---| -| `POST /v1/query` | Run SQL. Body: `{"sql", "async": true, "async_after_ms", "dialect"?}` | +| `POST /v1/query` | Run SQL. Body: `{"sql", "async": true, "async_after_ms", "dialect"?}`; `Accept: application/vnd.apache.arrow.stream` | | `GET /v1/query-runs/{qrun}` | Poll an async run | | `GET /v1/results/{rslt}?format=arrow` | Fetch full result as Arrow IPC stream | -| `GET /v1/databases` / `GET /v1/databases/{id}` | List databases / resolve `default_connection_id` | -| `GET /v1/information_schema?connection_id=…` | Tables; `&schema=&table=&include_columns=true` adds columns with Arrow types | +| `GET /v1/databases` / `GET /v1/databases/{id}` | List databases (paginated) / resolve `default_connection_id` | +| `GET /v1/information_schema?connection_id=…` | Tables (paginated); `&schema=&table=&include_columns=true` adds columns with Arrow types | | `GET /v1/workspaces` | Workspace list (config editor) | -**Fast query → `200`** (completed within `async_after_ms`): - -```json -{"query_run_id":"qrun…","result_id":"rslt…", - "columns":["one","f","s","ts","b"],"nullable":[false,true,false,false,false], - "rows":[[1,1.5,"x","2026-09-04T19:55:56.368620650",true]], - "row_count":1,"preview_row_count":1,"total_row_count":1, - "truncated":false,"execution_time_ms":2} -``` +**Fast query → `200`** (completed within `async_after_ms`): an Arrow IPC stream body holding the whole result. The run metadata comes as headers: `X-Query-Run-Id`, `X-Result-Id`, `X-Total-Row-Count`, `X-Execution-Time-Ms`. **Slow query → `202`:** @@ -51,21 +43,28 @@ All requests: `Authorization: Bearer `, `x-workspace-id: work…`. Quer {"query_run_id":"qrun…","status":"running","status_url":"/v1/query-runs/qrun…"} ``` -**Poll → `200`:** `{"id","status":"running"|"succeeded"|"failed","result_id","row_count","execution_time_ms","trace_id",…}` +**Result too large to inline → `202`** with the same body plus an `X-Result-Id` header: read it from the Results API, no run polling. + +**Poll → `200`:** `{"id","status":"running"|"succeeded"|"failed","result_id","row_count","execution_time_ms","error_message","trace_id",…}` + +**Results:** `200` Arrow when ready; `202` `{"result_id","status":"processing"}` + `Retry-After` while persisting; `409` `{"status":"failed","error_message"}` on failure. + +**Listings** (`/v1/databases`, `/v1/information_schema`) return one page: follow `next_cursor` (as `?cursor=`) while `has_more` is true. **Schema discovery** (`include_columns=true`): ```json {"tables":[{"schema":"cities","table":"global","synced":true, "columns":[{"name":"city","data_type":"Utf8View","nullable":true}, - {"name":"latitude","data_type":"Float64","nullable":true}]}]} + {"name":"latitude","data_type":"Float64","nullable":true}]}], + "count":1,"has_more":false} ``` **Operational behavior:** - `429` with `OVERLOADED` code + `Retry-After` header under load — retry with honor, cap 3 attempts. - Cold start: first query against an idle workspace blocks ~10–20s while the worker wakes. Not an error. - Dialects: `hotsql` (default), `postgres`, `duckdb`, `snowflake` — non-hotsql is transpiled server-side, read-only. -- Errors carry stable codes alongside messages — branch on the code, not the sentence. +- Errors carry stable codes alongside messages — branch on the code, not the sentence. Body shape: `{"error":{"code","message"}}` (a few endpoints send `{"error":"code"}`). ## 4. Query execution flow (backend) @@ -76,17 +75,20 @@ sequenceDiagram participant H as Hotdata API G->>P: QueryData(rawSql, timeRange) P->>P: interpolate macros (sqlutil) - P->>H: POST /v1/query {sql, async:true, async_after_ms:1000} + P->>H: POST /v1/query {sql, async:true, async_after_ms:1000} (Accept: Arrow) alt completed fast (200) - H-->>P: result_id + inline preview + H-->>P: Arrow IPC stream (+ X-Query-Run-Id, X-Result-Id headers) + else too large to inline (202 + X-Result-Id) + P->>H: GET /v1/results/{result_id}?format=arrow + H-->>P: Arrow IPC stream else still running (202) H-->>P: query_run_id + status_url loop 500ms → 2s backoff, until terminal or ctx.Done() P->>H: GET /v1/query-runs/{id} end + P->>H: GET /v1/results/{result_id}?format=arrow + H-->>P: Arrow IPC stream end - P->>H: GET /v1/results/{result_id}?format=arrow - H-->>P: Arrow IPC stream P->>P: record batches → data.Frame (+LongToWide if timeseries) P-->>G: frames ``` @@ -94,7 +96,7 @@ sequenceDiagram Details: - Set `async_after_ms: 1000` — the server-enforced minimum. - Poll at 500ms, ×1.5 backoff, 2s cap. Respect `ctx.Done()` — Grafana cancels on panel abort/timeout. -- Always fetch via the Arrow results endpoint, even on the fast path — the inline JSON is an untyped preview; `truncated: true` never applies to the Arrow fetch. +- Always read Arrow, never the JSON form — JSON is untyped. An inline Arrow response always holds the whole result; the server answers 202 instead of truncating. - Overall deadline comes from the query context; surface cold-start waits as a frame notice ("waking workspace worker"), not an error. ## 5. Configuration diff --git a/pkg/cmd/mockapi/main.go b/pkg/cmd/mockapi/main.go index 23e7e2a..de1c1da 100644 --- a/pkg/cmd/mockapi/main.go +++ b/pkg/cmd/mockapi/main.go @@ -65,7 +65,18 @@ func main() { _, _ = fmt.Fprintf(w, `{"query_run_id":%q,"status":"running","status_url":"/v1/query-runs/%s"}`, runID, runID) return } - _, _ = fmt.Fprintf(w, `{"query_run_id":%q,"result_id":%q,"columns":["ts","service","value"],"rows":[],"row_count":180,"truncated":false,"execution_time_ms":3}`, runID, resultID) + if r.Header.Get("Accept") != "application/vnd.apache.arrow.stream" { + _, _ = fmt.Fprintf(w, `{"query_run_id":%q,"result_id":%q,"columns":["ts","service","value"],"rows":[],"preview_row_count":180,"total_row_count":180,"truncated":false,"execution_time_ms":3}`, runID, resultID) + return + } + w.Header().Set("Content-Type", "application/vnd.apache.arrow.stream") + w.Header().Set("X-Query-Run-Id", runID) + w.Header().Set("X-Result-Id", resultID) + w.Header().Set("X-Total-Row-Count", "180") + w.Header().Set("X-Execution-Time-Ms", "3") + if err := writeSeries(w); err != nil { + log.Printf("writing arrow stream: %v", err) + } }) // Production contract: query-runs and results are database-scoped. @@ -114,10 +125,12 @@ func main() { _, _ = fmt.Fprint(w, `{"error":{"message":"result not found","code":"NOT_FOUND"}}`) return } - // Model the fast-path race: first fetch returns a JSON placeholder - // before the Arrow stream is materialized. + // First fetch: the result is still being persisted, answered with 202 + // and a JSON placeholder before the Arrow stream is served. if fetches == 0 { w.Header().Set("Content-Type", "application/json") + w.Header().Set("Retry-After", "1") + w.WriteHeader(http.StatusAccepted) _, _ = fmt.Fprintf(w, `{"result_id":%q,"status":"processing"}`, resultID) return } @@ -128,7 +141,7 @@ func main() { }) mux.HandleFunc("GET /v1/databases", func(w http.ResponseWriter, r *http.Request) { - _, _ = fmt.Fprint(w, `{"databases":[{"id":"dbmock","name":"mock-db","default_catalog":"default","default_schema":"public","created_at":"2026-09-04T00:00:00Z"}]}`) + _, _ = fmt.Fprint(w, `{"databases":[{"id":"dbmock","name":"mock-db","default_catalog":"default","default_schema":"public","created_at":"2026-09-04T00:00:00Z"}],"count":1,"limit":100,"has_more":false}`) }) mux.HandleFunc("GET /v1/databases/dbmock", func(w http.ResponseWriter, r *http.Request) { @@ -140,7 +153,7 @@ func main() { if r.URL.Query().Get("include_columns") == "true" { cols = `,"columns":[{"name":"ts","data_type":"Timestamp(Microsecond, Some(\"UTC\"))","nullable":true},{"name":"service","data_type":"Utf8View","nullable":true},{"name":"value","data_type":"Float64","nullable":true}]` } - _, _ = fmt.Fprintf(w, `{"tables":[{"connection":"__db_mock","schema":"public","table":"metrics","synced":true,"last_sync":"2026-09-04 00:00:00"%s,"partition_by":[],"sorted_by":[]}]}`, cols) + _, _ = fmt.Fprintf(w, `{"tables":[{"connection":"__db_mock","schema":"public","table":"metrics","synced":true,"last_sync":"2026-09-04 00:00:00"%s,"partition_by":[],"sorted_by":[],"constant_per_key":[]}],"count":1,"limit":100,"has_more":false}`, cols) }) mux.HandleFunc("GET /v1/workspaces", func(w http.ResponseWriter, r *http.Request) { diff --git a/pkg/hotdata/client.go b/pkg/hotdata/client.go index 3007dcf..e5f301e 100644 --- a/pkg/hotdata/client.go +++ b/pkg/hotdata/client.go @@ -74,10 +74,13 @@ type Workspace struct { type RunMeta struct { QueryRunID string ResultID string - RowCount int64 ExecutionTimeMS int64 } +// arrowStreamType is the media type Hotdata uses for Arrow IPC streams, both +// as the Accept value on POST /v1/query and the result Content-Type. +const arrowStreamType = "application/vnd.apache.arrow.stream" + // queryRequest is the POST /v1/query body. async_after_ms is the server-side // budget before the request degrades to a 202 + poll flow. type queryRequest struct { @@ -87,53 +90,76 @@ type queryRequest struct { Dialect string `json:"dialect,omitempty"` } +// queryResponse is the JSON body of POST /v1/query: a 202 run ticket, or a +// 200 from a server that answered JSON despite the Arrow Accept header. type queryResponse struct { QueryRunID string `json:"query_run_id"` ResultID string `json:"result_id"` - RowCount int64 `json:"row_count"` ExecutionTimeMS int64 `json:"execution_time_ms"` - // 202 fields - Status string `json:"status"` - StatusURL string `json:"status_url"` + Warning string `json:"warning"` } type queryRunStatus struct { ID string `json:"id"` Status string `json:"status"` ResultID string `json:"result_id"` - RowCount int64 `json:"row_count"` ExecutionTimeMS int64 `json:"execution_time_ms"` - Error string `json:"error"` - Message string `json:"message"` + ErrorMessage string `json:"error_message"` } -// RunQuery executes SQL in the given database and blocks until the run is -// terminal, following the async 202 + poll flow when the server chooses it. -// async_after_ms has a server-enforced minimum of 1000. -func (c *Client) RunQuery(ctx context.Context, databaseID, sql, dialect string) (RunMeta, error) { +// Query executes SQL in the given database and returns the complete result +// as an Arrow IPC stream. The caller must close the reader. +// +// The request asks for Arrow directly, so a query that finishes within +// async_after_ms comes back in one round trip. Otherwise the server answers +// 202: with an X-Result-Id header when the result was too large to inline, +// or with only a run ticket when the query went async, in which case the run +// is polled until it has a result. async_after_ms has a server minimum of 1000. +func (c *Client) Query(ctx context.Context, databaseID, sql, dialect string) (RunMeta, io.ReadCloser, error) { body := queryRequest{SQL: sql, Async: true, AsyncAfterMS: 1000, Dialect: dialect} - resp, err := c.do(ctx, http.MethodPost, "/v1/query", databaseID, body) + resp, err := c.doAccept(ctx, http.MethodPost, "/v1/query", databaseID, arrowStreamType, body) if err != nil { - return RunMeta{}, err + return RunMeta{}, nil, err + } + + if resp.StatusCode == http.StatusOK && isArrow(resp.Header.Get("Content-Type")) { + execMS, _ := strconv.ParseInt(resp.Header.Get("X-Execution-Time-Ms"), 10, 64) + return RunMeta{ + QueryRunID: resp.Header.Get("X-Query-Run-Id"), + ResultID: resp.Header.Get("X-Result-Id"), + ExecutionTimeMS: execMS, + }, resp.Body, nil } - defer func() { _ = resp.Body.Close() }() var qr queryResponse - if err := json.NewDecoder(resp.Body).Decode(&qr); err != nil { - return RunMeta{}, fmt.Errorf("decoding query response: %w", err) + err = json.NewDecoder(resp.Body).Decode(&qr) + _ = resp.Body.Close() + if err != nil { + return RunMeta{}, nil, fmt.Errorf("decoding query response: %w", err) + } + meta := RunMeta{QueryRunID: qr.QueryRunID, ResultID: qr.ResultID, ExecutionTimeMS: qr.ExecutionTimeMS} + if id := resp.Header.Get("X-Result-Id"); id != "" { + meta.ResultID = id } - if resp.StatusCode == http.StatusOK { - return RunMeta{ - QueryRunID: qr.QueryRunID, - ResultID: qr.ResultID, - RowCount: qr.RowCount, - ExecutionTimeMS: qr.ExecutionTimeMS, - }, nil + if resp.StatusCode == http.StatusAccepted && meta.ResultID == "" { + if meta, err = c.pollRun(ctx, qr.QueryRunID, databaseID); err != nil { + return RunMeta{}, nil, err + } + } + if meta.ResultID == "" { + msg := qr.Warning + if msg == "" { + msg = "no result id returned" + } + return RunMeta{}, nil, fmt.Errorf("query result unavailable: %s", msg) } - // 202: poll /v1/query-runs/{id} until terminal. - return c.pollRun(ctx, qr.QueryRunID, databaseID) + rc, err := c.FetchResultArrow(ctx, meta.ResultID, databaseID) + if err != nil { + return RunMeta{}, nil, err + } + return meta, rc, nil } // pollRun polls a query run until terminal. The endpoint is database-scoped, @@ -166,14 +192,10 @@ func (c *Client) pollRun(ctx context.Context, runID, databaseID string) (RunMeta return RunMeta{ QueryRunID: st.ID, ResultID: st.ResultID, - RowCount: st.RowCount, ExecutionTimeMS: st.ExecutionTimeMS, }, nil case "failed", "cancelled": - msg := st.Error - if msg == "" { - msg = st.Message - } + msg := st.ErrorMessage if msg == "" { msg = "query run " + st.Status } @@ -191,16 +213,13 @@ func (c *Client) pollRun(ctx context.Context, runID, databaseID string) (RunMeta // FetchResultArrow streams the complete result set as Arrow IPC. The endpoint // is database-scoped. The caller must close the reader. // -// A freshly created result is not immediately materialized: the endpoint -// answers HTTP 200 with a JSON body {"result_id":…,"status":"processing"} -// rather than blocking. This happens on the fast/sync query path, where the -// fetch can land within milliseconds of the result id being minted. So poll +// A result that is still being persisted answers 202 with a JSON body +// {"result_id":…,"status":"processing"} and a Retry-After header, so poll // until the response is an actual Arrow stream, backing off up to PollMax and -// honoring ctx. +// honoring ctx. A failed result answers 409, which surfaces as an APIError. func (c *Client) FetchResultArrow(ctx context.Context, resultID, databaseID string) (io.ReadCloser, error) { - // The placeholder race lasts milliseconds on the sync path, so start with a - // much shorter backoff than the run-poll cadence — waiting PollInitial here - // would add half a second to every fast query that hits the race. + // Persistence of a just-minted result usually takes milliseconds, so start + // with a much shorter backoff than the run-poll cadence. interval := 50 * time.Millisecond if interval > c.PollInitial { interval = c.PollInitial @@ -220,9 +239,8 @@ func (c *Client) FetchResultArrow(ctx context.Context, resultID, databaseID stri raw, _ := io.ReadAll(io.LimitReader(resp.Body, 8192)) _ = resp.Body.Close() var st struct { - Status string `json:"status"` - Error string `json:"error"` - Message string `json:"message"` + Status string `json:"status"` + ErrorMessage string `json:"error_message"` } if err := json.Unmarshal(raw, &st); err != nil { return nil, fmt.Errorf("unexpected non-arrow result response: %s", strings.TrimSpace(string(raw))) @@ -233,20 +251,21 @@ func (c *Client) FetchResultArrow(ctx context.Context, resultID, databaseID stri case "processing", "pending", "running": // keep polling default: - msg := st.Error - if msg == "" { - msg = st.Message - } + msg := st.ErrorMessage if msg == "" { msg = strings.TrimSpace(string(raw)) } return nil, fmt.Errorf("result not available (status %q): %s", st.Status, msg) } + wait := retryAfterDelay(resp.Header.Get("Retry-After"), interval) + if wait > c.PollMax { + wait = c.PollMax + } select { case <-ctx.Done(): return nil, ctx.Err() - case <-time.After(interval): + case <-time.After(wait): } if interval = interval * 3 / 2; interval > c.PollMax { interval = c.PollMax @@ -258,19 +277,62 @@ func isArrow(contentType string) bool { return strings.Contains(contentType, "arrow") } +// maxPages bounds every paginated listing so a server that keeps answering +// has_more can never loop forever. +const maxPages = 100 + +// page is the cursor envelope shared by the paginated listings. +type page struct { + HasMore bool `json:"has_more"` + NextCursor string `json:"next_cursor"` +} + +// listPages fetches path page by page, following next_cursor. decode reads +// one page's items and returns the page envelope. +func (c *Client) listPages(ctx context.Context, path string, q url.Values, decode func(io.Reader) (page, error)) error { + for i := 0; i < maxPages; i++ { + resp, err := c.do(ctx, http.MethodGet, path+"?"+q.Encode(), "", nil) + if err != nil { + return err + } + p, err := decode(resp.Body) + _ = resp.Body.Close() + if err != nil { + return err + } + if !p.HasMore || p.NextCursor == "" { + return nil + } + q.Set("cursor", p.NextCursor) + } + return nil +} + +// ListDatabases returns every database in the workspace, following pages. func (c *Client) ListDatabases(ctx context.Context) ([]Database, error) { - resp, err := c.do(ctx, http.MethodGet, "/v1/databases", "", nil) + var all []Database + err := c.listPages(ctx, "/v1/databases", url.Values{"limit": {"100"}}, func(r io.Reader) (page, error) { + var out struct { + page + Databases []Database `json:"databases"` + } + if err := json.NewDecoder(r).Decode(&out); err != nil { + return page{}, fmt.Errorf("decoding databases: %w", err) + } + all = append(all, out.Databases...) + return out.page, nil + }) + return all, err +} + +// Ping makes the cheapest workspace-scoped call, validating the API key and +// workspace id. +func (c *Client) Ping(ctx context.Context) error { + resp, err := c.do(ctx, http.MethodGet, "/v1/databases?limit=1", "", nil) if err != nil { - return nil, err - } - defer func() { _ = resp.Body.Close() }() - var out struct { - Databases []Database `json:"databases"` - } - if err := json.NewDecoder(resp.Body).Decode(&out); err != nil { - return nil, fmt.Errorf("decoding databases: %w", err) + return err } - return out.Databases, nil + return resp.Body.Close() } // GetDatabase resolves one database, including the default_connection_id that @@ -288,8 +350,9 @@ func (c *Client) GetDatabase(ctx context.Context, id string) (DatabaseDetail, er return out, nil } -// InformationSchema lists tables for a connection; schema/table narrow the -// listing and includeColumns adds column definitions with Arrow type names. +// InformationSchema lists tables for a connection, following pages; +// schema/table narrow the listing and includeColumns adds column definitions +// with Arrow type names. func (c *Client) InformationSchema(ctx context.Context, connectionID, schema, table string, includeColumns bool) ([]Table, error) { q := url.Values{"connection_id": {connectionID}} if schema != "" { @@ -301,18 +364,19 @@ func (c *Client) InformationSchema(ctx context.Context, connectionID, schema, ta if includeColumns { q.Set("include_columns", "true") } - resp, err := c.do(ctx, http.MethodGet, "/v1/information_schema?"+q.Encode(), "", nil) - if err != nil { - return nil, err - } - defer func() { _ = resp.Body.Close() }() - var out struct { - Tables []Table `json:"tables"` - } - if err := json.NewDecoder(resp.Body).Decode(&out); err != nil { - return nil, fmt.Errorf("decoding information schema: %w", err) - } - return out.Tables, nil + var all []Table + err := c.listPages(ctx, "/v1/information_schema", q, func(r io.Reader) (page, error) { + var out struct { + page + Tables []Table `json:"tables"` + } + if err := json.NewDecoder(r).Decode(&out); err != nil { + return page{}, fmt.Errorf("decoding information schema: %w", err) + } + all = append(all, out.Tables...) + return out.page, nil + }) + return all, err } func (c *Client) ListWorkspaces(ctx context.Context) ([]Workspace, error) { @@ -330,10 +394,14 @@ func (c *Client) ListWorkspaces(ctx context.Context) ([]Workspace, error) { return out.Workspaces, nil } -// do issues one request with auth headers, retrying on 429 (honoring +func (c *Client) do(ctx context.Context, method, path, databaseID string, body any) (*http.Response, error) { + return c.doAccept(ctx, method, path, databaseID, "", body) +} + +// doAccept issues one request with auth headers, retrying on 429 (honoring // Retry-After, capped at 3 attempts). Non-2xx responses become errors with a // body snippet, since Hotdata error payloads carry stable codes worth showing. -func (c *Client) do(ctx context.Context, method, path, databaseID string, body any) (*http.Response, error) { +func (c *Client) doAccept(ctx context.Context, method, path, databaseID, accept string, body any) (*http.Response, error) { var payload []byte if body != nil { var err error @@ -359,6 +427,9 @@ func (c *Client) do(ctx context.Context, method, path, databaseID string, body a if payload != nil { req.Header.Set("Content-Type", "application/json") } + if accept != "" { + req.Header.Set("Accept", accept) + } resp, err := c.hc.Do(req) if err != nil { @@ -422,5 +493,35 @@ type APIError struct { } func (e *APIError) Error() string { - return fmt.Sprintf("hotdata API: HTTP %d: %s", e.StatusCode, e.Body) + return fmt.Sprintf("hotdata API: HTTP %d: %s", e.StatusCode, e.message()) +} + +// message extracts the readable part of a Hotdata error body, which is one of +// {"error":{"code","message"}}, {"error":"code"} or, for a failed result, +// {"error_message":…}. Anything else is shown raw. +func (e *APIError) message() string { + var b struct { + Error json.RawMessage `json:"error"` + ErrorMessage string `json:"error_message"` + } + if json.Unmarshal([]byte(e.Body), &b) == nil { + var detail struct { + Code string `json:"code"` + Message string `json:"message"` + } + if json.Unmarshal(b.Error, &detail) == nil && detail.Message != "" { + if detail.Code != "" { + return detail.Message + " (" + detail.Code + ")" + } + return detail.Message + } + var code string + if json.Unmarshal(b.Error, &code) == nil && code != "" { + return code + } + if b.ErrorMessage != "" { + return b.ErrorMessage + } + } + return strings.TrimSpace(e.Body) } diff --git a/pkg/plugin/arrow.go b/pkg/plugin/arrow.go index c6d27ad..1559f2a 100644 --- a/pkg/plugin/arrow.go +++ b/pkg/plugin/arrow.go @@ -19,14 +19,13 @@ import ( // truncation to the user. const MaxResultRows = 1_000_000 -// FrameFromArrowStream decodes an Arrow IPC stream (as served by -// GET /v1/results/{id}?format=arrow) into a single Grafana data frame. The +// FrameFromArrowStream decodes an Arrow IPC stream (as served by POST /v1/query +// with an Arrow Accept header, or GET /v1/results/{id}?format=arrow) into a single Grafana data frame. The // boolean reports whether the result was truncated at MaxResultRows. // // Hotdata's engine reports Arrow types directly (Utf8View, Int64, Float64, // Timestamp, ...), so this mapping is the plugin's single source of type -// truth — the inline JSON preview from POST /v1/query is untyped and never -// used for data. +// truth — the JSON form of a result is untyped and never used for data. func FrameFromArrowStream(r io.Reader) (*data.Frame, bool, error) { rdr, err := ipc.NewReader(r) if err != nil { diff --git a/pkg/plugin/datasource.go b/pkg/plugin/datasource.go index ad54193..5e5e332 100644 --- a/pkg/plugin/datasource.go +++ b/pkg/plugin/datasource.go @@ -164,12 +164,7 @@ func (d *Datasource) query(ctx context.Context, query backend.DataQuery) backend sql := interpolateMacros(qm.RawSql, query.TimeRange, query.Interval) - meta, err := d.client.RunQuery(ctx, dbID, sql, dialect) - if err != nil { - return errResponse(err) - } - - body, err := d.client.FetchResultArrow(ctx, meta.ResultID, dbID) + meta, body, err := d.client.Query(ctx, dbID, sql, dialect) if err != nil { return errResponse(err) } @@ -246,7 +241,7 @@ func (d *Datasource) CheckHealth(ctx context.Context, req *backend.CheckHealthRe }, nil } - if _, err := d.client.ListDatabases(ctx); err != nil { + if err := d.client.Ping(ctx); err != nil { var apiErr *hotdata.APIError msg := fmt.Sprintf("Could not reach Hotdata: %v", err) if errors.As(err, &apiErr) && (apiErr.StatusCode == http.StatusUnauthorized || apiErr.StatusCode == http.StatusForbidden) { @@ -268,12 +263,14 @@ func (d *Datasource) CheckHealth(ctx context.Context, req *backend.CheckHealthRe qctx, cancel := context.WithTimeout(ctx, 35*time.Second) defer cancel() start := time.Now() - if _, err := d.client.RunQuery(qctx, d.settings.DefaultDatabaseID, "SELECT 1", ""); err != nil { + _, body, err := d.client.Query(qctx, d.settings.DefaultDatabaseID, "SELECT 1", "") + if err != nil { return &backend.CheckHealthResult{ Status: backend.HealthStatusError, Message: fmt.Sprintf("Connected, but test query failed: %v", err), }, nil } + _ = body.Close() msg := "Connected to Hotdata; test query succeeded." if elapsed := time.Since(start); elapsed > 5*time.Second { diff --git a/pkg/plugin/datasource_test.go b/pkg/plugin/datasource_test.go index 1e3b81a..bf1d374 100644 --- a/pkg/plugin/datasource_test.go +++ b/pkg/plugin/datasource_test.go @@ -76,8 +76,11 @@ func arrowFixture(t *testing.T) []byte { type serverOpts struct { async bool pollsNeeded int32 + runFails bool // the async run ends "failed" with an error_message + tooLarge bool // POST answers 202 + X-Result-Id: result too large to inline + jsonOnly bool // POST ignores the Arrow Accept header and answers JSON rateLimit bool - resultNotReady int32 // fetches that return a JSON "processing" placeholder first + resultNotReady int32 // fetches that return a 202 "processing" placeholder first postUnavailable bool // POST /v1/query answers 503 (never retryable: could re-run SQL) postCount *atomic.Int32 // when set, counts POST /v1/query requests } @@ -124,12 +127,27 @@ func newContractServer(t *testing.T, opts serverOpts) (*httptest.Server, *atomic _, _ = w.Write([]byte(`{"error":{"message":"async_after_ms must be at least 1000","code":"BAD_REQUEST"}}`)) return } + if opts.tooLarge { + w.Header().Set("X-Result-Id", "rslt1") + w.WriteHeader(http.StatusAccepted) + _, _ = w.Write([]byte(`{"query_run_id":"qrun1","status":"running","status_url":"/v1/query-runs/qrun1"}`)) + return + } if opts.async { w.WriteHeader(http.StatusAccepted) _, _ = w.Write([]byte(`{"query_run_id":"qrun1","status":"running","status_url":"/v1/query-runs/qrun1"}`)) return } - _, _ = w.Write([]byte(`{"query_run_id":"qrun1","result_id":"rslt1","columns":["ts","n","f","s","b"],"rows":[],"row_count":2,"truncated":false,"execution_time_ms":2}`)) + if opts.jsonOnly || r.Header.Get("Accept") != "application/vnd.apache.arrow.stream" { + _, _ = w.Write([]byte(`{"query_run_id":"qrun1","result_id":"rslt1","columns":["ts","n","f","s","b"],"rows":[],"preview_row_count":2,"total_row_count":2,"truncated":false,"execution_time_ms":2}`)) + return + } + w.Header().Set("Content-Type", "application/vnd.apache.arrow.stream") + w.Header().Set("X-Query-Run-Id", "qrun1") + w.Header().Set("X-Result-Id", "rslt1") + w.Header().Set("X-Total-Row-Count", "2") + w.Header().Set("X-Execution-Time-Ms", "2") + _, _ = w.Write(arrowBody) }) // Production contract: query-runs and results are database-scoped. requireDB := func(w http.ResponseWriter, r *http.Request) bool { @@ -148,6 +166,10 @@ func newContractServer(t *testing.T, opts serverOpts) (*httptest.Server, *atomic _, _ = w.Write([]byte(`{"id":"qrun1","status":"running"}`)) return } + if opts.runFails { + _, _ = w.Write([]byte(`{"id":"qrun1","status":"failed","error_message":"table 'default.public.t' not found"}`)) + return + } _, _ = w.Write([]byte(`{"id":"qrun1","status":"succeeded","result_id":"rslt1","row_count":2,"execution_time_ms":1500}`)) }) resultFetches := &atomic.Int32{} @@ -159,18 +181,25 @@ func newContractServer(t *testing.T, opts serverOpts) (*httptest.Server, *atomic w.WriteHeader(http.StatusBadRequest) return } - // Model the fast-path race: the result is not materialized yet, so the - // endpoint answers 200 with a JSON placeholder before the Arrow stream. + // The result is not persisted yet: 202 with a JSON placeholder and + // Retry-After before the Arrow stream is served. if resultFetches.Add(1) <= opts.resultNotReady { w.Header().Set("Content-Type", "application/json") + w.Header().Set("Retry-After", "0") + w.WriteHeader(http.StatusAccepted) _, _ = w.Write([]byte(`{"result_id":"rslt1","status":"processing"}`)) return } w.Header().Set("Content-Type", "application/vnd.apache.arrow.stream") _, _ = w.Write(arrowBody) }) + // Listings are paginated: the first page points at a second via next_cursor. mux.HandleFunc("GET /v1/databases", func(w http.ResponseWriter, r *http.Request) { - _, _ = w.Write([]byte(`{"databases":[{"id":"dbid123","name":"test","default_catalog":"default","default_schema":"public"}]}`)) + if r.URL.Query().Get("cursor") == "" { + _, _ = w.Write([]byte(`{"databases":[{"id":"dbid123","name":"test","default_catalog":"default","default_schema":"public"}],"count":1,"has_more":true,"next_cursor":"p2"}`)) + return + } + _, _ = w.Write([]byte(`{"databases":[{"id":"dbid456","name":"other","default_catalog":"default","default_schema":"main"}],"count":1,"has_more":false}`)) }) mux.HandleFunc("GET /v1/databases/dbid123", func(w http.ResponseWriter, r *http.Request) { _, _ = w.Write([]byte(`{"id":"dbid123","name":"test","default_connection_id":"conn123","default_catalog":"default","default_schema":"public"}`)) @@ -184,7 +213,15 @@ func newContractServer(t *testing.T, opts serverOpts) (*httptest.Server, *atomic if r.URL.Query().Get("include_columns") == "true" { cols = `,"columns":[{"name":"city","data_type":"Utf8View","nullable":true}]` } - _, _ = w.Write([]byte(`{"tables":[{"schema":"cities","table":"global","synced":true` + cols + `}]}`)) + if r.URL.Query().Get("table") == "" && r.URL.Query().Get("cursor") == "" { + _, _ = w.Write([]byte(`{"tables":[{"schema":"cities","table":"global","synced":true` + cols + `}],"count":1,"has_more":true,"next_cursor":"p2"}`)) + return + } + if r.URL.Query().Get("cursor") == "p2" { + _, _ = w.Write([]byte(`{"tables":[{"schema":"sales","table":"orders","synced":true}],"count":1,"has_more":false}`)) + return + } + _, _ = w.Write([]byte(`{"tables":[{"schema":"cities","table":"global","synced":true` + cols + `}],"count":1,"has_more":false}`)) }) srv := httptest.NewServer(mux) @@ -310,15 +347,53 @@ func TestQueryDataRetriesOn429(t *testing.T) { assertFixtureFrame(t, res) } -// The fast/sync path can fetch the result before it is materialized; the -// endpoint returns a JSON "processing" placeholder until the Arrow stream is -// ready. The client must poll through it. -func TestQueryDataResultNotReady(t *testing.T) { - srv, _ := newContractServer(t, serverOpts{resultNotReady: 2}) +// A server that answers JSON despite the Arrow Accept header falls back to +// fetching the result, polling through 202 "processing" placeholders until +// the Arrow stream is ready. +func TestQueryDataJSONFallback(t *testing.T) { + srv, _ := newContractServer(t, serverOpts{jsonOnly: true, resultNotReady: 2}) + ds := newTestDatasource(srv) + + res := runQuery(t, ds, `{"rawSql":"SELECT * FROM t","format":"table"}`) + assertFixtureFrame(t, res) +} + +// A result too large to inline answers 202 with X-Result-Id; the client reads +// it from the Results API without polling the run. +func TestQueryDataTooLarge(t *testing.T) { + srv, polls := newContractServer(t, serverOpts{tooLarge: true, resultNotReady: 1}) ds := newTestDatasource(srv) res := runQuery(t, ds, `{"rawSql":"SELECT * FROM t","format":"table"}`) assertFixtureFrame(t, res) + if polls.Load() != 0 { + t.Errorf("too-large flow should not poll the run, polled %d times", polls.Load()) + } +} + +// A failed async run surfaces the server's error_message. +func TestQueryDataRunFailed(t *testing.T) { + srv, _ := newContractServer(t, serverOpts{async: true, runFails: true}) + ds := newTestDatasource(srv) + + res := runQuery(t, ds, `{"rawSql":"SELECT * FROM t","format":"table"}`) + if res.Error == nil || !strings.Contains(res.Error.Error(), "table 'default.public.t' not found") { + t.Fatalf("expected the run's error_message, got %v", res.Error) + } +} + +func TestAPIErrorMessage(t *testing.T) { + for body, want := range map[string]string{ + `{"error":{"code":"NOT_FOUND","message":"Result 'x' not found"}}`: "Result 'x' not found (NOT_FOUND)", + `{"error":"missing_authorization"}`: "missing_authorization", + `{"status":"failed","error_message":"out of memory"}`: "out of memory", + `upstream timeout`: "upstream timeout", + } { + err := &hotdata.APIError{StatusCode: 400, Body: body} + if got := err.Error(); got != "hotdata API: HTTP 400: "+want { + t.Errorf("body %s: got %q, want message %q", body, got, want) + } + } } func TestQueryDataNoDatabase(t *testing.T) { @@ -424,8 +499,23 @@ func TestResources(t *testing.T) { if err := json.Unmarshal(rec.body, &tables); err != nil { t.Fatal(err) } - if len(tables) != 1 || tables[0].Schema != "cities" || tables[0].Table != "global" { - t.Fatalf("unexpected tables: %+v", tables) + if len(tables) != 2 || tables[0].Table != "global" || tables[1].Table != "orders" { + t.Fatalf("expected both pages of tables, got %+v", tables) + } + }) + + t.Run("databases", func(t *testing.T) { + rec := &resourceRecorder{} + err := ds.CallResource(context.Background(), &backend.CallResourceRequest{Path: "databases", URL: "databases"}, rec) + if err != nil || rec.status != 200 { + t.Fatalf("databases: err=%v status=%d body=%s", err, rec.status, rec.body) + } + var dbs []hotdata.Database + if err := json.Unmarshal(rec.body, &dbs); err != nil { + t.Fatal(err) + } + if len(dbs) != 2 || dbs[1].ID != "dbid456" { + t.Fatalf("expected both pages of databases, got %+v", dbs) } }) @@ -450,7 +540,7 @@ func TestResources(t *testing.T) { if err != nil || rec.status != 200 { t.Fatalf("schemas: err=%v status=%d body=%s", err, rec.status, rec.body) } - if string(rec.body) != `["cities"]` { + if string(rec.body) != `["cities","sales"]` { t.Fatalf("unexpected schemas: %s", rec.body) } })