feat(volume): enable native file transfers through ExecContext

.github/workflows/go.yml

@@ -105,6 +105,31 @@ jobs:
       - name: Build
         run: make linux
 
+  build-published-kernel:
+    name: Build (published kernel)
+    timeout-minutes: 15
+    runs-on:
+      group: databricks-protected-runner-group
+      labels: linux-ubuntu-latest
+    steps:
+      - name: Check out code into the Go module directory
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+
+      - name: Setup JFrog
+        uses: ./.github/actions/setup-jfrog
+
+      - name: Set up Go Toolchain
+        uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6.5.0
+        with:
+          go-version: '1.27.3'
+          cache: false
+
+      - name: Link against published native modules
+        run: go test -mod=readonly -tags databricks_kernel -run '^$' ./...
+        env:
+          CGO_ENABLED: 1
+          GOWORK: off
+
   build-and-test-kernel:
     name: Test (kernel backend)
     # Exercises the opt-in SEA-via-kernel backend: builds the Rust kernel static

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Release History
 
+## Unreleased
+- Support SQL PUT, GET, and REMOVE through kernel-backed `ExecContext` with a staging context. Kernel-backed `QueryContext` now returns raw transfer instructions, matching Thrift, instead of `ErrNotSupportedByKernel`; it does not transfer files. Callers should review fallback handling that depends on that error and use `ExecContext` to complete transfers.
+
 ## v1.19.0 (2026-11-04)
 - Upgrade the SQL kernel bindings to v1.4.0.
 - Upgrade go-oidc to v3.21.0 (go-jose v4.1.6) and x/oauth2 to v0.36.0 to address go-jose/v3 security findings (databricks/databricks-sql-go#526).

CONNECTION_PARAMETERS.md

@@ -237,3 +237,12 @@ unset), kernel telemetry follows its server feature flag, and `telemetry_batch_s
 The Go wrapper telemetry interceptor is skipped on the kernel path so it does not
 duplicate kernel-owned telemetry for the same connection and statements. See
 [`telemetry/DESIGN.md`](./telemetry/DESIGN.md).
+
+## Volume file transfers
+
+Both backends support SQL PUT, GET, and REMOVE through `ExecContext` with
+`driverctx.NewContextWithStagingInfo(ctx, allowedLocalPaths)`. The context must
+contain at least one allowed path, including for REMOVE. `QueryContext` returns
+the transfer instructions for callers that handle file transfers themselves.
+
+The credential-gated `TestKernelE2EVolumeFileLifecycle` uses `DATABRICKS_VOLUME_TEST_CATALOG` as a writable catalog and creates and removes an isolated schema and Volume. It covers ExecContext transfers and QueryContext raw instructions.

KERNEL_REV

@@ -1 +1 @@
-6f4b6d607a61bdde61dcda2cf82bc8e80fd32cb9
+37c7a3ec1b54f4690d9325d65f2220e0bb64b83e

connection.go

@@ -148,7 +148,7 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name
 
 	// Capture execution start time for telemetry before running the query
 	executeStart := time.Now()
-	op, err := c.runQuery(ctx, query, args)
+	op, err := c.runQuery(ctx, query, args, true)
 	// A nil op means execution never reached the backend (parameter conversion
 	// failed): no server statement to close or measure. Return the already-wrapped
 	// error directly.
@@ -278,7 +278,7 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam
 
 	// Capture execution start time for telemetry before running the query
 	executeStart := time.Now()
-	op, err := c.runQuery(ctx, query, args)
+	op, err := c.runQuery(ctx, query, args, false)
 	// A nil op means execution never reached the backend (parameter conversion
 	// failed): no statement id, no telemetry to emit. Log and return the
 	// already-wrapped error directly.
@@ -437,12 +437,12 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam
 // which case the wrapped error is returned directly. Once the backend is called,
 // backend.Execute guarantees a non-nil Operation, so callers guard only for the
 // pre-backend nil.
-func (c *conn) runQuery(ctx context.Context, query string, args []driver.NamedValue) (backend.Operation, error) {
+func (c *conn) runQuery(ctx context.Context, query string, args []driver.NamedValue, transferFiles bool) (backend.Operation, error) {
 	params, err := convertNamedValuesToParams(args)
 	if err != nil {
 		return nil, dbsqlerrint.NewExecutionError(ctx, dbsqlerr.ErrQueryExecution, err, nil)
 	}
-	return c.backend.Execute(ctx, backend.ExecRequest{Query: query, Params: params})
+	return c.backend.Execute(ctx, backend.ExecRequest{Query: query, Params: params, TransferFiles: transferFiles})
 }
 
 // enrichQueryId sets the operation's statement id as the context queryId. A

connection_test.go

@@ -18,6 +18,7 @@ import (
 	"github.com/pkg/errors"
 
 	dbsqlerr "github.com/databricks/databricks-sql-go/errors"
+	"github.com/databricks/databricks-sql-go/internal/backend"
 	thriftbackend "github.com/databricks/databricks-sql-go/internal/backend/thrift"
 	"github.com/databricks/databricks-sql-go/internal/cli_service"
 	"github.com/databricks/databricks-sql-go/internal/client"
@@ -974,3 +975,28 @@ func TestConn_handleStagingRetry(t *testing.T) {
 		assert.Less(t, elapsed, 2*time.Second, "context cancel should abort PUT retry backoff promptly")
 	})
 }
+
+type transferCaptureBackend struct {
+	backend.Backend
+	transferFiles []bool
+}
+
+func (b *transferCaptureBackend) Execute(ctx context.Context, req backend.ExecRequest) (backend.Operation, error) {
+	b.transferFiles = append(b.transferFiles, req.TransferFiles)
+	return b.Backend.Execute(ctx, req)
+}
+
+func TestConnectionTransferMode(t *testing.T) {
+	testClient := &client.TestClient{
+		FnExecuteStatement: func(context.Context, *cli_service.TExecuteStatementReq) (*cli_service.TExecuteStatementResp, error) {
+			return &cli_service.TExecuteStatementResp{}, fmt.Errorf("capture request")
+		},
+	}
+	capture := &transferCaptureBackend{Backend: thriftbackend.NewForTest(testClient, getTestSession(), config.WithDefaults())}
+	connection := &conn{cfg: config.WithDefaults(), backend: capture}
+	_, execErr := connection.ExecContext(context.Background(), "REMOVE '/Volumes/c/s/v/file'", nil)
+	_, queryErr := connection.QueryContext(context.Background(), "GET '/Volumes/c/s/v/file' TO '/tmp/file'", nil)
+	assert.Error(t, execErr)
+	assert.Error(t, queryErr)
+	assert.Equal(t, []bool{true, false}, capture.transferFiles)
+}

doc.go

@@ -220,16 +220,17 @@ RUST_LOG. Filter on the target databricks::sql::kernel (note the colons):
 Supported on the kernel backend: PAT and OAuth (M2M via WithClientCredentials, U2M
 via the authType=oauthU2M DSN param); reading scalar, nested, and complex-typed
 results (CloudFetch is transparent); bound query parameters (positional and named);
-context cancellation during execute; the initial namespace (WithInitialNamespace,
+context cancellation during execute; SQL PUT/GET/REMOVE on Unity Catalog Volumes
+through ExecContext with a staging context (QueryContext returns raw instructions);
+the initial namespace (WithInitialNamespace,
 applied post-connect via USE CATALOG / USE SCHEMA); metric-view metadata
 (WithEnableMetricViewMetadata); the retry / backoff policy (WithRetries:
 RetryWaitMin / RetryWaitMax / RetryMax, including the disable form, forwarded to the
 kernel's HTTP retry config); and the TLS, proxy, and session-conf (query tags,
 statement timeout, time zone) options. Federated token providers use one token snapshot
 during setup. Nothing is silently ignored: WithTimeout and token-provider / external /
-static authenticators are rejected at connect; staging (PUT/GET/REMOVE on a
-Unity Catalog volume) is rejected at execute. WithMaxRows is accepted but inert (the
-kernel manages fetching below the C ABI).
+static authenticators are rejected at connect. WithMaxRows is accepted but inert
+(the kernel manages fetching below the C ABI).
 
 Kernel HTTP requests inherit the driver's ClientTimeout (900s by default). If
 that value is zero, the kernel uses its own 120s default; zero is neither

examples/kernel/main.go

@@ -87,14 +87,4 @@ func main() {
 		log.Fatalf("query: %v", err)
 	}
 	fmt.Printf("server version: %s\n", version)
-
-	// Options the kernel can't honor are rejected with an error wrapping
-	// ErrNotSupportedByKernel, never silently ignored. Staging (PUT/GET/REMOVE on a
-	// Unity Catalog volume) is one such feature, rejected at execute.
-	_, err = db.ExecContext(ctx, "PUT '/tmp/local.csv' INTO '/Volumes/main/default/vol/f.csv' OVERWRITE")
-	if errors.Is(err, dbsqlerr.ErrNotSupportedByKernel) {
-		fmt.Println("staging is not supported on the kernel backend (use Thrift for staging)")
-	} else if err != nil {
-		fmt.Printf("staging: %v\n", err)
-	}
 }

internal/backend/backend.go

@@ -119,6 +119,9 @@ type ExecRequest struct {
 	// dbsql.Parameter/SqlType types; each backend maps these neutral params to its
 	// own wire type.
 	Params []Param
+	// TransferFiles requests native completion where the backend supports it.
+	// Other backends leave transfers to the connection layer.
+	TransferFiles bool
 }
 
 // Param is one bound query parameter in backend-neutral form: the name (empty

internal/backend/kernel/backend.go

@@ -71,8 +71,9 @@ import (
 	"sync/atomic"
 	"time"
 
-	dbsqlerr "github.com/databricks/databricks-sql-go/errors"
+	"github.com/databricks/databricks-sql-go/driverctx"
 	"github.com/databricks/databricks-sql-go/internal/backend"
+	dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors"
 )
 
 // kernelSessionSeq is a process-wide monotonic counter for kernel session ids.
@@ -890,15 +891,8 @@ func (k *KernelBackend) SessionID() string { return k.sessionID }
 // parameters (req.Params) are bound onto the statement in execute (see
 // bindParams).
 func (k *KernelBackend) Execute(ctx context.Context, req backend.ExecRequest) (backend.Operation, error) {
-	// Staging (Unity Catalog volume PUT/GET/REMOVE) needs a local file transfer the
-	// kernel path can't perform, and the kernel C ABI surfaces no IsStagingOperation
-	// signal to drive conn.execStagingOperation. Reject it here rather than let
-	// IsStaging return false and report success with no file moved (silent data loss).
-	// (Bound parameters, by contrast, ARE now supported — bound in execute via the
-	// kernel's raw-param C ABI.)
-	if isStagingStatement(req.Query) {
-		return &kernelOp{}, fmt.Errorf("databricks: staging operations (PUT/GET/REMOVE on a volume) are %w; "+
-			"use the default (Thrift) backend", dbsqlerr.ErrNotSupportedByKernel)
+	if req.TransferFiles && isStagingStatement(req.Query) && len(driverctx.StagingPathsFromContext(ctx)) == 0 {
+		return &kernelOp{}, dbsqlerrint.NewDriverError(ctx, "staging ctx must be provided.", nil)
 	}
 	return k.execute(ctx, req)
 }

internal/backend/kernel/include/databricks_kernel.h

@@ -688,6 +688,11 @@ KernelStatusCode kernel_session_rollback(kernel_session_t* session);
  * so a reused statement handle starts each query with a fresh parameter set. */
 KernelStatusCode kernel_statement_set_sql(kernel_statement_t* stmt, const char* sql);
 
+/** Configure allowed local paths; null restores caller-managed transfers.
+ * retain_result keeps instructions after transfer. */
+KernelStatusCode kernel_statement_set_staging_allowed_local_paths(
+    kernel_statement_t* stmt, const char* paths_json, bool retain_result);
+
 /* Typed positional parameter binding (1-based ordinal); the kernel infers the
  * SEA wire type from the C type. These are the "typed params" the
  * kernel_statement_bind_parameter contract refers to: within one statement they

internal/backend/kernel/kernel_test.go

@@ -531,30 +531,22 @@ func TestExecuteHandleLessOpContract(t *testing.T) {
 	}
 }
 
-// TestExecuteRejectsStaging drives Execute with a staging statement (not just the
-// isStagingStatement detector in isolation) to pin the detector→Execute wiring: a
-// refactor that dropped or reordered the check would silently reopen the
-// silent-no-op data-loss path. Mirrors TestExecuteRejectsParams.
-func TestExecuteRejectsStaging(t *testing.T) {
-	k := &KernelBackend{}
-	op, err := k.Execute(context.Background(), backend.ExecRequest{
-		Query: "PUT '/tmp/f' INTO '/Volumes/main/s/e/f.csv'",
-	})
-	if err == nil {
-		t.Fatal("expected an error for a staging statement, got nil")
-	}
-	if !errors.Is(err, dbsqlerr.ErrNotSupportedByKernel) {
-		t.Errorf("staging rejection should wrap ErrNotSupportedByKernel, got %v", err)
-	}
-	if op == nil {
-		t.Fatal("Execute must return a non-nil Operation per the Backend contract")
-	}
-	closed, closeErr := op.Close(context.Background())
-	if closeErr != nil {
-		t.Errorf("Close error = %v, want nil", closeErr)
-	}
-	if closed {
-		t.Error("Close on a handle-less op must report closed=false (no CLOSE_STATEMENT)")
+func TestExecuteRequiresStagingContext(t *testing.T) {
+	for _, query := range []string{"GET '/Volumes/main/s/v/f' TO '/tmp/f'", "PUT '/tmp/f' INTO '/Volumes/main/s/v/f'", "REMOVE '/Volumes/main/s/v/f'", "/* upload */ PUT '/tmp/f' INTO '/Volumes/main/s/v/f'"} {
+		t.Run(query, func(t *testing.T) {
+			k := &KernelBackend{}
+			op, err := k.Execute(context.Background(), backend.ExecRequest{Query: query, TransferFiles: true})
+			if err == nil || !strings.Contains(err.Error(), "staging ctx must be provided") {
+				t.Fatalf("expected missing staging context error, got %v", err)
+			}
+			if op == nil {
+				t.Fatal("expected non-nil operation")
+			}
+			closed, closeErr := op.Close(context.Background())
+			if closed || closeErr != nil {
+				t.Fatalf("Close = %v, %v", closed, closeErr)
+			}
+		})
 	}
 }
 

internal/backend/kernel/operation.go

@@ -11,6 +11,7 @@ import "C"
 import (
 	"context"
 	"database/sql/driver"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"sync"
@@ -35,8 +36,7 @@ import (
 //     client-timeout path, let an in-flight cancel finish in the background so
 //     it cannot extend the client deadline
 //
-// Binds any query parameters (bindParams) before executing; staging statements
-// are rejected up front by Execute, so none reach here.
+// Binds query parameters and configures local file access before executing.
 func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (backend.Operation, error) {
 	// Log the SQL length, not the text: query bodies can carry PII/secrets in
 	// WHERE/INSERT/SET, and this goes to stderr. Matches the driver's own
@@ -78,6 +78,30 @@ func (k *KernelBackend) execute(ctx context.Context, req backend.ExecRequest) (b
 	}
 	sql.free()
 
+	// Every ExecContext opts into completion, including transfers missed by the SQL recognizer.
+	// An empty list denies local transfers instead of reporting raw instructions as success.
+	if req.TransferFiles {
+		paths := driverctx.StagingPathsFromContext(ctx)
+		if paths == nil {
+			paths = []string{}
+		}
+		encoded, err := json.Marshal(paths)
+		if err != nil {
+			C.kernel_statement_close(stmt)
+			return &kernelOp{}, fmt.Errorf("kernel: encode allowed local paths: %w", err)
+		}
+		value := newCStr(string(encoded))
+		err = call(func() C.KernelStatusCode {
+			return C.kernel_statement_set_staging_allowed_local_paths(stmt, value.c, false)
+		})
+		value.free()
+		if err != nil {
+			C.kernel_statement_close(stmt)
+			k.evictIfSessionFatal(err)
+			return &kernelOp{}, fmt.Errorf("kernel: allowed local paths: %w", toStatementError(err))
+		}
+	}
+
 	// Bind parameters. The driver hands us backend.Param{Name, Type, Value}: the
 	// value is already stringified, Type is the Databricks SQL type name, and a nil
 	// Value is SQL NULL (Type "VOID"). Each maps 1:1 onto
@@ -384,10 +408,8 @@ func (o *kernelOp) Results(ctx context.Context, callbacks *dbsqlrows.TelemetryCa
 	return newKernelRows(ctx, o, stream, callbacks)
 }
 
-// IsStaging reports whether this is a staging (PUT/GET/REMOVE) operation. Always
-// false on the kernel path: staging statements are rejected up front in
-// Execute (isStagingStatement), so no kernelOp is ever produced for one — there
-// is nothing for conn.execStagingOperation to act on here.
+// IsStaging is false because the kernel completes configured file transfers
+// before returning the result.
 func (o *kernelOp) IsStaging(ctx context.Context) (bool, error) { return false, nil }
 
 // Close best-effort closes the executed statement and its statement handle. It is