dialect/sql/schema: added support for ydb schema migrations

.github/workflows/ci.yml

@@ -26,7 +26,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        go: ['1.23', '1.24']
+        go: ['1.23', '1.24.9']
     steps:
       - uses: actions/checkout@v4
       - uses: actions/setup-go@v5

dialect/sql/schema/atlas.go

@@ -1126,6 +1126,8 @@ func (a *Atlas) entDialect(ctx context.Context, drv dialect.Driver) (sqlDialect,
 		d = &SQLite{Driver: drv, WithForeignKeys: a.withForeignKeys}
 	case dialect.Postgres:
 		d = &Postgres{Driver: drv}
+	case dialect.YDB:
+		d = &YDB{Driver: drv}
 	default:
 		return nil, fmt.Errorf("sql/schema: unsupported dialect %q", a.dialect)
 	}

dialect/sql/schema/schema.go

@@ -17,6 +17,7 @@ import (
 	"ariga.io/atlas/sql/postgres"
 	"ariga.io/atlas/sql/schema"
 	"ariga.io/atlas/sql/sqlite"
+	"ariga.io/atlas/sql/ydb"
 	entdialect "entgo.io/ent/dialect"
 	"entgo.io/ent/dialect/entsql"
 	"entgo.io/ent/dialect/sql"
@@ -584,6 +585,14 @@ var drivers = func(v string) map[string]driver {
 			postgres.DefaultDiff,
 			postgres.DefaultPlan,
 		},
+		entdialect.YDB: {
+			&YDB{
+				version: v,
+				Driver:  nopDriver{dialect: entdialect.YDB},
+			},
+			ydb.DefaultDiff,
+			ydb.DefaultPlan,
+		},
 	}
 }
 

dialect/sql/schema/ydb.go

@@ -0,0 +1,255 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package schema
+
+import (
+	"context"
+	"database/sql"
+	"errors"
+	"fmt"
+	"strings"
+
+	"entgo.io/ent/dialect"
+	entsql "entgo.io/ent/dialect/sql"
+	entdriver "entgo.io/ent/dialect/ydb"
+	"entgo.io/ent/schema/field"
+
+	"ariga.io/atlas/sql/migrate"
+	"ariga.io/atlas/sql/schema"
+	atlas "ariga.io/atlas/sql/ydb"
+)
+
+// YDB adapter for Atlas migration engine.
+type YDB struct {
+	dialect.Driver
+
+	version string
+}
+
+// init loads the YDB version from the database for later use in the migration process.
+func (d *YDB) init(ctx context.Context) error {
+	if d.version != "" {
+		return nil // already initialized.
+	}
+
+	rows := &sql.Rows{}
+	if err := d.Driver.Query(ctx, "SELECT version()", nil, rows); err != nil {
+		return fmt.Errorf("ydb: failed to query version: %w", err)
+	}
+	defer rows.Close()
+
+	if !rows.Next() {
+		if err := rows.Err(); err != nil {
+			return err
+		}
+		return fmt.Errorf("ydb: version was not found")
+	}
+
+	var version string
+	if err := rows.Scan(&version); err != nil {
+		return fmt.Errorf("ydb: failed to scan version: %w", err)
+	}
+
+	d.version = version
+	return nil
+}
+
+// tableExist checks if a table exists in the database by querying the .sys/tables system table.
+func (d *YDB) tableExist(ctx context.Context, conn dialect.ExecQuerier, name string) (bool, error) {
+	query, args := entsql.Dialect(dialect.YDB).
+		Select(entsql.Count("*")).
+		From(entsql.Table(".sys/tables")).
+		Where(entsql.EQ("table_name", name)).
+		Query()
+
+	return exist(ctx, conn, query, args...)
+}
+
+// atOpen returns a custom Atlas migrate.Driver for YDB.
+func (d *YDB) atOpen(conn dialect.ExecQuerier) (migrate.Driver, error) {
+	ydbDriver, ok := conn.(*entdriver.YDBDriver)
+	if !ok {
+		return nil, fmt.Errorf("expected dialect/ydb.YDBDriver, but got %T", conn)
+	}
+
+	return atlas.Open(
+		ydbDriver.NativeDriver(),
+		ydbDriver.DB(),
+	)
+}
+
+func (d *YDB) atTable(table1 *Table, table2 *schema.Table) {
+	if table1.Annotation != nil {
+		setAtChecks(table1, table2)
+	}
+}
+
+// supportsDefault returns whether YDB supports DEFAULT values for the given column type.
+func (d *YDB) supportsDefault(column *Column) bool {
+	switch column.Default.(type) {
+	case Expr, map[string]Expr:
+		// Expression defaults are not well supported in YDB
+		return false
+	default:
+		// Simple literal defaults should work for basic types
+		return column.supportDefault()
+	}
+}
+
+// atTypeC converts an Ent column type to a YDB Atlas schema type.
+func (d *YDB) atTypeC(column1 *Column, column2 *schema.Column) error {
+	// Check for custom schema type override.
+	if column1.SchemaType != nil && column1.SchemaType[dialect.YDB] != "" {
+		typ, err := atlas.ParseType(
+			column1.SchemaType[dialect.YDB],
+		)
+		if err != nil {
+			return err
+		}
+		column2.Type.Type = typ
+		return nil
+	}
+
+	var (
+		typ schema.Type
+		err error
+	)
+
+	switch column1.Type {
+	case field.TypeBool:
+		typ = &schema.BoolType{T: atlas.TypeBool}
+	case field.TypeInt8:
+		typ = &schema.IntegerType{T: atlas.TypeInt8}
+	case field.TypeInt16:
+		typ = &schema.IntegerType{T: atlas.TypeInt16}
+	case field.TypeInt32:
+		typ = &schema.IntegerType{T: atlas.TypeInt32}
+	case field.TypeInt, field.TypeInt64:
+		typ = &schema.IntegerType{T: atlas.TypeInt64}
+	case field.TypeUint8:
+		typ = &schema.IntegerType{T: atlas.TypeUint8, Unsigned: true}
+	case field.TypeUint16:
+		typ = &schema.IntegerType{T: atlas.TypeUint16, Unsigned: true}
+	case field.TypeUint32:
+		typ = &schema.IntegerType{T: atlas.TypeUint32, Unsigned: true}
+	case field.TypeUint, field.TypeUint64:
+		typ = &schema.IntegerType{T: atlas.TypeUint64, Unsigned: true}
+	case field.TypeFloat32:
+		typ = &schema.FloatType{T: atlas.TypeFloat}
+	case field.TypeFloat64:
+		typ = &schema.FloatType{T: atlas.TypeDouble}
+	case field.TypeBytes:
+		typ = &schema.BinaryType{T: atlas.TypeString}
+	case field.TypeString:
+		typ = &schema.StringType{T: atlas.TypeUtf8}
+	case field.TypeJSON:
+		typ = &schema.JSONType{T: atlas.TypeJson}
+	case field.TypeTime:
+		typ = &schema.TimeType{T: atlas.TypeTimestamp}
+	case field.TypeUUID:
+		typ = &schema.UUIDType{T: atlas.TypeUuid}
+	case field.TypeEnum:
+		err = errors.New("ydb: Enum can't be used as column data type for tables")
+	case field.TypeOther:
+		typ = &schema.UnsupportedType{T: column1.typ}
+	default:
+		typ, err = atlas.ParseType(column1.typ)
+	}
+
+	if err != nil {
+		return err
+	}
+
+	column2.Type.Type = typ
+	return nil
+}
+
+// atUniqueC adds a unique constraint for a column.
+// In YDB, unique constraints are implemented as GLOBAL UNIQUE SYNC indexes.
+func (d *YDB) atUniqueC(
+	table1 *Table,
+	column1 *Column,
+	table2 *schema.Table,
+	column2 *schema.Column,
+) {
+	// Check if there's already an explicit unique index defined for this column.
+	for _, idx := range table1.Indexes {
+		if idx.Unique && len(idx.Columns) == 1 && idx.Columns[0].Name == column1.Name {
+			// Index already defined explicitly, will be added in atIndexes.
+			return
+		}
+	}
+	// Create a unique index for this column.
+	idxName := fmt.Sprintf("%s_%s_index", table1.Name, column1.Name)
+	index := schema.NewUniqueIndex(idxName).AddColumns(column2)
+
+	// Add YDB-specific attribute for GLOBAL SYNC index type.
+	index.AddAttrs(&atlas.YDBIndexAttributes{Global: true, Sync: true})
+
+	table2.AddIndexes(index)
+}
+
+// atIncrementC configures auto-increment for a column.
+// YDB uses Serial types for auto-increment.
+func (d *YDB) atIncrementC(table *schema.Table, column *schema.Column) {
+	if intType, ok := column.Type.Type.(*schema.IntegerType); ok {
+		column.Type.Type = atlas.SerialFromInt(intType)
+	}
+}
+
+// atIncrementT sets the table-level auto-increment starting value.
+func (d *YDB) atIncrementT(table *schema.Table, v int64) {
+	// not implemented
+}
+
+// atIndex configures an index for ydb.
+func (d *YDB) atIndex(
+	index1 *Index,
+	table2 *schema.Table,
+	index2 *schema.Index,
+) error {
+	for _, column1 := range index1.Columns {
+		column2, ok := table2.Column(column1.Name)
+		if !ok {
+			return fmt.Errorf("unexpected index %q column: %q", index1.Name, column1.Name)
+		}
+		index2.AddParts(&schema.IndexPart{C: column2})
+	}
+
+	// Set YDB-specific index attributes.
+	// By default, use GLOBAL SYNC for consistency.
+	idxType := &atlas.YDBIndexAttributes{Global: true, Sync: true}
+
+	// Check for annotation overrides.
+	if index1.Annotation != nil {
+		if indexType, ok := indexType(index1, dialect.YDB); ok {
+			// Parse YDB-specific index type from annotation.
+			switch strings.ToUpper(indexType) {
+			case "GLOBAL ASYNC", "ASYNC":
+				idxType.Sync = false
+			case "LOCAL":
+				idxType.Global = false
+			case "LOCAL ASYNC":
+				idxType.Global = false
+				idxType.Sync = false
+			}
+		}
+	}
+	index2.AddAttrs(idxType)
+	return nil
+}
+
+// atTypeRangeSQL returns the SQL statement to insert type ranges for global unique IDs.
+func (*YDB) atTypeRangeSQL(ts ...string) string {
+	values := make([]string, len(ts))
+	for i, t := range ts {
+		values[i] = fmt.Sprintf("('%s')", t)
+	}
+	return fmt.Sprintf(
+		"UPSERT INTO `%s` (`type`) VALUES %s",
+		TypeTable,
+		strings.Join(values, ", "),
+	)
+}

dialect/ydb/driver.go

@@ -32,7 +32,7 @@ func Open(ctx context.Context, dsn string) (*YDBDriver, error) {
 		ydb.WithTablePathPrefix(nativeDriver.Name()),
 	)
 	if err != nil {
-		panic(err)
+		return nil, err
 	}
 
 	dbSQLDriver := sql.OpenDB(conn)

entc/integration/go.mod

@@ -1,6 +1,8 @@
 module entgo.io/ent/entc/integration
 
-go 1.24.0
+go 1.24.9
+
+toolchain go1.24.11
 
 replace entgo.io/ent => ../../
 
@@ -12,8 +14,8 @@ require (
 	github.com/google/uuid v1.6.0
 	github.com/lib/pq v1.10.7
 	github.com/mattn/go-sqlite3 v1.14.28
-	github.com/stretchr/testify v1.10.0
-	golang.org/x/sync v0.17.0
+	github.com/stretchr/testify v1.11.1
+	golang.org/x/sync v0.19.0
 )
 
 require (
@@ -22,8 +24,10 @@ require (
 	github.com/bmatcuk/doublestar v1.3.4 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/go-openapi/inflect v0.19.0 // indirect
+	github.com/golang-jwt/jwt/v4 v4.5.2 // indirect
 	github.com/google/go-cmp v0.7.0 // indirect
 	github.com/hashicorp/hcl/v2 v2.18.1 // indirect
+	github.com/jonboulle/clockwork v0.5.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
 	github.com/mitchellh/go-wordwrap v1.0.1 // indirect
 	github.com/mitchellh/mapstructure v1.5.0 // indirect
@@ -32,9 +36,18 @@ require (
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/rogpeppe/go-internal v1.11.0 // indirect
 	github.com/stretchr/objx v0.5.2 // indirect
+	github.com/ydb-platform/ydb-go-genproto v0.0.0-20251222105147-0bf751469a4a // indirect
+	github.com/ydb-platform/ydb-go-sdk/v3 v3.125.0 // indirect
 	github.com/zclconf/go-cty v1.14.4 // indirect
 	github.com/zclconf/go-cty-yaml v1.1.0 // indirect
-	golang.org/x/mod v0.28.0 // indirect
-	golang.org/x/text v0.30.0 // indirect
+	golang.org/x/mod v0.30.0 // indirect
+	golang.org/x/net v0.48.0 // indirect
+	golang.org/x/sys v0.39.0 // indirect
+	golang.org/x/text v0.32.0 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20251222181119-0a764e51fe1b // indirect
+	google.golang.org/grpc v1.78.0 // indirect
+	google.golang.org/protobuf v1.36.11 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
+
+replace ariga.io/atlas => github.com/LostImagin4tion/atlas v0.0.5