dialect/sql/schema: univerdal id support for postgres

Summary: Pull Request resolved: https://github.com/facebookincubator/ent/pull/122

Reviewed By: alexsn

Differential Revision: D18113754

fbshipit-source-id: 1050f0c2b823db6ba1bc37c36f6749cfcffa982d
This commit is contained in:
Ariel Mashraki
2019-10-24 08:43:51 -07:00
committed by Facebook Github Bot
parent 7f598f34a2
commit d5c111b35a
4 changed files with 135 additions and 2 deletions

View File

@@ -335,7 +335,8 @@ func (m *Migrate) types(ctx context.Context, tx dialect.Tx) error {
return nil
}
rows := &sql.Rows{}
query, args := sql.Select("type").From(sql.Table(TypeTable)).OrderBy(sql.Asc("id")).Query()
query, args := sql.Dialect(m.Dialect()).
Select("type").From(sql.Table(TypeTable)).OrderBy(sql.Asc("id")).Query()
if err := tx.Query(ctx, query, args, rows); err != nil {
return fmt.Errorf("query types table: %v", err)
}
@@ -357,7 +358,8 @@ func (m *Migrate) allocPKRange(ctx context.Context, tx dialect.Tx, t *Table) err
if len(m.typeRanges) > MaxTypes {
return fmt.Errorf("max number of types exceeded: %d", MaxTypes)
}
query, args := sql.Insert(TypeTable).Columns("type").Values(t.Name).Query()
query, args := sql.Dialect(m.Dialect()).
Insert(TypeTable).Columns("type").Values(t.Name).Query()
if err := tx.Exec(ctx, query, args, new(sql.Result)); err != nil {
return fmt.Errorf("insert into type: %v", err)
}

View File

@@ -790,6 +790,7 @@ func TestMySQL_Create(t *testing.T) {
WithArgs("users").
WillReturnRows(sqlmock.NewRows([]string{"index_name", "column_name", "non_unique", "seq_in_index"}).
AddRow("PRIMARY", "id", "0", "1"))
// query groups table.
mock.ExpectQuery(escape("SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES WHERE `TABLE_SCHEMA` = (SELECT DATABASE()) AND `TABLE_NAME` = ?")).
WithArgs("groups").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))

View File

@@ -63,6 +63,9 @@ func (d *Postgres) fkExist(ctx context.Context, tx dialect.Tx, name string) (boo
// setRange sets restart the identity column to the given offset. Used by the universal-id option.
func (d *Postgres) setRange(ctx context.Context, tx dialect.Tx, name string, value int) error {
if value == 0 {
value = 1 // RESTART value cannot be < 1.
}
return tx.Exec(ctx, fmt.Sprintf("ALTER TABLE %s ALTER COLUMN id RESTART WITH %d", name, value), []interface{}{}, new(sql.Result))
}

View File

@@ -567,6 +567,133 @@ func TestPostgres_Create(t *testing.T) {
mock.ExpectCommit()
},
},
{
name: "universal id for all tables",
tables: []*Table{
NewTable("users").AddPrimary(&Column{Name: "id", Type: field.TypeInt, Increment: true}),
NewTable("groups").AddPrimary(&Column{Name: "id", Type: field.TypeInt, Increment: true}),
},
options: []MigrateOption{WithGlobalUniqueID(true)},
before: func(mock sqlmock.Sqlmock) {
mock.ExpectBegin()
mock.ExpectQuery(escape("SHOW server_version_num")).
WillReturnRows(sqlmock.NewRows([]string{"server_version_num"}).AddRow("120000"))
mock.ExpectQuery(escape(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES WHERE "table_schema" = CURRENT_SCHEMA() AND "table_name" = $1`)).
WithArgs("ent_types").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
// create ent_types table.
mock.ExpectExec(escape(`CREATE TABLE IF NOT EXISTS "ent_types"("id" bigint GENERATED BY DEFAULT AS IDENTITY NOT NULL, "type" varchar UNIQUE NOT NULL, PRIMARY KEY("id"))`)).
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectQuery(escape(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES WHERE "table_schema" = CURRENT_SCHEMA() AND "table_name" = $1`)).
WithArgs("users").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
mock.ExpectExec(escape(`CREATE TABLE IF NOT EXISTS "users"("id" bigint GENERATED BY DEFAULT AS IDENTITY NOT NULL, PRIMARY KEY("id"))`)).
WillReturnResult(sqlmock.NewResult(0, 1))
// set users id range.
mock.ExpectExec(escape(`INSERT INTO "ent_types" ("type") VALUES ($1)`)).
WithArgs("users").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec("ALTER TABLE users ALTER COLUMN id RESTART WITH 1").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectQuery(escape(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES WHERE "table_schema" = CURRENT_SCHEMA() AND "table_name" = $1`)).
WithArgs("groups").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
mock.ExpectExec(escape(`CREATE TABLE IF NOT EXISTS "groups"("id" bigint GENERATED BY DEFAULT AS IDENTITY NOT NULL, PRIMARY KEY("id"))`)).
WillReturnResult(sqlmock.NewResult(0, 1))
// set groups id range.
mock.ExpectExec(escape(`INSERT INTO "ent_types" ("type") VALUES ($1)`)).
WithArgs("groups").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec("ALTER TABLE groups ALTER COLUMN id RESTART WITH 4294967296").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectCommit()
},
},
{
name: "universal id for new tables",
tables: []*Table{
NewTable("users").AddPrimary(&Column{Name: "id", Type: field.TypeInt, Increment: true}),
NewTable("groups").AddPrimary(&Column{Name: "id", Type: field.TypeInt, Increment: true}),
},
options: []MigrateOption{WithGlobalUniqueID(true)},
before: func(mock sqlmock.Sqlmock) {
mock.ExpectBegin()
mock.ExpectQuery(escape("SHOW server_version_num")).
WillReturnRows(sqlmock.NewRows([]string{"server_version_num"}).AddRow("120000"))
mock.ExpectQuery(escape(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES WHERE "table_schema" = CURRENT_SCHEMA() AND "table_name" = $1`)).
WithArgs("ent_types").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1))
// query ent_types table.
mock.ExpectQuery(`SELECT "type" FROM "ent_types" ORDER BY "id" ASC`).
WillReturnRows(sqlmock.NewRows([]string{"type"}).AddRow("users"))
// query users table.
mock.ExpectQuery(escape(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES WHERE "table_schema" = CURRENT_SCHEMA() AND "table_name" = $1`)).
WithArgs("users").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1))
// users table has no changes.
mock.ExpectQuery(escape(`SELECT "column_name", "data_type", "is_nullable", "column_default" FROM INFORMATION_SCHEMA.COLUMNS WHERE "table_schema" = CURRENT_SCHEMA() AND "table_name" = $1`)).
WithArgs("users").
WillReturnRows(sqlmock.NewRows([]string{"column_name", "data_type", "is_nullable", "column_default"}).
AddRow("id", "bigint", "YES", "NULL"))
mock.ExpectQuery(escape(fmt.Sprintf(indexesQuery, "users"))).
WillReturnRows(sqlmock.NewRows([]string{"index_name", "column_name", "primary", "unique"}).
AddRow("users_pkey", "id", "t", "t"))
// query groups table.
mock.ExpectQuery(escape(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES WHERE "table_schema" = CURRENT_SCHEMA() AND "table_name" = $1`)).
WithArgs("groups").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
mock.ExpectExec(escape(`CREATE TABLE IF NOT EXISTS "groups"("id" bigint GENERATED BY DEFAULT AS IDENTITY NOT NULL, PRIMARY KEY("id"))`)).
WillReturnResult(sqlmock.NewResult(0, 1))
// set groups id range.
mock.ExpectExec(escape(`INSERT INTO "ent_types" ("type") VALUES ($1)`)).
WithArgs("groups").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec("ALTER TABLE groups ALTER COLUMN id RESTART WITH 4294967296").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectCommit()
},
},
{
name: "universal id for restored tables",
tables: []*Table{
NewTable("users").AddPrimary(&Column{Name: "id", Type: field.TypeInt, Increment: true}),
NewTable("groups").AddPrimary(&Column{Name: "id", Type: field.TypeInt, Increment: true}),
},
options: []MigrateOption{WithGlobalUniqueID(true)},
before: func(mock sqlmock.Sqlmock) {
mock.ExpectBegin()
mock.ExpectQuery(escape("SHOW server_version_num")).
WillReturnRows(sqlmock.NewRows([]string{"server_version_num"}).AddRow("120000"))
mock.ExpectQuery(escape(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES WHERE "table_schema" = CURRENT_SCHEMA() AND "table_name" = $1`)).
WithArgs("ent_types").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(1))
// query ent_types table.
mock.ExpectQuery(`SELECT "type" FROM "ent_types" ORDER BY "id" ASC`).
WillReturnRows(sqlmock.NewRows([]string{"type"}).AddRow("users"))
// query and create users (restored table).
mock.ExpectQuery(escape(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES WHERE "table_schema" = CURRENT_SCHEMA() AND "table_name" = $1`)).
WithArgs("users").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
mock.ExpectExec(escape(`CREATE TABLE IF NOT EXISTS "users"("id" bigint GENERATED BY DEFAULT AS IDENTITY NOT NULL, PRIMARY KEY("id"))`)).
WillReturnResult(sqlmock.NewResult(0, 1))
// set users id range (without inserting to ent_types).
mock.ExpectExec("ALTER TABLE users ALTER COLUMN id RESTART WITH 1").
WillReturnResult(sqlmock.NewResult(0, 1))
// query groups table.
mock.ExpectQuery(escape(`SELECT COUNT(*) FROM INFORMATION_SCHEMA.TABLES WHERE "table_schema" = CURRENT_SCHEMA() AND "table_name" = $1`)).
WithArgs("groups").
WillReturnRows(sqlmock.NewRows([]string{"count"}).AddRow(0))
mock.ExpectExec(escape(`CREATE TABLE IF NOT EXISTS "groups"("id" bigint GENERATED BY DEFAULT AS IDENTITY NOT NULL, PRIMARY KEY("id"))`)).
WillReturnResult(sqlmock.NewResult(0, 1))
// set groups id range.
mock.ExpectExec(escape(`INSERT INTO "ent_types" ("type") VALUES ($1)`)).
WithArgs("groups").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec("ALTER TABLE groups ALTER COLUMN id RESTART WITH 4294967296").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectCommit()
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {