Skip to content

Commit e2d7653

Browse files
asmyasnikovasmyasnikov
authored
Merge pull request #4 from art22m/refactor-ydb-lock
YDB Reimplement Lock Mechanics
2 parents 8d20ea1 + fa12828 commit e2d7653

File tree

3 files changed

+143
-37
lines changed

3 files changed

+143
-37
lines changed

‎database/ydb/README.md‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
|:----------------------------:|:--------------------------------------------------------------------------------------------:|
1515
| `x-auth-token` | Authentication token. |
1616
| `x-migrations-table` | Name of the migrations table (default `schema_migrations`). |
17+
| `x-lock-table` | Name of the table which maintains the migration lock (default `schema_lock`). |
1718
| `x-use-grpcs` | Enables gRPCS protocol for YDB connections (default grpc). |
1819
| `x-tls-ca` | The location of the CA (certificate authority) file. |
1920
| `x-tls-insecure-skip-verify` | Controls whether a client verifies the server's certificate chain and host name. |
@@ -37,4 +38,11 @@ Through the url query, you can change the default behavior:
3738
`ydb://user:password@host:port/database`
3839
- To connect to YDB using [token](https://ydb.tech/docs/en/recipes/ydb-sdk/auth-access-token) you need to specify token
3940
as query parameter:
40-
`ydb://host:port/database?x-auth-token=<YDB_TOKEN>`
41+
`ydb://host:port/database?x-auth-token=<YDB_TOKEN>`
42+
43+
### Locks
44+
45+
If golang-migrate fails to acquire the lock and no migrations are currently running.
46+
This may indicate that one of the migrations did not complete successfully.
47+
In this case, you need to analyze the previous migrations, rollback if necessary, and manually remove the lock from the
48+
`x-lock-table`.

‎database/ydb/ydb.go‎

Lines changed: 84 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,12 @@ import (
77
"fmt"
88
"io"
99
"net/url"
10-
"sync/atomic"
1110

1211
"github.com/hashicorp/go-multierror"
1312
"github.com/ydb-platform/ydb-go-sdk/v3"
13+
"github.com/ydb-platform/ydb-go-sdk/v3/balancers"
14+
"github.com/ydb-platform/ydb-go-sdk/v3/retry"
15+
"go.uber.org/atomic"
1416

1517
"github.com/golang-migrate/migrate/v4"
1618
"github.com/golang-migrate/migrate/v4/database"
@@ -22,9 +24,11 @@ func init() {
2224

2325
const (
2426
defaultMigrationsTable = "schema_migrations"
27+
defaultLockTable = "schema_lock"
2528

2629
queryParamAuthToken = "x-auth-token"
2730
queryParamMigrationsTable = "x-migrations-table"
31+
queryParamLockTable = "x-lock-table"
2832
queryParamUseGRPCS = "x-use-grpcs"
2933
queryParamTLSCertificateAuthorities = "x-tls-ca"
3034
queryParamTLSInsecureSkipVerify = "x-tls-insecure-skip-verify"
@@ -39,6 +43,8 @@ var (
3943

4044
type Config struct {
4145
MigrationsTable string
46+
LockTable string
47+
DatabaseName string
4248
}
4349

4450
type YDB struct {
@@ -63,6 +69,10 @@ func WithInstance(instance *sql.DB, config *Config) (database.Driver, error) {
6369
config.MigrationsTable = defaultMigrationsTable
6470
}
6571

72+
if len(config.LockTable) == 0 {
73+
config.LockTable = defaultLockTable
74+
}
75+
6676
conn, err := instance.Conn(context.TODO())
6777
if err != nil {
6878
return nil, err
@@ -73,6 +83,9 @@ func WithInstance(instance *sql.DB, config *Config) (database.Driver, error) {
7383
db: instance,
7484
config: config,
7585
}
86+
if err = db.ensureLockTable(); err != nil {
87+
return nil, err
88+
}
7689
if err = db.ensureVersionTable(); err != nil {
7790
return nil, err
7891
}
@@ -109,7 +122,11 @@ func (y *YDB) Open(dsn string) (database.Driver, error) {
109122
return nil, err
110123
}
111124

112-
nativeDriver, err := ydb.Open(context.TODO(), purl.String(), append(tlsOptions, credentials)...)
125+
nativeDriver, err := ydb.Open(
126+
context.TODO(),
127+
purl.String(),
128+
append(tlsOptions, credentials, ydb.WithBalancer(balancers.SingleConn()))...,
129+
)
113130
if err != nil {
114131
return nil, err
115132
}
@@ -123,6 +140,8 @@ func (y *YDB) Open(dsn string) (database.Driver, error) {
123140

124141
db, err := WithInstance(sql.OpenDB(connector), &Config{
125142
MigrationsTable: pquery.Get(queryParamMigrationsTable),
143+
LockTable: pquery.Get(queryParamLockTable),
144+
DatabaseName: purl.Path,
126145
})
127146
if err != nil {
128147
return nil, err
@@ -188,7 +207,7 @@ func (y *YDB) Run(migration io.Reader) error {
188207
return err
189208
}
190209

191-
if _, err = y.conn.ExecContext(ydb.WithQueryMode(context.TODO(), ydb.SchemeQueryMode), string(rawMigrations)); err != nil {
210+
if _, err = y.conn.ExecContext(context.Background(), string(rawMigrations)); err != nil {
192211
return database.Error{OrigErr: err, Err: "migration failed", Query: rawMigrations}
193212
}
194213
return nil
@@ -278,23 +297,77 @@ func (y *YDB) Drop() (err error) {
278297

279298
for _, path := range paths {
280299
dropQuery := fmt.Sprintf("DROP TABLE IF EXISTS `%s`", path)
281-
if _, err = y.conn.ExecContext(ydb.WithQueryMode(context.TODO(), ydb.SchemeQueryMode), dropQuery); err != nil {
300+
if _, err = y.conn.ExecContext(context.Background(), dropQuery); err != nil {
282301
return &database.Error{OrigErr: err, Query: []byte(dropQuery)}
283302
}
284303
}
285304
return nil
286305
}
287306

288307
func (y *YDB) Lock() error {
289-
if !y.isLocked.CompareAndSwap(false, true) {
290-
return database.ErrLocked
291-
}
292-
return nil
308+
return database.CasRestoreOnErr(&y.isLocked, false, true, database.ErrLocked, func() (err error) {
309+
return retry.DoTx(context.TODO(), y.db, func(ctx context.Context, tx *sql.Tx) (err error) {
310+
aid, err := database.GenerateAdvisoryLockId(y.config.DatabaseName)
311+
if err != nil {
312+
return err
313+
}
314+
315+
getLockQuery := fmt.Sprintf("SELECT * FROM %s WHERE lock_id = '%s'", y.config.LockTable, aid)
316+
rows, err := tx.Query(getLockQuery, aid)
317+
if err != nil {
318+
return database.Error{OrigErr: err, Err: "failed to fetch migration lock", Query: []byte(getLockQuery)}
319+
}
320+
defer func() {
321+
if errClose := rows.Close(); errClose != nil {
322+
err = multierror.Append(err, errClose)
323+
}
324+
}()
325+
326+
// If row exists at all, lock is present
327+
locked := rows.Next()
328+
if locked {
329+
return database.ErrLocked
330+
}
331+
332+
setLockQuery := fmt.Sprintf("INSERT INTO %s (lock_id) VALUES ('%s')", y.config.LockTable, aid)
333+
if _, err = tx.Exec(setLockQuery); err != nil {
334+
return database.Error{OrigErr: err, Err: "failed to set migration lock", Query: []byte(setLockQuery)}
335+
}
336+
return nil
337+
}, retry.WithTxOptions(&sql.TxOptions{Isolation: sql.LevelSerializable}))
338+
})
293339
}
294340

295341
func (y *YDB) Unlock() error {
296-
if !y.isLocked.CompareAndSwap(true, false) {
297-
return database.ErrNotLocked
342+
return database.CasRestoreOnErr(&y.isLocked, true, false, database.ErrNotLocked, func() (err error) {
343+
aid, err := database.GenerateAdvisoryLockId(y.config.DatabaseName)
344+
if err != nil {
345+
return err
346+
}
347+
348+
releaseLockQuery := fmt.Sprintf("DELETE FROM %s WHERE lock_id = '%s'", y.config.LockTable, aid)
349+
if _, err = y.conn.ExecContext(context.TODO(), releaseLockQuery); err != nil {
350+
// On drops, the lock table is fully removed; This is fine, and is a valid "unlocked" state for the schema.
351+
if ydb.IsOperationErrorSchemeError(err) {
352+
return nil
353+
}
354+
return database.Error{OrigErr: err, Err: "failed to release migration lock", Query: []byte(releaseLockQuery)}
355+
}
356+
357+
return nil
358+
})
359+
}
360+
361+
// ensureLockTable checks if lock table exists and, if not, creates it.
362+
func (y *YDB) ensureLockTable() (err error) {
363+
createLockTableQuery := fmt.Sprintf(`
364+
CREATE TABLE IF NOT EXISTS %s (
365+
lock_id String NOT NULL,
366+
PRIMARY KEY(lock_id)
367+
)
368+
`, y.config.LockTable)
369+
if _, err = y.conn.ExecContext(context.Background(), createLockTableQuery); err != nil {
370+
return &database.Error{OrigErr: err, Query: []byte(createLockTableQuery)}
298371
}
299372
return nil
300373
}
@@ -323,7 +396,7 @@ func (y *YDB) ensureVersionTable() (err error) {
323396
PRIMARY KEY(version)
324397
)
325398
`, y.config.MigrationsTable)
326-
if _, err = y.conn.ExecContext(ydb.WithQueryMode(context.TODO(), ydb.SchemeQueryMode), createVersionTableQuery); err != nil {
399+
if _, err = y.conn.ExecContext(context.Background(), createVersionTableQuery); err != nil {
327400
return &database.Error{OrigErr: err, Query: []byte(createVersionTableQuery)}
328401
}
329402
return nil

‎database/ydb/ydb_test.go‎

Lines changed: 50 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -3,55 +3,70 @@ package ydb
33
import (
44
"context"
55
"fmt"
6+
"log"
67
"strings"
78
"testing"
89
"time"
910

1011
"github.com/dhui/dktest"
1112
"github.com/docker/go-connections/nat"
1213
"github.com/ydb-platform/ydb-go-sdk/v3"
14+
"github.com/ydb-platform/ydb-go-sdk/v3/balancers"
1315

1416
"github.com/golang-migrate/migrate/v4"
1517
dt "github.com/golang-migrate/migrate/v4/database/testing"
18+
"github.com/golang-migrate/migrate/v4/dktesting"
1619
_ "github.com/golang-migrate/migrate/v4/source/file"
1720
)
1821

1922
const (
20-
image = "ydbplatform/local-ydb:latest"
21-
host = "localhost"
22-
port = "2136"
23+
defaultPort = 2136
2324
databaseName = "local"
2425
)
2526

2627
var (
27-
opts = dktest.Options{
28-
Env: map[string]string{
29-
"GRPC_TLS_PORT": "2135",
30-
"GRPC_PORT": "2136",
31-
"MON_PORT": "8765",
32-
},
33-
34-
PortBindings: nat.PortMap{
35-
nat.Port("2136/tcp"): []nat.PortBinding{
36-
{
37-
HostIP: "0.0.0.0",
38-
HostPort: port,
28+
getOptions = func(port string) dktest.Options {
29+
return dktest.Options{
30+
Env: map[string]string{
31+
"GRPC_TLS_PORT": "2135",
32+
"GRPC_PORT": "2136",
33+
"MON_PORT": "8765",
34+
},
35+
PortBindings: nat.PortMap{
36+
nat.Port(fmt.Sprintf("%d/tcp", defaultPort)): []nat.PortBinding{
37+
{
38+
HostIP: "0.0.0.0",
39+
HostPort: port,
40+
},
3941
},
4042
},
41-
},
43+
PortRequired: true,
44+
Hostname: "127.0.0.1",
45+
ReadyTimeout: 15 * time.Second,
46+
ReadyFunc: isReady,
47+
}
48+
}
4249

43-
Hostname: host,
44-
ReadyTimeout: 15 * time.Second,
45-
ReadyFunc: isReady,
50+
// Released version: https://ydb.tech/docs/downloads/#ydb-server
51+
specs = []dktesting.ContainerSpec{
52+
{ImageName: "ydbplatform/local-ydb:latest", Options: getOptions("22000")},
53+
{ImageName: "ydbplatform/local-ydb:24.3", Options: getOptions("22001")},
54+
{ImageName: "ydbplatform/local-ydb:24.2", Options: getOptions("22002")},
4655
}
4756
)
4857

49-
func connectionString(options ...string) string {
58+
func connectionString(host, port string, options ...string) string {
5059
return fmt.Sprintf("ydb://%s:%s/%s?%s", host, port, databaseName, strings.Join(options, "&"))
5160
}
5261

5362
func isReady(ctx context.Context, c dktest.ContainerInfo) bool {
54-
d, err := ydb.Open(ctx, fmt.Sprintf("grpc://%s:%s/%s", host, port, databaseName))
63+
ip, port, err := c.Port(defaultPort)
64+
if err != nil {
65+
log.Println("port error:", err)
66+
return false
67+
}
68+
69+
d, err := ydb.Open(ctx, fmt.Sprintf("grpc://%s:%s/%s", ip, port, databaseName), ydb.WithBalancer(balancers.SingleConn()))
5570
if err != nil {
5671
return false
5772
}
@@ -67,9 +82,14 @@ func isReady(ctx context.Context, c dktest.ContainerInfo) bool {
6782
}
6883

6984
func Test(t *testing.T) {
70-
dktest.Run(t, image, opts, func(t *testing.T, c dktest.ContainerInfo) {
85+
dktesting.ParallelTest(t, specs, func(t *testing.T, c dktest.ContainerInfo) {
86+
ip, port, err := c.Port(defaultPort)
87+
if err != nil {
88+
t.Fatal(err)
89+
}
90+
7191
db := &YDB{}
72-
d, err := db.Open(connectionString())
92+
d, err := db.Open(connectionString(ip, port))
7393
if err != nil {
7494
t.Fatal(err)
7595
}
@@ -86,9 +106,14 @@ func Test(t *testing.T) {
86106
}
87107

88108
func TestMigrate(t *testing.T) {
89-
dktest.Run(t, image, opts, func(t *testing.T, c dktest.ContainerInfo) {
109+
dktesting.ParallelTest(t, specs, func(t *testing.T, c dktest.ContainerInfo) {
110+
ip, port, err := c.Port(defaultPort)
111+
if err != nil {
112+
t.Fatal(err)
113+
}
114+
90115
db := &YDB{}
91-
d, err := db.Open(connectionString())
116+
d, err := db.Open(connectionString(ip, port))
92117
if err != nil {
93118
t.Fatal(err)
94119
}

0 commit comments

Comments
 (0)