golang/ent
1
0
Fork 0
mirror of https://github.com/ent/ent.git synced 2026-05-24 09:31:56 +03:00
Code Issues Packages Projects Releases Wiki Activity

entc/gen: support for upsert with client generated ids

Browse Source 
Fixed #1826
This commit is contained in:
Ariel Mashraki
2021-08-14 07:03:30 +03:00
committed by Ariel Mashraki
parent f25e0c17ea
commit b8532f87a6
55 changed files with 1471 additions and 387 deletions
Download Patch File Download Diff File Expand all files Collapse all files

132
dialect/sql/sqlgraph/graph.go
View File

@@ -982,15 +982,14 @@ func (c *creator) insert(ctx context.Context, tx dialect.ExecQuerier, insert *sq
// If the id field was provided by the user.
if c.ID.Value != nil {
insert.Set(c.ID.Column, c.ID.Value)
query, args := insert.Query()
return tx.Exec(ctx, query, args, &res)
// In case of "ON CONFLICT", the record may exists in the
// database, and we need to get back the database id field.
if len(c.CreateSpec.OnConflict) == 0 {
query, args := insert.Query()
return tx.Exec(ctx, query, args, &res)
}
}
id, err := insertLastID(ctx, tx, insert.Returning(c.ID.Column))
if err != nil {
return err
}
c.ID.Value = id
return nil
return c.insertLastID(ctx, tx, insert.Returning(c.ID.Column))
}
// ensureLastInsertID ensures the LAST_INSERT_ID was added to the
@@ -1014,18 +1013,7 @@ func (c *creator) batchInsert(ctx context.Context, tx dialect.ExecQuerier, inser
if opts := c.BatchCreateSpec.OnConflict; len(opts) > 0 {
insert.OnConflict(opts...)
}
ids, err := insertLastIDs(ctx, tx, insert.Returning(c.Nodes[0].ID.Column))
if err != nil {
return err
}
for i, node := range c.Nodes {
// If the ID field was not provided by the user,
// but was returned by the `RETURNING` clause.
if node.ID.Value == nil && i < len(ids) {
node.ID.Value = ids[i]
}
}
return nil
return c.insertLastIDs(ctx, tx, insert.Returning(c.Nodes[0].ID.Column))
}
// GroupRel groups edges by their relation type.
@@ -1280,63 +1268,99 @@ func setTableColumns(fields []*FieldSpec, edges map[Rel][]*EdgeSpec, set func(st
}
// insertLastID invokes the insert query on the transaction and returns the LastInsertID.
func insertLastID(ctx context.Context, tx dialect.ExecQuerier, insert *sql.InsertBuilder) (driver.Value, error) {
func (c *creator) insertLastID(ctx context.Context, tx dialect.ExecQuerier, insert *sql.InsertBuilder) error {
query, args := insert.Query()
// PostgreSQL does not support the LastInsertId() method of sql.Result
// on Exec, and should be extracted manually using the `RETURNING` clause.
if insert.Dialect() == dialect.Postgres {
if err := insert.Err(); err != nil {
return err
}
// MySQL does not support the "RETURNING" clause.
if insert.Dialect() != dialect.MySQL {
rows := &sql.Rows{}
if err := tx.Query(ctx, query, args, rows); err != nil {
return 0, err
return err
}
defer rows.Close()
return sql.ScanValue(rows)
if !c.ID.Type.Numeric() {
return sql.ScanOne(rows, &c.ID.Value)
}
// Normalize the type to int64 to make it looks
// like LastInsertId.
id, err := sql.ScanInt64(rows)
if err != nil {
return err
}
c.ID.Value = id
return nil
}
// MySQL, SQLite, etc.
// MySQL.
var res sql.Result
if err := tx.Exec(ctx, query, args, &res); err != nil {
return 0, err
return err
}
return res.LastInsertId()
// If the ID field is not numeric (e.g. string),
// there is no way to scan the LAST_INSERT_ID.
if c.ID.Type.Numeric() {
id, err := res.LastInsertId()
if err != nil {
return err
}
c.ID.Value = id
}
return nil
}
// insertLastIDs invokes the batch insert query on the transaction and returns the LastInsertID of all entities.
func insertLastIDs(ctx context.Context, tx dialect.ExecQuerier, insert *sql.InsertBuilder) (ids []driver.Value, err error) {
func (c *creator) insertLastIDs(ctx context.Context, tx dialect.ExecQuerier, insert *sql.InsertBuilder) error {
query, args := insert.Query()
// PostgreSQL does not support the LastInsertId() method of sql.Result
// on Exec, and should be extracted manually using the `RETURNING` clause.
if insert.Dialect() == dialect.Postgres {
if err := insert.Err(); err != nil {
return err
}
// MySQL does not support the "RETURNING" clause.
if insert.Dialect() != dialect.MySQL {
rows := &sql.Rows{}
if err := tx.Query(ctx, query, args, rows); err != nil {
return nil, err
return err
}
defer rows.Close()
return ids, sql.ScanSlice(rows, &ids)
for i := 0; rows.Next(); i++ {
node := c.Nodes[i]
if node.ID.Type.Numeric() {
// Normalize the type to int64 to make it looks
// like LastInsertId.
var id int64
if err := rows.Scan(&id); err != nil {
return err
}
node.ID.Value = id
} else if err := rows.Scan(&node.ID.Value); err != nil {
return err
}
}
return nil
}
// MySQL, SQLite, etc.
// MySQL.
var res sql.Result
if err := tx.Exec(ctx, query, args, &res); err != nil {
return nil, err
return err
}
id, err := res.LastInsertId()
if err != nil {
return nil, err
}
affected, err := res.RowsAffected()
if err != nil {
return nil, err
}
ids = make([]driver.Value, 0, affected)
switch insert.Dialect() {
case dialect.SQLite:
id -= affected - 1
fallthrough
case dialect.MySQL:
for i := int64(0); i < affected; i++ {
ids = append(ids, id+i)
// If the ID field is not numeric (e.g. string),
// there is no way to scan the LAST_INSERT_ID.
if len(c.Nodes) > 0 && c.Nodes[0].ID.Type.Numeric() {
id, err := res.LastInsertId()
if err != nil {
return err
}
affected, err := res.RowsAffected()
if err != nil {
return err
}
// Assume the ID field is AUTO_INCREMENT
// if its type is numeric.
for i := 0; int64(i) < affected && i < len(c.Nodes); i++ {
c.Nodes[i].ID.Value = id + int64(i)
}
}
return ids, nil
return nil
}
// rollback calls to tx.Rollback and wraps the given error with the rollback error if occurred.

56
dialect/sql/sqlgraph/graph_test.go
View File

@@ -844,7 +844,7 @@ func TestCreateNode(t *testing.T) {
name: "fields",
spec: &CreateSpec{
Table: "users",
ID: &FieldSpec{Column: "id"},
ID: &FieldSpec{Column: "id", Type: field.TypeInt},
Fields: []*FieldSpec{
{Column: "age", Type: field.TypeInt, Value: 30},
{Column: "name", Type: field.TypeString, Value: "a8m"},
@@ -862,7 +862,7 @@ func TestCreateNode(t *testing.T) {
name: "modifiers",
spec: &CreateSpec{
Table: "users",
ID: &FieldSpec{Column: "id"},
ID: &FieldSpec{Column: "id", Type: field.TypeInt},
Fields: []*FieldSpec{
{Column: "age", Type: field.TypeInt, Value: 30},
{Column: "name", Type: field.TypeString, Value: "a8m"},
@@ -873,7 +873,7 @@ func TestCreateNode(t *testing.T) {
},
expect: func(m sqlmock.Sqlmock) {
m.ExpectBegin()
m.ExpectExec(escape("INSERT INTO `users` (`age`, `name`) VALUES (?, ?) ON DUPLICATE KEY UPDATE `age` = VALUES(`age`), `name` = VALUES(`name`)")).
m.ExpectExec(escape("INSERT INTO `users` (`age`, `name`) VALUES (?, ?) ON DUPLICATE KEY UPDATE `age` = VALUES(`age`), `name` = VALUES(`name`), `id` = LAST_INSERT_ID(`users`.`id`)")).
WithArgs(30, "a8m").
WillReturnResult(sqlmock.NewResult(1, 1))
m.ExpectCommit()
@@ -901,7 +901,7 @@ func TestCreateNode(t *testing.T) {
name: "fields/json",
spec: &CreateSpec{
Table: "users",
ID: &FieldSpec{Column: "id"},
ID: &FieldSpec{Column: "id", Type: field.TypeInt},
Fields: []*FieldSpec{
{Column: "json", Type: field.TypeJSON, Value: struct{}{}},
},
@@ -918,7 +918,7 @@ func TestCreateNode(t *testing.T) {
name: "edges/m2o",
spec: &CreateSpec{
Table: "pets",
ID: &FieldSpec{Column: "id"},
ID: &FieldSpec{Column: "id", Type: field.TypeInt},
Fields: []*FieldSpec{
{Column: "name", Type: field.TypeString, Value: "pedro"},
},
@@ -938,7 +938,7 @@ func TestCreateNode(t *testing.T) {
name: "edges/o2o/inverse",
spec: &CreateSpec{
Table: "cards",
ID: &FieldSpec{Column: "id"},
ID: &FieldSpec{Column: "id", Type: field.TypeInt},
Fields: []*FieldSpec{
{Column: "number", Type: field.TypeString, Value: "0001"},
},
@@ -958,7 +958,7 @@ func TestCreateNode(t *testing.T) {
name: "edges/o2m",
spec: &CreateSpec{
Table: "users",
ID: &FieldSpec{Column: "id"},
ID: &FieldSpec{Column: "id", Type: field.TypeInt},
Fields: []*FieldSpec{
{Column: "name", Type: field.TypeString, Value: "a8m"},
},
@@ -981,7 +981,7 @@ func TestCreateNode(t *testing.T) {
name: "edges/o2m",
spec: &CreateSpec{
Table: "users",
ID: &FieldSpec{Column: "id"},
ID: &FieldSpec{Column: "id", Type: field.TypeInt},
Fields: []*FieldSpec{
{Column: "name", Type: field.TypeString, Value: "a8m"},
},
@@ -1004,7 +1004,7 @@ func TestCreateNode(t *testing.T) {
name: "edges/o2o",
spec: &CreateSpec{
Table: "users",
ID: &FieldSpec{Column: "id"},
ID: &FieldSpec{Column: "id", Type: field.TypeInt},
Fields: []*FieldSpec{
{Column: "name", Type: field.TypeString, Value: "a8m"},
},
@@ -1027,7 +1027,7 @@ func TestCreateNode(t *testing.T) {
name: "edges/o2o/bidi",
spec: &CreateSpec{
Table: "users",
ID: &FieldSpec{Column: "id"},
ID: &FieldSpec{Column: "id", Type: field.TypeInt},
Fields: []*FieldSpec{
{Column: "name", Type: field.TypeString, Value: "a8m"},
},
@@ -1050,7 +1050,7 @@ func TestCreateNode(t *testing.T) {
name: "edges/m2m",
spec: &CreateSpec{
Table: "groups",
ID: &FieldSpec{Column: "id"},
ID: &FieldSpec{Column: "id", Type: field.TypeInt},
Fields: []*FieldSpec{
{Column: "name", Type: field.TypeString, Value: "GitHub"},
},
@@ -1073,7 +1073,7 @@ func TestCreateNode(t *testing.T) {
name: "edges/m2m/inverse",
spec: &CreateSpec{
Table: "users",
ID: &FieldSpec{Column: "id"},
ID: &FieldSpec{Column: "id", Type: field.TypeInt},
Fields: []*FieldSpec{
{Column: "name", Type: field.TypeString, Value: "mashraki"},
},
@@ -1096,7 +1096,7 @@ func TestCreateNode(t *testing.T) {
name: "edges/m2m/bidi",
spec: &CreateSpec{
Table: "users",
ID: &FieldSpec{Column: "id"},
ID: &FieldSpec{Column: "id", Type: field.TypeInt},
Fields: []*FieldSpec{
{Column: "name", Type: field.TypeString, Value: "mashraki"},
},
@@ -1119,7 +1119,7 @@ func TestCreateNode(t *testing.T) {
name: "edges/m2m/bidi/batch",
spec: &CreateSpec{
Table: "users",
ID: &FieldSpec{Column: "id"},
ID: &FieldSpec{Column: "id", Type: field.TypeInt},
Fields: []*FieldSpec{
{Column: "name", Type: field.TypeString, Value: "mashraki"},
},
@@ -1149,7 +1149,7 @@ func TestCreateNode(t *testing.T) {
spec: &CreateSpec{
Table: "users",
Schema: "mydb",
ID: &FieldSpec{Column: "id"},
ID: &FieldSpec{Column: "id", Type: field.TypeInt},
Fields: []*FieldSpec{
{Column: "age", Type: field.TypeInt, Value: 30},
{Column: "name", Type: field.TypeString, Value: "a8m"},
@@ -1196,7 +1196,7 @@ func TestBatchCreate(t *testing.T) {
Nodes: []*CreateSpec{
{
Table: "users",
ID: &FieldSpec{Column: "id"},
ID: &FieldSpec{Column: "id", Type: field.TypeInt},
Fields: []*FieldSpec{
{Column: "age", Type: field.TypeInt, Value: 32},
{Column: "name", Type: field.TypeString, Value: "a8m"},
@@ -1205,7 +1205,7 @@ func TestBatchCreate(t *testing.T) {
},
{
Table: "users",
ID: &FieldSpec{Column: "id"},
ID: &FieldSpec{Column: "id", Type: field.TypeInt},
Fields: []*FieldSpec{
{Column: "age", Type: field.TypeInt, Value: 30},
{Column: "name", Type: field.TypeString, Value: "nati"},
@@ -1231,7 +1231,7 @@ func TestBatchCreate(t *testing.T) {
Nodes: []*CreateSpec{
{
Table: "users",
ID: &FieldSpec{Column: "id"},
ID: &FieldSpec{Column: "id", Type: field.TypeInt},
Fields: []*FieldSpec{
{Column: "age", Type: field.TypeInt, Value: 32},
{Column: "name", Type: field.TypeString, Value: "a8m"},
@@ -1240,14 +1240,14 @@ func TestBatchCreate(t *testing.T) {
Edges: []*EdgeSpec{
{Rel: M2M, Inverse: true, Table: "group_users", Columns: []string{"group_id", "user_id"}, Target: &EdgeTarget{Nodes: []driver.Value{2}, IDSpec: &FieldSpec{Column: "id"}}},
{Rel: M2M, Table: "user_products", Columns: []string{"user_id", "product_id"}, Target: &EdgeTarget{Nodes: []driver.Value{2}, IDSpec: &FieldSpec{Column: "id"}}},
{Rel: M2M, Table: "user_friends", Bidi: true, Columns: []string{"user_id", "friend_id"}, Target: &EdgeTarget{IDSpec: &FieldSpec{Column: "id"}, Nodes: []driver.Value{2}}},
{Rel: M2M, Table: "user_friends", Bidi: true, Columns: []string{"user_id", "friend_id"}, Target: &EdgeTarget{IDSpec: &FieldSpec{Column: "id", Type: field.TypeInt}, Nodes: []driver.Value{2}}},
{Rel: M2O, Table: "company", Columns: []string{"workplace_id"}, Target: &EdgeTarget{Nodes: []driver.Value{2}}},
{Rel: O2M, Table: "pets", Columns: []string{"owner_id"}, Target: &EdgeTarget{Nodes: []driver.Value{2}, IDSpec: &FieldSpec{Column: "id"}}},
},
},
{
Table: "users",
ID: &FieldSpec{Column: "id"},
ID: &FieldSpec{Column: "id", Type: field.TypeInt},
Fields: []*FieldSpec{
{Column: "age", Type: field.TypeInt, Value: 30},
{Column: "name", Type: field.TypeString, Value: "nati"},
@@ -1255,7 +1255,7 @@ func TestBatchCreate(t *testing.T) {
Edges: []*EdgeSpec{
{Rel: M2M, Inverse: true, Table: "group_users", Columns: []string{"group_id", "user_id"}, Target: &EdgeTarget{Nodes: []driver.Value{2}, IDSpec: &FieldSpec{Column: "id"}}},
{Rel: M2M, Table: "user_products", Columns: []string{"user_id", "product_id"}, Target: &EdgeTarget{Nodes: []driver.Value{2}, IDSpec: &FieldSpec{Column: "id"}}},
{Rel: M2M, Table: "user_friends", Bidi: true, Columns: []string{"user_id", "friend_id"}, Target: &EdgeTarget{IDSpec: &FieldSpec{Column: "id"}, Nodes: []driver.Value{2}}},
{Rel: M2M, Table: "user_friends", Bidi: true, Columns: []string{"user_id", "friend_id"}, Target: &EdgeTarget{IDSpec: &FieldSpec{Column: "id", Type: field.TypeInt}, Nodes: []driver.Value{2}}},
{Rel: O2M, Table: "pets", Columns: []string{"owner_id"}, Target: &EdgeTarget{Nodes: []driver.Value{3}, IDSpec: &FieldSpec{Column: "id"}}},
},
},
@@ -1461,10 +1461,10 @@ func TestUpdateNode(t *testing.T) {
Edges: EdgeMut{
Clear: []*EdgeSpec{
{Rel: O2O, Table: "users", Bidi: true, Columns: []string{"partner_id"}, Target: &EdgeTarget{IDSpec: &FieldSpec{Column: "id"}}},
{Rel: O2O, Table: "users", Bidi: true, Columns: []string{"spouse_id"}, Target: &EdgeTarget{IDSpec: &FieldSpec{Column: "id"}, Nodes: []driver.Value{2}}},
{Rel: O2O, Table: "users", Bidi: true, Columns: []string{"spouse_id"}, Target: &EdgeTarget{IDSpec: &FieldSpec{Column: "id", Type: field.TypeInt}, Nodes: []driver.Value{2}}},
},
Add: []*EdgeSpec{
{Rel: O2O, Table: "users", Bidi: true, Columns: []string{"spouse_id"}, Target: &EdgeTarget{IDSpec: &FieldSpec{Column: "id"}, Nodes: []driver.Value{3}}},
{Rel: O2O, Table: "users", Bidi: true, Columns: []string{"spouse_id"}, Target: &EdgeTarget{IDSpec: &FieldSpec{Column: "id", Type: field.TypeInt}, Nodes: []driver.Value{3}}},
},
},
},
@@ -1505,8 +1505,8 @@ func TestUpdateNode(t *testing.T) {
},
Edges: EdgeMut{
Clear: []*EdgeSpec{
{Rel: M2M, Table: "user_friends", Bidi: true, Columns: []string{"user_id", "friend_id"}, Target: &EdgeTarget{IDSpec: &FieldSpec{Column: "id"}, Nodes: []driver.Value{2}}},
{Rel: M2M, Inverse: true, Table: "group_users", Columns: []string{"group_id", "user_id"}, Target: &EdgeTarget{IDSpec: &FieldSpec{Column: "id"}, Nodes: []driver.Value{3, 7}}},
{Rel: M2M, Table: "user_friends", Bidi: true, Columns: []string{"user_id", "friend_id"}, Target: &EdgeTarget{IDSpec: &FieldSpec{Column: "id", Type: field.TypeInt}, Nodes: []driver.Value{2}}},
{Rel: M2M, Inverse: true, Table: "group_users", Columns: []string{"group_id", "user_id"}, Target: &EdgeTarget{IDSpec: &FieldSpec{Column: "id", Type: field.TypeInt}, Nodes: []driver.Value{3, 7}}},
// Clear all "following" edges (and their inverse).
{Rel: M2M, Table: "user_following", Bidi: true, Columns: []string{"following_id", "follower_id"}, Target: &EdgeTarget{IDSpec: &FieldSpec{Column: "id"}}},
// Clear all "user_blocked" edges.
@@ -1515,9 +1515,9 @@ func TestUpdateNode(t *testing.T) {
{Rel: M2M, Inverse: true, Table: "comment_responders", Columns: []string{"comment_id", "responder_id"}, Target: &EdgeTarget{IDSpec: &FieldSpec{Column: "id"}}},
},
Add: []*EdgeSpec{
{Rel: M2M, Table: "user_friends", Bidi: true, Columns: []string{"user_id", "friend_id"}, Target: &EdgeTarget{IDSpec: &FieldSpec{Column: "id"}, Nodes: []driver.Value{4}}},
{Rel: M2M, Inverse: true, Table: "group_users", Columns: []string{"group_id", "user_id"}, Target: &EdgeTarget{IDSpec: &FieldSpec{Column: "id"}, Nodes: []driver.Value{5}}},
{Rel: M2M, Inverse: true, Table: "group_users", Columns: []string{"group_id", "user_id"}, Target: &EdgeTarget{IDSpec: &FieldSpec{Column: "id"}, Nodes: []driver.Value{6, 8}}},
{Rel: M2M, Table: "user_friends", Bidi: true, Columns: []string{"user_id", "friend_id"}, Target: &EdgeTarget{IDSpec: &FieldSpec{Column: "id", Type: field.TypeInt}, Nodes: []driver.Value{4}}},
{Rel: M2M, Inverse: true, Table: "group_users", Columns: []string{"group_id", "user_id"}, Target: &EdgeTarget{IDSpec: &FieldSpec{Column: "id", Type: field.TypeInt}, Nodes: []driver.Value{5}}},
{Rel: M2M, Inverse: true, Table: "group_users", Columns: []string{"group_id", "user_id"}, Target: &EdgeTarget{IDSpec: &FieldSpec{Column: "id", Type: field.TypeInt}, Nodes: []driver.Value{6, 8}}},
},
},
},