feat(kernel): send per-statement query tags on the kernel backend

KERNEL_REV

@@ -1 +1 @@
-167651ecc67143ef258ad70fd2682ced00beaa99
+21504eae074c7f5b11888a4ebcc7468780b2f844

internal/backend/kernel/include/databricks_kernel.h

@@ -688,6 +688,22 @@ 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 (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);
+
 /*
  * Wait-for-result execution. On success, `*out` holds an executed handle
  * released with `kernel_executed_statement_close`.

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.

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
+}

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)
+	}
+}