From b4d7a9d967cd3626b31de807c15ec39ad9681805 Mon Sep 17 00:00:00 2001 From: Rahul Singhal Date: Thu, 3 Sep 2026 22:09:58 +0000 Subject: [PATCH 1/2] feat(kernel): send per-statement query tags on the kernel backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Thrift backend already forwards per-statement query tags (from driverctx.QueryTagsFromContext) via confOverlay["query_tags"]; the kernel (SEA) backend did not, because the kernel C ABI had no per-statement tags setter. It now does (kernel_statement_set_query_tags), so wire it up: - internal/backend/kernel/operation.go: after binding params, read the context query tags, serialize with the shared querytags.Serialize, and pass the wire string to kernel_statement_set_query_tags before execute. Reuses the same user-facing API and serializer as the Thrift path. - internal/backend/kernel/querytags_nul.go: checkQueryTags/errQueryTagsNUL guards the serialized string against an interior NUL before the length-less C setter silently truncates it — the query-tags counterpart to checkQueryText (set_sql) and checkParamValue (bind_parameter), keeping the kernel path at parity with Thrift (which sends tags whole). Pure Go, tested under CGO_ENABLED=0 (querytags_nul_test.go). - internal/backend/kernel/include/databricks_kernel.h: add the declaration. NOTE (do not merge until unblocked): the databricks_kernel-tagged CI leg can only link once the databricks-sql-kernel-bindings modules are rebuilt at a kernel rev that includes kernel_statement_set_query_tags (merged as 21504eae074c7f5b11888a4ebcc7468780b2f844), and KERNEL_REV + the go.mod bindings versions are bumped to that rev in a follow-up commit. Validated locally by building/linking/vetting the kernel-tagged package + running the guard test against a kernel lib built from that rev. Co-authored-by: Isaac Signed-off-by: Rahul Singhal --- .../kernel/include/databricks_kernel.h | 14 +++++++++ internal/backend/kernel/operation.go | 26 ++++++++++++++++ internal/backend/kernel/querytags_nul.go | 30 +++++++++++++++++++ internal/backend/kernel/querytags_nul_test.go | 27 +++++++++++++++++ 4 files changed, 97 insertions(+) create mode 100644 internal/backend/kernel/querytags_nul.go create mode 100644 internal/backend/kernel/querytags_nul_test.go diff --git a/internal/backend/kernel/include/databricks_kernel.h b/internal/backend/kernel/include/databricks_kernel.h index 6ea9bccb..c00e1f97 100644 --- a/internal/backend/kernel/include/databricks_kernel.h +++ b/internal/backend/kernel/include/databricks_kernel.h @@ -688,6 +688,20 @@ KernelStatusCode kernel_statement_bind_parameter(kernel_statement_t* stmt, const char* sql_type, const char* value); +/* + * Set per-statement query tags — key/value annotations sent with this + * statement's execution for cost attribution and tracking. `query_tags` is the + * serialized wire form: comma-separated `key:value` pairs, with a bare `key` + * (no colon) for a valueless tag, e.g. "team:eng,job:etl,production"; backslash + * escapes the separators in a value ("\:", "\,", "\\"). This is the same string + * a host already builds for the Thrift `confOverlay` path, so it is passed + * through unchanged and re-parsed into the native SEA query_tags array at + * execute time. An empty string is a no-op; calling again replaces the + * statement's tags. `query_tags` must be a NUL-terminated UTF-8 string. + */ +KernelStatusCode kernel_statement_set_query_tags(kernel_statement_t* stmt, + const char* query_tags); + /* * Wait-for-result execution. On success, `*out` holds an executed handle * released with `kernel_executed_statement_close`. diff --git a/internal/backend/kernel/operation.go b/internal/backend/kernel/operation.go index ef63c902..e59bbe90 100644 --- a/internal/backend/kernel/operation.go +++ b/internal/backend/kernel/operation.go @@ -16,9 +16,11 @@ import ( "sync" "time" + "github.com/databricks/databricks-sql-go/driverctx" dbsqlerr "github.com/databricks/databricks-sql-go/errors" "github.com/databricks/databricks-sql-go/internal/backend" dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors" + "github.com/databricks/databricks-sql-go/internal/querytags" dbsqlrows "github.com/databricks/databricks-sql-go/internal/rows" ) @@ -78,6 +80,30 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b return &kernelOp{}, fmt.Errorf("kernel: bind params: %w", toStatementError(err)) } + // Per-statement query tags from context — the same source and serializer the + // Thrift backend uses; here the serialized wire string goes to the kernel's + // per-statement setter instead of Thrift confOverlay. Empty/absent → skip. + if queryTags := driverctx.QueryTagsFromContext(ctx); len(queryTags) > 0 { + if serialized := querytags.Serialize(queryTags); serialized != "" { + // Guard against an interior NUL before the length-less C setter + // silently truncates it — same parity guard as set_sql/bind_parameter. + if err := checkQueryTags(serialized); err != nil { + C.kernel_statement_close(stmt) + return &kernelOp{}, fmt.Errorf("kernel: %w", err) + } + tags := newCStr(serialized) + err := call(func() C.KernelStatusCode { + return C.kernel_statement_set_query_tags(stmt, tags.c) + }) + tags.free() + if err != nil { + C.kernel_statement_close(stmt) + k.evictIfSessionFatal(err) + return &kernelOp{}, fmt.Errorf("kernel: set_query_tags: %w", toStatementError(err)) + } + } + } + // Detached canceller, obtained before execute so it observes the server // statement id the moment execute publishes it. Non-fatal on failure: proceed // without cancellation rather than failing the query. diff --git a/internal/backend/kernel/querytags_nul.go b/internal/backend/kernel/querytags_nul.go new file mode 100644 index 00000000..7e28c752 --- /dev/null +++ b/internal/backend/kernel/querytags_nul.go @@ -0,0 +1,30 @@ +package kernel + +import ( + "errors" + "strings" +) + +// This file is intentionally NOT behind the `cgo && databricks_kernel` build tag: +// the interior-NUL guard on the serialized query-tags string is pure Go, so it is +// unit-tested under CGO_ENABLED=0 (see querytags_nul_test.go). The tagged execute +// path (operation.go) calls checkQueryTags before newCStr(serialized). + +// errQueryTagsNUL rejects a serialized query-tags string containing an interior +// NUL. The kernel's set_query_tags C ABI takes it as a NUL-terminated C string +// with no length, so a NUL would silently truncate it — sending fewer/corrupted +// tags than intended — whereas the Thrift path puts the same tags in +// confOverlay["query_tags"] and transmits them whole. This is the query-tags +// counterpart to errQueryNUL (sqltext.go) and errParamNUL (bindparams.go), which +// guard SQL text and bound values for the identical reason. Fail loudly rather +// than diverge from Thrift. +var errQueryTagsNUL = errors.New("query tags contain a NUL byte, which the kernel set_query_tags ABI cannot carry") + +// checkQueryTags validates the serialized query-tags string before the cgo layer +// C-string-marshals it. Returns errQueryTagsNUL when it contains an interior NUL. +func checkQueryTags(serialized string) error { + if strings.IndexByte(serialized, 0) >= 0 { + return errQueryTagsNUL + } + return nil +} diff --git a/internal/backend/kernel/querytags_nul_test.go b/internal/backend/kernel/querytags_nul_test.go new file mode 100644 index 00000000..602a5cc1 --- /dev/null +++ b/internal/backend/kernel/querytags_nul_test.go @@ -0,0 +1,27 @@ +package kernel + +import ( + "errors" + "testing" +) + +// checkQueryTags must reject a serialized query-tags string with an interior NUL +// (the kernel's set_query_tags ABI would truncate it, sending fewer tags than +// Thrift) but accept ordinary serialized tags and an empty string. Runs under +// CGO_ENABLED=0. +func TestCheckQueryTags(t *testing.T) { + if err := checkQueryTags("team:eng,job:etl,production"); err != nil { + t.Errorf("plain tags: got %v, want nil", err) + } + if err := checkQueryTags(""); err != nil { + t.Errorf("empty: got %v, want nil", err) + } + // NUL inside a value (querytags.Serialize escapes only \\, :, and , — never NUL). + if err := checkQueryTags("team:e\x00ng"); !errors.Is(err, errQueryTagsNUL) { + t.Errorf("NUL in value: got %v, want errQueryTagsNUL", err) + } + // NUL inside a key must be caught too. + if err := checkQueryTags("te\x00am:eng"); !errors.Is(err, errQueryTagsNUL) { + t.Errorf("NUL in key: got %v, want errQueryTagsNUL", err) + } +} From 949edc8e18bf114bf775c00179c47e9d606f1eea Mon Sep 17 00:00:00 2001 From: Rahul Singhal Date: Fri, 4 Sep 2026 19:08:31 +0000 Subject: [PATCH 2/2] chore(kernel): bump KERNEL_REV to include the query-tags C ABI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Advance the pinned kernel to 21504eae074c7f5b11888a4ebcc7468780b2f844 (databricks-sql-kernel main, which merged kernel_statement_set_query_tags via #314) and sync the committed C-ABI header to that rev. The databricks_kernel CI leg source-builds the kernel .a from KERNEL_REV (make kernel-lib), so this is what lets it link the new setter and run the tagged tests. The header sync also pulls in the other C-ABI declarations that landed since the previous pin (max_connections, telemetry circuit breaker, a result-stream canceller type, and the kernel_session_test timeout param); the Go backend calls none of them, so no binding code changes are needed. Note: the published databricks-sql-kernel-bindings modules (go.mod, still v1.0.0) are NOT bumped here — a matching bindings release is a fast-follow for `go get -tags databricks_kernel` consumers. CI and this repo's tagged tests build from KERNEL_REV and are unaffected. Co-authored-by: Isaac Signed-off-by: Rahul Singhal --- KERNEL_REV | 2 +- internal/backend/kernel/include/databricks_kernel.h | 10 ++++++---- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/KERNEL_REV b/KERNEL_REV index 110c974a..b914d0a7 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -167651ecc67143ef258ad70fd2682ced00beaa99 +21504eae074c7f5b11888a4ebcc7468780b2f844 diff --git a/internal/backend/kernel/include/databricks_kernel.h b/internal/backend/kernel/include/databricks_kernel.h index c00e1f97..525f2f8b 100644 --- a/internal/backend/kernel/include/databricks_kernel.h +++ b/internal/backend/kernel/include/databricks_kernel.h @@ -694,10 +694,12 @@ KernelStatusCode kernel_statement_bind_parameter(kernel_statement_t* stmt, * serialized wire form: comma-separated `key:value` pairs, with a bare `key` * (no colon) for a valueless tag, e.g. "team:eng,job:etl,production"; backslash * escapes the separators in a value ("\:", "\,", "\\"). This is the same string - * a host already builds for the Thrift `confOverlay` path, so it is passed - * through unchanged and re-parsed into the native SEA query_tags array at - * execute time. An empty string is a no-op; calling again replaces the - * statement's tags. `query_tags` must be a NUL-terminated UTF-8 string. + * a host already builds for the Thrift `confOverlay` path (Go/Node connectors), + * so it is passed through unchanged and re-parsed into the native SEA query_tags + * array at execute time. A non-empty value replaces the statement's tags; an + * empty string clears any previously-set tags (they otherwise persist across + * kernel_statement_set_sql, so pass "" to reset a reused statement handle). + * `query_tags` must be a NUL-terminated UTF-8 string. */ KernelStatusCode kernel_statement_set_query_tags(kernel_statement_t* stmt, const char* query_tags);