Add support for ScyllaDB (remove cql counter use)

bodyguard-ai · Dec 19, 2016 · 9e11d70 · 9e11d70
1 parent ce1b59b
commit 9e11d70
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 9 deletions.
diff --git a/driver/cassandra/cassandra.go b/driver/cassandra/cassandra.go
@@ -25,14 +25,19 @@ const (
 
 type counterStmt bool
 
-func (c counterStmt) String() string {
-	sign := ""
+func (c counterStmt) Exec(session *gocql.Session) error {
+	var version int64
+	if err := session.Query("SELECT version FROM "+tableName+" WHERE versionRow = ?", versionRow).Scan(&version); err != nil {
+		return err
+	}
+
 	if bool(c) {
-		sign = "+"
+		version++
 	} else {
-		sign = "-"
+		version--
 	}
-	return "UPDATE " + tableName + " SET version = version " + sign + " 1 where versionRow = ?"
+
+	return session.Query("UPDATE "+tableName+" SET version = ? WHERE versionRow = ?", version, versionRow).Exec()
 }
 
 const (
@@ -95,14 +100,17 @@ func (driver *Driver) Close() error {
 }
 
 func (driver *Driver) ensureVersionTableExists() error {
-	err := driver.session.Query("CREATE TABLE IF NOT EXISTS " + tableName + " (version counter, versionRow bigint primary key);").Exec()
+	err := driver.session.Query("CREATE TABLE IF NOT EXISTS " + tableName + " (version int, versionRow bigint primary key);").Exec()
 	if err != nil {
 		return err
 	}
 
 	_, err = driver.Version()
 	if err != nil {
-		driver.session.Query(up.String(), versionRow).Exec()
+		if err.Error() == "not found" {
+			return driver.session.Query("UPDATE "+tableName+" SET version = ? WHERE versionRow = ?", 1, versionRow).Exec()
+		}
+		return err
 	}
 
 	return nil
@@ -123,7 +131,7 @@ func (driver *Driver) version(d direction.Direction, invert bool) error {
 	if invert {
 		stmt = !stmt
 	}
-	return driver.session.Query(stmt.String(), versionRow).Exec()
+	return stmt.Exec(driver.session)
 }
 
 func (driver *Driver) Migrate(f file.File, pipe chan interface{}) {

diff --git a/driver/cassandra/cassandra_test.go b/driver/cassandra/cassandra_test.go
@@ -36,6 +36,9 @@ func TestMigrate(t *testing.T) {
 		t.Fatal(err)
 	}
 
+	if err := session.Query(`DROP KEYSPACE IF EXISTS migrate;`).Exec(); err != nil {
+		t.Fatal(err)
+	}
 	if err := session.Query(`CREATE KEYSPACE IF NOT EXISTS migrate WITH REPLICATION = {'class': 'SimpleStrategy', 'replication_factor': 1};`).Exec(); err != nil {
 		t.Fatal(err)
 	}
@@ -61,7 +64,6 @@ func TestMigrate(t *testing.T) {
                     msg text
                 );
 
-				CREATE INDEX ON yolo (msg);
             `),
 		},
 		{