feat: allow DDL statements with QueryContext (#515)

Allow DDL statements to be used with the QueryContext function. This makes the
function a 'catch-all' function that can be used for all types of statements,
without having to worry beforehand what the type of statement is. The driver
returns an empty row iterator for DDL statements.

Author: olavloite

conn.go

@@ -836,9 +836,13 @@ func (c *conn) queryContext(ctx context.Context, query string, execOptions *Exec
 		return nil, err
 	}
 	statementType := c.parser.DetectStatementType(query)
-	// DDL statements are not supported in QueryContext so fail early.
+	// DDL statements are not supported in QueryContext so use the execContext method for the execution.
 	if statementType.StatementType == parser.StatementTypeDdl {
-		return nil, spanner.ToSpannerError(status.Errorf(codes.FailedPrecondition, "QueryContext does not support DDL statements, use ExecContext instead"))
+		res, err := c.execContext(ctx, query, execOptions, args)
+		if err != nil {
+			return nil, err
+		}
+		return createDriverResultRows(res, execOptions), nil
 	}
 	var iter rowIterator
 	if c.tx == nil {

conn_with_mockserver_test.go

@@ -567,17 +567,26 @@ func TestDropDatabase(t *testing.T) {
 func TestDDLUsingQueryContext(t *testing.T) {
 	t.Parallel()
 
-	db, _, teardown := setupTestDBConnection(t)
+	db, server, teardown := setupTestDBConnection(t)
 	defer teardown()
+	var expectedResponse = &emptypb.Empty{}
+	anyMsg, _ := anypb.New(expectedResponse)
+	server.TestDatabaseAdmin.SetResps([]proto.Message{
+		&longrunningpb.Operation{
+			Done:   true,
+			Result: &longrunningpb.Operation_Response{Response: anyMsg},
+			Name:   "test-operation",
+		},
+	})
 	ctx := context.Background()
 
-	// DDL statements should not use the query context.
-	_, err := db.QueryContext(ctx, "CREATE TABLE Foo (Bar STRING(100))")
-	if err == nil {
-		t.Fatal("expected error for DDL statement using QueryContext, got nil")
-	}
-	if g, w := err.Error(), `spanner: code = "FailedPrecondition", desc = "QueryContext does not support DDL statements, use ExecContext instead"`; g != w {
-		t.Fatalf("error mismatch\n Got: %v\nWant: %v", g, w)
+	// DDL statements should be able to use QueryContext.
+	if it, err := db.QueryContext(ctx, "CREATE TABLE Foo (Bar STRING(100))"); err != nil {
+		t.Fatal(err)
+	} else {
+		if it.Next() {
+			t.Fatalf("DDL should not return any rows")
+		}
 	}
 }
 
@@ -598,7 +607,7 @@ func TestDDLUsingQueryContextInReadOnlyTx(t *testing.T) {
 	if err == nil {
 		t.Fatal("expected error for DDL statement using QueryContext in read-only transaction, got nil")
 	}
-	if g, w := err.Error(), `spanner: code = "FailedPrecondition", desc = "QueryContext does not support DDL statements, use ExecContext instead"`; g != w {
+	if g, w := err.Error(), `spanner: code = "FailedPrecondition", desc = "cannot execute DDL as part of a transaction"`; g != w {
 		t.Fatalf("error mismatch\n Got: %v\nWant: %v", g, w)
 	}
 }
@@ -621,7 +630,7 @@ func TestDDLUsingQueryContextInReadWriteTransaction(t *testing.T) {
 	if err == nil {
 		t.Fatal("expected error for DDL statement using QueryContext in read-write transaction, got nil")
 	}
-	if g, w := err.Error(), `spanner: code = "FailedPrecondition", desc = "QueryContext does not support DDL statements, use ExecContext instead"`; g != w {
+	if g, w := err.Error(), `spanner: code = "FailedPrecondition", desc = "cannot execute DDL as part of a transaction"`; g != w {
 		t.Fatalf("error mismatch\n Got: %v\nWant: %v", g, w)
 	}
 }

rows.go

@@ -476,3 +476,97 @@ func (r *rows) nextStats(dest []driver.Value) error {
 	dest[0] = r.it.ResultSetStats()
 	return nil
 }
+
+var _ driver.Rows = (*emptyRows)(nil)
+var _ driver.RowsNextResultSet = (*emptyRows)(nil)
+var emptyRowsMetadata = &sppb.ResultSetMetadata{
+	RowType: &sppb.StructType{
+		Fields: []*sppb.StructType_Field{{Name: "affected_rows", Type: &sppb.Type{Code: sppb.TypeCode_INT64}}},
+	},
+}
+var emptyRowsStats = &sppb.ResultSetStats{}
+
+type emptyRows struct {
+	currentResultSetType    resultSetType
+	returnResultSetMetadata bool
+	returnResultSetStats    bool
+
+	hasReturnedResultSetMetadata bool
+	hasReturnedResultSetStats    bool
+}
+
+func createDriverResultRows(_ driver.Result, opts *ExecOptions) *emptyRows {
+	res := &emptyRows{
+		returnResultSetMetadata: opts.ReturnResultSetMetadata,
+		returnResultSetStats:    opts.ReturnResultSetStats,
+	}
+	if !opts.ReturnResultSetMetadata {
+		res.currentResultSetType = resultSetTypeResults
+	}
+	return res
+}
+
+func (e *emptyRows) HasNextResultSet() bool {
+	if e.currentResultSetType == resultSetTypeMetadata && e.returnResultSetMetadata {
+		return true
+	}
+	if e.currentResultSetType == resultSetTypeResults && e.returnResultSetStats {
+		return true
+	}
+	return false
+}
+
+func (e *emptyRows) NextResultSet() error {
+	if !e.HasNextResultSet() {
+		return io.EOF
+	}
+	e.currentResultSetType++
+	return nil
+}
+
+func (e *emptyRows) Columns() []string {
+	switch e.currentResultSetType {
+	case resultSetTypeMetadata:
+		return []string{"metadata"}
+	case resultSetTypeResults:
+		return []string{"affected_rows"}
+	case resultSetTypeStats:
+		return []string{"stats"}
+	case resultSetTypeNoMoreResults:
+		return []string{}
+	}
+	return []string{}
+}
+
+func (e *emptyRows) Close() error {
+	return nil
+}
+
+func (e *emptyRows) Next(dest []driver.Value) error {
+	if e.currentResultSetType == resultSetTypeMetadata {
+		return e.nextMetadata(dest)
+	}
+	if e.currentResultSetType == resultSetTypeStats {
+		return e.nextStats(dest)
+	}
+
+	return io.EOF
+}
+
+func (e *emptyRows) nextMetadata(dest []driver.Value) error {
+	if e.hasReturnedResultSetMetadata {
+		return io.EOF
+	}
+	e.hasReturnedResultSetMetadata = true
+	dest[0] = emptyRowsMetadata
+	return nil
+}
+
+func (e *emptyRows) nextStats(dest []driver.Value) error {
+	if e.hasReturnedResultSetStats {
+		return io.EOF
+	}
+	e.hasReturnedResultSetStats = true
+	dest[0] = emptyRowsStats
+	return nil
+}

rows_test.go

@@ -16,12 +16,15 @@ package spannerdriver
 
 import (
 	"database/sql/driver"
+	"errors"
 	"fmt"
 	"io"
+	"reflect"
 	"testing"
 
 	"cloud.google.com/go/spanner"
 	sppb "cloud.google.com/go/spanner/apiv1/spannerpb"
+	"github.com/google/go-cmp/cmp"
 	"google.golang.org/protobuf/types/known/structpb"
 )
 
@@ -153,3 +156,70 @@ func TestRows_Next_Unsupported(t *testing.T) {
 		t.Fatalf("expected error %q, but got %q", expectedError, err.Error())
 	}
 }
+
+func TestEmptyRows(t *testing.T) {
+	r := createDriverResultRows(&result{}, &ExecOptions{})
+
+	if g, w := r.Columns(), []string{"affected_rows"}; !cmp.Equal(g, w) {
+		t.Fatalf("columns mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	if r.HasNextResultSet() {
+		t.Fatalf("unexpected next result set available")
+	}
+}
+
+func TestEmptyRowsWithMetadataAndStats(t *testing.T) {
+	r := createDriverResultRows(&result{}, &ExecOptions{ReturnResultSetMetadata: true, ReturnResultSetStats: true})
+
+	// The first result set should contain ResultSetMetadata.
+	if g, w := r.Columns(), []string{"metadata"}; !cmp.Equal(g, w) {
+		t.Fatalf("columns mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	values := make([]driver.Value, 1)
+	if err := r.Next(values); err != nil {
+		t.Fatalf("unexpected error from Next: %v", err)
+	}
+	if g, w := reflect.TypeOf(values[0]), reflect.TypeOf(&sppb.ResultSetMetadata{}); g != w {
+		t.Fatalf("result set metadata type mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	if g, w := r.Next(values), io.EOF; !errors.Is(g, w) {
+		t.Fatalf("next result set mismatch\n Got: %v\nWant: %v", g, w)
+	}
+
+	// The second result set should contain the actual data (which is empty).
+	if !r.HasNextResultSet() {
+		t.Fatalf("missing next result set")
+	}
+	if err := r.NextResultSet(); err != nil {
+		t.Fatalf("unexpected error from NextResultSet: %v", err)
+	}
+	if g, w := r.Columns(), []string{"affected_rows"}; !cmp.Equal(g, w) {
+		t.Fatalf("columns mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	// There should be no data.
+	if g, w := r.Next(values), io.EOF; !errors.Is(g, w) {
+		t.Fatalf("next result set mismatch\n Got: %v\nWant: %v", g, w)
+	}
+
+	// The third result set should contain ResultSetStats.
+	if !r.HasNextResultSet() {
+		t.Fatalf("missing next result set")
+	}
+	if err := r.NextResultSet(); err != nil {
+		t.Fatalf("unexpected error from NextResultSet: %v", err)
+	}
+	if err := r.Next(values); err != nil {
+		t.Fatalf("unexpected error from Next: %v", err)
+	}
+	if g, w := reflect.TypeOf(values[0]), reflect.TypeOf(&sppb.ResultSetStats{}); g != w {
+		t.Fatalf("result set stats type mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	if g, w := r.Next(values), io.EOF; !errors.Is(g, w) {
+		t.Fatalf("next result set mismatch\n Got: %v\nWant: %v", g, w)
+	}
+
+	// There should be no more result sets.
+	if r.HasNextResultSet() {
+		t.Fatalf("unexpected next result set available")
+	}
+}