YDB Support

Makefile

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

README.md

@@ -44,6 +44,7 @@ Database drivers run migrations. [Add a new database?](database/driver.go)
 * [Firebird](database/firebird)
 * [MS SQL Server](database/sqlserver)
 * [rqlite](database/rqlite)
+* [YDB](database/ydb)
 
 ### Database URLs
 

database/ydb/README.md

@@ -0,0 +1,30 @@
+# [YDB](https://ydb.tech/docs/)
+
+`ydb://[user:password@]host:port/database?QUERY_PARAMS`
+
+| URL Query  |               Description               |
+|:----------:|:---------------------------------------:|
+|   `user`   |         The user to sign in as.         |
+| `password` |          The user's password.           | 
+|   `host`   |         The host to connect to.         |
+|   `port`   |          The port to bind to.           |                                     |
+| `database` | The name of the database to connect to. |
+
+|   URL Query Params   |                         Description                         |
+|:--------------------:|:-----------------------------------------------------------:|
+|    `x-auth-token`    |                    Authentication token.                    |
+| `x-migrations-table` | Name of the migrations table (default `schema_migrations`). |
+|    `x-use-grpcs`     | Enables gRPCS protocol for YDB connections (default grpc).  |
+
+### Authentication
+
+By default, golang-migrate connects to YDB
+using [anonymous credentials](https://ydb.tech/docs/en/recipes/ydb-sdk/auth-anonymous). \
+Through the url query, you can change the default behavior:
+
+- To connect to YDB using [static credentials](https://ydb.tech/docs/en/recipes/ydb-sdk/auth-static) you need to specify
+  username and password:
+  `ydb://user:password@host:port/database`
+- To connect to YDB using [token](https://ydb.tech/docs/en/recipes/ydb-sdk/auth-access-token) you need to specify token
+  as query parameter:
+  `ydb://host:port/database?x-auth-token=<YDB_TOKEN>`

database/ydb/examples/migrations/001_create_users.down.yql

@@ -0,0 +1 @@
+DROP TABLE `test/users`;

database/ydb/examples/migrations/001_create_users.up.yql

@@ -0,0 +1,6 @@
+CREATE TABLE `test/users` (
+    id Uint64,
+    name String,
+    email String,
+    PRIMARY KEY (id)
+);

database/ydb/examples/migrations/002_add_city_to_users.down.yql

@@ -0,0 +1,4 @@
+DROP TABLE `test/cities`;
+
+ALTER TABLE `test/users`
+    DROP COLUMN city;

database/ydb/examples/migrations/002_add_city_to_users.up.yql

@@ -0,0 +1,8 @@
+CREATE TABLE `test/cities` (
+    id Uint64,
+    name String,
+    PRIMARY KEY (id)
+);
+
+ALTER TABLE `test/users`
+    ADD COLUMN city Uint64;

database/ydb/examples/migrations/003_create_topic.down.yql

@@ -0,0 +1 @@
+DROP TOPIC `test/topic`;

database/ydb/examples/migrations/003_create_topic.up.yql

@@ -0,0 +1 @@
+CREATE TOPIC `test/topic`;

database/ydb/ydb.go

@@ -0,0 +1,276 @@
+package ydb
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net/url"
+	"sync/atomic"
+
+	"github.com/hashicorp/go-multierror"
+	"github.com/ydb-platform/ydb-go-sdk/v3"
+	"github.com/ydb-platform/ydb-go-sdk/v3/query"
+
+	"github.com/golang-migrate/migrate/v4"
+	"github.com/golang-migrate/migrate/v4/database"
+)
+
+func init() {
+	db := YDB{}
+	database.Register("ydb", &db)
+}
+
+const (
+	defaultMigrationsTable = "schema_migrations"
+
+	queryParamAuthToken       = "x-auth-token"
+	queryParamMigrationsTable = "x-migrations-table"
+	queryParamUseGRPCS        = "x-use-grpcs"
+)
+
+var (
+	ErrNilConfig      = fmt.Errorf("no config")
+	ErrNoDatabaseName = fmt.Errorf("no database name")
+)
+
+type Config struct {
+	MigrationsTable string
+}
+
+type YDB struct {
+	driver *ydb.Driver
+	config *Config
+
+	isLocked atomic.Bool
+}
+
+func WithInstance(driver *ydb.Driver, config *Config) (database.Driver, error) {
+	if config == nil {
+		return nil, ErrNilConfig
+	}
+
+	if len(config.MigrationsTable) == 0 {
+		config.MigrationsTable = defaultMigrationsTable
+	}
+
+	db := &YDB{
+		driver: driver,
+		config: config,
+	}
+	if err := db.ensureVersionTable(); err != nil {
+		return nil, err
+	}
+	return db, nil
+}
+
+func (db *YDB) Open(dsn string) (database.Driver, error) {
+	purl, err := url.Parse(dsn)
+	if err != nil {
+		return nil, err
+	}
+
+	if len(purl.Path) == 0 {
+		return nil, ErrNoDatabaseName
+	}
+
+	pquery, err := url.ParseQuery(purl.RawQuery)
+	if err != nil {
+		return nil, err
+	}
+
+	switch {
+	case pquery.Has(queryParamUseGRPCS):
+		purl.Scheme = "grpcs"
+	default:
+		purl.Scheme = "grpc"
+	}
+
+	purl = migrate.FilterCustomQuery(purl)
+	credentials := db.parseCredentials(purl, pquery)
+	driver, err := ydb.Open(context.TODO(), purl.String(), credentials)
+	if err != nil {
+		return nil, err
+	}
+
+	px, err := WithInstance(driver, &Config{
+		MigrationsTable: pquery.Get(queryParamMigrationsTable),
+	})
+	if err != nil {
+		return nil, err
+	}
+	return px, nil
+}
+
+func (db *YDB) parseCredentials(url *url.URL, query url.Values) (credentials ydb.Option) {
+	switch {
+	case query.Has(queryParamAuthToken):
+		credentials = ydb.WithAccessTokenCredentials(query.Get(queryParamAuthToken))
+	case url.User != nil:
+		user := url.User.Username()
+		password, _ := url.User.Password()
+		credentials = ydb.WithStaticCredentials(user, password)
+	default:
+		credentials = ydb.WithAnonymousCredentials()
+	}
+	url.User = nil
+	return credentials
+}
+
+func (db *YDB) Close() error {
+	return db.driver.Close(context.TODO())
+}
+
+func (db *YDB) Run(migration io.Reader) error {
+	rawMigrations, err := io.ReadAll(migration)
+	if err != nil {
+		return err
+	}
+
+	res, err := db.driver.Scripting().Execute(context.TODO(), string(rawMigrations), nil)
+	if err != nil {
+		return err
+	}
+	return res.Close()
+}
+
+func (db *YDB) SetVersion(version int, dirty bool) error {
+	deleteVersionQuery := fmt.Sprintf(`
+		DELETE FROM %s 
+	`, db.config.MigrationsTable)
+
+	insertVersionQuery := fmt.Sprintf(`
+		INSERT INTO %s (version, dirty, created) VALUES (%d, %t, CurrentUtcTimestamp())
+	`, db.config.MigrationsTable, version, dirty)
+
+	ctx := context.TODO()
+	err := db.driver.Query().DoTx(ctx, func(ctx context.Context, tx query.TxActor) error {
+		if err := tx.Exec(ctx, deleteVersionQuery); err != nil {
+			return err
+		}
+		// Also re-write the schema version for nil dirty versions to prevent
+		// empty schema version for failed down migration on the first migration
+		// See: https://github.com/golang-migrate/migrate/issues/330
+		if version >= 0 || (version == database.NilVersion && dirty) {
+			if err := tx.Exec(ctx, insertVersionQuery); err != nil {
+				return err
+			}
+		}
+		return nil
+	}, query.WithTxSettings(query.TxSettings(query.WithSerializableReadWrite())))
+	return err
+}
+
+func (db *YDB) Version() (version int, dirty bool, err error) {
+	ctx := context.TODO()
+
+	getVersionQuery := fmt.Sprintf(`
+		SELECT version, dirty FROM %s ORDER BY version DESC LIMIT 1
+	`, db.config.MigrationsTable)
+
+	rs, err := db.driver.Query().QueryResultSet(ctx, getVersionQuery)
+	if err != nil {
+		return 0, false, &database.Error{OrigErr: err, Query: []byte(getVersionQuery)}
+	}
+	defer func() {
+		if closeErr := rs.Close(ctx); closeErr != nil {
+			err = multierror.Append(err, closeErr)
+		}
+	}()
+
+	row, err := rs.NextRow(ctx)
+	if err != nil {
+		if errors.Is(err, io.EOF) {
+			return database.NilVersion, false, nil
+		}
+		return 0, false, err
+	}
+
+	var v uint64
+	if err = row.Scan(&v, &dirty); err != nil {
+		return 0, false, &database.Error{OrigErr: err, Query: []byte(getVersionQuery)}
+	}
+	return int(v), dirty, err
+}
+
+func (db *YDB) Drop() (err error) {
+	ctx := context.TODO()
+
+	listQuery := "SELECT DISTINCT Path FROM `.sys/partition_stats` WHERE Path NOT LIKE '%/.sys%'"
+	rs, err := db.driver.Query().QueryResultSet(context.TODO(), listQuery)
+	if err != nil {
+		return &database.Error{OrigErr: err, Query: []byte(listQuery)}
+	}
+	defer func() {
+		if closeErr := rs.Close(ctx); closeErr != nil {
+			err = multierror.Append(err, closeErr)
+		}
+	}()
+
+	for {
+		var row query.Row
+		if row, err = rs.NextRow(ctx); err != nil {
+			if errors.Is(err, io.EOF) {
+				err = nil
+				break
+			}
+			return err
+		}
+
+		var table string
+		if err = row.Scan(&table); err != nil {
+			return err
+		}
+
+		dropQuery := fmt.Sprintf("DROP TABLE `%s`", table)
+		if err = db.driver.Query().Exec(ctx, dropQuery); err != nil {
+			return &database.Error{OrigErr: err, Query: []byte(dropQuery)}
+		}
+	}
+
+	return err
+}
+func (db *YDB) Lock() error {
+	if !db.isLocked.CompareAndSwap(false, true) {
+		return database.ErrLocked
+	}
+	return nil
+}
+
+func (db *YDB) Unlock() error {
+	if !db.isLocked.CompareAndSwap(true, false) {
+		return database.ErrNotLocked
+	}
+	return nil
+}
+
+// ensureVersionTable checks if versions table exists and, if not, creates it.
+func (db *YDB) ensureVersionTable() (err error) {
+	if err = db.Lock(); err != nil {
+		return err
+	}
+
+	defer func() {
+		if unlockErr := db.Unlock(); unlockErr != nil {
+			if err == nil {
+				err = unlockErr
+			} else {
+				err = multierror.Append(err, unlockErr)
+			}
+		}
+	}()
+
+	createVersionTableQuery := fmt.Sprintf(`
+		CREATE TABLE IF NOT EXISTS %s (
+			version Uint64,
+			dirty Bool,
+			created Timestamp,
+			PRIMARY KEY(version)
+		)
+	`, db.config.MigrationsTable)
+	err = db.driver.Query().Exec(context.TODO(), createVersionTableQuery)
+	if err != nil {
+		return err
+	}
+	return err
+}