Skip to content

Commit 2ef588c

Browse files
wizardishungrywizardishungry
committed
Add AstraDB support
1 parent c378583 commit 2ef588c

File tree

10 files changed

+371
-22
lines changed

10 files changed

+371
-22
lines changed

Diff for: ‎Makefile

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
SOURCE ?= file go_bindata github github_ee bitbucket aws_s3 google_cloud_storage godoc_vfs gitlab
2-
DATABASE ?= postgres mysql redshift cassandra spanner cockroachdb yugabytedb clickhouse mongodb sqlserver firebird neo4j pgx pgx5 rqlite
2+
DATABASE ?= postgres mysql redshift cassandra spanner cockroachdb yugabytedb clickhouse mongodb sqlserver firebird neo4j pgx pgx5 rqlite astra
33
DATABASE_TEST ?= $(DATABASE) sqlite sqlite3 sqlcipher
44
VERSION ?= $(shell git describe --tags 2>/dev/null | cut -c 2-)
55
TEST_FLAGS ?=

Diff for: ‎README.md

+1-1
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ Database drivers run migrations. [Add a new database?](database/driver.go)
2828
* [PGX v5](database/pgx/v5)
2929
* [Redshift](database/redshift)
3030
* [Ql](database/ql)
31-
* [Cassandra / ScyllaDB](database/cassandra)
31+
* [Cassandra / ScyllaDB / AstraDB](database/cassandra)
3232
* [SQLite](database/sqlite)
3333
* [SQLite3](database/sqlite3) ([todo #165](https://github.com/mattes/migrate/issues/165))
3434
* [SQLCipher](database/sqlcipher)

Diff for: ‎database/cassandra/README.md

+22-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Cassandra / ScyllaDB
1+
# Cassandra / ScyllaDB / AstraDB
22

33
* `Drop()` method will not work on Cassandra 2.X because it rely on
44
system_schema table which comes with 3.X
@@ -13,6 +13,10 @@ system_schema table which comes with 3.X
1313
* The `Drop()` method` works for ScyllaDB 5.1
1414

1515

16+
**AstraDB**
17+
18+
* Astra uses different parameters for authentication. See below.
19+
1620
## Usage
1721
`cassandra://host:port/keyspace?param1=value&param2=value2`
1822

@@ -36,6 +40,23 @@ system_schema table which comes with 3.X
3640

3741
`timeout` is parsed using [time.ParseDuration(s string)](https://golang.org/pkg/time/#ParseDuration)
3842

43+
### [AstraDB](https://docs.datastax.com/)
44+
45+
`astra:///keyspace?bundle=bundle.zip&token=token` or
46+
`astra:///keyspace?token=token&database_id=database_id`. *Note the triple slash.*
47+
48+
Astra supports two authentication schemes;
49+
[bundle](https://pkg.go.dev/github.com/datastax/gocql-astra#NewClusterFromURL) and
50+
[token](https://pkg.go.dev/github.com/datastax/gocql-astra#NewClusterFromURL).
51+
The additional parameters are:
52+
53+
54+
| URL Query | Default value | Description |
55+
|------------|-------------|-----------|
56+
| `token` | | Astra Bearer Token (beginning with AstraCS) |
57+
| `database_id` | | Database ID |
58+
| `bundle` | | Path to secure connect bundle |
59+
| `api_url` | `https://api.astra.datastax.com` | Custom Astra Endpoint |
3960

4061
## Upgrading from v1
4162

Diff for: ‎database/cassandra/astra.go

+18
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
package cassandra
2+
3+
import (
4+
"time"
5+
6+
"github.com/gocql/gocql"
7+
)
8+
9+
// These are stubs to keep from depending on the astra driver.
10+
// The astra package assigns to these from the astra driver.
11+
var (
12+
GocqlastraNewClusterFromURL = func(url string, databaseID string, token string, timeout time.Duration) (*gocql.ClusterConfig, error) {
13+
panic("should not be used for cassandra")
14+
}
15+
GocqlastraNewClusterFromBundle = func(path string, username string, password string, timeout time.Duration) (*gocql.ClusterConfig, error) {
16+
panic("should not be used for cassandra")
17+
}
18+
)

Diff for: ‎database/cassandra/astra/astra.go

+19
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
package cassandra
2+
3+
import (
4+
gocqlastra "github.com/datastax/gocql-astra"
5+
"github.com/golang-migrate/migrate/v4/database"
6+
"github.com/golang-migrate/migrate/v4/database/cassandra"
7+
)
8+
9+
func init() {
10+
db := new(Astra)
11+
database.Register("astra", db)
12+
}
13+
14+
type Astra = cassandra.Cassandra
15+
16+
func init() {
17+
cassandra.GocqlastraNewClusterFromURL = gocqlastra.NewClusterFromURL
18+
cassandra.GocqlastraNewClusterFromBundle = gocqlastra.NewClusterFromBundle
19+
}

Diff for: ‎database/cassandra/astra/astra_test.go

+148
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
1+
package cassandra
2+
3+
import (
4+
"errors"
5+
"net/url"
6+
"testing"
7+
"time"
8+
9+
"github.com/gocql/gocql"
10+
cas "github.com/golang-migrate/migrate/v4/database/cassandra"
11+
)
12+
13+
func TestAstra(t *testing.T) {
14+
type mockResult struct {
15+
timeout time.Duration
16+
17+
// NewClusterFromBundle
18+
path string
19+
username string
20+
password string
21+
22+
// NewClusterFromURL
23+
apiUrl string
24+
databaseID string
25+
token string
26+
}
27+
28+
var (
29+
errNewClusterFromBundle = errors.New("NewClusterFromBundle")
30+
errNewClusterFromURL = errors.New("NewClusterFromURL")
31+
)
32+
33+
test := func(t *testing.T, url string) (mockResult, error) {
34+
t.Helper()
35+
36+
var mr mockResult
37+
38+
// Since we can't actually call the Astra API, we mock the calls and return an error so we never dial.
39+
cas.GocqlastraNewClusterFromBundle = func(path string, username string, password string, timeout time.Duration) (*gocql.ClusterConfig, error) {
40+
mr.path = path
41+
mr.username = username
42+
mr.password = password
43+
mr.timeout = timeout
44+
return nil, errNewClusterFromBundle
45+
}
46+
47+
cas.GocqlastraNewClusterFromURL = func(apiUrl string, databaseID string, token string, timeout time.Duration) (*gocql.ClusterConfig, error) {
48+
mr.apiUrl = apiUrl
49+
mr.databaseID = databaseID
50+
mr.token = token
51+
mr.timeout = timeout
52+
return nil, errNewClusterFromURL
53+
}
54+
55+
astra := &Astra{}
56+
57+
_, err := astra.Open(url)
58+
return mr, err
59+
}
60+
61+
t.Run("Token", func(t *testing.T) {
62+
mr, err := test(t, "astra:///testks?token=token&database_id=database_id")
63+
if err != errNewClusterFromURL {
64+
t.Error("Expected", errNewClusterFromURL, "but got", err)
65+
}
66+
if mr.token != "token" {
67+
t.Error("Expected token to be 'token' but got", mr.token)
68+
}
69+
if mr.databaseID != "database_id" {
70+
t.Error("Expected database_id to be 'database_id' but got", mr.databaseID)
71+
}
72+
})
73+
t.Run("Bundle", func(t *testing.T) {
74+
mr, err := test(t, "astra:///testks?bundle=bundle.zip&token=AstraCS:password")
75+
if err != errNewClusterFromBundle {
76+
t.Error("Expected", errNewClusterFromBundle, "but got", err)
77+
}
78+
if mr.path != "bundle.zip" {
79+
t.Error("Expected path to be 'bundle.zip' but got", mr.path)
80+
}
81+
if mr.username != "token" {
82+
t.Error("Expected username to be 'token' but got", mr.username)
83+
}
84+
if mr.password != "AstraCS:password" {
85+
t.Error("Expected password to be 'AstraCS:password' but got", mr.password)
86+
}
87+
})
88+
89+
t.Run("No Keyspace", func(t *testing.T) {
90+
astra := &Astra{}
91+
_, err := astra.Open("astra://")
92+
if err != cas.ErrNoKeyspace {
93+
t.Error("Expected", cas.ErrNoKeyspace, "but got", err)
94+
}
95+
})
96+
97+
t.Run("AstraMissing", func(t *testing.T) {
98+
astra := &Astra{}
99+
_, err := astra.Open("astra:///testks")
100+
if err != cas.ErrAstraMissing {
101+
t.Error("Expected", cas.ErrAstraMissing, "but got", err)
102+
}
103+
})
104+
t.Run("No Token", func(t *testing.T) {
105+
astra := &Astra{}
106+
_, err := astra.Open("astra:///testks?database_id=database_id")
107+
if err != cas.ErrAstraMissing {
108+
t.Error("Expected", cas.ErrAstraMissing, "but got", err)
109+
}
110+
})
111+
t.Run("No DatabaseID", func(t *testing.T) {
112+
astra := &Astra{}
113+
_, err := astra.Open("astra:///testks?token=AstraCS:password")
114+
if err != cas.ErrAstraMissing {
115+
t.Error("Expected", cas.ErrAstraMissing, "but got", err)
116+
}
117+
})
118+
t.Run("No Bundle", func(t *testing.T) {
119+
astra := &Astra{}
120+
_, err := astra.Open("astra:///testks?token=AstraCS:password")
121+
if err != cas.ErrAstraMissing {
122+
t.Error("Expected", cas.ErrAstraMissing, "but got", err)
123+
}
124+
})
125+
t.Run("Custom API URL", func(t *testing.T) {
126+
mr, err := test(t, "astra:///testks?token=token&database_id=database_id&api_url=api_url")
127+
if err != errNewClusterFromURL {
128+
t.Error("Expected", errNewClusterFromURL, "but got", err)
129+
}
130+
if mr.apiUrl != "api_url" {
131+
t.Error("Expected api_url to be 'api_url' but got", mr.apiUrl)
132+
}
133+
})
134+
}
135+
136+
func TestTripleSlashInURLMeansNoHost(t *testing.T) {
137+
const str = "astra:///testks?token=token&database_id=database_id"
138+
u, err := url.Parse(str)
139+
if err != nil {
140+
t.Fatal(err)
141+
}
142+
if u.Host != "" {
143+
t.Error("Expected host to be empty but got", u.Host)
144+
}
145+
if u.Path != "/testks" {
146+
t.Error("Expected path to be '/testks' but got", u.Path)
147+
}
148+
}

Diff for: ‎database/cassandra/cassandra.go

+46-5
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import (
1111

1212
"go.uber.org/atomic"
1313

14+
gocqlastra "github.com/datastax/gocql-astra"
1415
"github.com/gocql/gocql"
1516
"github.com/golang-migrate/migrate/v4/database"
1617
"github.com/golang-migrate/migrate/v4/database/multistmt"
@@ -35,6 +36,7 @@ var (
3536
ErrNoKeyspace = errors.New("no keyspace provided")
3637
ErrDatabaseDirty = errors.New("database is dirty")
3738
ErrClosedSession = errors.New("session is closed")
39+
ErrAstraMissing = errors.New("missing required parameters for Astra connection")
3840
)
3941

4042
type Config struct {
@@ -84,25 +86,64 @@ func WithInstance(session *gocql.Session, config *Config) (database.Driver, erro
8486
}
8587

8688
func (c *Cassandra) Open(url string) (database.Driver, error) {
89+
const (
90+
timeout = 1 * time.Minute
91+
)
8792
u, err := nurl.Parse(url)
8893
if err != nil {
8994
return nil, err
9095
}
9196

97+
isAstra := u.Scheme == "astra"
98+
99+
username := u.Query().Get("username")
100+
password := u.Query().Get("password")
101+
102+
if isAstra {
103+
username = "token"
104+
}
105+
92106
// Check for missing mandatory attributes
93107
if len(u.Path) == 0 {
94108
return nil, ErrNoKeyspace
95109
}
96110

97-
cluster := gocql.NewCluster(u.Host)
111+
var cluster *gocql.ClusterConfig
112+
113+
if isAstra {
114+
bundle := u.Query().Get("bundle")
115+
databaseID := u.Query().Get("database_id")
116+
token := u.Query().Get("token")
117+
apiUrl := u.Query().Get("api_url")
118+
if apiUrl == "" {
119+
apiUrl = gocqlastra.AstraAPIURL
120+
}
121+
122+
if bundle == "" && databaseID != "" && token != "" {
123+
cluster, err = GocqlastraNewClusterFromURL(apiUrl, databaseID, token, timeout)
124+
if err != nil {
125+
return nil, err
126+
}
127+
} else if bundle != "" && token != "" {
128+
cluster, err = GocqlastraNewClusterFromBundle(bundle, username, token, timeout)
129+
if err != nil {
130+
return nil, err
131+
}
132+
} else {
133+
return nil, ErrAstraMissing
134+
}
135+
} else {
136+
cluster = gocql.NewCluster(u.Host)
137+
}
138+
98139
cluster.Keyspace = strings.TrimPrefix(u.Path, "/")
99140
cluster.Consistency = gocql.All
100-
cluster.Timeout = 1 * time.Minute
141+
cluster.Timeout = timeout
101142

102-
if len(u.Query().Get("username")) > 0 && len(u.Query().Get("password")) > 0 {
143+
if !isAstra && len(username) > 0 && len(password) > 0 {
103144
authenticator := gocql.PasswordAuthenticator{
104-
Username: u.Query().Get("username"),
105-
Password: u.Query().Get("password"),
145+
Username: username,
146+
Password: password,
106147
}
107148
cluster.Authenticator = authenticator
108149
}

0 commit comments

Comments
 (0)