Skip to content

Commit bd4e2b4

Browse files
olavloiteolavloite
committed
Merge branch 'create-drop-database' into spanner-lib
2 parents abc6e5b + 88f2dbf commit bd4e2b4

File tree

4 files changed

+261
-7
lines changed

4 files changed

+261
-7
lines changed

‎conn_with_mockserver_test.go‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -334,6 +334,78 @@ func TestIsolationLevelAutoCommit(t *testing.T) {
334334
}
335335
}
336336

337+
func TestCreateDatabase(t *testing.T) {
338+
t.Parallel()
339+
340+
ctx := context.Background()
341+
db, server, teardown := setupTestDBConnection(t)
342+
defer teardown()
343+
344+
var expectedResponse = &databasepb.Database{}
345+
anyMsg, _ := anypb.New(expectedResponse)
346+
server.TestDatabaseAdmin.SetResps([]proto.Message{
347+
&longrunningpb.Operation{
348+
Done: true,
349+
Result: &longrunningpb.Operation_Response{Response: anyMsg},
350+
Name: "test-operation",
351+
},
352+
})
353+
354+
conn, err := db.Conn(ctx)
355+
if err != nil {
356+
t.Fatal(err)
357+
}
358+
defer silentClose(conn)
359+
360+
if _, err = conn.ExecContext(ctx, "create database foo"); err != nil {
361+
t.Fatalf("failed to execute CREATE DATABASE: %v", err)
362+
}
363+
364+
requests := server.TestDatabaseAdmin.Reqs()
365+
if g, w := len(requests), 1; g != w {
366+
t.Fatalf("requests count mismatch\nGot: %v\nWant: %v", g, w)
367+
}
368+
if req, ok := requests[0].(*databasepb.CreateDatabaseRequest); ok {
369+
if g, w := req.Parent, "projects/p/instances/i"; g != w {
370+
t.Fatalf("parent mismatch\n Got: %v\nWant: %v", g, w)
371+
}
372+
} else {
373+
t.Fatalf("request type mismatch, got %v", requests[0])
374+
}
375+
}
376+
377+
func TestDropDatabase(t *testing.T) {
378+
t.Parallel()
379+
380+
ctx := context.Background()
381+
db, server, teardown := setupTestDBConnection(t)
382+
defer teardown()
383+
384+
server.TestDatabaseAdmin.SetResps([]proto.Message{&emptypb.Empty{}})
385+
386+
conn, err := db.Conn(ctx)
387+
if err != nil {
388+
t.Fatal(err)
389+
}
390+
defer silentClose(conn)
391+
392+
if _, err = conn.ExecContext(ctx, "drop database foo"); err != nil {
393+
t.Fatalf("failed to execute DROP DATABASE: %v", err)
394+
}
395+
396+
requests := server.TestDatabaseAdmin.Reqs()
397+
if g, w := len(requests), 1; g != w {
398+
t.Fatalf("requests count mismatch\nGot: %v\nWant: %v", g, w)
399+
}
400+
if req, ok := requests[0].(*databasepb.DropDatabaseRequest); ok {
401+
if g, w := req.Database, "projects/p/instances/i/databases/foo"; g != w {
402+
t.Fatalf("database name mismatch\n Got: %v\nWant: %v", g, w)
403+
}
404+
} else {
405+
t.Fatalf("request type mismatch, got %v", requests[0])
406+
}
407+
}
408+
337409
func TestDDLUsingQueryContext(t *testing.T) {
338410
t.Parallel()
339411

‎statement_parser.go‎

Lines changed: 19 additions & 7 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}
@@ -951,6 +955,14 @@ func (p *statementParser) isQuery(query string) bool {
951955
return info.statementType == StatementTypeQuery
952956
}
953957

958+
func isCreateKeyword(keyword string) bool {
959+
return isStatementKeyword(keyword, createStatements)
960+
}
961+
962+
func isDropKeyword(keyword string) bool {
963+
return isStatementKeyword(keyword, dropStatements)
964+
}
965+
954966
func isQueryKeyword(keyword string) bool {
955967
return isStatementKeyword(keyword, selectStatements)
956968
}

‎statements.go‎

Lines changed: 145 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,10 @@ func parseStatement(parser *statementParser, keyword, query string) (parsedState
2525
stmt = &parsedSetStatement{}
2626
} else if isResetStatementKeyword(keyword) {
2727
stmt = &parsedResetStatement{}
28+
} else if isCreateKeyword(keyword) && isCreateDatabase(parser, query) {
29+
stmt = &parsedCreateDatabaseStatement{}
30+
} else if isDropKeyword(keyword) && isDropDatabase(parser, query) {
31+
stmt = &parsedDropDatabaseStatement{}
2832
} else {
2933
return nil, nil
3034
}
@@ -34,6 +38,28 @@ func parseStatement(parser *statementParser, keyword, query string) (parsedState
3438
return stmt, nil
3539
}
3640

41+
func isCreateDatabase(parser *statementParser, query string) bool {
42+
sp := &simpleParser{sql: []byte(query), statementParser: parser}
43+
if _, ok := sp.eatKeyword("create"); !ok {
44+
return false
45+
}
46+
if _, ok := sp.eatKeyword("database"); !ok {
47+
return false
48+
}
49+
return true
50+
}
51+
52+
func isDropDatabase(parser *statementParser, query string) bool {
53+
sp := &simpleParser{sql: []byte(query), statementParser: parser}
54+
if _, ok := sp.eatKeyword("drop"); !ok {
55+
return false
56+
}
57+
if _, ok := sp.eatKeyword("database"); !ok {
58+
return false
59+
}
60+
return true
61+
}
62+
3763
// SHOW [VARIABLE] [my_extension.]my_property
3864
type parsedShowStatement struct {
3965
query string
@@ -250,3 +276,122 @@ func createEmptyRows(opts *ExecOptions) *rows {
250276
returnResultSetStats: opts.ReturnResultSetStats,
251277
}
252278
}
279+
280+
type parsedCreateDatabaseStatement struct {
281+
query string
282+
identifier identifier
283+
}
284+
285+
func (s *parsedCreateDatabaseStatement) parse(parser *statementParser, query string) error {
286+
// Parse a statement of the form
287+
// CREATE DATABASE <database-name>
288+
sp := &simpleParser{sql: []byte(query), statementParser: parser}
289+
if _, ok := sp.eatKeyword("CREATE"); !ok {
290+
return status.Error(codes.InvalidArgument, "statement does not start with CREATE DATABASE")
291+
}
292+
if _, ok := sp.eatKeyword("DATABASE"); !ok {
293+
return status.Error(codes.InvalidArgument, "statement does not start with CREATE DATABASE")
294+
}
295+
identifier, err := sp.eatIdentifier()
296+
if err != nil {
297+
return err
298+
}
299+
s.query = query
300+
s.identifier = identifier
301+
return nil
302+
}
303+
304+
func (s *parsedCreateDatabaseStatement) execContext(ctx context.Context, c *conn, params string, opts *ExecOptions, args []driver.NamedValue) (driver.Result, error) {
305+
instance := fmt.Sprintf("projects/%s/instances/%s", c.connector.connectorConfig.Project, c.connector.connectorConfig.Instance)
306+
request := &databasepb.CreateDatabaseRequest{
307+
CreateStatement: s.query,
308+
Parent: instance,
309+
DatabaseDialect: c.parser.dialect,
310+
}
311+
op, err := c.adminClient.CreateDatabase(ctx, request)
312+
if err != nil {
313+
return nil, err
314+
}
315+
if _, err := op.Wait(ctx); err != nil {
316+
return nil, err
317+
}
318+
return driver.ResultNoRows, nil
319+
}
320+
321+
func (s *parsedCreateDatabaseStatement) queryContext(ctx context.Context, c *conn, params string, opts *ExecOptions, args []driver.NamedValue) (driver.Rows, error) {
322+
if _, err := s.execContext(ctx, c, params, opts, args); err != nil {
323+
return nil, err
324+
}
325+
return createEmptyRows(opts), nil
326+
}
327+
328+
func (s *parsedCreateDatabaseStatement) executableStatement(c *conn) *executableClientSideStatement {
329+
return &executableClientSideStatement{
330+
conn: c,
331+
query: s.query,
332+
clientSideStatement: &clientSideStatement{
333+
Name: "CREATE DATABASE",
334+
execContext: s.execContext,
335+
queryContext: s.queryContext,
336+
},
337+
}
338+
}
339+
340+
type parsedDropDatabaseStatement struct {
341+
query string
342+
identifier identifier
343+
}
344+
345+
func (s *parsedDropDatabaseStatement) parse(parser *statementParser, query string) error {
346+
// Parse a statement of the form
347+
// DROP DATABASE <database-name>
348+
sp := &simpleParser{sql: []byte(query), statementParser: parser}
349+
if _, ok := sp.eatKeyword("DROP"); !ok {
350+
return status.Error(codes.InvalidArgument, "statement does not start with DROP DATABASE")
351+
}
352+
if _, ok := sp.eatKeyword("DATABASE"); !ok {
353+
return status.Error(codes.InvalidArgument, "statement does not start with DROP DATABASE")
354+
}
355+
identifier, err := sp.eatIdentifier()
356+
if err != nil {
357+
return err
358+
}
359+
s.query = query
360+
s.identifier = identifier
361+
return nil
362+
}
363+
364+
func (s *parsedDropDatabaseStatement) execContext(ctx context.Context, c *conn, params string, opts *ExecOptions, args []driver.NamedValue) (driver.Result, error) {
365+
database := fmt.Sprintf(
366+
"projects/%s/instances/%s/databases/%s",
367+
c.connector.connectorConfig.Project,
368+
c.connector.connectorConfig.Instance,
369+
s.identifier.String())
370+
request := &databasepb.DropDatabaseRequest{
371+
Database: database,
372+
}
373+
err := c.adminClient.DropDatabase(ctx, request)
374+
if err != nil {
375+
return nil, err
376+
}
377+
return driver.ResultNoRows, nil
378+
}
379+
380+
func (s *parsedDropDatabaseStatement) queryContext(ctx context.Context, c *conn, params string, opts *ExecOptions, args []driver.NamedValue) (driver.Rows, error) {
381+
if _, err := s.execContext(ctx, c, params, opts, args); err != nil {
382+
return nil, err
383+
}
384+
return createEmptyRows(opts), nil
385+
}
386+
387+
func (s *parsedDropDatabaseStatement) executableStatement(c *conn) *executableClientSideStatement {
388+
return &executableClientSideStatement{
389+
conn: c,
390+
query: s.query,
391+
clientSideStatement: &clientSideStatement{
392+
Name: "DROP DATABASE",
393+
execContext: s.execContext,
394+
queryContext: s.queryContext,
395+
},
396+
}
397+
}

‎testutil/inmem_database_admin_server.go‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ import (
2323
databasepb "cloud.google.com/go/spanner/admin/database/apiv1/databasepb"
2424
"google.golang.org/grpc/metadata"
2525
"google.golang.org/protobuf/proto"
26+
"google.golang.org/protobuf/types/known/emptypb"
2627
)
2728

2829
// InMemDatabaseAdminServer contains the DatabaseAdminServer interface plus a couple
@@ -55,6 +56,30 @@ func NewInMemDatabaseAdminServer() InMemDatabaseAdminServer {
5556
return res
5657
}
5758

59+
func (s *inMemDatabaseAdminServer) CreateDatabase(ctx context.Context, req *databasepb.CreateDatabaseRequest) (*longrunningpb.Operation, error) {
60+
md, _ := metadata.FromIncomingContext(ctx)
61+
if xg := md["x-goog-api-client"]; len(xg) == 0 || !strings.Contains(xg[0], "gl-go/") {
62+
return nil, fmt.Errorf("x-goog-api-client = %v, expected gl-go key", xg)
63+
}
64+
s.reqs = append(s.reqs, req)
65+
if s.err != nil {
66+
return nil, s.err
67+
}
68+
return s.resps[0].(*longrunningpb.Operation), nil
69+
}
70+
71+
func (s *inMemDatabaseAdminServer) DropDatabase(ctx context.Context, req *databasepb.DropDatabaseRequest) (*emptypb.Empty, error) {
72+
md, _ := metadata.FromIncomingContext(ctx)
73+
if xg := md["x-goog-api-client"]; len(xg) == 0 || !strings.Contains(xg[0], "gl-go/") {
74+
return nil, fmt.Errorf("x-goog-api-client = %v, expected gl-go key", xg)
75+
}
76+
s.reqs = append(s.reqs, req)
77+
if s.err != nil {
78+
return nil, s.err
79+
}
80+
return s.resps[0].(*emptypb.Empty), nil
81+
}
82+
5883
func (s *inMemDatabaseAdminServer) UpdateDatabaseDdl(ctx context.Context, req *databasepb.UpdateDatabaseDdlRequest) (*longrunningpb.Operation, error) {
5984
md, _ := metadata.FromIncomingContext(ctx)
6085
if xg := md["x-goog-api-client"]; len(xg) == 0 || !strings.Contains(xg[0], "gl-go/") {

0 commit comments

Comments
 (0)