diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml index 92e32bb6..ed9c95e1 100644 --- a/.github/workflows/go.yml +++ b/.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 diff --git a/CHANGELOG.md b/CHANGELOG.md index 0b736e9a..f4bd5356 100644 --- a/CHANGELOG.md +++ b/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). diff --git a/CONNECTION_PARAMETERS.md b/CONNECTION_PARAMETERS.md index 35a4ad0e..9469e29a 100644 --- a/CONNECTION_PARAMETERS.md +++ b/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. diff --git a/KERNEL_REV b/KERNEL_REV index d025be44..b1078d37 100644 --- a/KERNEL_REV +++ b/KERNEL_REV @@ -1 +1 @@ -6f4b6d607a61bdde61dcda2cf82bc8e80fd32cb9 +37c7a3ec1b54f4690d9325d65f2220e0bb64b83e diff --git a/connection.go b/connection.go index 19de7f76..548fd589 100644 --- a/connection.go +++ b/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 diff --git a/connection_test.go b/connection_test.go index 59be9150..8a9f9fe1 100644 --- a/connection_test.go +++ b/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) +} diff --git a/doc.go b/doc.go index e78f5a61..9d43ddac 100644 --- a/doc.go +++ b/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 diff --git a/examples/kernel/main.go b/examples/kernel/main.go index 51d99c31..e3f19775 100644 --- a/examples/kernel/main.go +++ b/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) - } } diff --git a/internal/backend/backend.go b/internal/backend/backend.go index 37fa1a43..7bfc239c 100644 --- a/internal/backend/backend.go +++ b/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 diff --git a/internal/backend/kernel/backend.go b/internal/backend/kernel/backend.go index f8e6ebb7..c82a3f00 100644 --- a/internal/backend/kernel/backend.go +++ b/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) } diff --git a/internal/backend/kernel/include/databricks_kernel.h b/internal/backend/kernel/include/databricks_kernel.h index 965c09b6..0fe272f7 100644 --- a/internal/backend/kernel/include/databricks_kernel.h +++ b/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 diff --git a/internal/backend/kernel/kernel_test.go b/internal/backend/kernel/kernel_test.go index cd379412..959b87e9 100644 --- a/internal/backend/kernel/kernel_test.go +++ b/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) + } + }) } } diff --git a/internal/backend/kernel/operation.go b/internal/backend/kernel/operation.go index e2f355f0..ce4efa0d 100644 --- a/internal/backend/kernel/operation.go +++ b/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,26 @@ 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, _ := json.Marshal(paths) + 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 +404,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 diff --git a/internal/backend/kernel/staging.go b/internal/backend/kernel/staging.go index c614169c..52cc59f7 100644 --- a/internal/backend/kernel/staging.go +++ b/internal/backend/kernel/staging.go @@ -6,24 +6,14 @@ import "strings" // Staging detection is pure Go (string inspection of the SQL), so keeping it // untagged lets its test run under CGO_ENABLED=0. -// isStagingStatement reports whether sql is a Unity Catalog volume staging -// command (PUT / GET / REMOVE). The Thrift backend learns a statement is staging -// from server result metadata (IsStagingOperation); the kernel C ABI exposes no -// such signal, and the kernel path cannot perform the local file transfer these -// commands require. We therefore detect them from the SQL and reject at execute -// time (see Execute) rather than silently returning success with no file moved — -// which is what returning false from IsStaging would do. Case-insensitive; skips -// leading whitespace AND leading SQL comments (`-- ...` and `/* ... */`), so a -// comment-prefixed staging command is still caught rather than slipping through -// as a silent no-op. Errs toward matching only a clear leading keyword — a false -// negative just yields the kernel's normal "unsupported syntax" server error, -// whereas a false positive would reject a legitimate query. +// isStagingStatement detects PUT, GET, or REMOVE after leading SQL comments, +// ignoring keyword case. func isStagingStatement(sql string) bool { s := stripLeadingSQLComments(sql) for _, kw := range []string{"PUT", "GET", "REMOVE"} { if len(s) > len(kw) && strings.EqualFold(s[:len(kw)], kw) && - isSQLTokenBreak(s[len(kw)]) { + (isSQLTokenBreak(s[len(kw)]) || strings.HasPrefix(s[len(kw):], "/*") || strings.HasPrefix(s[len(kw):], "--")) { return true } } @@ -33,9 +23,9 @@ func isStagingStatement(sql string) bool { // stripLeadingSQLComments removes leading whitespace and leading SQL comments // (both `-- ...` line comments and `/* ... */` block comments), repeatedly, so // the first real token is exposed. An unterminated block comment consumes the -// rest of the string (yielding ""). Comment bodies are not scanned for nested -// markers — SQL block comments do not nest — and this only touches the LEADING -// run, so a comment later in the statement is untouched. +// rest of the string (yielding ""). Nested block comments must be balanced so +// their trailing text cannot hide a staging keyword. Only the leading run is +// removed; comments later in the statement are untouched. func stripLeadingSQLComments(sql string) string { s := strings.TrimSpace(sql) for { @@ -47,11 +37,23 @@ func stripLeadingSQLComments(sql string) string { return "" // line comment to end of input } case strings.HasPrefix(s, "/*"): - if i := strings.Index(s[2:], "*/"); i >= 0 { - s = s[2+i+2:] - } else { + depth, i := 1, 2 + for i < len(s) && depth > 0 { + switch { + case strings.HasPrefix(s[i:], "/*"): + depth++ + i += 2 + case strings.HasPrefix(s[i:], "*/"): + depth-- + i += 2 + default: + i++ + } + } + if depth != 0 { return "" // unterminated block comment } + s = s[i:] default: return s } @@ -59,9 +61,9 @@ func stripLeadingSQLComments(sql string) string { } } -// isSQLTokenBreak reports whether b ends a leading keyword — a space or a quote +// isSQLTokenBreak reports whether b ends a leading keyword, whitespace or a quote // (staging commands take a quoted path, e.g. PUT '' INTO ''), so // "PUT '..." matches but an identifier like "PUTS" or "GETTER" does not. func isSQLTokenBreak(b byte) bool { - return b == ' ' || b == '\t' || b == '\n' || b == '\r' || b == '\'' || b == '"' + return b == ' ' || b == '\t' || b == '\n' || b == '\r' || b == '\f' || b == '\v' || b == '\'' || b == '"' } diff --git a/internal/backend/kernel/staging_cancel_test.go b/internal/backend/kernel/staging_cancel_test.go new file mode 100644 index 00000000..c4c174c8 --- /dev/null +++ b/internal/backend/kernel/staging_cancel_test.go @@ -0,0 +1,182 @@ +//go:build cgo && databricks_kernel + +package kernel + +import ( + "bytes" + "context" + "encoding/base64" + "encoding/json" + "errors" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + "github.com/apache/arrow/go/v12/arrow" + "github.com/apache/arrow/go/v12/arrow/array" + "github.com/apache/arrow/go/v12/arrow/ipc" + "github.com/apache/arrow/go/v12/arrow/memory" + "github.com/databricks/databricks-sql-go/driverctx" + "github.com/databricks/databricks-sql-go/internal/backend" +) + +func TestExecuteCancelsFileTransfer(t *testing.T) { + for _, verb := range []string{"PUT", "GET"} { + t.Run(verb, func(t *testing.T) { + directory := t.TempDir() + local := filepath.Join(directory, "file") + original := bytes.Repeat([]byte("a"), 8*1024*1024) + if err := os.WriteFile(local, original, 0600); err != nil { + t.Fatal(err) + } + started := make(chan struct{}) + var startedOnce sync.Once + release := make(chan struct{}) + var instruction []byte + instructionReady := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch { + case r.URL.Path == "/storage": + if verb == "PUT" { + if _, err := io.CopyN(io.Discard, r.Body, 1); err != nil { + return + } + } else { + _, _ = w.Write([]byte("partial")) + w.(http.Flusher).Flush() + } + startedOnce.Do(func() { close(started) }) + select { + case <-r.Context().Done(): + case <-release: + } + case r.Method == "POST" && r.URL.Path == "/api/2.0/sql/statements": + <-instructionReady + _, _ = w.Write(instruction) + case strings.HasSuffix(r.URL.Path, "/sessions"): + _, _ = io.WriteString(w, `{"session_id":"session","warehouse_id":"test"}`) + default: + _, _ = io.WriteString(w, `{}`) + } + })) + defer server.Close() + instruction = fileTransferInstruction(t, verb, server.URL+"/storage", local) + close(instructionReady) + k := New(Config{Host: server.URL, WarehouseID: "test", Auth: Auth{Mode: AuthPAT, Token: "test"}}) + if err := k.OpenSession(context.Background()); err != nil { + t.Fatal(err) + } + defer k.CloseSession(context.Background()) + for _, rawCtx := range []context.Context{ + context.Background(), + driverctx.NewContextWithStagingInfo(context.Background(), []string{directory}), + } { + raw, err := k.Execute(rawCtx, backend.ExecRequest{Query: "get '/Volumes/c/s/v/file' to '" + local + "'"}) + if err != nil { + t.Fatalf("raw execution: %v", err) + } + if _, err := raw.Close(context.Background()); err != nil { + t.Fatal(err) + } + } + select { + case <-started: + t.Fatal("raw execution started a file transfer") + default: + } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + ctx = driverctx.NewContextWithStagingInfo(ctx, []string{directory}) + finished := make(chan error, 1) + stopped := make(chan struct{}) + query := "PUT '" + local + "' INTO '/Volumes/c/s/v/file' OVERWRITE" + if verb == "GET" { + query = "GET '/Volumes/c/s/v/file' TO '" + local + "'" + } + go func() { + defer close(stopped) + op, err := k.Execute(ctx, backend.ExecRequest{Query: query, TransferFiles: true}) + if op != nil { + _, _ = op.Close(context.Background()) + } + finished <- err + }() + defer func() { + close(release) + <-stopped + }() + select { + case <-started: + case err := <-finished: + t.Fatalf("execute returned before transfer: %v", err) + case <-time.After(5 * time.Second): + t.Fatal("transfer did not start") + } + cancel() + select { + case err := <-finished: + if !errors.Is(err, context.Canceled) { + t.Fatalf("execute error = %v, want context.Canceled", err) + } + case <-time.After(2 * time.Second): + t.Fatal("cancelled execute is still waiting for the transfer") + } + if verb == "GET" { + data, err := os.ReadFile(local) + if err != nil || !bytes.Equal(data, original) { + t.Fatalf("cancelled GET changed the destination: %v", err) + } + } + }) + } +} + +func fileTransferInstruction(t *testing.T, verb, url, local string) []byte { + t.Helper() + names := []string{"operation", "url", "headers", "local"} + fields := make([]arrow.Field, len(names)) + columns := make([]map[string]any, len(names)) + for i, name := range names { + fields[i] = arrow.Field{Name: name, Type: arrow.BinaryTypes.String} + columns[i] = map[string]any{"name": name, "type_name": "STRING", "type_text": "STRING", "position": i} + } + schema := arrow.NewSchema(fields, nil) + builder := array.NewRecordBuilder(memory.NewGoAllocator(), schema) + defer builder.Release() + for i, value := range []string{verb, url, "{}", local} { + builder.Field(i).(*array.StringBuilder).Append(value) + } + record := builder.NewRecord() + defer record.Release() + var buffer bytes.Buffer + writer := ipc.NewWriter(&buffer, ipc.WithSchema(schema)) + if err := writer.Write(record); err != nil { + t.Fatal(err) + } + if err := writer.Close(); err != nil { + t.Fatal(err) + } + body, err := json.Marshal(map[string]any{ + "statement_id": "transfer", + "status": map[string]any{"state": "SUCCEEDED"}, + "manifest": map[string]any{ + "format": "ARROW_STREAM", "is_volume_operation": true, + "schema": map[string]any{"column_count": len(names), "columns": columns}, + "total_chunk_count": 1, "total_row_count": 1, "total_byte_count": buffer.Len(), "truncated": false, + }, + "result": map[string]any{ + "chunk_index": 0, "row_offset": 0, "row_count": 1, "byte_count": buffer.Len(), + "attachment": base64.StdEncoding.EncodeToString(buffer.Bytes()), + }, + }) + if err != nil { + t.Fatal(err) + } + return body +} diff --git a/internal/backend/kernel/staging_test.go b/internal/backend/kernel/staging_test.go index 91d69794..8a29843c 100644 --- a/internal/backend/kernel/staging_test.go +++ b/internal/backend/kernel/staging_test.go @@ -13,6 +13,11 @@ func TestIsStagingStatement(t *testing.T) { "/* a */ /* b */ PUT '/l' INTO '/Volumes/x'", // multiple block comments "-- upload the file\nPUT '/l' INTO '/Volumes/x'", // leading line comment, then staging " /* c */\n -- d\n REMOVE '/Volumes/x'", // mixed comments + whitespace + "/* outer /* inner */ outer */ REMOVE '/Volumes/x'", + "REMOVE/* note */ '/Volumes/x'", + "REMOVE-- note\n'/Volumes/x'", + "REMOVE\f'/Volumes/x'", + "REMOVE\v'/Volumes/x'", } for _, sql := range staging { if !isStagingStatement(sql) { @@ -32,6 +37,7 @@ func TestIsStagingStatement(t *testing.T) { "SELECT 'PUT ' AS col", // keyword only inside a literal "/* unterminated PUT '/x'", // unterminated block comment consumes the rest "/* c */ SELECT 1", // real leading comment, but not a staging command + "/* outer /* inner */ REMOVE '/Volumes/x'", } for _, sql := range notStaging { if isStagingStatement(sql) { diff --git a/kernel_e2e_test.go b/kernel_e2e_test.go index 55feacae..51c3cfdd 100644 --- a/kernel_e2e_test.go +++ b/kernel_e2e_test.go @@ -8,8 +8,11 @@ import ( "database/sql" "database/sql/driver" "errors" + "fmt" "io" "os" + "path/filepath" + "strings" "sync" "sync/atomic" "testing" @@ -17,6 +20,7 @@ import ( "github.com/databricks/databricks-sql-go/driverctx" dbsqlerr "github.com/databricks/databricks-sql-go/errors" + dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors" dbsqlrows "github.com/databricks/databricks-sql-go/rows" ) @@ -671,3 +675,105 @@ func TestKernelE2EGetArrowBatches(t *testing.T) { t.Errorf("GetArrowBatches row count = %d, want %d", rowCount, want) } } + +func TestKernelE2EVolumeFileLifecycle(t *testing.T) { + catalog := os.Getenv("DATABRICKS_VOLUME_TEST_CATALOG") + if catalog == "" { + t.Skip("set DATABRICKS_VOLUME_TEST_CATALOG to a writable catalog") + } + db := kernelTestDB(t) + defer db.Close() + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute) + defer cancel() + schema := fmt.Sprintf("driver_volume_%d", time.Now().UnixNano()) + identifier := func(value string) string { return "`" + strings.ReplaceAll(value, "`", "``") + "`" } + literal := func(value string) string { return "'" + strings.NewReplacer(`\`, `\\`, `'`, `\'`).Replace(value) + "'" } + namespace := identifier(catalog) + "." + identifier(schema) + if _, err := db.ExecContext(ctx, "CREATE SCHEMA "+namespace); err != nil { + t.Fatal(err) + } + defer func() { + cleanupCtx, cancelCleanup := context.WithTimeout(context.Background(), time.Minute) + defer cancelCleanup() + if _, err := db.ExecContext(cleanupCtx, "DROP SCHEMA "+namespace+" CASCADE"); err != nil { + t.Error(err) + } + }() + if _, err := db.ExecContext(ctx, "CREATE VOLUME "+namespace+".files"); err != nil { + t.Fatal(err) + } + directory := t.TempDir() + source, destination := filepath.Join(directory, "input.bin"), filepath.Join(directory, "output.bin") + payload := bytes.Repeat([]byte{0, 255, 1, 128}, 262144) + if err := os.WriteFile(source, payload, 0600); err != nil { + t.Fatal(err) + } + remote := "/Volumes/" + catalog + "/" + schema + "/files/input.bin" + stagingCtx := driverctx.NewContextWithStagingInfo(ctx, []string{directory}) + if _, err := db.ExecContext(stagingCtx, "PUT "+literal(source)+" INTO "+literal(remote)+" OVERWRITE"); err != nil { + t.Fatal(err) + } + if _, err := db.ExecContext(ctx, "REMOVE "+literal(remote)); err == nil || !strings.Contains(err.Error(), "staging ctx must be provided") { + t.Fatalf("REMOVE without staging context: %v", err) + } + get := "GET " + literal(remote) + " TO " + literal(destination) + rows, err := db.QueryContext(stagingCtx, get) + if err != nil { + t.Fatal(err) + } + defer rows.Close() + columns, err := rows.Columns() + if err != nil || len(columns) != 4 { + t.Fatalf("raw instruction columns: %v, %v", columns, err) + } + if !rows.Next() { + t.Fatalf("missing raw instruction: %v", rows.Err()) + } + var operation, storageURL, local string + var headers any + if err := rows.Scan(&operation, &storageURL, &headers, &local); err != nil { + t.Fatal(err) + } + if operation != "GET" || local != destination || storageURL == "" { + t.Fatalf("unexpected raw instruction: operation=%q local=%q", operation, local) + } + if rows.Next() || rows.Err() != nil { + t.Fatalf("expected one raw instruction: %v", rows.Err()) + } + if err := rows.Close(); err != nil { + t.Fatal(err) + } + if _, err := os.Stat(destination); !os.IsNotExist(err) { + t.Fatalf("raw instructions changed local destination: %v", err) + } + if _, err := db.ExecContext(stagingCtx, get); err != nil { + t.Fatal(err) + } + actual, err := os.ReadFile(destination) + if err != nil || !bytes.Equal(actual, payload) { + t.Fatalf("download differs: %v", err) + } + denied := filepath.Join(t.TempDir(), "denied.bin") + if err := os.WriteFile(denied, payload, 0600); err != nil { + t.Fatal(err) + } + for _, test := range []struct { + local string + category dbsqlerrint.ErrorCategory + }{ + {denied, dbsqlerrint.CategoryPermissionError}, + {filepath.Join(directory, "missing.bin"), dbsqlerrint.CategoryNotFound}, + } { + if _, err := db.ExecContext(stagingCtx, "PUT "+literal(test.local)+" INTO "+literal(remote)+" OVERWRITE"); err == nil || dbsqlerrint.CategoryFromError(err) != test.category { + t.Fatalf("PUT %q: expected category %q, got %v", test.local, test.category, err) + } + } + for range 2 { + if _, err := db.ExecContext(stagingCtx, "REMOVE "+literal(remote)); err != nil { + t.Fatal(err) + } + } + if _, err := db.ExecContext(stagingCtx, get); err == nil { + t.Fatal("GET succeeded after REMOVE") + } +}