Skip to content

Commit a9c4c8a

Browse files
sagebeesagebee
committed
fix: added ddl support to driver, changed tests to call driver rather than api directly
1 parent fd70120 commit a9c4c8a

2 files changed

Lines changed: 294 additions & 16 deletions

File tree

‎driver.go‎

Lines changed: 71 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,17 @@ import (
1919
"database/sql"
2020
"database/sql/driver"
2121
"errors"
22+
"os"
23+
"regexp"
2224
"time"
2325

2426
"cloud.google.com/go/spanner"
2527
"github.com/rakyll/go-sql-driver-spanner/internal"
2628
"google.golang.org/api/option"
29+
30+
adminapi "cloud.google.com/go/spanner/admin/database/apiv1"
31+
adminpb "google.golang.org/genproto/googleapis/spanner/admin/database/v1"
32+
"google.golang.org/grpc"
2733
)
2834

2935
const userAgent = "go-sql-driver-spanner/0.1"
@@ -78,17 +84,47 @@ func openDriverConn(ctx context.Context, d *Driver, name string) (driver.Conn, e
7884
if err != nil {
7985
return nil, err
8086
}
81-
return &conn{client: client}, nil
87+
88+
adminClient, err := CreateAdminClient(ctx)
89+
return &conn{client: client, adminClient: adminClient, name: name}, nil
90+
}
91+
92+
func CreateAdminClient(ctx context.Context) (*adminapi.DatabaseAdminClient, error) {
93+
94+
var adminClient *adminapi.DatabaseAdminClient
95+
var err error
96+
97+
// Admin client will connect tp emulator if SPANNER_EMULATOR_HOST
98+
// is set in the environment.
99+
if spannerHost, ok := os.LookupEnv("SPANNER_EMULATOR_HOST"); ok {
100+
adminClient, err = adminapi.NewDatabaseAdminClient(
101+
ctx,
102+
option.WithoutAuthentication(),
103+
option.WithEndpoint(spannerHost),
104+
option.WithGRPCDialOption(grpc.WithInsecure()))
105+
if err != nil {
106+
return nil, err
107+
}
108+
} else {
109+
adminClient, err = adminapi.NewDatabaseAdminClient(ctx)
110+
if err != nil {
111+
return nil, err
112+
}
113+
}
114+
115+
return adminClient, nil
82116
}
83117

84118
func (c *connector) Driver() driver.Driver {
85119
return &Driver{}
86120
}
87121

88122
type conn struct {
89-
client *spanner.Client
90-
roTx *spanner.ReadOnlyTransaction
91-
rwTx *rwTx
123+
client *spanner.Client
124+
adminClient *adminapi.DatabaseAdminClient
125+
roTx *spanner.ReadOnlyTransaction
126+
rwTx *rwTx
127+
name string
92128
}
93129

94130
func (c *conn) Prepare(query string) (driver.Stmt, error) {
@@ -105,6 +141,27 @@ func (c *conn) PrepareContext(ctx context.Context, query string) (driver.Stmt, e
105141
}
106142

107143
func (c *conn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
144+
145+
// Use admin API if DDL statement is provided.
146+
ddl, err := IsDdlStatement(query)
147+
if err != nil {
148+
return nil, err
149+
}
150+
151+
if ddl {
152+
op, err := c.adminClient.UpdateDatabaseDdl(ctx, &adminpb.UpdateDatabaseDdlRequest{
153+
Database: c.name,
154+
Statements: []string{query},
155+
})
156+
if err != nil {
157+
return nil, err
158+
}
159+
if err := op.Wait(ctx); err != nil {
160+
return nil, err
161+
}
162+
return &result{rowsAffected: 0}, nil
163+
}
164+
108165
if c.roTx != nil {
109166
return nil, errors.New("cannot write in read-only transaction")
110167
}
@@ -125,6 +182,16 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name
125182
return &result{rowsAffected: rowsAffected}, nil
126183
}
127184

185+
func IsDdlStatement(query string) (bool, error) {
186+
187+
matchddl, err := regexp.MatchString(`(?is)^\n*\s*(CREATE|DROP|ALTER)\s+.+$`, query)
188+
if err != nil {
189+
return false, err
190+
}
191+
192+
return matchddl, nil
193+
}
194+
128195
func (c *conn) Close() error {
129196
c.client.Close()
130197
return nil

‎driver_test.go‎

Lines changed: 223 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -12,26 +12,237 @@
1212
// See the License for the specific language governing permissions and
1313
// limitations under the License.
1414

15-
package spannerdriver_test
15+
package spannerdriver
1616

1717
import (
18+
"cloud.google.com/go/spanner"
19+
"context"
1820
"database/sql"
19-
"log"
21+
"os"
22+
"reflect"
23+
"testing"
24+
)
2025

21-
spannerdriver "github.com/rakyll/go-sql-driver-spanner"
22-
"google.golang.org/api/option"
26+
var (
27+
dsn string
2328
)
2429

25-
func ExampleDriver() {
26-
driver := &spannerdriver.Driver{
27-
Options: []option.ClientOption{
28-
option.WithCredentialsFile("/path/to/service-account-key.json"),
30+
type Connector struct {
31+
ctx context.Context
32+
client *spanner.Client
33+
}
34+
35+
func NewConnector() (*Connector, error) {
36+
37+
ctx := context.Background()
38+
39+
dataClient, err := spanner.NewClient(ctx, dsn)
40+
if err != nil {
41+
return nil, err
42+
}
43+
44+
conn := &Connector{
45+
ctx: ctx,
46+
client: dataClient,
47+
}
48+
return conn, nil
49+
}
50+
51+
func (c *Connector) Close() {
52+
c.client.Close()
53+
}
54+
55+
func init() {
56+
57+
var projectId, instanceId, databaseId string
58+
var ok bool
59+
60+
// Get environment variables or set to default.
61+
if instanceId, ok = os.LookupEnv("SPANNER_TEST_INSTANCE"); !ok {
62+
instanceId = "test-instance"
63+
}
64+
if projectId, ok = os.LookupEnv("SPANNER_TEST_PROJECT"); !ok {
65+
projectId = "test-project"
66+
}
67+
if databaseId, ok = os.LookupEnv("SPANNER_TEST_DBID"); !ok {
68+
databaseId = "gotest"
69+
}
70+
71+
// Derive data source name.
72+
dsn = "projects/" + projectId + "/instances/" + instanceId + "/databases/" + databaseId
73+
}
74+
75+
// Executes DML using the client library.
76+
func ExecuteDMLClientLib(dml []string) error {
77+
78+
// Open client/
79+
ctx := context.Background()
80+
client, err := spanner.NewClient(ctx, dsn)
81+
if err != nil {
82+
return err
83+
}
84+
defer client.Close()
85+
86+
// Put strings into spanner.Statement structure.
87+
var stmts []spanner.Statement
88+
for _, line := range dml {
89+
stmts = append(stmts, spanner.NewStatement(line))
90+
}
91+
92+
// Execute statements.
93+
_, err = client.ReadWriteTransaction(ctx, func(ctx context.Context, txn *spanner.ReadWriteTransaction) error {
94+
_, err := txn.BatchUpdate(ctx, stmts)
95+
if err != nil {
96+
return err
97+
}
98+
return nil
99+
})
100+
101+
return err
102+
}
103+
104+
func TestQueryContext(t *testing.T) {
105+
106+
// Open db.
107+
ctx := context.Background()
108+
db, err := sql.Open("spanner", dsn)
109+
if err != nil {
110+
t.Fatal(err)
111+
}
112+
defer db.Close()
113+
114+
// Set up test table.
115+
_, err = db.ExecContext(ctx, `CREATE TABLE TestQueryContext (
116+
A STRING(1024),
117+
B STRING(1024),
118+
C STRING(1024)
119+
) PRIMARY KEY (A)`)
120+
if err != nil {
121+
t.Fatal(err)
122+
}
123+
124+
conn, err := NewConnector()
125+
if err != nil {
126+
t.Fatal(err)
127+
}
128+
defer conn.Close()
129+
130+
err = ExecuteDMLClientLib([]string{`INSERT INTO TestQueryContext (A, B, C)
131+
VALUES ("a1", "b1", "c1"), ("a2", "b2", "c2") , ("a3", "b3", "c3") `})
132+
if err != nil {
133+
t.Fatal(err)
134+
}
135+
136+
type testQueryContextRow struct {
137+
A, B, C string
138+
}
139+
140+
tests := []struct {
141+
name string
142+
input string
143+
want []testQueryContextRow
144+
wantErrorQuery bool
145+
wantErrorScan bool
146+
wantErrorClose bool
147+
}{
148+
{
149+
name: "empty query",
150+
wantErrorClose: true,
151+
input: "",
152+
want: []testQueryContextRow{},
153+
},
154+
{
155+
name: "syntax error",
156+
wantErrorClose: true,
157+
input: "SELECT SELECT * FROM TestQueryContext",
158+
want: []testQueryContextRow{},
159+
},
160+
{
161+
name: "return nothing",
162+
input: "SELECT * FROM TestQueryContext WHERE A = \"hihihi\"",
163+
want: []testQueryContextRow{},
164+
},
165+
{
166+
name: "select one tuple",
167+
input: "SELECT * FROM TestQueryContext WHERE A = \"a1\"",
168+
want: []testQueryContextRow{
169+
{A: "a1", B: "b1", C: "c1"},
170+
},
171+
},
172+
{
173+
name: "select subset of tuples",
174+
input: "SELECT * FROM TestQueryContext WHERE A = \"a1\" OR A = \"a2\"",
175+
want: []testQueryContextRow{
176+
{A: "a1", B: "b1", C: "c1"},
177+
{A: "a2", B: "b2", C: "c2"},
178+
},
179+
},
180+
{
181+
name: "select subset of tuples with !=",
182+
input: "SELECT * FROM TestQueryContext WHERE A != \"a3\"",
183+
want: []testQueryContextRow{
184+
{A: "a1", B: "b1", C: "c1"},
185+
{A: "a2", B: "b2", C: "c2"},
186+
},
187+
},
188+
{
189+
name: "select entire table",
190+
input: "SELECT * FROM TestQueryContext ORDER BY A",
191+
want: []testQueryContextRow{
192+
{A: "a1", B: "b1", C: "c1"},
193+
{A: "a2", B: "b2", C: "c2"},
194+
{A: "a3", B: "b3", C: "c3"},
195+
},
196+
},
197+
{
198+
name: "query non existant table",
199+
wantErrorClose: true,
200+
input: "SELECT * FROM TestQueryContexta", want: []testQueryContextRow{},
29201
},
30202
}
31-
connector, err := driver.OpenConnector("projects/$PROJECT/instances/$INSTANCE/databases/$DATABASE")
203+
204+
// Run tests
205+
for _, tc := range tests {
206+
207+
rows, err := db.QueryContext(ctx, tc.input)
208+
if (err != nil) && (!tc.wantErrorQuery) {
209+
t.Errorf("%s: unexpected query error: %v", tc.name, err)
210+
}
211+
if (err == nil) && (tc.wantErrorQuery) {
212+
t.Errorf("%s: expected query error but error was %v", tc.name, err)
213+
}
214+
215+
got := []testQueryContextRow{}
216+
for rows.Next() {
217+
var curr testQueryContextRow
218+
err := rows.Scan(&curr.A, &curr.B, &curr.C)
219+
if (err != nil) && (!tc.wantErrorScan) {
220+
t.Errorf("%s: unexpected query error: %v", tc.name, err)
221+
}
222+
if (err == nil) && (tc.wantErrorScan) {
223+
t.Errorf("%s: expected query error but error was %v", tc.name, err)
224+
}
225+
226+
got = append(got, curr)
227+
}
228+
229+
rows.Close()
230+
err = rows.Err()
231+
if (err != nil) && (!tc.wantErrorClose) {
232+
t.Errorf("%s: unexpected query error: %v", tc.name, err)
233+
}
234+
if (err == nil) && (tc.wantErrorClose) {
235+
t.Errorf("%s: expected query error but error was %v", tc.name, err)
236+
}
237+
if !reflect.DeepEqual(tc.want, got) {
238+
t.Errorf("Test failed: %s. want: %v, got: %v", tc.name, tc.want, got)
239+
}
240+
241+
}
242+
243+
// Drop table.
244+
_, err = db.ExecContext(ctx, `DROP TABLE TestQueryContext`)
32245
if err != nil {
33-
log.Fatal(err)
246+
t.Error(err)
34247
}
35-
db := sql.OpenDB(connector)
36-
_ = db // Use db.
37248
}

0 commit comments

Comments
 (0)