From 8d6607be335211d73a54a48f1fc389b2fef684a4 Mon Sep 17 00:00:00 2001 From: TonyStack Date: Tue, 22 Sep 2026 09:39:35 +0800 Subject: [PATCH 1/2] feat: download an MCQA sub query's result over the instance tunnel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit M-B's B3: `SQLExecutor(INTERACTIVE)` reads a sub query through the instance tunnel unless the client opts out, and that is the path JDBC's MaxQA mode takes. Without it an interactive consumer silently re-runs its statement as an offline instance, so the session plane could not be accepted end to end — the previous increment had to set `useInstanceTunnel(false)` in its own acceptance test to avoid being fooled by that fallback. - `GET .../instances/{id}?data&cached&taskname=..[&queryid=N]` answers one sub query as a Protobuf record stream. There is no download session behind it, so the stream carries its own schema frame (field 1 JSON, then SCHEMA_END_TAG 33553920 with the CRC32C of the field number and payload); the records after it are byte-identical to the ones the offline download path writes. - `odps-tunnel-record-count` reports the rows still available from the read's start offset, not the size of the result: `SessionRecordSetIterator` keeps the first answer's count and pages `openRecordReader(offset+cursor, fetchSize)` while the cursor is below it, so reporting the total sends a client reading past the end. - `rowrange` windows, `sizelimit` shortens a response by halving (at least one row, so a reader always advances) and `instance_tunnel_limit_enabled` applies READ_TABLE_MAX_ROW to that read only. - A sub query that is still running makes the read wait — long polling is what `cached` asks for — and a wait past the bound answers with the SDK's own retryable tunnel timeout instead of an empty stream. Failed, cancelled, unknown and dropped sub queries, and one-shot instances, each get their own structured JSON error; a statement with no result set is refused with `InstanceTypeNotSupported`, which is how the service sends the client to the session API instead of an empty download. - Sub queries now retain their typed result next to the CSV text, inside the same retention window and the same instance byte budget, so both read paths of one statement serve the same rows and the same types. Evidence: 10 new Go cases (framing and CRC checked by decoding the response with our own record decoder, paging, size limit, the 10k cap, non-select, structured failures, long poll, idle-TTL activity, retention) — `go test -race ./internal/...` green. Against a locally built binary `McqaSessionTest` is 6/6, including two new SDK cases: the default fetch, and an offset read that pages the same result; both fail on the previous build with `NoSuchObject: unsupported endpoint`, so they are not passing by falling back. `EmulatorTest` 26/26 and `UdfMetadataTest` 4/4 stay green, and the PyODPS contract probe stays at 30/0 on a fresh process. --- CHANGELOG.md | 1 + README.md | 4 +- docs/data-transfer.md | 1 + docs/protocol.md | 3 + internal/server/mcqa.go | 13 +- internal/server/mcqa_download.go | 257 +++++++++ internal/server/mcqa_download_test.go | 499 ++++++++++++++++++ internal/server/server.go | 13 +- internal/wire/rows.go | 46 ++ tests/java/src/test/java/McqaSessionTest.java | 78 ++- 10 files changed, 904 insertions(+), 11 deletions(-) create mode 100644 internal/server/mcqa_download.go create mode 100644 internal/server/mcqa_download_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index f0d5ca9..93e2f3f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,7 @@ - Require write grants to create or update resources and functions under strict authentication; the Tunnel read-session `create` exception no longer applies to metadata verbs. - Add MCQA / SQLRT interactive sessions: an `SQLRT` task instance stays `Running` until the client stops it or it idles out, and its statements run as sub queries over the instance information KV (`?info`) using the Java SDK's own object status codes, with `query`/`cancel` writes and `status`/`progress`/`result`/`result_` reads. - Declare `mcqa` in `/capabilities`, and name the remaining session gaps (named-session attach, MaxQA v2) separately from the surface that works. +- Download an MCQA sub query's result over the instance tunnel (`?data&cached&taskname=..&queryid=..`): a record stream that carries its own schema, the `odps-tunnel-record-count` the SDK pages on, `rowrange` paging, `sizelimit` truncation and the `READ_TABLE_MAX_ROW` cap, with `InstanceTypeNotSupported` for statements that have no result set. This is the Java SDK's default interactive fetch and JDBC MaxQA's read path. ## 1.1.0 — 2026-09-17 diff --git a/README.md b/README.md index aa07f61..f85de89 100644 --- a/README.md +++ b/README.md @@ -88,7 +88,7 @@ session plane; `interactiveMode=maxqa` is refused by name. | --- | --- | | SQL | CREATE/DROP/TRUNCATE TABLE, INSERT INTO/OVERWRITE, SELECT/WITH; ODPS grammar validation with DuckDB execution | | Metadata | Table identity, schema, primary keys and properties; table listing; STRING partition create/delete/exists/list/pagination; persistent empty partitions | -| MCQA session (SQLRT) | `Instance/Job/Tasks/SQLRT` creates a session instance that stays `Running` until the client stops it or it idles out; statements arrive as sub queries over the instance information KV (`?info&taskname=...`), with `status` / `progress` / `result_` reads, `query` / `cancel` writes, and the object status codes the Java SDK polls on | +| MCQA session (SQLRT) | `Instance/Job/Tasks/SQLRT` creates a session instance that stays `Running` until the client stops it or it idles out; statements arrive as sub queries over the instance information KV (`?info&taskname=...`), with `status` / `progress` / `result_` reads, `query` / `cancel` writes, and the object status codes the Java SDK polls on. A sub query's result is also downloadable over the instance tunnel (`?data&cached&taskname=...&queryid=N`), which is the Java SDK's default fetch and JDBC MaxQA's read path | | Resources and functions | File-like resource upload (single payload or Java SDK part + merge), metadata read, download with `rOffset`/`rSize`, update, delete, prefix/paginated listing; TABLE resource metadata; Java/SQL/embedded function registration referencing existing resources | | Tunnel | Protobuf and Arrow batch upload/download; stream upload; Upsert/delete/partial updates; instance result download | | Storage API v2 | Arrow read/write, Batch/BatchCompatible/Streaming/StreamingRealtime sessions, commit/abort, projections and partition selection | @@ -101,7 +101,7 @@ See [data transfer](docs/data-transfer.md) and [protocol details](docs/protocol. This is a local/CI test service. By default authentication signatures and permissions are **not validated**; optional strict mode uses local test credentials; keep it on a trusted test network. It is not a replacement for real MaxCompute acceptance tests. -MCQA sessions cover the Java SDK's information-channel path (`SQLExecutor` with `useInstanceTunnel(false)`): sub-query answers are CSV whose first line is the column-name header, and session-scoped state (variables, temp objects) is not isolated — statements run against the shared engine. Not simulated: named-session attach (a session name is metadata; equal names create separate sessions), MaxQA v2 (`/mcqa` request prefix), instance-tunnel download of sub-query results, and per-statement statistics. +MCQA sessions cover both of the Java SDK's read paths: the information channel (`SQLExecutor` with `useInstanceTunnel(false)`) answers with CSV whose first line is the column-name header, and the instance tunnel (`?data&cached&taskname=..&queryid=..`, the default fetch and the one JDBC MaxQA uses) answers with a record stream that carries its own schema, so a sub-query result arrives with the same rows and the same types either way. `instance_tunnel_limit_enabled` applies the service's `READ_TABLE_MAX_ROW` (10,000 rows) to that read, `rowrange` pages through it, and `sizelimit` shortens a response instead of failing it. Statements that produce no result set are refused on the tunnel with `InstanceTypeNotSupported`, like the service, so a client confirms them through the session API rather than reading an empty stream. Session-scoped state (variables, temp objects) is not isolated — statements run against the shared engine. Not simulated: named-session attach (a session name is metadata; equal names create separate sessions), MaxQA v2 (`/mcqa` request prefix), compressed or Arrow-encoded sub-query downloads (Protobuf records only), and per-statement statistics. Unsupported: Storage v1, Volume/Blob, CDC/incremental reads, filter predicate pushdown, explicit Schema management, column schema evolution, UDF **execution** (functions are registered metadata; SQL `CREATE FUNCTION`/`DROP FUNCTION` and calling a UDF return `UnsupportedFeature`), volume-backed resources, distributed scheduling, and full ODPS SQL semantics. Resource payloads are capped at 64 MiB each and 512 MiB per project/schema. Resource and function names resolve case-insensitively while the uploaded spelling is what listings return. Unsupported operations return errors rather than cloud behavior being assumed. diff --git a/docs/data-transfer.md b/docs/data-transfer.md index 0db55cd..e692ece 100644 --- a/docs/data-transfer.md +++ b/docs/data-transfer.md @@ -11,6 +11,7 @@ | StreamUploadSession | Protobuf pack | flush 即可见、动态/静态分区、retry trace-id 去重 | | UpsertSession | Protobuf U/D | PRIMARY KEY、按键替换/删除、部分列更新、提交/中止、重试去重 | | InstanceTunnel | Protobuf；HTTP Arrow 格式复用 | SQL 结果快照；Java InstanceTunnel 的 RecordReader 验收 | +| InstanceTunnel MCQA 直读 | Protobuf（schema 随流） | 子查询按 queryid 读取、`rowrange` 分页、`sizelimit` 裁剪、`instance_tunnel_limit_enabled` 应用 10000 行上限、非 select 返回 InstanceTypeNotSupported；不做压缩与 Arrow | | Storage 表读取 | Arrow IPC | read session/get、Index 或 Offset/Count、MaxBatchRows、投影、分区列表、preview | | Storage 实例读取 | Arrow IPC | create/get session、Offset/Count 结果读取 | | Storage Batch | Arrow IPC | stream 创建/恢复/version、flush 暂存、close、选择 streams commit、abort、静态分区 overwrite | diff --git a/docs/protocol.md b/docs/protocol.md index 28abfad..c5f47b2 100644 --- a/docs/protocol.md +++ b/docs/protocol.md @@ -33,6 +33,8 @@ JDBC 驱动对每条语句都会先签一个 logview token，再为实例建立 ## 编码 +MCQA 子查询直读没有 download session 可查，schema 只能随流走：帧头是 field 1 的 JSON schema（TunnelTableSchema 的 `columns`/`partitionKeys` 形状），后接 SCHEMA_END_TAG 33553920 与该帧自己的 CRC32C（覆盖 field number 的 4 字节小端值与 JSON 字节，不含 tag 与长度 varint），校验和后 CRC 归零，其后的记录部分与下面的 Protobuf 完全一致。响应头 `odps-tunnel-record-count` 是从本次 `rowrange` 起点起还能读到的行数，不是全量结果：Java 的 SessionRecordSetIterator 只取第一次响应的这个数并据此翻页，报大全量会让客户端越过结果末尾继续读取。查询仍在执行时该读取阻塞等待（长轮询），超时返回 Java SDK 认得的 `OdpsTaskTimeout`/“Wait for cache data timeout”，而不是先给一个空流。 + Protobuf 按 public SDK 的列序编码，NULL 不写字段。记录结束 tag 33553408 后写 CRC32C，流尾 meta-count 33554430 与 checksum-of-checksums 33554431；CRC 数值按小端逻辑值更新，不包含传输长度和 NULL 标记。DATE 为天数，DATETIME 为毫秒，TIMESTAMP 为秒+纳秒，DECIMAL 为十进制文本。MAP 在引擎边界从 DuckDB OrderedMap 转换，ARRAY/STRUCT 递归编码。 Tunnel Arrow 不是普通带 Schema 消息的 IPC stream：客户端从会话元数据构造 schema，服务端只发送 RecordBatch payload。帧头为大端 chunk size（65536），每满块追加 CRC32C，最后追加全流 CRC32C。ZSTD 和 LZ4 frame 压缩 Arrow IPC buffer，外层 Tunnel CRC 帧保持未压缩；deflate 仍包裹整个 Tunnel 帧。空表和 count=0 返回合法空数据帧。常规 Arrow 请求每批最多4096行；带 raw_size 的请求最多一批，CPP BufferArrowRecordReader 按实际行数推进；raw_size 过小时仍返回至少一行，以保证进度。 @@ -58,6 +60,7 @@ Tunnel Arrow 不是普通带 Schema 消息的 IPC stream：客户端从会话元 | POST/GET/PUT `/projects/p/tables/t/streams` | 创建/恢复/写 pack；flush 后可见；dynamic_partition 支持逐 pack 分区 | | POST/GET/PUT/POST/DELETE `/projects/p/tables/t/upserts` | 创建/恢复/暂存 U/D/提交/中止；支持部分列更新 | | POST/GET `/projects/p/instances/i?downloads` / `?downloadid=id` | SQL 实例结果快照与 Tunnel 下载，复用表读取格式 | +| GET `/projects/p/instances/i?data&cached&taskname=t[&queryid=n][&rowrange=(s,c)][&sizelimit=n][&instance_tunnel_limit_enabled]` | MCQA 子查询结果直读（Java SDK 默认交互取数、JDBC MaxQA 路径）：自带 schema 的 Protobuf 记录流，只支持未压缩 | | POST `/api/storage/v2` 或 `/api/storage/v3` | Action + Target 协议；Java 0.61.2-public 实际调用 v3 | Storage 表 Action：TableCreateReadSession、TableGetReadSession、TableRead、TablePreview、TableCreateWriteSession、TableGetWriteSession、TableCreateWriteStream、TableGetWriteStream、TableWrite、TableCloseWriteStream、TableCommitWriteSession、TableAbortWriteSession。实例 Action：InstanceCreateReadSession、InstanceGetReadSession（GET）、InstanceRead。 diff --git a/internal/server/mcqa.go b/internal/server/mcqa.go index 0fe711e..0dd38f8 100644 --- a/internal/server/mcqa.go +++ b/internal/server/mcqa.go @@ -59,6 +59,12 @@ type mcqaQuery struct { Warnings string ID int Created time.Time + // Data is the typed result behind the CSV text of the information channel. The + // instance-tunnel read serves it as records, so a consumer gets the same values and + // the same types through either path. It is retained for the same window as the CSV + // text — the mcqaRetainedQueries entries dropped by mcqaSession.drop, and counted by + // bytes into the instance budget. + Data engine.Result } // mcqaSession is the state behind a SQLRT instance. Fields are guarded by Server.mu; @@ -115,7 +121,7 @@ func (m *mcqaSession) bytes() int64 { } var total int64 for _, q := range m.queries { - total += int64(len(q.Result) + len(q.Query)) + total += q.Data.Bytes + int64(len(q.Result)+len(q.Query)) } return total } @@ -360,13 +366,16 @@ func (s *Server) runSubQuery(ctx context.Context, i *instance, id int, query str q.Result = "Sub query cancelled" default: q.Status = mcqaStatusTerminated + q.Data = res q.Result = sessionResultCSV(res) } } // sessionResultCSV renders a result the way the SQLRT information channel does: the // first line is the column-name header, because CSVRecordParser#parse reads that line -// to build the schema. Values are formatted like the offline instance result. +// to build the schema. Values are formatted like the offline instance result, and every +// row of the result is rendered — the tunnel read of the same sub query serves the same +// rows, so neither path can look truncated relative to the other. func sessionResultCSV(res engine.Result) string { header := make([]string, len(res.Columns)) for n, c := range res.Columns { diff --git a/internal/server/mcqa_download.go b/internal/server/mcqa_download.go new file mode 100644 index 0000000..ed66fd1 --- /dev/null +++ b/internal/server/mcqa_download.go @@ -0,0 +1,257 @@ +package server + +// MCQA sub-query results are also downloadable through the instance tunnel, which is the +// default way the Java SDK's interactive executor reads them: +// +// SQLExecutorImpl#getSessionResultSetByInstanceTunnel +// -> InstanceTunnel#createDirectDownloadSession(project, instance, taskname, queryid, limit) +// -> TunnelRecordReader -> GET /projects/{p}/instances/{id} +// ?data&cached&schema_in_stream&taskname=..[&queryid=N] +// [&rowrange=(start,count)][&sizelimit=bytes][&instance_tunnel_limit_enabled] +// +// Two things differ from an offline instance download: +// +// - There is no download session to create first, so nothing else carries the schema. +// The body starts with the in-stream JSON schema frame that +// ProtobufRecordStreamReader#readSchema reads (the Java SDK labels that method "for +// MCQA direct download"). +// - `odps-tunnel-record-count` is the number of rows the read makes available from its +// start offset, because that is what SessionRecordSetIterator iterates: it remembers +// the first answer's count and keeps calling openRecordReader(offset+cursor, +// fetchSize) while cursor < it. Reporting the whole result there makes a client that +// asked for an offset read past the end of it. + +import ( + "context" + "fmt" + "net/http" + "strconv" + "strings" + "time" + + "github.com/dingxin-tech/maxcompute-emulator/internal/engine" + "github.com/dingxin-tech/maxcompute-emulator/internal/wire" +) + +const ( + // mcqaLimitRows is the service's READ_TABLE_MAX_ROW, which a client asks the tunnel + // to apply by naming the download limited (instance_tunnel_limit_enabled). + mcqaLimitRows = 10000 + // mcqaDirectWait bounds how long a direct read waits for a sub query that is still + // running. Clients normally poll the information channel first; one that reads right + // away is long polling on purpose, so the answer blocks instead of serving an empty + // stream. The code and message are the pair the Java SDK treats as a retryable + // tunnel timeout (SQLExecutorConstants.sessionTunnelTimeout*). + mcqaDirectWait = 30 * time.Second + // mcqaDirectPoll is how often a waiting read re-checks the sub query. Statements run + // with the server lock released, so the wait must not hold it. + mcqaDirectPoll = 5 * time.Millisecond +) + +// mcqaRead is the outcome of one attempt to resolve a direct-download request. +type mcqaRead struct { + Result engine.Result + Status int + Code string + Message string + Pending bool // sub query still running: wait, then retry or time out +} + +// failed reports a structured answer rather than a result stream. +func (m mcqaRead) failed() bool { return m.Code != "" } + +// mcqaDirectDownload answers a session sub-query download with a record stream. +func (s *Server) mcqaDirectDownload(w http.ResponseWriter, r *http.Request, p, sc, iid string) { + q := r.URL.Query() + task := q.Get("taskname") + if task == "" { + fail(w, r, 400, "InvalidParameter", fmt.Errorf("taskname is required to download a session sub query result")) + return + } + subQueryID := -1 + if raw := q.Get("queryid"); raw != "" { + v, e := strconv.Atoi(raw) + if e != nil { + fail(w, r, 400, "InvalidParameter", fmt.Errorf("invalid queryid %q", raw)) + return + } + subQueryID = v + } + start, count, e := parseRowRange(q.Get("rowrange")) + if e != nil { + fail(w, r, 400, "InvalidParameter", e) + return + } + sizeLimit := int64(0) + if raw := q.Get("sizelimit"); raw != "" { + sizeLimit, e = strconv.ParseInt(raw, 10, 64) + if e != nil || sizeLimit < 0 { + fail(w, r, 400, "InvalidParameter", fmt.Errorf("invalid sizelimit %q", raw)) + return + } + } + limited := q.Has("instance_tunnel_limit_enabled") + query := "the current sub query" + if subQueryID >= 0 { + query = "sub query " + strconv.Itoa(subQueryID) + } + deadline := time.Now().Add(mcqaDirectWait) + for { + res := s.mcqaResolveRead(p, sc, iid, task, subQueryID, limited) + switch { + case !res.Pending && !res.failed(): + s.writeRecordStream(w, r, res.Result, start, count, sizeLimit) + return + case !res.Pending: + fail(w, r, res.Status, res.Code, fmt.Errorf("%s", res.Message)) + return + } + if time.Now().Add(mcqaDirectPoll).After(deadline) { + fail(w, r, 500, "OdpsTaskTimeout", fmt.Errorf("Wait for cache data timeout: %s of session %s is still running", query, iid)) + return + } + if !sleepUntil(r.Context(), mcqaDirectPoll) { + return // the client went away; there is nobody left to answer + } + } +} + +// mcqaResolveRead looks up the sub query a download names and returns the rows it may +// serve. It runs under the server lock; statement execution takes the same lock to store +// its answer, so callers never hold it across a wait. +func (s *Server) mcqaResolveRead(p, sc, iid, task string, subQueryID int, limited bool) mcqaRead { + s.mu.Lock() + defer s.mu.Unlock() + i := s.instances[iid] + if i == nil || i.Project != p || i.Schema != sc { + return mcqaRead{Status: 404, Code: "NoSuchInstance", Message: "unknown instance"} + } + m := i.MCQA + if m == nil { + // An offline instance has no sub queries. Saying so is what lets the SDK use its + // documented fallback instead of parsing an error out of a result stream. + return mcqaRead{Status: 409, Code: "InvalidState", Message: "instance is not an MCQA session; download its result with ?downloads"} + } + if s.mcqaIdle(i) { + delete(s.instances, i.ID) + return mcqaRead{Status: 404, Code: "NoSuchInstance", Message: "session " + iid + " has expired"} + } + if task != i.Name { + return mcqaRead{Status: 400, Code: "InvalidParameter", Message: fmt.Sprintf("session %s has no task %s", iid, task)} + } + // Reading a result is session activity, like the information calls are. + m.lastActive = time.Now() + id := m.lastID // no queryid means "the current sub query" + if subQueryID >= 0 { + id = subQueryID + } + q := m.queries[id] + switch { + case q == nil: // never submitted, or dropped by the retention window + return mcqaRead{Status: 404, Code: "NoSuchInstance", Message: fmt.Sprintf("session %s has no sub query %d", iid, id)} + case q.Status == mcqaStatusRunning: + return mcqaRead{Pending: true} + case q.Status == mcqaStatusCancelled: + return mcqaRead{Status: 409, Code: "OdpsJobCancelledException", Message: fmt.Sprintf("Job is cancelled: session %s sub query %d", iid, id)} + case q.Status == mcqaStatusFailed: + // The failure text is the statement error. The SDK answers a tunnel failure for a + // select by re-running it offline, where the same error surfaces again — so this + // must not be dressed up as an empty result. + return mcqaRead{Status: 409, Code: "InvalidState", Message: fmt.Sprintf("session %s sub query %d failed: %s", iid, id, q.Result)} + case q.Status != mcqaStatusTerminated: + return mcqaRead{Status: 409, Code: "InvalidState", Message: fmt.Sprintf("session %s sub query %d is not terminated", iid, id)} + case len(q.Data.Columns) == 0: + // The service will not tunnel a non-select statement. These two strings are the + // ones SQLExecutorImpl#checkIsSelect looks for, after which it confirms the + // outcome through the session API rather than reading records. + return mcqaRead{Status: 400, Code: "InstanceTypeNotSupported", Message: "Non select query not supported by session result download"} + } + count := len(q.Data.Rows) + if limited && count > mcqaLimitRows { + // READ_TABLE_MAX_ROW applies to this read only; the sub query keeps its full + // result, and an unlimited read of it still returns every row. + count = mcqaLimitRows + } + return mcqaRead{Result: engine.Result{Columns: q.Data.Columns, Rows: q.Data.Rows[:count]}} +} + +// writeRecordStream sends the rows of res that a read starting at `start` may take, as a +// schema-first record stream, honouring a byte limit by dropping trailing rows. At least +// one row is sent when the limit is smaller than a single record: a reader that cannot +// advance would otherwise stall. +func (s *Server) writeRecordStream(w http.ResponseWriter, r *http.Request, res engine.Result, start, count, sizeLimit int64) { + total := int64(len(res.Rows)) + if start > total { + fail(w, r, 400, "InvalidParameter", fmt.Errorf("rowrange starts at %d, the result has %d rows", start, total)) + return + } + // The count the client gets is what it iterates, so it has to be the rows this read + // makes available from `start` — the remainder of the result, cut short by the + // requested count. SessionRecordSetIterator keeps that first answer's number and + // pages until it reaches it, so a total here would send it past the end of the + // result looking for rows it was promised. + available := total - start + if count >= 0 && count < available { + available = count + } + out := engine.Result{Columns: res.Columns, Rows: res.Rows[start : start+available]} + data, e := wire.ProtobufSchema(out) + if e != nil { + fail(w, r, 400, "SerializationError", e) + return + } + // Halving keeps the re-encodes logarithmic; dropping one row at a time would let a + // one-byte limit turn a wide result into thousands of encodes. + for n := len(out.Rows); sizeLimit > 0 && int64(len(data)) > sizeLimit && n > 1; { + n = max(1, n/2) + out.Rows = out.Rows[:n] + if data, e = wire.ProtobufSchema(out); e != nil { + fail(w, r, 400, "SerializationError", e) + return + } + } + w.Header().Set("Content-Type", "application/octet-stream") + w.Header().Set("Content-Length", strconv.Itoa(len(data))) + w.Header().Set("odps-tunnel-record-count", strconv.FormatInt(available, 10)) + w.WriteHeader(200) + w.Write(data) +} + +// parseRowRange reads the tunnel rowrange parameter. An empty value means "from the start, +// all rows"; the SDK sends a negative count for "no limit" and Long.MAX_VALUE when it was +// given only a start offset. +func parseRowRange(text string) (start, count int64, e error) { + if text == "" { + return 0, -1, nil + } + if !strings.HasPrefix(text, "(") || !strings.HasSuffix(text, ")") { + return 0, 0, fmt.Errorf("rowrange=(start,count) required") + } + a := strings.Split(strings.TrimSuffix(strings.TrimPrefix(text, "("), ")"), ",") + if len(a) != 2 { + return 0, 0, fmt.Errorf("invalid rowrange %q", text) + } + if start, e = strconv.ParseInt(a[0], 10, 64); e != nil || start < 0 { + return 0, 0, fmt.Errorf("invalid rowrange %q", text) + } + if count, e = strconv.ParseInt(a[1], 10, 64); e != nil { + return 0, 0, fmt.Errorf("invalid rowrange %q", text) + } + if count < 0 { + count = -1 + } + return start, count, nil +} + +// sleepUntil waits for d and reports whether it finished, bailing out early when the +// request is cancelled. +func sleepUntil(ctx context.Context, d time.Duration) bool { + t := time.NewTimer(d) + defer t.Stop() + select { + case <-ctx.Done(): + return false + case <-t.C: + return true + } +} diff --git a/internal/server/mcqa_download_test.go b/internal/server/mcqa_download_test.go new file mode 100644 index 0000000..e8e3ec6 --- /dev/null +++ b/internal/server/mcqa_download_test.go @@ -0,0 +1,499 @@ +package server + +import ( + "bytes" + "encoding/binary" + "encoding/json" + "fmt" + "hash/crc32" + "io" + "net/http" + "strconv" + "strings" + "testing" + "time" + + "github.com/dingxin-tech/maxcompute-emulator/internal/engine" + "github.com/dingxin-tech/maxcompute-emulator/internal/wire" +) + +// downloadPath is the request com.aliyun.odps.tunnel.InstanceTunnel makes for one sub +// query when an interactive executor reads its result: no download id, because there is +// no download session to create first. +func downloadPath(id string, queryID int, extra string) string { + path := "/projects/p/instances/" + id + "?data&cached&schema_in_stream&taskname=" + sessionTask + if queryID >= 0 { + path += "&queryid=" + strconv.Itoa(queryID) + } + if extra != "" { + path += "&" + extra + } + return path +} + +var castagnoli = crc32.MakeTable(crc32.Castagnoli) + +// readSchemaFrame parses the in-stream schema frame the way +// ProtobufRecordStreamReader#readSchema does, checksum included, and returns the columns +// it names plus the record stream that follows. The tag and magic numbers are spelled out +// from the SDK (field 1 / SCHEMA_END_TAG 2^25-512) rather than imported from our own +// encoder, so this fails if the implementation drifts. +func readSchemaFrame(t *testing.T, b []byte) ([]engine.Column, []byte) { + t.Helper() + tag, used := binary.Uvarint(b) + if tag != 1<<3|2 { + t.Fatalf("stream must open with the schema field (tag 0x0a), got tag %d in % x", tag, b[:min(len(b), 6)]) + } + n, nUsed := binary.Uvarint(b[used:]) + raw := b[used+nUsed : used+nUsed+int(n)] + end, endUsed := binary.Uvarint(b[used+nUsed+int(n):]) + if end>>3 != 33553920 { + t.Fatalf("schema end tag=%d want 33553920", end>>3) + } + sum, sumUsed := binary.Uvarint(b[used+nUsed+int(n)+endUsed:]) + // The reader CRCs the field number as a 4-byte little-endian int, then the payload; + // the tag and the length varint are outside the checksum. + crc := crc32.Update(0, castagnoli, binary.LittleEndian.AppendUint32(nil, 1)) + crc = crc32.Update(crc, castagnoli, raw) + if uint32(sum) != crc { + t.Fatalf("schema frame checksum=%d want %d", uint32(sum), crc) + } + var envelope struct { + Columns []engine.Column `json:"columns"` + PartitionKeys []engine.Column `json:"partitionKeys"` + } + if e := json.Unmarshal(raw, &envelope); e != nil { + t.Fatalf("schema frame is not the JSON the SDK parses: %v / %s", e, raw) + } + if len(envelope.PartitionKeys) != 0 { + t.Fatalf("a sub query result carries no partition keys: %s", raw) + } + for i := range envelope.Columns { + parsed, e := engine.ParseType(envelope.Columns[i].Type) + if e != nil { + t.Fatalf("schema column %s has type %q: %v", envelope.Columns[i].Name, envelope.Columns[i].Type, e) + } + envelope.Columns[i].Parsed = parsed + } + return envelope.Columns, b[used+nUsed+int(n)+endUsed+sumUsed:] +} + +// readRecords decodes a direct-download response into rows. wire.DecodeProtobuf re-checks +// every record CRC and the stream footer, so a header that is off by a byte shows up here +// instead of as a silently short result. +func readRecords(t *testing.T, body []byte) (engine.Result, []engine.Column) { + t.Helper() + cols, rest := readSchemaFrame(t, body) + res, e := wire.DecodeProtobuf(rest, cols) + if e != nil { + t.Fatalf("decode records: %v (%d bytes of stream)", e, len(body)) + } + res.Columns = cols + return res, cols +} + +func directDownloadJSONError(t *testing.T, code int, body []byte) (string, string) { + t.Helper() + // A tunnel client parses a JSON error body; the instance XML error shape would leave + // it with no error code to route on. + var e struct{ Code, Message string } + if u := json.Unmarshal(body, &e); u != nil { + t.Fatalf("download error is not tunnel JSON (code=%d): %v / %s", code, u, body) + } + if e.Code == "" { + t.Fatalf("download error has no Code: %s", body) + } + return e.Code, e.Message +} + +// TestMCQADirectDownloadServesTypedRows is the acceptance point for M-B's B3: the same +// statement answered by the session must come back over the instance tunnel with real +// types. The information channel can only ever hand the client strings, so a read that +// falls back to CSV would fail the int64 assertion rather than pass by accident. +func TestMCQADirectDownloadServesTypedRows(t *testing.T) { + _, h := fixture(t, Config{SessionTTL: time.Hour}) + id := createSession(t, h, "") + queryID, _ := submitStatement(t, h, id, "select id, s from t order by id limit 3") + code, header, body := mcqaCall(t, h, "GET", downloadPath(id, queryID, ""), nil) + if code != 200 { + t.Fatalf("direct download: code=%d body=%s", code, body) + } + if got := header.Get("odps-tunnel-record-count"); got != "3" { + t.Fatalf("record count header=%q want 3", got) + } + res, cols := readRecords(t, body) + if len(cols) != 2 || cols[0].Name != "id" || cols[1].Name != "s" { + t.Fatalf("schema columns: %+v", cols) + } + if len(res.Rows) != 3 { + t.Fatalf("rows: %d", len(res.Rows)) + } + if v, ok := res.Rows[0][0].(int64); !ok || v != 0 { + t.Fatalf("id must arrive as a typed bigint, got %#v", res.Rows[0][0]) + } + if v, ok := res.Rows[2][1].(string); !ok || !strings.HasPrefix(v, "valuevalue") { + t.Fatalf("s must arrive as a string, got %#v", res.Rows[2][1]) + } + // Both reads describe the same result, so a client that switches paths cannot see a + // different row count. + csv := getInfo(t, h, infoPath(id, sessionTask, "result_"+strconv.Itoa(queryID))) + if lines := strings.Split(strings.TrimRight(csv.Result, "\r\n"), "\r\n"); len(lines) != 4 { + t.Fatalf("information channel rows and tunnel rows disagree: %q", csv.Result) + } +} + +// TestMCQADirectDownloadPagesWithRowRange pins the paging contract: rowrange selects a +// window, and the count the header reports is what is left to read from that window's start +// (the requested count at most). Those two numbers are all SessionRecordSetIterator has to +// walk a result: report the total for an offset read and it keeps asking for rows that are +// not there, report the window and it stops early. +func TestMCQADirectDownloadPagesWithRowRange(t *testing.T) { + _, h := fixture(t, Config{SessionTTL: time.Hour}) + id := createSession(t, h, "") + queryID, _ := submitStatement(t, h, id, "select id, s from t order by id limit 3") + for _, tc := range []struct { + rangeText string + want []int64 + count string + }{ + {"", []int64{0, 1, 2}, "3"}, + {"rowrange=(0,1)", []int64{0}, "1"}, + {"rowrange=(1,-1)", []int64{1, 2}, "2"}, + {"rowrange=(2,10)", []int64{2}, "1"}, + {"rowrange=(3,10)", []int64{}, "0"}, + } { + code, header, body := mcqaCall(t, h, "GET", downloadPath(id, queryID, tc.rangeText), nil) + if code != 200 { + t.Fatalf("%s: code=%d body=%s", tc.rangeText, code, body) + } + if got := header.Get("odps-tunnel-record-count"); got != tc.count { + t.Fatalf("%s: record count header=%q want %q", tc.rangeText, got, tc.count) + } + res, _ := readRecords(t, body) + var got []int64 + for _, row := range res.Rows { + v, _ := row[0].(int64) + got = append(got, v) + } + if fmt.Sprint(got) != fmt.Sprint(tc.want) { + t.Fatalf("%s: rows=%v want %v", tc.rangeText, got, tc.want) + } + } + // A range that starts past the end is a client error, and a malformed one is too — + // neither may come back as a valid empty stream. + for _, bad := range []string{"rowrange=(4,10)", "rowrange=3,10", "rowrange=(x,1)", "rowrange=(-1,1)"} { + code, body := func() (int, []byte) { + c, _, b := mcqaCall(t, h, "GET", downloadPath(id, queryID, bad), nil) + return c, b + }() + if code != 400 { + t.Fatalf("%s: code=%d body=%s", bad, code, body) + } + if got, _ := directDownloadJSONError(t, code, body); got != "InvalidParameter" { + t.Fatalf("%s: error code=%s want InvalidParameter", bad, got) + } + } +} + +// TestMCQADirectDownloadHonoursSizeLimit keeps a byte-limited read inside the limit: the +// SDK turns an oversized response into a client-side error, so the server has to stop +// early — but never at zero rows, or a reader could not advance. +func TestMCQADirectDownloadHonoursSizeLimit(t *testing.T) { + _, h := fixture(t, Config{SessionTTL: time.Hour}) + id := createSession(t, h, "") + queryID, _ := submitStatement(t, h, id, "select id, s from t order by id limit 50") + _, _, full := mcqaCall(t, h, "GET", downloadPath(id, queryID, ""), nil) + for _, tc := range []struct{ limit, maxRows, minRows int }{ + {len(full), 50, 50}, // the whole stream fits, so nothing is dropped + {len(full) / 2, 49, 1}, // half of it: fewer rows, still progress + {1, 1, 1}, // smaller than one record: one row is sent anyway + } { + code, header, body := mcqaCall(t, h, "GET", downloadPath(id, queryID, "sizelimit="+strconv.Itoa(tc.limit)), nil) + if code != 200 { + t.Fatalf("sizelimit=%d: code=%d body=%s", tc.limit, code, body) + } + // A byte limit shortens this response, not the result: the client is told how many + // rows are still to come and keeps reading. + if got := header.Get("odps-tunnel-record-count"); got != "50" { + t.Fatalf("sizelimit=%d: record count header=%q want the 50 rows from offset 0", tc.limit, got) + } + res, _ := readRecords(t, body) + if len(res.Rows) > tc.maxRows || len(res.Rows) < tc.minRows { + t.Fatalf("sizelimit=%d sent %d rows, want between %d and %d (body=%d bytes)", + tc.limit, len(res.Rows), tc.minRows, tc.maxRows, len(body)) + } + } +} + +// TestMCQADirectDownloadLimitedFlagCapsRows covers READ_TABLE_MAX_ROW. The cap belongs to +// the read that asked for it: the sub query keeps its full result, so an unlimited read of +// the same sub query still returns every row. +func TestMCQADirectDownloadLimitedFlagCapsRows(t *testing.T) { + _, h := fixture(t, Config{SessionTTL: time.Hour}) + id := createSession(t, h, "") + // 2048 fixture rows, doubled three times: 16384, i.e. past the 10000 cap. + for n := 0; n < 3; n++ { + if _, ir := submitStatement(t, h, id, fmt.Sprintf("insert into t select id+%d, s from t", n+1)); ir.Status != mcqaInfoOK { + t.Fatalf("grow the table: %+v", ir) + } + } + queryID, _ := submitStatement(t, h, id, "select id from t") + for _, tc := range []struct { + extra string + want string + rows int + }{{"", "16384", 16384}, {"instance_tunnel_limit_enabled", "10000", 10000}} { //nolint:dupl + code, header, body := mcqaCall(t, h, "GET", downloadPath(id, queryID, tc.extra), nil) + if code != 200 { + t.Fatalf("%q: code=%d body=%s", tc.extra, code, body) + } + if got := header.Get("odps-tunnel-record-count"); got != tc.want { + t.Fatalf("%q: record count header=%q want %q", tc.extra, got, tc.want) + } + res, _ := readRecords(t, body) + if len(res.Rows) != tc.rows { + t.Fatalf("%q: %d rows want %d", tc.extra, len(res.Rows), tc.rows) + } + } + // The paging loop must stop at the cap instead of walking the rest of the result. + code, header, body := mcqaCall(t, h, "GET", downloadPath(id, queryID, "instance_tunnel_limit_enabled&rowrange=(9999,1)"), nil) + if code != 200 { + t.Fatalf("limited read at the cap: code=%d body=%s", code, body) + } + if got := header.Get("odps-tunnel-record-count"); got != "1" { + t.Fatalf("a limited read one row before the cap reports %q rows left", got) + } + res, _ := readRecords(t, body) + if len(res.Rows) != 1 { + t.Fatalf("limited read at row 9999 sent %d rows", len(res.Rows)) + } + // The cap is a hard boundary: one row past it is the empty stream a paging client + // stops on, and a start beyond it is the range error the offline download path gives. + if code, header, body = mcqaCall(t, h, "GET", downloadPath(id, queryID, "instance_tunnel_limit_enabled&rowrange=(10000,1)"), nil); code != 200 { + t.Fatalf("a limited read at the cap: code=%d body=%s", code, body) + } else if got := header.Get("odps-tunnel-record-count"); got != "0" { + t.Fatalf("a limited read at the cap reports %q rows left", got) + } else if res, _ := readRecords(t, body); len(res.Rows) != 0 { + t.Fatalf("a limited read served %d rows past the cap", len(res.Rows)) + } + if code, _, _ = mcqaCall(t, h, "GET", downloadPath(id, queryID, "instance_tunnel_limit_enabled&rowrange=(10001,1)"), nil); code != 400 { + t.Fatalf("a limited read starting past the cap must be a range error") + } +} + +// TestMCQADirectDownloadRefusesNonSelect matches the service, which will not tunnel a +// statement without a result set. The code and message are the pair +// SQLExecutorImpl#checkIsSelect recognises; anything else would be read as a select that +// happened to come back empty. +func TestMCQADirectDownloadRefusesNonSelect(t *testing.T) { + _, h := fixture(t, Config{SessionTTL: time.Hour}) + id := createSession(t, h, "") + for _, statement := range []string{"create table ns(a bigint)", "insert into t values (9001,'x')"} { + queryID, ir := submitStatement(t, h, id, statement) + if ir.Status != mcqaInfoOK { + t.Fatalf("%s: %+v", statement, ir) + } + code, _, body := mcqaCall(t, h, "GET", downloadPath(id, queryID, ""), nil) + if code != 400 { + t.Fatalf("%s: code=%d body=%s", statement, code, body) + } + got, message := directDownloadJSONError(t, code, body) + if got != "InstanceTypeNotSupported" || !strings.Contains(message, "Non select query not supported") { + t.Fatalf("%s: code=%s message=%s", statement, got, message) + } + } + // A select with no rows is not a non-select: it answers with a schema and zero records. + queryID, _ := submitStatement(t, h, id, "select id, s from t where id = 99999") + code, header, body := mcqaCall(t, h, "GET", downloadPath(id, queryID, ""), nil) + if code != 200 || header.Get("odps-tunnel-record-count") != "0" { + t.Fatalf("empty select: code=%d count=%q body=%s", code, header.Get("odps-tunnel-record-count"), body) + } + res, cols := readRecords(t, body) + if len(res.Rows) != 0 || len(cols) != 2 { + t.Fatalf("empty select lost its schema: %d rows, %d columns", len(res.Rows), len(cols)) + } +} + +// TestMCQADirectDownloadFailuresAreStructured keeps every rejection distinguishable: a +// wrong task, an unknown sub query and a one-shot instance are three different problems, +// and an empty 200 would look like "the query returned nothing" to all three. +func TestMCQADirectDownloadFailuresAreStructured(t *testing.T) { + _, h := fixture(t, Config{SessionTTL: time.Hour}) + id := createSession(t, h, "") + queryID, _ := submitStatement(t, h, id, "select id from t limit 1") + wrongTask := "/projects/p/instances/" + id + "?data&cached&taskname=other_task&queryid=" + strconv.Itoa(queryID) + if code, _, body := mcqaCall(t, h, "GET", wrongTask, nil); code != 400 { + t.Fatalf("unknown task name: code=%d body=%s", code, body) + } else if got, _ := directDownloadJSONError(t, code, body); got != "InvalidParameter" { + t.Fatalf("unknown task name code=%s", got) + } + if code, _, body := mcqaCall(t, h, "GET", downloadPath(id, 999, ""), nil); code != 404 { + t.Fatalf("unknown sub query: code=%d body=%s", code, body) + } else if got, message := directDownloadJSONError(t, code, body); got != "NoSuchInstance" || !strings.Contains(message, "999") { + t.Fatalf("unknown sub query: code=%s message=%s", got, message) + } + // A failed statement is not an empty result either: the SDK re-runs a select whose + // tunnel read failed, and the error text is what the user ends up seeing. + failed, _ := submitStatement(t, h, id, "select * from missing_table") + code, _, body := mcqaCall(t, h, "GET", downloadPath(id, failed, ""), nil) + if code != 409 { + t.Fatalf("failed sub query: code=%d body=%s", code, body) + } + if got, message := directDownloadJSONError(t, code, body); got != "InvalidState" || !strings.Contains(message, "missing_table") { + t.Fatalf("failed sub query: code=%s message=%s", got, message) + } + // A one-shot instance has no sub queries; saying InvalidState is what lets the client + // use its documented download path instead. + code, header, body := mcqaCall(t, h, "POST", "/projects/p/instances", + []byte(`oneshotselect 1`)) + if code != 201 { + t.Fatalf("one-shot create: code=%d body=%s", code, body) + } + oneShot := strings.TrimPrefix(header.Get("Location"), "/projects/p/instances/") + path := "/projects/p/instances/" + oneShot + "?data&cached&taskname=oneshot" + if code, _, body := mcqaCall(t, h, "GET", path, nil); code != 409 { + t.Fatalf("direct download of a one-shot instance: code=%d body=%s", code, body) + } else if got, message := directDownloadJSONError(t, code, body); got != "InvalidState" || !strings.Contains(message, "downloads") { + t.Fatalf("one-shot instance download: code=%s message=%s", got, message) + } + // Unknown instances stay the plain not-found they have always been. + if code, _, _ := mcqaCall(t, h, "GET", downloadPath("missing", 1, ""), nil); code != 404 { + t.Fatalf("missing instance: code=%d", code) + } +} + +// TestMCQADirectDownloadLongPollsARunningSubQuery is why the read must not answer an empty +// stream while a statement is in flight: the SDK's "cached" read is a long poll, and a +// premature empty answer is indistinguishable from a query that returned nothing. +func TestMCQADirectDownloadLongPollsARunningSubQuery(t *testing.T) { + s, h := fixture(t, Config{SessionTTL: time.Hour}) + id := createSession(t, h, "") + queryID, _ := submitStatement(t, h, id, "select id from t limit 2") + setStatus := func(status int) { + s.mu.Lock() + defer s.mu.Unlock() + q := s.instances[id].MCQA.queries[queryID] + if q == nil { + t.Fatalf("sub query %d is gone from the session", queryID) + } + q.Status = status + } + // The emulator finishes a statement before it answers the submission, so an in-flight + // sub query is only reachable by a client that reads without polling first — or by a + // test that puts the session back into that state. + setStatus(mcqaStatusRunning) + done := make(chan *http.Response, 1) + go func() { + res, e := http.Get(h.URL + downloadPath(id, queryID, "")) + if e != nil { + t.Errorf("long poll: %v", e) + close(done) + return + } + done <- res + }() + select { + case res := <-done: + body, _ := io.ReadAll(res.Body) + t.Fatalf("a running sub query answered immediately: code=%d body=%s", res.StatusCode, body) + case <-time.After(100 * time.Millisecond): + } + setStatus(mcqaStatusTerminated) + var res *http.Response + select { + case res = <-done: + case <-time.After(5 * time.Second): + t.Fatal("the read never came back after the sub query finished") + } + defer res.Body.Close() + if res.StatusCode != 200 { + t.Fatalf("long poll after completion: code=%d", res.StatusCode) + } + body, e := io.ReadAll(res.Body) + if e != nil { + t.Fatal(e) + } + if got := res.Header.Get("odps-tunnel-record-count"); got != "2" { + t.Fatalf("record count header=%q want 2", got) + } + if decoded, _ := readRecords(t, body); len(decoded.Rows) != 2 { + t.Fatalf("long poll rows: %d", len(decoded.Rows)) + } + // A cancelled sub query answers with the flag the SDK rethrows instead of re-running. + setStatus(mcqaStatusCancelled) + code, _, body := mcqaCall(t, h, "GET", downloadPath(id, queryID, ""), nil) + if code != 409 { + t.Fatalf("cancelled sub query: code=%d body=%s", code, body) + } + if _, message := directDownloadJSONError(t, code, body); !strings.Contains(message, "Job is cancelled") { + t.Fatalf("cancelled sub query message=%s", message) + } +} + +// TestMCQADirectDownloadKeepsTheSessionAlive: a result read is session activity, the same +// way an information call is. Otherwise a client that polls results through the tunnel and +// only submits every now and then would have its session dropped underneath it. +func TestMCQADirectDownloadKeepsTheSessionAlive(t *testing.T) { + _, h := fixture(t, Config{SessionTTL: 250 * time.Millisecond}) + id := createSession(t, h, "") + queryID, _ := submitStatement(t, h, id, "select id from t limit 1") + for n := 0; n < 2; n++ { + time.Sleep(150 * time.Millisecond) + code, _, body := mcqaCall(t, h, "GET", downloadPath(id, queryID, ""), nil) + if code != 200 { + t.Fatalf("read %d after %dms of tunnel activity: code=%d body=%s", n, 150*(n+1), code, body) + } + } + time.Sleep(300 * time.Millisecond) + code, _, body := mcqaCall(t, h, "GET", downloadPath(id, queryID, ""), nil) + if code != 404 { + t.Fatalf("an idle session must still expire: code=%d body=%s", code, body) + } +} + +// TestMCQADirectDownloadFollowsRetention: the tunnel must not keep serving a result the +// session has already dropped, or a client would read rows the information channel has +// reported as gone. +func TestMCQADirectDownloadFollowsRetention(t *testing.T) { + _, h := fixture(t, Config{SessionTTL: time.Hour}) + id := createSession(t, h, "") + first, _ := submitStatement(t, h, id, "select id from t limit 2") + for n := 0; n < mcqaRetainedQueries; n++ { + if _, ir := submitStatement(t, h, id, "select 1"); ir.Status != mcqaInfoOK { + t.Fatalf("fill the retention window: %+v", ir) + } + } + if r := getInfo(t, h, infoPath(id, sessionTask, "result_"+strconv.Itoa(first))); r.Status != mcqaStatusFailed { + t.Fatalf("the information channel kept a dropped sub query: %+v", r) + } + code, _, body := mcqaCall(t, h, "GET", downloadPath(id, first, ""), nil) + if code != 404 { + t.Fatalf("a dropped sub query is still downloadable: code=%d body=%s", code, body) + } +} + +// TestProtobufSchemaTailMatchesPlainStream pins the claim in the encoder's comment: only +// the schema frame is added, so the record path stays byte-identical to the one the +// download sessions already serve. +func TestProtobufSchemaTailMatchesPlainStream(t *testing.T) { + id, e := engine.ParseType("bigint") + if e != nil { + t.Fatal(e) + } + res := engine.Result{Columns: []engine.Column{{Name: "id", Type: "bigint", Nullable: true, Parsed: id}}, Rows: [][]any{{int64(1)}, {nil}, {int64(3)}}} + withSchema, e := wire.ProtobufSchema(res) + if e != nil { + t.Fatal(e) + } + plain, e := wire.Protobuf(res) + if e != nil { + t.Fatal(e) + } + cols, rest := readSchemaFrame(t, withSchema) + if cols[0].Name != "id" || cols[0].Type != "bigint" { + t.Fatalf("schema frame columns: %+v", cols) + } + if !bytes.Equal(rest, plain) { + t.Fatalf("record stream after the schema frame differs from wire.Protobuf") + } +} diff --git a/internal/server/server.go b/internal/server/server.go index fc669c5..018a909 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -108,7 +108,9 @@ func fail(w http.ResponseWriter, r *http.Request, status int, code string, e err if len(message) > 600 { message = message[:600] } - if r.URL.Query().Has("downloads") || r.URL.Query().Has("downloadid") || r.URL.Query().Has("uploads") || r.URL.Query().Has("uploadid") || strings.Contains(r.URL.Path, "/upserts") || strings.Contains(r.URL.Path, "/streams") || strings.Contains(r.URL.Path, "storage") { + // `cached` is the direct session-result download, whose client parses a JSON tunnel + // error rather than an XML one. + if r.URL.Query().Has("downloads") || r.URL.Query().Has("downloadid") || r.URL.Query().Has("uploads") || r.URL.Query().Has("uploadid") || r.URL.Query().Has("cached") || strings.Contains(r.URL.Path, "/upserts") || strings.Contains(r.URL.Path, "/streams") || strings.Contains(r.URL.Path, "storage") { jsonResponse(w, status, map[string]string{"Code": code, "Message": message, "RequestId": w.Header().Get("x-odps-request-id")}) } else { w.Header().Set("Content-Type", "application/xml") @@ -198,7 +200,7 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { return } if path == "capabilities" { - jsonResponse(w, 200, map[string]any{"version": Version, "tunnel_download": []string{"create", "reload", "protobuf", "arrow", "complete"}, "compression": []string{"identity", "deflate", "zstd", "lz4_frame"}, "storage_v2": true, "storage_paths": []string{"/api/storage/v2", "/api/storage/v3"}, "tunnel_upload": []string{"protobuf", "arrow", "blocks", "stream", "upsert"}, "storage_write_modes": []string{"Batch", "BatchCompatible", "Streaming", "StreamingRealtime"}, "sql": "CREATE/DROP/INSERT/SELECT; static partitions; ODPS2 subset", "resources": []string{"file", "jar", "py", "archive", "table-metadata"}, "functions": []string{"java-udf-metadata", "sql-function-metadata", "embedded-function-metadata"}, "mcqa": []string{"sqlrt-session", "subquery-query", "subquery-result", "subquery-cancel", "session-stop"}, "unsupported": []string{"volume-resources", "sql-udf-execution", "volumes", "mcqa-named-session-attach", "mcqa-v2-maxqa", "catalogapi"}, "auth": s.authDescription(), "test_faults": s.cfg.TestMode, "quotas": append([]string{"default"}, s.cfg.Quotas...)}) + jsonResponse(w, 200, map[string]any{"version": Version, "tunnel_download": []string{"create", "reload", "protobuf", "arrow", "complete"}, "compression": []string{"identity", "deflate", "zstd", "lz4_frame"}, "storage_v2": true, "storage_paths": []string{"/api/storage/v2", "/api/storage/v3"}, "tunnel_upload": []string{"protobuf", "arrow", "blocks", "stream", "upsert"}, "storage_write_modes": []string{"Batch", "BatchCompatible", "Streaming", "StreamingRealtime"}, "sql": "CREATE/DROP/INSERT/SELECT; static partitions; ODPS2 subset", "resources": []string{"file", "jar", "py", "archive", "table-metadata"}, "functions": []string{"java-udf-metadata", "sql-function-metadata", "embedded-function-metadata"}, "mcqa": []string{"sqlrt-session", "subquery-query", "subquery-result", "subquery-cancel", "session-stop", "subquery-instance-tunnel"}, "unsupported": []string{"volume-resources", "sql-udf-execution", "volumes", "mcqa-named-session-attach", "mcqa-v2-maxqa", "catalogapi"}, "auth": s.authDescription(), "test_faults": s.cfg.TestMode, "quotas": append([]string{"default"}, s.cfg.Quotas...)}) return } if path == "logview/host" { @@ -416,6 +418,13 @@ func (s *Server) instance(w http.ResponseWriter, r *http.Request, p, sc string, s.instanceDownload(w, r, p, sc, rest[0]) return } + // A session sub query is read straight off the instance: the Java SDK's direct + // download arrives as GET ?data&cached&taskname=.., with no download id to look up. + if len(rest) == 1 && r.Method == http.MethodGet && r.URL.Query().Has("data") && + (r.URL.Query().Has("cached") || r.URL.Query().Has("taskname")) { + s.mcqaDirectDownload(w, r, p, sc, rest[0]) + return + } if len(rest) != 1 || r.Method != "GET" && r.Method != "PUT" { fail(w, r, 400, "UnsupportedOperation", fmt.Errorf("unsupported instance operation")) return diff --git a/internal/wire/rows.go b/internal/wire/rows.go index 61856c6..7914807 100644 --- a/internal/wire/rows.go +++ b/internal/wire/rows.go @@ -5,6 +5,7 @@ package wire import ( "bytes" "encoding/binary" + "encoding/json" "fmt" "hash/crc32" "math" @@ -151,6 +152,51 @@ func (e *encoder) value(t engine.Type, v any) error { } func Protobuf(r engine.Result) ([]byte, error) { return protobuf(r, true) } +// schemaEndTag closes the in-stream schema frame; ProtoWireConstant.SCHEMA_END_TAG in +// the Java SDK, magic number 2^25-512. +const schemaEndTag = 33553920 + +// ProtobufSchema encodes a result as a Tunnel stream whose first frame carries the +// schema as JSON. It is what an MCQA sub-query download needs: that read has no +// download session to ask for a schema, so the stream has to bring its own +// (ProtobufRecordStreamReader#readSchema, which the Java SDK calls the MCQA direct +// download interface). The schema frame is checksummed on its own and the reader +// resets before the first record, so the record stream after it is byte-identical to +// Protobuf's. +func ProtobufSchema(r engine.Result) ([]byte, error) { + head, e := schemaFrame(r.Columns) + if e != nil { + return nil, e + } + body, e := protobuf(r, true) + if e != nil { + return nil, e + } + return append(head, body...), nil +} + +// schemaFrame renders the in-stream schema frame for cols. The reader CRCs the field +// number as a 4-byte little-endian int and then the payload bytes, exactly as the +// record path CRCs each field; the tag and the length varint are outside the checksum. +func schemaFrame(cols []engine.Column) ([]byte, error) { + if cols == nil { + cols = []engine.Column{} + } + s, e := json.Marshal(map[string]any{"columns": cols, "partitionKeys": []engine.Column{}}) + if e != nil { + return nil, e + } + w := &encoder{} + w.i32(1) + w.u(1<<3 | 2) + w.u(uint64(len(s))) + w.b = append(w.b, s...) + w.sum(s) + w.u(schemaEndTag << 3) + w.u(uint64(w.crc)) + return w.b, nil +} + // ProtobufPrefix omits the stream trailer for row-boundary disconnect tests. func ProtobufPrefix(r engine.Result) ([]byte, error) { return protobuf(r, false) } func protobuf(r engine.Result, trailer bool) (out []byte, err error) { diff --git a/tests/java/src/test/java/McqaSessionTest.java b/tests/java/src/test/java/McqaSessionTest.java index 9eb5c05..f3c0849 100644 --- a/tests/java/src/test/java/McqaSessionTest.java +++ b/tests/java/src/test/java/McqaSessionTest.java @@ -67,11 +67,9 @@ void client() throws Exception { /** * An interactive executor that reads results through the session information channel. * - *

{@code useInstanceTunnel(false)} is deliberate, and it is a scope boundary rather - * than a convenience: by default the SDK fetches sub-query results through the instance - * tunnel, which the emulator does not serve for sessions yet, and it then silently - * re-runs the statement as an offline instance — an acceptance test written that way - * would pass even with nothing implemented on the session plane. + *

{@code useInstanceTunnel(false)} is deliberate: it pins one read path down so a + * regression on the other one cannot hide behind it. With the tunnel enabled the SDK + * never asks the session for CSV text at all — see {@link #tunnelExecutor()}. */ private SQLExecutor executor() throws OdpsException { return SQLExecutorBuilder.builder() @@ -81,6 +79,76 @@ private SQLExecutor executor() throws OdpsException { .build(); } + /** + * The fetch path the SDK uses by default, and the one JDBC's MaxQA mode takes: a + * sub-query result is downloaded from the instance tunnel, and because a direct + * download has no download session to ask, the schema has to travel in the stream. + */ + private SQLExecutor tunnelExecutor() throws OdpsException { + return SQLExecutorBuilder.builder() + .odps(odps) + .executeMode(ExecuteMode.INTERACTIVE) + .tunnelEndpoint(endpoint) + .build(); + } + + @Test + void defaultFetchDownloadsTheSubQueryResult() throws Exception { + SQLExecutor executor = tunnelExecutor(); + try { + executor.run("select id, s from " + table + " order by id", new HashMap<>()); + String sessionId = executor.getInstance().getId(); + java.util.List records = executor.getResult(); + assertEquals(2, records.size(), "rows over the instance tunnel"); + // Typed values, not the strings the information channel's CSV hands back: the + // column types came from the in-stream schema, which only this read has. + assertEquals(1L, ((Number) records.get(0).get(0)).longValue(), "bigint id"); + assertEquals("two", records.get(1).getString(1), "string s"); + assertEquals(2, records.get(0).getColumns().length, "schema from the stream"); + // The observable that separates "the tunnel answered" from "the SDK gave up": + // a failed session download re-runs the statement offline, which replaces the + // session instance with a new one and logs why. + assertEquals(sessionId, executor.getInstance().getId(), + "an offline fallback would replace the session instance"); + assertTrue(executor.isRunningInInteractiveMode(), "still interactive after the download"); + // QueryInfo writes "Will fallback to offline mode" when a read leaves the session, + // and "Running in interactive mode" when it stays on one. getExecutionLog() drains, + // so the list has to be read once. + String logText = String.join("\n", executor.getExecutionLog()); + assertTrue(logText.contains("Running in interactive mode"), + "the executor should report the session it ran on: " + logText); + assertFalse(logText.contains("fallback"), "the download fell back: " + logText); + } finally { + executor.close(); + } + } + + /** + * The row count a direct download reports is what the SDK iterates, so it has to be the + * rows still available from the requested offset — the total would send it reading past + * the end of the result. An offset read is the only way to exercise that without a + * ten-thousand-row statement in the acceptance suite. + */ + @Test + void offsetReadPagesTheSubQueryResult() throws Exception { + SQLExecutor executor = tunnelExecutor(); + try { + executor.run("select id from " + table + " order by id", new HashMap<>()); + String sessionId = executor.getInstance().getId(); + java.util.List tail = executor.getResult(1L, null, null); + assertEquals(1, tail.size(), "rows from offset 1"); + assertEquals(2L, ((Number) tail.get(0).get(0)).longValue(), "the row after the offset"); + java.util.List window = executor.getResult(0L, 1L, null); + assertEquals(1, window.size(), "a count limit is a count limit"); + assertEquals(1L, ((Number) window.get(0).get(0)).longValue()); + java.util.List rest = executor.getResult(2L, 10L, null); + assertTrue(rest.isEmpty(), "reading past the end is empty, not an error"); + assertEquals(sessionId, executor.getInstance().getId(), "paging stayed on the session"); + } finally { + executor.close(); + } + } + @Test void selectRunsThroughSessionNotOffline() throws Exception { SQLExecutor executor = executor(); From ca2e4dddb5daad7f3a3def96888ae62f88c0f4e3 Mon Sep 17 00:00:00 2001 From: tony-stack Date: Tue, 22 Sep 2026 11:09:36 +0800 Subject: [PATCH 2/2] fix: a select-only session declines a non-select before running it Driving the published JDBC driver with interactiveMode=mcqa and the instance tunnel applied every write statement twice: the session executed the sub query, the tunnel read then answered InstanceTypeNotSupported, and the client's own reaction to that answer is to run the same statement offline again (SQLExecutorImpl NON_SELECT_QUERY -> runQueryInternal(OFFLINE, ..., true)). An insert of three rows landed six. The restriction is a session property, not an emulator opinion, and the client can lift it: odps.sql.session.select.only, in the session settings or per statement. The refusal is the pair the client already acts on - queryId -1, status ok, and a message carrying ODPS-185 plus the service's own sentence - so a declined submission becomes one offline run instead of an error. Multi-statement submissions are classified by every segment, because 'select 1; drop table t' must not pass for a select. --- CHANGELOG.md | 1 + README.md | 2 +- docs/protocol.md | 2 + internal/server/mcqa.go | 24 +++- internal/server/mcqa_download_test.go | 8 +- internal/server/mcqa_select_only.go | 133 +++++++++++++++++++++++ internal/server/mcqa_select_only_test.go | 120 ++++++++++++++++++++ internal/server/server.go | 2 +- 8 files changed, 284 insertions(+), 8 deletions(-) create mode 100644 internal/server/mcqa_select_only.go create mode 100644 internal/server/mcqa_select_only_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index 93e2f3f..73ac7ea 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ - Answer the three endpoints the official JDBC driver calls before any statement runs: `POST /projects/p/authorization?sign_bearer_token` (placeholder logview token), `GET /logview/host`, and `GET /connection/mcqa` (refused by name, so a MaxQA misconfiguration stops resurfacing as an unrelated endpoint error). - Emit schema `columns`/`partitionKeys` as arrays instead of JSON `null`, which the Java SDK's `TunnelTableSchema` rejects — every DDL and INSERT result has no columns, so the download session the driver opens for one was unparseable. - Add `tests/jdbc`: acceptance against the published `odps-jdbc` artifact (offline mode), covering a statement with no result set end to end. +- Decline non-select sub queries in a SQLRT session before executing them (`queryId: -1` with `ODPS-1850001 Non select query not supported.`, lift with `odps.sql.session.select.only=false`): a client that discovers a non-select at fetch time reruns the statement offline, so a session that had already executed an INSERT applied it twice. - Add resource CRUD: single-payload and Java SDK chunked part/merge uploads, `?meta` header metadata, ranged downloads, update, delete, prefix and paginated listing, plus TABLE resource metadata. - Add function registration (`/projects/p/registration/functions`) for Java, SQL and embedded functions with resource-reference validation; execution still returns `UnsupportedFeature`. - Declare `resources`, `functions` and `unsupported` surfaces in `/capabilities`. diff --git a/README.md b/README.md index f85de89..767a6c3 100644 --- a/README.md +++ b/README.md @@ -101,7 +101,7 @@ See [data transfer](docs/data-transfer.md) and [protocol details](docs/protocol. This is a local/CI test service. By default authentication signatures and permissions are **not validated**; optional strict mode uses local test credentials; keep it on a trusted test network. It is not a replacement for real MaxCompute acceptance tests. -MCQA sessions cover both of the Java SDK's read paths: the information channel (`SQLExecutor` with `useInstanceTunnel(false)`) answers with CSV whose first line is the column-name header, and the instance tunnel (`?data&cached&taskname=..&queryid=..`, the default fetch and the one JDBC MaxQA uses) answers with a record stream that carries its own schema, so a sub-query result arrives with the same rows and the same types either way. `instance_tunnel_limit_enabled` applies the service's `READ_TABLE_MAX_ROW` (10,000 rows) to that read, `rowrange` pages through it, and `sizelimit` shortens a response instead of failing it. Statements that produce no result set are refused on the tunnel with `InstanceTypeNotSupported`, like the service, so a client confirms them through the session API rather than reading an empty stream. Session-scoped state (variables, temp objects) is not isolated — statements run against the shared engine. Not simulated: named-session attach (a session name is metadata; equal names create separate sessions), MaxQA v2 (`/mcqa` request prefix), compressed or Arrow-encoded sub-query downloads (Protobuf records only), and per-statement statistics. +MCQA sessions cover both of the Java SDK's read paths: the information channel (`SQLExecutor` with `useInstanceTunnel(false)`) answers with CSV whose first line is the column-name header, and the instance tunnel (`?data&cached&taskname=..&queryid=..`, the default fetch and the one JDBC MaxQA uses) answers with a record stream that carries its own schema, so a sub-query result arrives with the same rows and the same types either way. `instance_tunnel_limit_enabled` applies the service's `READ_TABLE_MAX_ROW` (10,000 rows) to that read, `rowrange` pages through it, and `sizelimit` shortens a response instead of failing it. Statements that produce no result set are refused on the tunnel with `InstanceTypeNotSupported`, like the service, so a client confirms them through the session API rather than reading an empty stream. A session is **select-only by default**: a `CREATE`/`INSERT`/`DROP` submitted as a sub query is declined before anything runs (`queryId: -1` plus `ODPS-1850001 Non select query not supported.`), because the Java client reacts to a non-select it learns about at fetch time by running the same statement offline again — and an INSERT that already ran in the session would land twice. Send `odps.sql.session.select.only=false`, in the session settings or per statement, to run them in the session anyway. Session-scoped state (variables, temp objects) is not isolated — statements run against the shared engine. Not simulated: named-session attach (a session name is metadata; equal names create separate sessions), MaxQA v2 (`/mcqa` request prefix), compressed or Arrow-encoded sub-query downloads (Protobuf records only), and per-statement statistics. Unsupported: Storage v1, Volume/Blob, CDC/incremental reads, filter predicate pushdown, explicit Schema management, column schema evolution, UDF **execution** (functions are registered metadata; SQL `CREATE FUNCTION`/`DROP FUNCTION` and calling a UDF return `UnsupportedFeature`), volume-backed resources, distributed scheduling, and full ODPS SQL semantics. Resource payloads are capped at 64 MiB each and 512 MiB per project/schema. Resource and function names resolve case-insensitively while the uploaded spelling is what listings return. Unsupported operations return errors rather than cloud behavior being assumed. diff --git a/docs/protocol.md b/docs/protocol.md index c5f47b2..0f55c0e 100644 --- a/docs/protocol.md +++ b/docs/protocol.md @@ -35,6 +35,8 @@ JDBC 驱动对每条语句都会先签一个 logview token，再为实例建立 MCQA 子查询直读没有 download session 可查，schema 只能随流走：帧头是 field 1 的 JSON schema（TunnelTableSchema 的 `columns`/`partitionKeys` 形状），后接 SCHEMA_END_TAG 33553920 与该帧自己的 CRC32C（覆盖 field number 的 4 字节小端值与 JSON 字节，不含 tag 与长度 varint），校验和后 CRC 归零，其后的记录部分与下面的 Protobuf 完全一致。响应头 `odps-tunnel-record-count` 是从本次 `rowrange` 起点起还能读到的行数，不是全量结果：Java 的 SessionRecordSetIterator 只取第一次响应的这个数并据此翻页，报大全量会让客户端越过结果末尾继续读取。查询仍在执行时该读取阻塞等待（长轮询），超时返回 Java SDK 认得的 `OdpsTaskTimeout`/“Wait for cache data timeout”，而不是先给一个空流。 +会话默认 select-only：`?info` 写入 `Key=query` 时，非 select 语句（含多语句提交里任何一段不是 select）在**执行前**就被拒绝，回 `{"queryId":-1,"status":"ok","result":"ODPS-1850001 Non select query not supported."}`——客户端对"取结果时才发现不是 select"的处置是 offline 重跑同一条语句，会话若已执行过，写语句就会落两次。`odps.sql.session.select.only=false`（会话设置或单条设置）可放开。 + Protobuf 按 public SDK 的列序编码，NULL 不写字段。记录结束 tag 33553408 后写 CRC32C，流尾 meta-count 33554430 与 checksum-of-checksums 33554431；CRC 数值按小端逻辑值更新，不包含传输长度和 NULL 标记。DATE 为天数，DATETIME 为毫秒，TIMESTAMP 为秒+纳秒，DECIMAL 为十进制文本。MAP 在引擎边界从 DuckDB OrderedMap 转换，ARRAY/STRUCT 递归编码。 Tunnel Arrow 不是普通带 Schema 消息的 IPC stream：客户端从会话元数据构造 schema，服务端只发送 RecordBatch payload。帧头为大端 chunk size（65536），每满块追加 CRC32C，最后追加全流 CRC32C。ZSTD 和 LZ4 frame 压缩 Arrow IPC buffer，外层 Tunnel CRC 帧保持未压缩；deflate 仍包裹整个 Tunnel 帧。空表和 count=0 返回合法空数据帧。常规 Arrow 请求每批最多4096行；带 raw_size 的请求最多一批，CPP BufferArrowRecordReader 按实际行数推进；raw_size 过小时仍返回至少一行，以保证进度。 diff --git a/internal/server/mcqa.go b/internal/server/mcqa.go index 0dd38f8..848bbe0 100644 --- a/internal/server/mcqa.go +++ b/internal/server/mcqa.go @@ -70,7 +70,11 @@ type mcqaQuery struct { // mcqaSession is the state behind a SQLRT instance. Fields are guarded by Server.mu; // statement execution runs outside that lock (see runSubQuery). type mcqaSession struct { - name string + name string + // selectOnly is the session's promise that nothing but a whole select runs here. + // The service defaults it to true; a client that wants DDL or DML inside the session + // says so with odps.sql.session.select.only. + selectOnly bool live bool nextID int lastID int @@ -79,8 +83,8 @@ type mcqaSession struct { lastActive time.Time } -func newMCQASession(name string) *mcqaSession { - return &mcqaSession{name: name, live: true, queries: map[int]*mcqaQuery{}, lastActive: time.Now()} +func newMCQASession(name string, selectOnly bool) *mcqaSession { + return &mcqaSession{name: name, selectOnly: selectOnly, live: true, queries: map[int]*mcqaQuery{}, lastActive: time.Now()} } // expired reports whether the session has been idle past the configured TTL. Idle, not @@ -174,7 +178,7 @@ func (s *Server) mcqaCreate(w http.ResponseWriter, r *http.Request, p, sc, taskN i := &instance{ Project: p, Schema: sc, ID: id(), Name: taskName, TaskType: "SQLRT", Status: "Running", Created: time.Now(), - MCQA: newMCQASession(name), + MCQA: newMCQASession(name, !sessionRunsNonSelect(mcqaSettingsMap(settings))), } if !s.reserveInstance(i) { fail(w, r, 429, "ResourceLimit", fmt.Errorf("instance limit")) @@ -300,6 +304,18 @@ func (s *Server) mcqaSetInfo(w http.ResponseWriter, r *http.Request, i *instance s.informationReply(w, mcqaInfoFailed, "Invalid sub query payload") return } + if m.selectOnly && !sessionRunsNonSelect(sub.Settings) && !isSelectStatement(sub.Query) { + // Refuse before running anything. A client that only learns the statement + // was not a select when it comes to fetch the result reruns it offline — + // and an INSERT that already ran in the session then applies twice. queryId + // -1 with status ok is how the service reports a sub query it declined to + // start, and the message is the pair FallbackPolicy matches on ("ODPS-185"), + // which is what turns the refusal into an offline rerun rather than an error. + s.mu.Unlock() + payload, _ := json.Marshal(map[string]any{"queryId": -1, "status": mcqaInfoOK, "result": mcqaSelectOnlyRefusal}) + s.informationReply(w, mcqaInfoOK, string(payload)) + return + } m.nextID++ id := m.nextID m.drop() diff --git a/internal/server/mcqa_download_test.go b/internal/server/mcqa_download_test.go index e8e3ec6..37be918 100644 --- a/internal/server/mcqa_download_test.go +++ b/internal/server/mcqa_download_test.go @@ -230,7 +230,8 @@ func TestMCQADirectDownloadHonoursSizeLimit(t *testing.T) { // the same sub query still returns every row. func TestMCQADirectDownloadLimitedFlagCapsRows(t *testing.T) { _, h := fixture(t, Config{SessionTTL: time.Hour}) - id := createSession(t, h, "") + // The table grows through the session, which a select-only session would decline. + id := createSession(t, h, `{"`+mcqaSelectOnlyKey+`":"false"}`) // 2048 fixture rows, doubled three times: 16384, i.e. past the 10000 cap. for n := 0; n < 3; n++ { if _, ir := submitStatement(t, h, id, fmt.Sprintf("insert into t select id+%d, s from t", n+1)); ir.Status != mcqaInfoOK { @@ -287,7 +288,10 @@ func TestMCQADirectDownloadLimitedFlagCapsRows(t *testing.T) { // happened to come back empty. func TestMCQADirectDownloadRefusesNonSelect(t *testing.T) { _, h := fixture(t, Config{SessionTTL: time.Hour}) - id := createSession(t, h, "") + // Select-only has to be lifted to get here: a default session declines a non-select + // before running it (mcqa_select_only_test.go), so reaching the download with one + // means the client asked for exactly this combination. + id := createSession(t, h, `{"`+mcqaSelectOnlyKey+`":"false"}`) for _, statement := range []string{"create table ns(a bigint)", "insert into t values (9001,'x')"} { queryID, ir := submitStatement(t, h, id, statement) if ir.Status != mcqaInfoOK { diff --git a/internal/server/mcqa_select_only.go b/internal/server/mcqa_select_only.go new file mode 100644 index 0000000..9b0f003 --- /dev/null +++ b/internal/server/mcqa_select_only.go @@ -0,0 +1,133 @@ +package server + +// A SQLRT session runs whole selects. Anything else — DDL, DML, anything with a side +// effect — the session declines to start, and the client that submitted it runs it +// somewhere else instead. +// +// The order matters, and it is the reason this file exists. A session that *executes* a +// non-select and only then says the result cannot be downloaded is not merely untidy: the +// Java SDK's tunnel read reacts to exactly that answer by re-running the statement offline +// (SQLExecutorImpl#getSessionResultSetByInstanceTunnel -> TunnelRetryStatus.NON_SELECT_QUERY +// -> runQueryInternal(OFFLINE, ..., true)). An INSERT applied that way lands twice. + +import ( + "encoding/json" + "strings" +) + +const ( + // mcqaSelectOnlyKey is the setting a client sends to lift the restriction, either in + // the session's task settings or in the per-statement settings of a sub query. + mcqaSelectOnlyKey = "odps.sql.session.select.only" + // mcqaSelectOnlyRefusal is what the service says, and the pair the client acts on: + // "ODPS-185" is the flag FallbackPolicy#shouldFallback reads to decide that an + // unsupported-in-session statement should be retried offline. Say only the sentence + // without the code and the same submission comes back as a hard error instead. + mcqaSelectOnlyRefusal = "ODPS-1850001 Non select query not supported." +) + +// sessionRunsNonSelect reports whether a settings block asks for non-select statements to +// run in the session. Absent or any other value keeps the select-only default. +func sessionRunsNonSelect(settings map[string]string) bool { + value, ok := settings[mcqaSelectOnlyKey] + return ok && strings.EqualFold(strings.TrimSpace(value), "false") +} + +// mcqaSettingsMap decodes the JSON settings property of a SQLRT task. The same helper +// reads it for the session name, so a settings block that is not a JSON object simply +// yields nothing and keeps the defaults. +func mcqaSettingsMap(settings string) map[string]string { + if strings.TrimSpace(settings) == "" { + return nil + } + var raw map[string]any + if e := json.Unmarshal([]byte(settings), &raw); e != nil { + return nil + } + out := map[string]string{} + for k, v := range raw { + out[k] = fmtSprint(v) + } + return out +} + +func fmtSprint(v any) string { + s, _ := json.Marshal(v) + return strings.Trim(string(s), `"`) +} + +// isSelectStatement is the emulator's stand-in for the planner's own select check: every +// statement in the submission has to start with SELECT or WITH, ignoring comments and the +// parentheses a client may wrap a query in. It is deliberately conservative in the +// direction that costs nothing to be wrong about — anything it does not recognise is run +// offline by the client instead of twice by the session. +func isSelectStatement(sql string) bool { + statements := sqlStatements(sql) + if len(statements) == 0 { + return false + } + for _, statement := range statements { + head := strings.TrimLeft(statement, "( \t\n") + if !strings.HasPrefix(head, "select") && !strings.HasPrefix(head, "with") { + return false + } + } + return true +} + +// sqlStatements returns the lower-cased text of each statement in a submission with +// comments removed, so that a `select 1; drop table t` cannot pass for a select. Quotes +// are respected: a semicolon inside a string literal does not end a statement. +func sqlStatements(sql string) []string { + var out []string + var b strings.Builder + quote := byte(0) + for n := 0; n < len(sql); { + c := sql[n] + switch { + case quote != 0 && c == '\\' && quote != '`': + b.WriteByte(c) + n++ + if n < len(sql) { + b.WriteByte(sql[n]) + n++ + } + case quote != 0: + if c == quote { + quote = 0 + } + b.WriteByte(c) + n++ + case c == '\'' || c == '"' || c == '`': + quote = c + b.WriteByte(c) + n++ + case c == '-' && n+1 < len(sql) && sql[n+1] == '-': + for n < len(sql) && sql[n] != '\n' { + n++ + } + case c == '/' && n+1 < len(sql) && sql[n+1] == '*': + end := strings.Index(sql[n+2:], "*/") + if end < 0 { + n = len(sql) + } else { + n += end + 4 + } + case c == ';': + out = addStatement(out, b.String()) + b.Reset() + n++ + default: + b.WriteByte(c) + n++ + } + } + return addStatement(out, b.String()) +} + +func addStatement(out []string, text string) []string { + if text = strings.TrimSpace(text); text == "" { + return out + } + return append(out, strings.ToLower(text)) +} diff --git a/internal/server/mcqa_select_only_test.go b/internal/server/mcqa_select_only_test.go new file mode 100644 index 0000000..014d760 --- /dev/null +++ b/internal/server/mcqa_select_only_test.go @@ -0,0 +1,120 @@ +package server + +import ( + "encoding/json" + "fmt" + "net/http/httptest" + "strings" + "testing" + "time" +) + +// submitStatementWithSettings is submitStatement plus the per-statement settings block, +// which is where a client says it wants non-select statements to run in the session. +func submitStatementWithSettings(t *testing.T, h *httptest.Server, id, query string, settings map[string]string) (int, informationResult) { + t.Helper() + payload, e := json.Marshal(map[string]any{"query": query, "settings": settings}) + if e != nil { + t.Fatal(e) + } + body := "query" + string(payload) + "" + code, _, b := mcqaCall(t, h, "PUT", infoPath(id, sessionTask, ""), []byte(body)) + if code != 200 { + t.Fatalf("setInformation code=%d body=%s", code, b) + } + var ir informationResult + if e := json.Unmarshal(b, &ir); e != nil { + t.Fatalf("setInformation body not JSON: %v / %s", e, b) + } + return parseQueryID(t, ir), ir +} + +func refusedResult(t *testing.T, ir informationResult) map[string]any { + t.Helper() + var sub map[string]any + if e := json.Unmarshal([]byte(ir.Result), &sub); e != nil { + t.Fatalf("sub query payload is not an object: %v / %s", e, ir.Result) + } + return sub +} + +func TestIsSelectStatement(t *testing.T) { + cases := []struct { + sql string + want bool + }{ + {"select 1", true}, + {" SELECT id FROM t", true}, + {"with x as (select 1 as a) select * from x", true}, + {"(select 1)", true}, + {"-- a leading comment\nselect 1", true}, + {"/* hint */ select 1", true}, + {"create table t(a int)", false}, + {"insert into t values (1)", false}, + {"truncate table t", false}, + {"select 1; insert into t values (1)", false}, + {"select ';' as x", true}, + {"", false}, + {" ", false}, + } + for _, c := range cases { + if got := isSelectStatement(c.sql); got != c.want { + t.Errorf("isSelectStatement(%q) = %v, want %v", c.sql, got, c.want) + } + } +} + +// TestSessionRefusesNonSelectWithoutRunningIt is the point of the file: the refusal has to +// come before execution, because the client's reaction to a refusal is to run the same +// statement again somewhere else. +func TestSessionRefusesNonSelectWithoutRunningIt(t *testing.T) { + _, h := fixture(t, Config{SessionTTL: time.Hour}) + id := createSession(t, h, "") + subQueryID, ir := submitStatementWithSettings(t, h, id, "create table nope(a bigint)", nil) + if ir.Status != mcqaInfoOK || subQueryID != -1 { + t.Fatalf("a select-only session must decline to start a DDL: %+v", ir) + } + result := fmt.Sprint(refusedResult(t, ir)["result"]) + for _, want := range []string{"ODPS-185", "Non select query not supported"} { + if !strings.Contains(result, want) { + t.Fatalf("refusal %q must carry %q: the first is what makes the client retry offline, the second is what a human reads", result, want) + } + } + if code, _, b := request(t, h, "GET", "/projects/p/tables/nope", ""); code == 200 { + t.Fatalf("the refused statement must not have run: table exists, body=%s", b) + } + // Nothing was allocated, so the next statement is still the session's first. + if next, _ := submitStatementWithSettings(t, h, id, "select 1", nil); next != 1 { + t.Fatalf("a declined submission must not consume a sub query id, got %d", next) + } +} + +// TestSessionRunsNonSelectWhenTheClientSaysSo covers the other half: the restriction is a +// default, not an opinion of the emulator's. +func TestSessionRunsNonSelectWhenTheClientSaysSo(t *testing.T) { + _, h := fixture(t, Config{SessionTTL: time.Hour}) + id := createSession(t, h, "") + subQueryID, ir := submitStatementWithSettings(t, h, id, "create table yes(a bigint)", + map[string]string{mcqaSelectOnlyKey: "false"}) + if ir.Status != mcqaInfoOK || subQueryID < 1 { + t.Fatalf("with %s=false the session must run the statement: %+v", mcqaSelectOnlyKey, ir) + } + if code, _, b := request(t, h, "GET", "/projects/p/tables/yes", ""); code != 200 { + t.Fatalf("the statement should have run inside the session: code=%d body=%s", code, b) + } +} + +func TestSessionCreateSettingLiftsSelectOnlyForEveryStatement(t *testing.T) { + _, h := fixture(t, Config{SessionTTL: time.Hour}) + id := createSession(t, h, `{"`+mcqaSelectOnlyKey+`":"false","odps.sql.session.name":"wide"}`) + if subQueryID, _ := submitStatementWithSettings(t, h, id, "create table wide_session(a bigint)", nil); subQueryID != 1 { + t.Fatalf("the session-level setting should apply without a per-statement one, got id=%d", subQueryID) + } +} + +func TestSQLStatementsKeepsSemicolonsInsideStrings(t *testing.T) { + got := sqlStatements("select 'a;b' as x; -- trailing\n insert into t values (1)") + if len(got) != 2 || !strings.HasPrefix(got[0], "select 'a;b'") || !strings.HasPrefix(got[1], "insert") { + t.Fatalf("statements=%q", got) + } +} diff --git a/internal/server/server.go b/internal/server/server.go index 018a909..52ce2b3 100644 --- a/internal/server/server.go +++ b/internal/server/server.go @@ -200,7 +200,7 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) { return } if path == "capabilities" { - jsonResponse(w, 200, map[string]any{"version": Version, "tunnel_download": []string{"create", "reload", "protobuf", "arrow", "complete"}, "compression": []string{"identity", "deflate", "zstd", "lz4_frame"}, "storage_v2": true, "storage_paths": []string{"/api/storage/v2", "/api/storage/v3"}, "tunnel_upload": []string{"protobuf", "arrow", "blocks", "stream", "upsert"}, "storage_write_modes": []string{"Batch", "BatchCompatible", "Streaming", "StreamingRealtime"}, "sql": "CREATE/DROP/INSERT/SELECT; static partitions; ODPS2 subset", "resources": []string{"file", "jar", "py", "archive", "table-metadata"}, "functions": []string{"java-udf-metadata", "sql-function-metadata", "embedded-function-metadata"}, "mcqa": []string{"sqlrt-session", "subquery-query", "subquery-result", "subquery-cancel", "session-stop", "subquery-instance-tunnel"}, "unsupported": []string{"volume-resources", "sql-udf-execution", "volumes", "mcqa-named-session-attach", "mcqa-v2-maxqa", "catalogapi"}, "auth": s.authDescription(), "test_faults": s.cfg.TestMode, "quotas": append([]string{"default"}, s.cfg.Quotas...)}) + jsonResponse(w, 200, map[string]any{"version": Version, "tunnel_download": []string{"create", "reload", "protobuf", "arrow", "complete"}, "compression": []string{"identity", "deflate", "zstd", "lz4_frame"}, "storage_v2": true, "storage_paths": []string{"/api/storage/v2", "/api/storage/v3"}, "tunnel_upload": []string{"protobuf", "arrow", "blocks", "stream", "upsert"}, "storage_write_modes": []string{"Batch", "BatchCompatible", "Streaming", "StreamingRealtime"}, "sql": "CREATE/DROP/INSERT/SELECT; static partitions; ODPS2 subset", "resources": []string{"file", "jar", "py", "archive", "table-metadata"}, "functions": []string{"java-udf-metadata", "sql-function-metadata", "embedded-function-metadata"}, "mcqa": []string{"sqlrt-session", "subquery-query", "subquery-result", "subquery-cancel", "session-stop", "subquery-instance-tunnel", "session-select-only"}, "unsupported": []string{"volume-resources", "sql-udf-execution", "volumes", "mcqa-named-session-attach", "mcqa-v2-maxqa", "catalogapi"}, "auth": s.authDescription(), "test_faults": s.cfg.TestMode, "quotas": append([]string{"default"}, s.cfg.Quotas...)}) return } if path == "logview/host" {