Skip to content

Commit d0a5fc3

Browse files
committed
feat: create/drop database statements
Support CREATE DATABASE and DROP DATABASE statements.
1 parent 72aa0c8 commit d0a5fc3

File tree

5 files changed

+329
-8
lines changed

5 files changed

+329
-8
lines changed

conn_with_mockserver_test.go

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,15 @@ import (
2121
"reflect"
2222
"testing"
2323

24+
"cloud.google.com/go/longrunning/autogen/longrunningpb"
2425
"cloud.google.com/go/spanner"
2526
"cloud.google.com/go/spanner/admin/database/apiv1/databasepb"
2627
"cloud.google.com/go/spanner/apiv1/spannerpb"
2728
"github.com/googleapis/go-sql-spanner/connectionstate"
2829
"github.com/googleapis/go-sql-spanner/testutil"
30+
"google.golang.org/protobuf/proto"
31+
"google.golang.org/protobuf/types/known/anypb"
32+
"google.golang.org/protobuf/types/known/emptypb"
2933
)
3034

3135
func TestBeginTx(t *testing.T) {
@@ -310,6 +314,78 @@ func TestIsolationLevelAutoCommit(t *testing.T) {
310314
}
311315
}
312316

317+
func TestCreateDatabase(t *testing.T) {
318+
t.Parallel()
319+
320+
ctx := context.Background()
321+
db, server, teardown := setupTestDBConnection(t)
322+
defer teardown()
323+
324+
var expectedResponse = &databasepb.Database{}
325+
anyMsg, _ := anypb.New(expectedResponse)
326+
server.TestDatabaseAdmin.SetResps([]proto.Message{
327+
&longrunningpb.Operation{
328+
Done: true,
329+
Result: &longrunningpb.Operation_Response{Response: anyMsg},
330+
Name: "test-operation",
331+
},
332+
})
333+
334+
conn, err := db.Conn(ctx)
335+
if err != nil {
336+
t.Fatal(err)
337+
}
338+
defer silentClose(conn)
339+
340+
if _, err = conn.ExecContext(ctx, "create database `foo`"); err != nil {
341+
t.Fatalf("failed to execute CREATE DATABASE: %v", err)
342+
}
343+
344+
requests := server.TestDatabaseAdmin.Reqs()
345+
if g, w := len(requests), 1; g != w {
346+
t.Fatalf("requests count mismatch\nGot: %v\nWant: %v", g, w)
347+
}
348+
if req, ok := requests[0].(*databasepb.CreateDatabaseRequest); ok {
349+
if g, w := req.Parent, "projects/p/instances/i"; g != w {
350+
t.Fatalf("parent mismatch\n Got: %v\nWant: %v", g, w)
351+
}
352+
} else {
353+
t.Fatalf("request type mismatch, got %v", requests[0])
354+
}
355+
}
356+
357+
func TestDropDatabase(t *testing.T) {
358+
t.Parallel()
359+
360+
ctx := context.Background()
361+
db, server, teardown := setupTestDBConnection(t)
362+
defer teardown()
363+
364+
server.TestDatabaseAdmin.SetResps([]proto.Message{&emptypb.Empty{}})
365+
366+
conn, err := db.Conn(ctx)
367+
if err != nil {
368+
t.Fatal(err)
369+
}
370+
defer silentClose(conn)
371+
372+
if _, err = conn.ExecContext(ctx, "drop database foo"); err != nil {
373+
t.Fatalf("failed to execute DROP DATABASE: %v", err)
374+
}
375+
376+
requests := server.TestDatabaseAdmin.Reqs()
377+
if g, w := len(requests), 1; g != w {
378+
t.Fatalf("requests count mismatch\nGot: %v\nWant: %v", g, w)
379+
}
380+
if req, ok := requests[0].(*databasepb.DropDatabaseRequest); ok {
381+
if g, w := req.Database, "projects/p/instances/i/databases/foo"; g != w {
382+
t.Fatalf("database name mismatch\n Got: %v\nWant: %v", g, w)
383+
}
384+
} else {
385+
t.Fatalf("request type mismatch, got %v", requests[0])
386+
}
387+
}
388+
313389
func TestDDLUsingQueryContext(t *testing.T) {
314390
t.Parallel()
315391

statement_parser.go

Lines changed: 49 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -41,13 +41,17 @@ var updateStatements = map[string]bool{"UPDATE": true}
4141
var deleteStatements = map[string]bool{"DELETE": true}
4242
var dmlStatements = union(insertStatements, union(updateStatements, deleteStatements))
4343
var clientSideKeywords = map[string]bool{
44-
"SHOW": true,
45-
"SET": true,
46-
"RESET": true,
47-
"START": true,
48-
"RUN": true,
49-
"ABORT": true,
50-
}
44+
"SHOW": true,
45+
"SET": true,
46+
"RESET": true,
47+
"START": true,
48+
"RUN": true,
49+
"ABORT": true,
50+
"CREATE": true, // CREATE DATABASE is handled as a client-side statement
51+
"DROP": true, // DROP DATABASE is handled as a client-side statement
52+
}
53+
var createStatements = map[string]bool{"CREATE": true}
54+
var dropStatements = map[string]bool{"DROP": true}
5155
var showStatements = map[string]bool{"SHOW": true}
5256
var setStatements = map[string]bool{"SET": true}
5357
var resetStatements = map[string]bool{"RESET": true}
@@ -201,18 +205,36 @@ func (i *identifier) String() string {
201205
// eatIdentifier reads the identifier at the current parser position, updates the parser position,
202206
// and returns the identifier.
203207
func (p *simpleParser) eatIdentifier() (identifier, error) {
204-
// TODO: Add support for quoted identifiers.
205208
p.skipWhitespacesAndComments()
206209
if p.pos >= len(p.sql) {
207210
return identifier{}, status.Errorf(codes.InvalidArgument, "no identifier found at position %d", p.pos)
208211
}
212+
209213
startPos := p.pos
210214
first := true
211215
result := identifier{parts: make([]string, 0, 1)}
212216
appendLastPart := true
213217
for p.pos < len(p.sql) {
214218
if first {
215219
first = false
220+
// Check if this is a quoted identifier.
221+
if p.sql[p.pos] == p.statementParser.identifierQuoteToken() {
222+
pos, quoteLen, err := p.statementParser.skipQuoted(p.sql, p.pos, p.sql[p.pos])
223+
if err != nil {
224+
return identifier{}, err
225+
}
226+
p.pos = pos
227+
result.parts = append(result.parts, string(p.sql[startPos+quoteLen:pos-quoteLen]))
228+
if p.eatToken('.') {
229+
p.skipWhitespacesAndComments()
230+
startPos = p.pos
231+
first = true
232+
continue
233+
} else {
234+
appendLastPart = false
235+
break
236+
}
237+
}
216238
if !p.isValidFirstIdentifierChar() {
217239
return identifier{}, status.Errorf(codes.InvalidArgument, "invalid first identifier character found at position %d: %s", p.pos, p.sql[p.pos:p.pos+1])
218240
}
@@ -446,6 +468,13 @@ func (p *statementParser) supportsNestedComments() bool {
446468
return p.dialect == databasepb.DatabaseDialect_POSTGRESQL
447469
}
448470

471+
func (p *statementParser) identifierQuoteToken() byte {
472+
if p.dialect == databasepb.DatabaseDialect_POSTGRESQL {
473+
return '"'
474+
}
475+
return '`'
476+
}
477+
449478
func (p *statementParser) supportsBacktickQuotes() bool {
450479
return p.dialect != databasepb.DatabaseDialect_POSTGRESQL
451480
}
@@ -653,6 +682,10 @@ func (p *statementParser) skipMultiLineComment(sql []byte, pos int) int {
653682
return pos
654683
}
655684

685+
// skipQuoted skips a quoted string at the given position in the sql string and
686+
// returns the new position, the quote length, or an error if the quoted string
687+
// could not be read.
688+
// The quote length is either 1 for normal quoted strings, and 3 for triple-quoted string.
656689
func (p *statementParser) skipQuoted(sql []byte, pos int, quote byte) (int, int, error) {
657690
isTripleQuoted := p.supportsTripleQuotedLiterals() && len(sql) > pos+2 && sql[pos+1] == quote && sql[pos+2] == quote
658691
if isTripleQuoted && (isMultibyte(sql[pos+1]) || isMultibyte(sql[pos+2])) {
@@ -951,6 +984,14 @@ func (p *statementParser) isQuery(query string) bool {
951984
return info.statementType == statementTypeQuery
952985
}
953986

987+
func isCreateKeyword(keyword string) bool {
988+
return isStatementKeyword(keyword, createStatements)
989+
}
990+
991+
func isDropKeyword(keyword string) bool {
992+
return isStatementKeyword(keyword, dropStatements)
993+
}
994+
954995
func isQueryKeyword(keyword string) bool {
955996
return isStatementKeyword(keyword, selectStatements)
956997
}

statement_parser_test.go

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2494,15 +2494,28 @@ func TestEatIdentifier(t *testing.T) {
24942494
input: "my_property",
24952495
want: identifier{parts: []string{"my_property"}},
24962496
},
2497+
{
2498+
input: "`my_property`",
2499+
want: identifier{parts: []string{"my_property"}},
2500+
},
24972501
{
24982502
input: "my_extension.my_property",
24992503
want: identifier{parts: []string{"my_extension", "my_property"}},
25002504
},
2505+
{
2506+
input: "`my_extension`.`my_property`",
2507+
want: identifier{parts: []string{"my_extension", "my_property"}},
2508+
},
25012509
{
25022510
// spaces are allowed
25032511
input: " \n my_extension . \t my_property ",
25042512
want: identifier{parts: []string{"my_extension", "my_property"}},
25052513
},
2514+
{
2515+
// spaces are allowed
2516+
input: " \n `my_extension` . \t `my_property` ",
2517+
want: identifier{parts: []string{"my_extension", "my_property"}},
2518+
},
25062519
{
25072520
// comments are treated the same as spaces and are allowed
25082521
input: " /* comment */ \n my_extension -- yet another comment\n. \t -- Also a comment \nmy_property ",
@@ -2512,6 +2525,14 @@ func TestEatIdentifier(t *testing.T) {
25122525
input: "p1.p2.p3.p4",
25132526
want: identifier{parts: []string{"p1", "p2", "p3", "p4"}},
25142527
},
2528+
{
2529+
input: "`p1`.`p2`.`p3`.`p4`",
2530+
want: identifier{parts: []string{"p1", "p2", "p3", "p4"}},
2531+
},
2532+
{
2533+
input: "`p1`.p2.`p3`.p4",
2534+
want: identifier{parts: []string{"p1", "p2", "p3", "p4"}},
2535+
},
25152536
{
25162537
input: "a.b.c",
25172538
want: identifier{parts: []string{"a", "b", "c"}},
@@ -2520,6 +2541,11 @@ func TestEatIdentifier(t *testing.T) {
25202541
input: "1a",
25212542
wantErr: true,
25222543
},
2544+
{
2545+
// Double-quotes are not valid around identifiers in GoogleSQL.
2546+
input: `"1a""`,
2547+
wantErr: true,
2548+
},
25232549
{
25242550
input: "my_extension.",
25252551
wantErr: true,
@@ -2537,6 +2563,14 @@ func TestEatIdentifier(t *testing.T) {
25372563
input: "a . 1a",
25382564
wantErr: true,
25392565
},
2566+
{
2567+
input: "`p1 /* looks like a comment */ `.`p2`",
2568+
want: identifier{parts: []string{"p1 /* looks like a comment */ ", "p2"}},
2569+
},
2570+
{
2571+
input: "```p1 -- looks like a comment\n ```.`p2`",
2572+
want: identifier{parts: []string{"p1 -- looks like a comment\n ", "p2"}},
2573+
},
25402574
}
25412575
for _, test := range tests {
25422576
sp := &simpleParser{sql: []byte(test.input), statementParser: parser}

0 commit comments

Comments
 (0)