testserver: register SQL exec routes as overridable defaults

Move the SQL Statement Execution lifecycle routes from testserver.New into
AddDefaultHandlers. Registering them in New made them permanent first-wins
registrations, so tests could not install raw handlers for those paths (needed
for malformed bodies, transport errors, or custom status codes the testsql fake
does not model). As default handlers they are overridable: a raw Server.Handle
for the same pattern registered before AddDefaultHandlers wins, as do test.toml
stubs in acceptance. The libs/sqlexec HTTP tests opt in via a newServer helper.

Co-authored-by: Isaac

Author: pietern

libs/sqlexec/sqlexec_http_test.go

@@ -18,6 +18,16 @@ import (
 // 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"})
@@ -27,7 +37,7 @@ func httpClient(t *testing.T, server *testserver.Server) *sqlexec.Client {
 }
 
 func TestHTTPExecuteSuccess(t *testing.T) {
-	server := testserver.New(t)
+	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"}}}
 	})
@@ -39,7 +49,7 @@ func TestHTTPExecuteSuccess(t *testing.T) {
 }
 
 func TestHTTPExecutePolls(t *testing.T) {
-	server := testserver.New(t)
+	server := newServer(t)
 	server.HandleSQL("SELECT 1", func(testsql.Request) testsql.Result {
 		return testsql.Result{Rows: [][]string{{"done"}}, Polls: 1}
 	})
@@ -50,7 +60,7 @@ func TestHTTPExecutePolls(t *testing.T) {
 }
 
 func TestHTTPExecutePaginatesChunks(t *testing.T) {
-	server := testserver.New(t)
+	server := newServer(t)
 	server.HandleSQL("SELECT * FROM big", func(testsql.Request) testsql.Result {
 		return testsql.Result{Rows: [][]string{{"0"}, {"1"}, {"2"}}, Chunks: 3}
 	})
@@ -61,7 +71,7 @@ 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.HandleSQL("SELECT * FROM nope", func(testsql.Request) testsql.Result {
@@ -78,7 +88,7 @@ func TestHTTPExecuteFailedReturns200(t *testing.T) {
 }
 
 func TestHTTPSubmitAndCancel(t *testing.T) {
-	server := testserver.New(t)
+	server := newServer(t)
 	server.HandleSQL("SELECT 1", func(testsql.Request) testsql.Result {
 		return testsql.Result{}
 	})

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

@@ -297,14 +297,6 @@ Response.Body = '<response body here>'
 		}
 	})
 
-	// Register the SQL Statement Execution lifecycle endpoints. Tests program
-	// these by registering matchers via Server.HandleSQL / HandleSQLPattern;
-	// see statements.go.
-	s.Handle("POST", "/api/2.0/sql/statements", s.sqlExecuteStatement)
-	s.Handle("GET", "/api/2.0/sql/statements/{statement_id}", s.sqlGetStatement)
-	s.Handle("GET", "/api/2.0/sql/statements/{statement_id}/result/chunks/{chunk_index}", s.sqlGetStatementResultChunk)
-	s.Handle("POST", "/api/2.0/sql/statements/{statement_id}/cancel", s.sqlCancelStatement)
-
 	return s
 }