Add a SQL Statement Execution fake to the testserver

integration/libs/sqlexec/sqlexec_test.go

@@ -70,6 +70,9 @@ func TestSQLExecFailedStatement(t *testing.T) {
 	assert.Equal(t, sql.StatementStateFailed, se.State)
 	assert.NotEmpty(t, se.Code)
 	assert.NotEmpty(t, se.Message)
+	// The API populates status.sql_state (SQLSTATE) on failures, e.g. 42P01 for
+	// a missing relation; assert it survives into the typed error.
+	assert.NotEmpty(t, se.SQLState)
 }
 
 func TestSQLExecSubmitAndCancel(t *testing.T) {

libs/sqlexec/sqlexec.go

@@ -180,7 +180,13 @@ func (e *StatementError) Error() string {
 // Columns and Rows.
 type Result struct {
 	Columns []string
-	Rows    [][]string
+	// Rows holds every result row with each cell as a string. A SQL NULL and an
+	// empty string are indistinguishable here: the API wire format encodes NULL
+	// as JSON null (distinct from ""), but the SDK models cells as [][]string, so
+	// JSON null is decoded as "". The distinction is recoverable only with a
+	// null-aware cell type (e.g. [][]*string) or the ARROW_STREAM format; no
+	// caller needs it today, so we keep the simpler string rows.
+	Rows [][]string
 }
 
 // Scalar returns the top-left cell of the result, or "" when there are no rows.

libs/sqlexec/sqlexec_http_test.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/databricks/cli/libs/sqlexec"
 	"github.com/databricks/cli/libs/testserver"
+	"github.com/databricks/cli/libs/testserver/testsql"
 	"github.com/databricks/databricks-sdk-go"
 	"github.com/databricks/databricks-sdk-go/service/sql"
 	"github.com/stretchr/testify/assert"
@@ -14,9 +15,19 @@ import (
 
 // These tests drive the engine through a real SDK client over HTTP against the
 // in-process testserver, with the statement-execution endpoints programmed per
-// test. Unlike the mock-interface unit tests they exercise the full
-// request/response JSON serialization, and unlike the integration tests they are
-// hermetic and run on every PR without a warehouse.
+// test via server.HandleSQL matchers. Unlike the mock-interface unit tests they
+// exercise the full request/response JSON serialization, and unlike the
+// integration tests they are hermetic and run on every PR without a warehouse.
+// newServer returns a testserver with the default handlers installed, which is
+// where the SQL Statement Execution routes live (they delegate to the matchers
+// registered via server.HandleSQL).
+func newServer(t *testing.T) *testserver.Server {
+	t.Helper()
+	server := testserver.New(t)
+	testserver.AddDefaultHandlers(server)
+	return server
+}
+
 func httpClient(t *testing.T, server *testserver.Server) *sqlexec.Client {
 	t.Helper()
 	w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"})
@@ -26,14 +37,9 @@ func httpClient(t *testing.T, server *testserver.Server) *sqlexec.Client {
 }
 
 func TestHTTPExecuteSuccess(t *testing.T) {
-	server := testserver.New(t)
-	server.Handle("POST", "/api/2.0/sql/statements", func(testserver.Request) any {
-		return sql.StatementResponse{
-			StatementId: "s1",
-			Status:      &sql.StatementStatus{State: sql.StatementStateSucceeded},
-			Manifest:    &sql.ResultManifest{Schema: &sql.ResultSchema{Columns: []sql.ColumnInfo{{Name: "a"}, {Name: "b"}}}, TotalChunkCount: 1},
-			Result:      &sql.ResultData{DataArray: [][]string{{"1", "2"}}},
-		}
+	server := newServer(t)
+	server.HandleSQL("SELECT 1 AS a, 2 AS b", func(testsql.Request) testsql.Result {
+		return testsql.Result{Columns: []string{"a", "b"}, Rows: [][]string{{"1", "2"}}}
 	})
 
 	r, err := httpClient(t, server).Execute(t.Context(), "SELECT 1 AS a, 2 AS b")
@@ -43,42 +49,20 @@ func TestHTTPExecuteSuccess(t *testing.T) {
 }
 
 func TestHTTPExecutePolls(t *testing.T) {
-	server := testserver.New(t)
-	server.Handle("POST", "/api/2.0/sql/statements", func(testserver.Request) any {
-		return sql.StatementResponse{StatementId: "s1", Status: &sql.StatementStatus{State: sql.StatementStatePending}}
-	})
-	polls := 0
-	server.Handle("GET", "/api/2.0/sql/statements/{statement_id}", func(req testserver.Request) any {
-		assert.Equal(t, "s1", req.Vars["statement_id"])
-		polls++
-		if polls < 2 {
-			return sql.StatementResponse{StatementId: "s1", Status: &sql.StatementStatus{State: sql.StatementStateRunning}}
-		}
-		return sql.StatementResponse{
-			StatementId: "s1",
-			Status:      &sql.StatementStatus{State: sql.StatementStateSucceeded},
-			Result:      &sql.ResultData{DataArray: [][]string{{"done"}}},
-		}
+	server := newServer(t)
+	server.HandleSQL("SELECT 1", func(testsql.Request) testsql.Result {
+		return testsql.Result{Rows: [][]string{{"done"}}, Polls: 1}
 	})
 
 	got, err := httpClient(t, server).ExecuteScalar(t.Context(), "SELECT 1")
 	require.NoError(t, err)
 	assert.Equal(t, "done", got)
-	assert.GreaterOrEqual(t, polls, 2)
 }
 
 func TestHTTPExecutePaginatesChunks(t *testing.T) {
-	server := testserver.New(t)
-	server.Handle("POST", "/api/2.0/sql/statements", func(testserver.Request) any {
-		return sql.StatementResponse{
-			StatementId: "s1",
-			Status:      &sql.StatementStatus{State: sql.StatementStateSucceeded},
-			Manifest:    &sql.ResultManifest{TotalChunkCount: 3},
-			Result:      &sql.ResultData{DataArray: [][]string{{"0"}}},
-		}
-	})
-	server.Handle("GET", "/api/2.0/sql/statements/{statement_id}/result/chunks/{chunk_index}", func(req testserver.Request) any {
-		return sql.ResultData{DataArray: [][]string{{req.Vars["chunk_index"]}}}
+	server := newServer(t)
+	server.HandleSQL("SELECT * FROM big", func(testsql.Request) testsql.Result {
+		return testsql.Result{Rows: [][]string{{"0"}, {"1"}, {"2"}}, Chunks: 3}
 	})
 
 	r, err := httpClient(t, server).Execute(t.Context(), "SELECT * FROM big")
@@ -87,18 +71,11 @@ func TestHTTPExecutePaginatesChunks(t *testing.T) {
 }
 
 func TestHTTPExecuteFailedReturns200(t *testing.T) {
-	server := testserver.New(t)
+	server := newServer(t)
 	// A failed statement comes back as HTTP 200 with state=FAILED, not an HTTP
 	// error; the engine must inspect the body and surface a *StatementError.
-	server.Handle("POST", "/api/2.0/sql/statements", func(testserver.Request) any {
-		return sql.StatementResponse{
-			StatementId: "s1",
-			Status: &sql.StatementStatus{
-				State:    sql.StatementStateFailed,
-				SqlState: "42P01",
-				Error:    &sql.ServiceError{ErrorCode: sql.ServiceErrorCodeBadRequest, Message: "no such table"},
-			},
-		}
+	server.HandleSQL("SELECT * FROM nope", func(testsql.Request) testsql.Result {
+		return testsql.Result{Error: &testsql.Error{Code: sql.ServiceErrorCodeBadRequest, Message: "no such table", SQLState: "42P01"}}
 	})
 
 	_, err := httpClient(t, server).Execute(t.Context(), "SELECT * FROM nope")
@@ -111,33 +88,19 @@ func TestHTTPExecuteFailedReturns200(t *testing.T) {
 }
 
 func TestHTTPSubmitAndCancel(t *testing.T) {
-	server := testserver.New(t)
-	server.Handle("POST", "/api/2.0/sql/statements", func(testserver.Request) any {
-		return sql.StatementResponse{StatementId: "s1", Status: &sql.StatementStatus{State: sql.StatementStatePending}}
-	})
-	canceled := false
-	server.Handle("POST", "/api/2.0/sql/statements/{statement_id}/cancel", func(req testserver.Request) any {
-		assert.Equal(t, "s1", req.Vars["statement_id"])
-		canceled = true
-		return map[string]string{}
-	})
-	server.Handle("GET", "/api/2.0/sql/statements/{statement_id}", func(testserver.Request) any {
-		state := sql.StatementStatePending
-		if canceled {
-			state = sql.StatementStateCanceled
-		}
-		return sql.StatementResponse{StatementId: "s1", Status: &sql.StatementStatus{State: state}}
+	server := newServer(t)
+	server.HandleSQL("SELECT 1", func(testsql.Request) testsql.Result {
+		return testsql.Result{}
 	})
 
 	c := httpClient(t, server)
 	ctx := t.Context()
 
 	stmt, err := c.Submit(ctx, "SELECT 1")
 	require.NoError(t, err)
-	assert.Equal(t, "s1", stmt.ID)
+	assert.Equal(t, "statement-1", stmt.ID)
 
 	require.NoError(t, c.Cancel(ctx, stmt.ID))
-	assert.True(t, canceled)
 
 	stmt, err = c.Poll(ctx, stmt)
 	require.NoError(t, err)

libs/testserver/handlers.go

@@ -567,6 +567,17 @@ func AddDefaultHandlers(server *Server) {
 		return MapDelete(req.Workspace, req.Workspace.SqlWarehouses, req.Vars["warehouse_id"])
 	})
 
+	// SQL Statement Execution lifecycle. Tests program these by registering
+	// matchers via Server.HandleSQL / HandleSQLPattern (see statements.go).
+	// They live here, not in New, so they act as overridable defaults: a test
+	// that needs raw control over a SQL endpoint (malformed body, transport
+	// error, custom status) can register its own handler for the same pattern
+	// before calling AddDefaultHandlers, or stub it via test.toml in acceptance.
+	server.Handle("POST", "/api/2.0/sql/statements", server.sqlExecuteStatement)
+	server.Handle("GET", "/api/2.0/sql/statements/{statement_id}", server.sqlGetStatement)
+	server.Handle("GET", "/api/2.0/sql/statements/{statement_id}/result/chunks/{chunk_index}", server.sqlGetStatementResultChunk)
+	server.Handle("POST", "/api/2.0/sql/statements/{statement_id}/cancel", server.sqlCancelStatement)
+
 	server.Handle("GET", "/api/2.0/preview/sql/data_sources", func(req Request) any {
 		return req.Workspace.SqlDataSourcesList(req)
 	})

libs/testserver/server.go

@@ -17,6 +17,7 @@ import (
 	"sync"
 
 	"github.com/databricks/cli/internal/testutil"
+	"github.com/databricks/cli/libs/testserver/testsql"
 )
 
 const testPidKey = "test-pid"
@@ -49,6 +50,8 @@ type Server struct {
 	kills  *killRules
 	faults *FaultRules
 
+	sqlHandler *testsql.Handler
+
 	RequestCallback  func(request *Request)
 	ResponseCallback func(request *Request, response *EncodedResponse)
 }
@@ -228,6 +231,7 @@ func New(t testutil.TestingT) *Server {
 		fakeOidc:       &FakeOidc{url: server.URL},
 		kills:          kills,
 		faults:         faults,
+		sqlHandler:     testsql.New(),
 	}
 	router.Dispatch = s.serve
 

libs/testserver/statements.go

@@ -0,0 +1,64 @@
+package testserver
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"regexp"
+	"strconv"
+
+	"github.com/databricks/cli/libs/testserver/testsql"
+	"github.com/databricks/databricks-sdk-go/service/sql"
+)
+
+// HandleSQL registers a matcher that runs fn when a submitted statement equals
+// statement exactly (after trimming).
+func (s *Server) HandleSQL(statement string, fn func(testsql.Request) testsql.Result) {
+	s.sqlHandler.Handle(statement, fn)
+}
+
+// HandleSQLPattern registers a matcher that runs fn when re matches a submitted
+// statement, passing the submatches through as Request.Match.
+func (s *Server) HandleSQLPattern(re *regexp.Regexp, fn func(testsql.Request) testsql.Result) {
+	s.sqlHandler.HandlePattern(re, fn)
+}
+
+// sqlExecuteStatement handles POST /api/2.0/sql/statements. A statement that
+// terminates as FAILED comes back as HTTP 200 with state=FAILED; the engine
+// builds that response and this HTTP layer is just transport.
+func (s *Server) sqlExecuteStatement(req Request) any {
+	var r sql.ExecuteStatementRequest
+	if err := json.Unmarshal(req.Body, &r); err != nil {
+		return Response{StatusCode: http.StatusBadRequest, Body: fmt.Sprintf("invalid execute statement request: %s", err)}
+	}
+	return s.sqlHandler.Submit(r.Statement, r.WaitTimeout, r.Parameters)
+}
+
+// sqlGetStatement handles GET /api/2.0/sql/statements/{statement_id}.
+func (s *Server) sqlGetStatement(req Request) any {
+	resp := s.sqlHandler.Get(req.Vars["statement_id"])
+	if resp == nil {
+		return Response{StatusCode: http.StatusNotFound}
+	}
+	return resp
+}
+
+// sqlGetStatementResultChunk handles GET
+// /api/2.0/sql/statements/{statement_id}/result/chunks/{chunk_index}.
+func (s *Server) sqlGetStatementResultChunk(req Request) any {
+	idx, err := strconv.Atoi(req.Vars["chunk_index"])
+	if err != nil {
+		return Response{StatusCode: http.StatusBadRequest, Body: fmt.Sprintf("invalid chunk index: %s", err)}
+	}
+	data := s.sqlHandler.Chunk(req.Vars["statement_id"], idx)
+	if data == nil {
+		return Response{StatusCode: http.StatusNotFound}
+	}
+	return data
+}
+
+// sqlCancelStatement handles POST /api/2.0/sql/statements/{statement_id}/cancel.
+func (s *Server) sqlCancelStatement(req Request) any {
+	s.sqlHandler.Cancel(req.Vars["statement_id"])
+	return map[string]any{}
+}

libs/testserver/statements_test.go

@@ -0,0 +1,51 @@
+package testserver
+
+import (
+	"encoding/json"
+	"regexp"
+	"testing"
+
+	"github.com/databricks/cli/libs/testserver/testsql"
+	"github.com/databricks/databricks-sdk-go/service/sql"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// submitSQL runs a statement through the POST handler and returns the response.
+func submitSQL(t *testing.T, server *Server, statement string) *sql.StatementResponse {
+	t.Helper()
+	body, err := json.Marshal(sql.ExecuteStatementRequest{Statement: statement, WaitTimeout: "10s"})
+	require.NoError(t, err)
+	resp, ok := server.sqlExecuteStatement(Request{Body: body}).(*sql.StatementResponse)
+	require.True(t, ok)
+	return resp
+}
+
+func TestHandleSQL(t *testing.T) {
+	server := New(t)
+	server.HandleSQL("SELECT 1 AS a", func(r testsql.Request) testsql.Result {
+		assert.Equal(t, []string{"SELECT 1 AS a"}, r.Match)
+		return testsql.Result{Columns: []string{"a"}, Rows: [][]string{{"1"}}}
+	})
+
+	resp := submitSQL(t, server, "SELECT 1 AS a")
+	require.NotNil(t, resp.Status)
+	assert.Equal(t, sql.StatementStateSucceeded, resp.Status.State)
+	require.NotNil(t, resp.Result)
+	assert.Equal(t, [][]string{{"1"}}, resp.Result.DataArray)
+}
+
+func TestHandleSQLPattern(t *testing.T) {
+	server := New(t)
+	// A regex matcher echoes back a captured submatch, exercising both the
+	// HandleSQLPattern registration and Request.Match.
+	server.HandleSQLPattern(regexp.MustCompile(`^SELECT (\d+)$`), func(r testsql.Request) testsql.Result {
+		return testsql.Result{Columns: []string{"n"}, Rows: [][]string{{r.Match[1]}}}
+	})
+
+	resp := submitSQL(t, server, "SELECT 42")
+	require.NotNil(t, resp.Status)
+	assert.Equal(t, sql.StatementStateSucceeded, resp.Status.State)
+	require.NotNil(t, resp.Result)
+	assert.Equal(t, [][]string{{"42"}}, resp.Result.DataArray)
+}