Skip to content

Commit 61944de

Browse files
negachovnegachov
committed
sql/mysql: add AUTO_RANDOM support for TiDB
1 parent 1a19b3d commit 61944de

9 files changed

Lines changed: 1071 additions & 1 deletion

File tree

‎internal/integration/docker-compose.yaml‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,13 @@ services:
170170
image: pingcap/tidb:v6.6.0
171171
ports:
172172
- "4310:4000"
173+
174+
tidb8:
175+
container_name: atlas-integration-tidb8
176+
platform: linux/amd64
177+
image: pingcap/tidb:v8.5.5
178+
ports:
179+
- "4311:4000"
173180

174181
cockroach21.2.11:
175182
container_name: atlas-integration-cockroach21.2.11

‎internal/integration/tidb_test.go‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import (
2121
var tidbTests = map[string]*myTest{
2222
"tidb5": {port: 4309},
2323
"tidb6": {port: 4310},
24+
"tidb8": {port: 4311},
2425
}
2526

2627
func tidbRun(t *testing.T, fn func(*myTest)) {
@@ -1038,3 +1039,74 @@ create table atlas_types_sanity
10381039
})
10391040
})
10401041
}
1042+
1043+
func TestTiDB_AutoRandom(t *testing.T) {
1044+
t.Run("CreateAndInspect", func(t *testing.T) {
1045+
tidbRun(t, func(t *myTest) {
1046+
t.dropTables("ar_test")
1047+
_, err := t.db.Exec("CREATE TABLE ar_test (id bigint NOT NULL AUTO_RANDOM(5), PRIMARY KEY (id) CLUSTERED)")
1048+
require.NoError(t, err)
1049+
tbl := t.loadTable("ar_test")
1050+
require.NotNil(t, tbl)
1051+
col, ok := tbl.Column("id")
1052+
require.True(t, ok)
1053+
ar := findAutoRandom(t, col)
1054+
require.Equal(t, 5, ar.ShardBits)
1055+
})
1056+
})
1057+
t.Run("CreateAndInspectWithRange", func(t *testing.T) {
1058+
tidbRun(t, func(t *myTest) {
1059+
t.dropTables("ar_range")
1060+
_, err := t.db.Exec("CREATE TABLE ar_range (id bigint NOT NULL AUTO_RANDOM(3, 32), PRIMARY KEY (id) CLUSTERED)")
1061+
require.NoError(t, err)
1062+
tbl := t.loadTable("ar_range")
1063+
require.NotNil(t, tbl)
1064+
col, ok := tbl.Column("id")
1065+
require.True(t, ok)
1066+
ar := findAutoRandom(t, col)
1067+
require.Equal(t, 3, ar.ShardBits)
1068+
require.Equal(t, 32, ar.RangeBits)
1069+
})
1070+
})
1071+
t.Run("NoDrift", func(t *testing.T) {
1072+
tidbRun(t, func(t *myTest) {
1073+
t.dropTables("ar_nodrift")
1074+
_, err := t.db.Exec("CREATE TABLE ar_nodrift (id bigint NOT NULL AUTO_RANDOM(5), PRIMARY KEY (id) CLUSTERED)")
1075+
require.NoError(t, err)
1076+
tbl := t.loadTable("ar_nodrift")
1077+
ensureNoChange(t, tbl)
1078+
})
1079+
})
1080+
t.Run("HCLRoundTrip", func(t *testing.T) {
1081+
tidbRun(t, func(t *myTest) {
1082+
t.dropTables("ar_hcl")
1083+
_, err := t.db.Exec("CREATE TABLE ar_hcl (id bigint NOT NULL AUTO_RANDOM(5), PRIMARY KEY (id) CLUSTERED)")
1084+
require.NoError(t, err)
1085+
realm := t.loadRealm()
1086+
spec, err := mysql.MarshalHCL(realm.Schemas[0])
1087+
require.NoError(t, err)
1088+
require.Contains(t, string(spec), "auto_random")
1089+
var s schema.Realm
1090+
err = mysql.EvalHCLBytes(spec, &s, nil)
1091+
require.NoError(t, err)
1092+
tbl, ok := s.Schemas[0].Table("ar_hcl")
1093+
require.True(t, ok)
1094+
col, ok := tbl.Column("id")
1095+
require.True(t, ok)
1096+
ar := findAutoRandom(t, col)
1097+
require.Equal(t, 5, ar.ShardBits)
1098+
})
1099+
})
1100+
}
1101+
1102+
// findAutoRandom returns the AutoRandom attribute from the column or fails the test.
1103+
func findAutoRandom(t testing.TB, col *schema.Column) *mysql.AutoRandom {
1104+
t.Helper()
1105+
for _, a := range col.Attrs {
1106+
if ar, ok := a.(*mysql.AutoRandom); ok {
1107+
return ar
1108+
}
1109+
}
1110+
t.Fatal("AutoRandom attribute not found on column")
1111+
return nil
1112+
}

‎sql/mysql/inspect_oss.go‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -276,6 +276,11 @@ func (i *inspect) addColumn(s *schema.Schema, rows *sql.Rows) error {
276276
}
277277
c.Attrs = append(c.Attrs, a)
278278
}
279+
if attr.autorandom {
280+
// Placeholder with zero values; actual ShardBits/RangeBits are
281+
// extracted from the CREATE TABLE statement in tinspect.setAutoRandom.
282+
c.Attrs = append(c.Attrs, &AutoRandom{})
283+
}
279284
if attr.onUpdate != "" {
280285
c.Attrs = append(c.Attrs, &OnUpdate{A: attr.onUpdate})
281286
}
@@ -476,6 +481,7 @@ func (i *inspect) indexQuery() string {
476481
// extraAttr is a parsed version of the information_schema EXTRA column.
477482
type extraAttr struct {
478483
autoinc bool
484+
autorandom bool
479485
onUpdate string
480486
generatedType string
481487
defaultGenerated bool
@@ -498,6 +504,11 @@ func parseExtra(extra string) (*extraAttr, error) {
498504
// and it is handled in Driver.addColumn.
499505
case el == autoIncrement:
500506
attr.autoinc = true
507+
case el == "auto_random" || strings.HasPrefix(el, "auto_random("):
508+
// TiDB returns "auto_random" (or "auto_random(5)" in v7+) in the EXTRA
509+
// column for columns with the AUTO_RANDOM attribute. Shard/range bits
510+
// are extracted from the CREATE TABLE statement in tinspect.setAutoRandom.
511+
attr.autorandom = true
501512
case reTimeOnUpdate.MatchString(extra):
502513
attr.onUpdate = reTimeOnUpdate.FindStringSubmatch(extra)[1]
503514
case reGenerateType.MatchString(extra):
@@ -800,6 +811,20 @@ type (
800811
V int64
801812
}
802813

814+
// AutoRandom is a TiDB-specific attribute for BIGINT primary key columns
815+
// that generates random unique IDs to avoid write hotspots.
816+
//
817+
// Restrictions:
818+
// - Only valid on BIGINT primary key columns with CLUSTERED index.
819+
// - Cannot be removed once set (TiDB limitation).
820+
// - ShardBits: 1-15 (default 5).
821+
// - RangeBits: 32-64 or 0 for default (64).
822+
AutoRandom struct {
823+
schema.Attr
824+
ShardBits int
825+
RangeBits int
826+
}
827+
803828
// CreateOptions attribute for describing extra options used with CREATE TABLE.
804829
CreateOptions struct {
805830
schema.Attr

‎sql/mysql/migrate_oss.go‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -576,6 +576,16 @@ func (s *state) column(b *sqlx.Builder, t *schema.Table, c *schema.Column) error
576576
if a.V > 0 && !sqlx.Has(t.Attrs, &AutoIncrement{}) {
577577
t.Attrs = append(t.Attrs, a)
578578
}
579+
case *AutoRandom:
580+
if a.ShardBits == 0 {
581+
// ShardBits=0 is a zero-value placeholder; skip SQL generation.
582+
break
583+
}
584+
if a.RangeBits > 0 && a.RangeBits != 64 {
585+
b.P(fmt.Sprintf("AUTO_RANDOM(%d, %d)", a.ShardBits, a.RangeBits))
586+
} else {
587+
b.P(fmt.Sprintf("AUTO_RANDOM(%d)", a.ShardBits))
588+
}
579589
default:
580590
s.attr(b, a)
581591
}

‎sql/mysql/migrate_oss_test.go‎

Lines changed: 183 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1146,6 +1146,189 @@ func TestPlanChanges(t *testing.T) {
11461146
},
11471147
wantErr: true,
11481148
},
1149+
// AUTO_RANDOM with shard bits.
1150+
{
1151+
version: "5.7.25-TiDB-v6.1.0",
1152+
changes: []schema.Change{
1153+
func() *schema.AddTable {
1154+
t := &schema.Table{
1155+
Name: "users",
1156+
Columns: []*schema.Column{
1157+
{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "bigint"}}, Attrs: []schema.Attr{&AutoRandom{ShardBits: 5}}},
1158+
{Name: "name", Type: &schema.ColumnType{Type: &schema.StringType{T: "varchar", Size: 255}}},
1159+
},
1160+
}
1161+
t.PrimaryKey = &schema.Index{Parts: []*schema.IndexPart{{C: t.Columns[0]}}}
1162+
return &schema.AddTable{T: t}
1163+
}(),
1164+
},
1165+
wantPlan: &migrate.Plan{
1166+
Reversible: true,
1167+
Changes: []*migrate.Change{{Cmd: "CREATE TABLE `users` (`id` bigint NOT NULL AUTO_RANDOM(5), `name` varchar(255) NOT NULL, PRIMARY KEY (`id`))", Reverse: "DROP TABLE `users`"}},
1168+
},
1169+
},
1170+
// AUTO_RANDOM with zero ShardBits (placeholder) should not emit AUTO_RANDOM in SQL.
1171+
{
1172+
version: "5.7.25-TiDB-v6.1.0",
1173+
changes: []schema.Change{
1174+
func() *schema.AddTable {
1175+
t := &schema.Table{
1176+
Name: "users",
1177+
Columns: []*schema.Column{
1178+
{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "bigint"}}, Attrs: []schema.Attr{&AutoRandom{}}},
1179+
},
1180+
}
1181+
t.PrimaryKey = &schema.Index{Parts: []*schema.IndexPart{{C: t.Columns[0]}}}
1182+
return &schema.AddTable{T: t}
1183+
}(),
1184+
},
1185+
wantPlan: &migrate.Plan{
1186+
Reversible: true,
1187+
Changes: []*migrate.Change{{Cmd: "CREATE TABLE `users` (`id` bigint NOT NULL, PRIMARY KEY (`id`))", Reverse: "DROP TABLE `users`"}},
1188+
},
1189+
},
1190+
// AUTO_RANDOM with min shard bits (boundary).
1191+
{
1192+
version: "5.7.25-TiDB-v6.1.0",
1193+
changes: []schema.Change{
1194+
func() *schema.AddTable {
1195+
t := &schema.Table{
1196+
Name: "users",
1197+
Columns: []*schema.Column{
1198+
{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "bigint"}}, Attrs: []schema.Attr{&AutoRandom{ShardBits: 1}}},
1199+
},
1200+
}
1201+
t.PrimaryKey = &schema.Index{Parts: []*schema.IndexPart{{C: t.Columns[0]}}}
1202+
return &schema.AddTable{T: t}
1203+
}(),
1204+
},
1205+
wantPlan: &migrate.Plan{
1206+
Reversible: true,
1207+
Changes: []*migrate.Change{{Cmd: "CREATE TABLE `users` (`id` bigint NOT NULL AUTO_RANDOM(1), PRIMARY KEY (`id`))", Reverse: "DROP TABLE `users`"}},
1208+
},
1209+
},
1210+
// AUTO_RANDOM with max shard bits (boundary).
1211+
{
1212+
version: "5.7.25-TiDB-v6.1.0",
1213+
changes: []schema.Change{
1214+
func() *schema.AddTable {
1215+
t := &schema.Table{
1216+
Name: "users",
1217+
Columns: []*schema.Column{
1218+
{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "bigint"}}, Attrs: []schema.Attr{&AutoRandom{ShardBits: 15}}},
1219+
},
1220+
}
1221+
t.PrimaryKey = &schema.Index{Parts: []*schema.IndexPart{{C: t.Columns[0]}}}
1222+
return &schema.AddTable{T: t}
1223+
}(),
1224+
},
1225+
wantPlan: &migrate.Plan{
1226+
Reversible: true,
1227+
Changes: []*migrate.Change{{Cmd: "CREATE TABLE `users` (`id` bigint NOT NULL AUTO_RANDOM(15), PRIMARY KEY (`id`))", Reverse: "DROP TABLE `users`"}},
1228+
},
1229+
},
1230+
// AUTO_RANDOM with shard and range bits.
1231+
{
1232+
version: "5.7.25-TiDB-v6.1.0",
1233+
changes: []schema.Change{
1234+
func() *schema.AddTable {
1235+
t := &schema.Table{
1236+
Name: "users",
1237+
Columns: []*schema.Column{
1238+
{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "bigint"}}, Attrs: []schema.Attr{&AutoRandom{ShardBits: 5, RangeBits: 32}}},
1239+
},
1240+
}
1241+
t.PrimaryKey = &schema.Index{Parts: []*schema.IndexPart{{C: t.Columns[0]}}}
1242+
return &schema.AddTable{T: t}
1243+
}(),
1244+
},
1245+
wantPlan: &migrate.Plan{
1246+
Reversible: true,
1247+
Changes: []*migrate.Change{{Cmd: "CREATE TABLE `users` (`id` bigint NOT NULL AUTO_RANDOM(5, 32), PRIMARY KEY (`id`))", Reverse: "DROP TABLE `users`"}},
1248+
},
1249+
},
1250+
// AUTO_RANDOM with range bits = 64 (default, should not emit range).
1251+
{
1252+
version: "5.7.25-TiDB-v6.1.0",
1253+
changes: []schema.Change{
1254+
func() *schema.AddTable {
1255+
t := &schema.Table{
1256+
Name: "users",
1257+
Columns: []*schema.Column{
1258+
{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "bigint"}}, Attrs: []schema.Attr{&AutoRandom{ShardBits: 5, RangeBits: 64}}},
1259+
},
1260+
}
1261+
t.PrimaryKey = &schema.Index{Parts: []*schema.IndexPart{{C: t.Columns[0]}}}
1262+
return &schema.AddTable{T: t}
1263+
}(),
1264+
},
1265+
wantPlan: &migrate.Plan{
1266+
Reversible: true,
1267+
Changes: []*migrate.Change{{Cmd: "CREATE TABLE `users` (`id` bigint NOT NULL AUTO_RANDOM(5), PRIMARY KEY (`id`))", Reverse: "DROP TABLE `users`"}},
1268+
},
1269+
},
1270+
// ALTER TABLE: add AUTO_RANDOM to an existing column.
1271+
{
1272+
version: "5.7.25-TiDB-v6.1.0",
1273+
changes: []schema.Change{
1274+
&schema.ModifyTable{
1275+
T: &schema.Table{
1276+
Name: "users",
1277+
Columns: []*schema.Column{
1278+
{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "bigint"}}, Attrs: []schema.Attr{&AutoRandom{ShardBits: 5}}},
1279+
},
1280+
},
1281+
Changes: []schema.Change{
1282+
&schema.ModifyColumn{
1283+
Change: schema.ChangeAttr,
1284+
From: &schema.Column{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "bigint"}}},
1285+
To: &schema.Column{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "bigint"}}, Attrs: []schema.Attr{&AutoRandom{ShardBits: 5}}},
1286+
},
1287+
},
1288+
},
1289+
},
1290+
wantPlan: &migrate.Plan{
1291+
Reversible: true,
1292+
Transactional: false,
1293+
Changes: []*migrate.Change{
1294+
{
1295+
Cmd: "ALTER TABLE `users` MODIFY COLUMN `id` bigint NOT NULL AUTO_RANDOM(5)",
1296+
Reverse: "ALTER TABLE `users` MODIFY COLUMN `id` bigint NOT NULL",
1297+
},
1298+
},
1299+
},
1300+
},
1301+
// ALTER TABLE: change AUTO_RANDOM shard bits.
1302+
{
1303+
version: "5.7.25-TiDB-v6.1.0",
1304+
changes: []schema.Change{
1305+
&schema.ModifyTable{
1306+
T: &schema.Table{
1307+
Name: "users",
1308+
Columns: []*schema.Column{
1309+
{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "bigint"}}, Attrs: []schema.Attr{&AutoRandom{ShardBits: 10}}},
1310+
},
1311+
},
1312+
Changes: []schema.Change{
1313+
&schema.ModifyColumn{
1314+
Change: schema.ChangeAttr,
1315+
From: &schema.Column{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "bigint"}}, Attrs: []schema.Attr{&AutoRandom{ShardBits: 5}}},
1316+
To: &schema.Column{Name: "id", Type: &schema.ColumnType{Type: &schema.IntegerType{T: "bigint"}}, Attrs: []schema.Attr{&AutoRandom{ShardBits: 10}}},
1317+
},
1318+
},
1319+
},
1320+
},
1321+
wantPlan: &migrate.Plan{
1322+
Reversible: true,
1323+
Transactional: false,
1324+
Changes: []*migrate.Change{
1325+
{
1326+
Cmd: "ALTER TABLE `users` MODIFY COLUMN `id` bigint NOT NULL AUTO_RANDOM(10)",
1327+
Reverse: "ALTER TABLE `users` MODIFY COLUMN `id` bigint NOT NULL AUTO_RANDOM(5)",
1328+
},
1329+
},
1330+
},
1331+
},
11491332
}
11501333
for i, tt := range tests {
11511334
t.Run(strconv.Itoa(i), func(t *testing.T) {

0 commit comments

Comments
 (0)