entc/gen: generate sql builders with dialect option

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

Reviewed By: alexsn

Differential Revision: D18164397

fbshipit-source-id: 2858d69d3ff85c06b51382c01c3d4369ee2c3bdb

examples/o2m2types/ent/client.go

@@ -170,11 +170,12 @@ func (c *PetClient) GetX(ctx context.Context, id int) *Pet {
 func (c *PetClient) QueryOwner(pe *Pet) *UserQuery {
 	query := &UserQuery{config: c.config}
 	id := pe.ID
-	t1 := sql.Table(user.Table)
-	t2 := sql.Select(pet.OwnerColumn).
-		From(sql.Table(pet.OwnerTable)).
+	builder := sql.Dialect(pe.driver.Dialect())
+	t1 := builder.Table(user.Table)
+	t2 := builder.Select(pet.OwnerColumn).
+		From(builder.Table(pet.OwnerTable)).
 		Where(sql.EQ(pet.FieldID, id))
-	query.sql = sql.Select().From(t1).Join(t2).On(t1.C(user.FieldID), t2.C(pet.OwnerColumn))
+	query.sql = builder.Select().From(t1).Join(t2).On(t1.C(user.FieldID), t2.C(pet.OwnerColumn))
 
 	return query
 }
@@ -247,7 +248,8 @@ func (c *UserClient) GetX(ctx context.Context, id int) *User {
 func (c *UserClient) QueryPets(u *User) *PetQuery {
 	query := &PetQuery{config: c.config}
 	id := u.ID
-	query.sql = sql.Select().From(sql.Table(pet.Table)).
+	builder := sql.Dialect(u.driver.Dialect())
+	query.sql = builder.Select().From(builder.Table(pet.Table)).
 		Where(sql.EQ(user.PetsColumn, id))
 
 	return query

examples/o2m2types/ent/ent.go

@@ -7,6 +7,7 @@
 package ent
 
 import (
+	"context"
 	"fmt"
 	"strings"
 
@@ -168,9 +169,19 @@ func IsConstraintFailure(err error) bool {
 }
 
 func isSQLConstraintError(err error) (*ErrConstraintFailed, bool) {
-	// Error number 1062 is ER_DUP_ENTRY in mysql, and "UNIQUE constraint failed" is SQLite prefix.
-	if msg := err.Error(); strings.HasPrefix(msg, "Error 1062") || strings.HasPrefix(msg, "UNIQUE constraint failed") {
-		return &ErrConstraintFailed{msg, err}, true
+	var (
+		msg = err.Error()
+		// error format per dialect.
+		errors = [...]string{
+			"Error 1062",               // MySQL 1062 error (ER_DUP_ENTRY).
+			"UNIQUE constraint failed", // SQLite.
+			"duplicate key value violates unique constraint", // PostgreSQL.
+		}
+	)
+	for i := range errors {
+		if strings.Contains(msg, errors[i]) {
+			return &ErrConstraintFailed{msg, err}, true
+		}
 	}
 	return nil, false
 }
@@ -186,6 +197,38 @@ func rollback(tx dialect.Tx, err error) error {
 	return err
 }
 
+// insertLastID invokes the insert query on the transaction and returns the LastInsertID.
+func insertLastID(ctx context.Context, tx dialect.Tx, insert *sql.InsertBuilder) (int64, 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 {
+		rows := &sql.Rows{}
+		if err := tx.Query(ctx, query, args, rows); err != nil {
+			return 0, err
+		}
+		defer rows.Close()
+		if !rows.Next() {
+			return 0, fmt.Errorf("no rows found for query: %v", query)
+		}
+		var id int64
+		if err := rows.Scan(&id); err != nil {
+			return 0, err
+		}
+		return id, nil
+	}
+	// MySQL, SQLite, etc.
+	var res sql.Result
+	if err := tx.Exec(ctx, query, args, &res); err != nil {
+		return 0, err
+	}
+	id, err := res.LastInsertId()
+	if err != nil {
+		return 0, err
+	}
+	return id, nil
+}
+
 // keys returns the keys/ids from the edge map.
 func keys(m map[int]struct{}) []int {
 	s := make([]int, 0, len(m))

examples/o2m2types/ent/pet/where.go

@@ -272,8 +272,9 @@ func HasOwner() predicate.Pet {
 func HasOwnerWith(preds ...predicate.User) predicate.Pet {
 	return predicate.Pet(
 		func(s *sql.Selector) {
+			builder := sql.Dialect(s.Dialect())
 			t1 := s.Table()
-			t2 := sql.Select(FieldID).From(sql.Table(OwnerInverseTable))
+			t2 := builder.Select(FieldID).From(builder.Table(OwnerInverseTable))
 			for _, p := range preds {
 				p(t2)
 			}

examples/o2m2types/ent/pet_create.go

@@ -71,32 +71,27 @@ func (pc *PetCreate) SaveX(ctx context.Context) *Pet {
 
 func (pc *PetCreate) sqlSave(ctx context.Context) (*Pet, error) {
 	var (
-		res sql.Result
-		pe  = &Pet{config: pc.config}
+		res     sql.Result
+		builder = sql.Dialect(pc.driver.Dialect())
+		pe      = &Pet{config: pc.config}
 	)
 	tx, err := pc.driver.Tx(ctx)
 	if err != nil {
 		return nil, err
 	}
-	builder := sql.Dialect(pc.driver.Dialect()).
-		Insert(pet.Table).
-		Default()
+	insert := builder.Insert(pet.Table).Default()
 	if value := pc.name; value != nil {
-		builder.Set(pet.FieldName, *value)
+		insert.Set(pet.FieldName, *value)
 		pe.Name = *value
 	}
-	query, args := builder.Query()
-	if err := tx.Exec(ctx, query, args, &res); err != nil {
-		return nil, rollback(tx, err)
-	}
-	id, err := res.LastInsertId()
+	id, err := insertLastID(ctx, tx, insert.Returning(pet.FieldID))
 	if err != nil {
 		return nil, rollback(tx, err)
 	}
 	pe.ID = int(id)
 	if len(pc.owner) > 0 {
 		for eid := range pc.owner {
-			query, args := sql.Update(pet.OwnerTable).
+			query, args := builder.Update(pet.OwnerTable).
 				Set(pet.OwnerColumn, eid).
 				Where(sql.EQ(pet.FieldID, id)).
 				Query()

examples/o2m2types/ent/pet_delete.go

@@ -41,12 +41,15 @@ func (pd *PetDelete) ExecX(ctx context.Context) int {
 }
 
 func (pd *PetDelete) sqlExec(ctx context.Context) (int, error) {
-	var res sql.Result
-	selector := sql.Select().From(sql.Table(pet.Table))
+	var (
+		res     sql.Result
+		builder = sql.Dialect(pd.driver.Dialect())
+	)
+	selector := builder.Select().From(sql.Table(pet.Table))
 	for _, p := range pd.predicates {
 		p(selector)
 	}
-	query, args := sql.Delete(pet.Table).FromSelect(selector).Query()
+	query, args := builder.Delete(pet.Table).FromSelect(selector).Query()
 	if err := pd.driver.Exec(ctx, query, args, &res); err != nil {
 		return 0, err
 	}

examples/o2m2types/ent/pet_query.go

@@ -57,10 +57,12 @@ func (pq *PetQuery) Order(o ...Order) *PetQuery {
 // QueryOwner chains the current query on the owner edge.
 func (pq *PetQuery) QueryOwner() *UserQuery {
 	query := &UserQuery{config: pq.config}
-	t1 := sql.Table(user.Table)
+
+	builder := sql.Dialect(pq.driver.Dialect())
+	t1 := builder.Table(user.Table)
 	t2 := pq.sqlQuery()
 	t2.Select(t2.C(pet.OwnerColumn))
-	query.sql = sql.Select(t1.Columns(user.Columns...)...).
+	query.sql = builder.Select(t1.Columns(user.Columns...)...).
 		From(t1).
 		Join(t2).
 		On(t1.C(user.FieldID), t2.C(pet.OwnerColumn))
@@ -328,8 +330,9 @@ func (pq *PetQuery) sqlExist(ctx context.Context) (bool, error) {
 }
 
 func (pq *PetQuery) sqlQuery() *sql.Selector {
-	t1 := sql.Table(pet.Table)
-	selector := sql.Select(t1.Columns(pet.Columns...)...).From(t1)
+	builder := sql.Dialect(pq.driver.Dialect())
+	t1 := builder.Table(pet.Table)
+	selector := builder.Select(t1.Columns(pet.Columns...)...).From(t1)
 	if pq.sql != nil {
 		selector = pq.sql
 		selector.Select(selector.Columns(pet.Columns...)...)
@@ -598,5 +601,6 @@ func (ps *PetSelect) sqlScan(ctx context.Context, v interface{}) error {
 
 func (ps *PetSelect) sqlQuery() sql.Querier {
 	view := "pet_view"
-	return sql.Select(ps.fields...).From(ps.sql.As(view))
+	return sql.Dialect(ps.driver.Dialect()).
+		Select(ps.fields...).From(ps.sql.As(view))
 }

examples/o2m2types/ent/pet_update.go

@@ -97,7 +97,10 @@ func (pu *PetUpdate) ExecX(ctx context.Context) {
 }
 
 func (pu *PetUpdate) sqlSave(ctx context.Context) (n int, err error) {
-	selector := sql.Select(pet.FieldID).From(sql.Table(pet.Table))
+	var (
+		builder  = sql.Dialect(pu.driver.Dialect())
+		selector = builder.Select(pet.FieldID).From(builder.Table(pet.Table))
+	)
 	for _, p := range pu.predicates {
 		p(selector)
 	}
@@ -125,19 +128,19 @@ func (pu *PetUpdate) sqlSave(ctx context.Context) (n int, err error) {
 	}
 	var (
 		res     sql.Result
-		builder = sql.Update(pet.Table).Where(sql.InInts(pet.FieldID, ids...))
+		updater = builder.Update(pet.Table).Where(sql.InInts(pet.FieldID, ids...))
 	)
 	if value := pu.name; value != nil {
-		builder.Set(pet.FieldName, *value)
+		updater.Set(pet.FieldName, *value)
 	}
-	if !builder.Empty() {
-		query, args := builder.Query()
+	if !updater.Empty() {
+		query, args := updater.Query()
 		if err := tx.Exec(ctx, query, args, &res); err != nil {
 			return 0, rollback(tx, err)
 		}
 	}
 	if pu.clearedOwner {
-		query, args := sql.Update(pet.OwnerTable).
+		query, args := builder.Update(pet.OwnerTable).
 			SetNull(pet.OwnerColumn).
 			Where(sql.InInts(user.FieldID, ids...)).
 			Query()
@@ -147,7 +150,7 @@ func (pu *PetUpdate) sqlSave(ctx context.Context) (n int, err error) {
 	}
 	if len(pu.owner) > 0 {
 		for eid := range pu.owner {
-			query, args := sql.Update(pet.OwnerTable).
+			query, args := builder.Update(pet.OwnerTable).
 				Set(pet.OwnerColumn, eid).
 				Where(sql.InInts(pet.FieldID, ids...)).
 				Query()
@@ -236,7 +239,10 @@ func (puo *PetUpdateOne) ExecX(ctx context.Context) {
 }
 
 func (puo *PetUpdateOne) sqlSave(ctx context.Context) (pe *Pet, err error) {
-	selector := sql.Select(pet.Columns...).From(sql.Table(pet.Table))
+	var (
+		builder  = sql.Dialect(puo.driver.Dialect())
+		selector = builder.Select(pet.Columns...).From(builder.Table(pet.Table))
+	)
 	pet.ID(puo.id)(selector)
 	rows := &sql.Rows{}
 	query, args := selector.Query()
@@ -267,20 +273,20 @@ func (puo *PetUpdateOne) sqlSave(ctx context.Context) (pe *Pet, err error) {
 	}
 	var (
 		res     sql.Result
-		builder = sql.Update(pet.Table).Where(sql.InInts(pet.FieldID, ids...))
+		updater = builder.Update(pet.Table).Where(sql.InInts(pet.FieldID, ids...))
 	)
 	if value := puo.name; value != nil {
-		builder.Set(pet.FieldName, *value)
+		updater.Set(pet.FieldName, *value)
 		pe.Name = *value
 	}
-	if !builder.Empty() {
-		query, args := builder.Query()
+	if !updater.Empty() {
+		query, args := updater.Query()
 		if err := tx.Exec(ctx, query, args, &res); err != nil {
 			return nil, rollback(tx, err)
 		}
 	}
 	if puo.clearedOwner {
-		query, args := sql.Update(pet.OwnerTable).
+		query, args := builder.Update(pet.OwnerTable).
 			SetNull(pet.OwnerColumn).
 			Where(sql.InInts(user.FieldID, ids...)).
 			Query()
@@ -290,7 +296,7 @@ func (puo *PetUpdateOne) sqlSave(ctx context.Context) (pe *Pet, err error) {
 	}
 	if len(puo.owner) > 0 {
 		for eid := range puo.owner {
-			query, args := sql.Update(pet.OwnerTable).
+			query, args := builder.Update(pet.OwnerTable).
 				Set(pet.OwnerColumn, eid).
 				Where(sql.InInts(pet.FieldID, ids...)).
 				Query()

examples/o2m2types/ent/user/where.go

@@ -364,11 +364,12 @@ func HasPets() predicate.User {
 	return predicate.User(
 		func(s *sql.Selector) {
 			t1 := s.Table()
+			builder := sql.Dialect(s.Dialect())
 			s.Where(
 				sql.In(
 					t1.C(FieldID),
-					sql.Select(PetsColumn).
-						From(sql.Table(PetsTable)).
+					builder.Select(PetsColumn).
+						From(builder.Table(PetsTable)).
 						Where(sql.NotNull(PetsColumn)),
 				),
 			)
@@ -380,8 +381,9 @@ func HasPets() predicate.User {
 func HasPetsWith(preds ...predicate.Pet) predicate.User {
 	return predicate.User(
 		func(s *sql.Selector) {
+			builder := sql.Dialect(s.Dialect())
 			t1 := s.Table()
-			t2 := sql.Select(PetsColumn).From(sql.Table(PetsTable))
+			t2 := builder.Select(PetsColumn).From(builder.Table(PetsTable))
 			for _, p := range preds {
 				p(t2)
 			}

examples/o2m2types/ent/user_create.go

@@ -78,29 +78,24 @@ func (uc *UserCreate) SaveX(ctx context.Context) *User {
 
 func (uc *UserCreate) sqlSave(ctx context.Context) (*User, error) {
 	var (
-		res sql.Result
-		u   = &User{config: uc.config}
+		res     sql.Result
+		builder = sql.Dialect(uc.driver.Dialect())
+		u       = &User{config: uc.config}
 	)
 	tx, err := uc.driver.Tx(ctx)
 	if err != nil {
 		return nil, err
 	}
-	builder := sql.Dialect(uc.driver.Dialect()).
-		Insert(user.Table).
-		Default()
+	insert := builder.Insert(user.Table).Default()
 	if value := uc.age; value != nil {
-		builder.Set(user.FieldAge, *value)
+		insert.Set(user.FieldAge, *value)
 		u.Age = *value
 	}
 	if value := uc.name; value != nil {
-		builder.Set(user.FieldName, *value)
+		insert.Set(user.FieldName, *value)
 		u.Name = *value
 	}
-	query, args := builder.Query()
-	if err := tx.Exec(ctx, query, args, &res); err != nil {
-		return nil, rollback(tx, err)
-	}
-	id, err := res.LastInsertId()
+	id, err := insertLastID(ctx, tx, insert.Returning(user.FieldID))
 	if err != nil {
 		return nil, rollback(tx, err)
 	}
@@ -110,7 +105,7 @@ func (uc *UserCreate) sqlSave(ctx context.Context) (*User, error) {
 		for eid := range uc.pets {
 			p.Or().EQ(pet.FieldID, eid)
 		}
-		query, args := sql.Update(user.PetsTable).
+		query, args := builder.Update(user.PetsTable).
 			Set(user.PetsColumn, id).
 			Where(sql.And(p, sql.IsNull(user.PetsColumn))).
 			Query()

examples/o2m2types/ent/user_delete.go

@@ -41,12 +41,15 @@ func (ud *UserDelete) ExecX(ctx context.Context) int {
 }
 
 func (ud *UserDelete) sqlExec(ctx context.Context) (int, error) {
-	var res sql.Result
-	selector := sql.Select().From(sql.Table(user.Table))
+	var (
+		res     sql.Result
+		builder = sql.Dialect(ud.driver.Dialect())
+	)
+	selector := builder.Select().From(sql.Table(user.Table))
 	for _, p := range ud.predicates {
 		p(selector)
 	}
-	query, args := sql.Delete(user.Table).FromSelect(selector).Query()
+	query, args := builder.Delete(user.Table).FromSelect(selector).Query()
 	if err := ud.driver.Exec(ctx, query, args, &res); err != nil {
 		return 0, err
 	}

examples/o2m2types/ent/user_query.go

@@ -57,10 +57,12 @@ func (uq *UserQuery) Order(o ...Order) *UserQuery {
 // QueryPets chains the current query on the pets edge.
 func (uq *UserQuery) QueryPets() *PetQuery {
 	query := &PetQuery{config: uq.config}
-	t1 := sql.Table(pet.Table)
+
+	builder := sql.Dialect(uq.driver.Dialect())
+	t1 := builder.Table(pet.Table)
 	t2 := uq.sqlQuery()
 	t2.Select(t2.C(user.FieldID))
-	query.sql = sql.Select().
+	query.sql = builder.Select().
 		From(t1).
 		Join(t2).
 		On(t1.C(user.PetsColumn), t2.C(user.FieldID))
@@ -328,8 +330,9 @@ func (uq *UserQuery) sqlExist(ctx context.Context) (bool, error) {
 }
 
 func (uq *UserQuery) sqlQuery() *sql.Selector {
-	t1 := sql.Table(user.Table)
-	selector := sql.Select(t1.Columns(user.Columns...)...).From(t1)
+	builder := sql.Dialect(uq.driver.Dialect())
+	t1 := builder.Table(user.Table)
+	selector := builder.Select(t1.Columns(user.Columns...)...).From(t1)
 	if uq.sql != nil {
 		selector = uq.sql
 		selector.Select(selector.Columns(user.Columns...)...)
@@ -598,5 +601,6 @@ func (us *UserSelect) sqlScan(ctx context.Context, v interface{}) error {
 
 func (us *UserSelect) sqlQuery() sql.Querier {
 	view := "user_view"
-	return sql.Select(us.fields...).From(us.sql.As(view))
+	return sql.Dialect(us.driver.Dialect()).
+		Select(us.fields...).From(us.sql.As(view))
 }

examples/o2m2types/ent/user_update.go

@@ -124,7 +124,10 @@ func (uu *UserUpdate) ExecX(ctx context.Context) {
 }
 
 func (uu *UserUpdate) sqlSave(ctx context.Context) (n int, err error) {
-	selector := sql.Select(user.FieldID).From(sql.Table(user.Table))
+	var (
+		builder  = sql.Dialect(uu.driver.Dialect())
+		selector = builder.Select(user.FieldID).From(builder.Table(user.Table))
+	)
 	for _, p := range uu.predicates {
 		p(selector)
 	}
@@ -152,19 +155,19 @@ func (uu *UserUpdate) sqlSave(ctx context.Context) (n int, err error) {
 	}
 	var (
 		res     sql.Result
-		builder = sql.Update(user.Table).Where(sql.InInts(user.FieldID, ids...))
+		updater = builder.Update(user.Table).Where(sql.InInts(user.FieldID, ids...))
 	)
 	if value := uu.age; value != nil {
-		builder.Set(user.FieldAge, *value)
+		updater.Set(user.FieldAge, *value)
 	}
 	if value := uu.addage; value != nil {
-		builder.Add(user.FieldAge, *value)
+		updater.Add(user.FieldAge, *value)
 	}
 	if value := uu.name; value != nil {
-		builder.Set(user.FieldName, *value)
+		updater.Set(user.FieldName, *value)
 	}
-	if !builder.Empty() {
-		query, args := builder.Query()
+	if !updater.Empty() {
+		query, args := updater.Query()
 		if err := tx.Exec(ctx, query, args, &res); err != nil {
 			return 0, rollback(tx, err)
 		}
@@ -174,7 +177,7 @@ func (uu *UserUpdate) sqlSave(ctx context.Context) (n int, err error) {
 		for eid := range uu.removedPets {
 			eids = append(eids, eid)
 		}
-		query, args := sql.Update(user.PetsTable).
+		query, args := builder.Update(user.PetsTable).
 			SetNull(user.PetsColumn).
 			Where(sql.InInts(user.PetsColumn, ids...)).
 			Where(sql.InInts(pet.FieldID, eids...)).
@@ -189,7 +192,7 @@ func (uu *UserUpdate) sqlSave(ctx context.Context) (n int, err error) {
 			for eid := range uu.pets {
 				p.Or().EQ(pet.FieldID, eid)
 			}
-			query, args := sql.Update(user.PetsTable).
+			query, args := builder.Update(user.PetsTable).
 				Set(user.PetsColumn, id).
 				Where(sql.And(p, sql.IsNull(user.PetsColumn))).
 				Query()
@@ -313,7 +316,10 @@ func (uuo *UserUpdateOne) ExecX(ctx context.Context) {
 }
 
 func (uuo *UserUpdateOne) sqlSave(ctx context.Context) (u *User, err error) {
-	selector := sql.Select(user.Columns...).From(sql.Table(user.Table))
+	var (
+		builder  = sql.Dialect(uuo.driver.Dialect())
+		selector = builder.Select(user.Columns...).From(builder.Table(user.Table))
+	)
 	user.ID(uuo.id)(selector)
 	rows := &sql.Rows{}
 	query, args := selector.Query()
@@ -344,22 +350,22 @@ func (uuo *UserUpdateOne) sqlSave(ctx context.Context) (u *User, err error) {
 	}
 	var (
 		res     sql.Result
-		builder = sql.Update(user.Table).Where(sql.InInts(user.FieldID, ids...))
+		updater = builder.Update(user.Table).Where(sql.InInts(user.FieldID, ids...))
 	)
 	if value := uuo.age; value != nil {
-		builder.Set(user.FieldAge, *value)
+		updater.Set(user.FieldAge, *value)
 		u.Age = *value
 	}
 	if value := uuo.addage; value != nil {
-		builder.Add(user.FieldAge, *value)
+		updater.Add(user.FieldAge, *value)
 		u.Age += *value
 	}
 	if value := uuo.name; value != nil {
-		builder.Set(user.FieldName, *value)
+		updater.Set(user.FieldName, *value)
 		u.Name = *value
 	}
-	if !builder.Empty() {
-		query, args := builder.Query()
+	if !updater.Empty() {
+		query, args := updater.Query()
 		if err := tx.Exec(ctx, query, args, &res); err != nil {
 			return nil, rollback(tx, err)
 		}
@@ -369,7 +375,7 @@ func (uuo *UserUpdateOne) sqlSave(ctx context.Context) (u *User, err error) {
 		for eid := range uuo.removedPets {
 			eids = append(eids, eid)
 		}
-		query, args := sql.Update(user.PetsTable).
+		query, args := builder.Update(user.PetsTable).
 			SetNull(user.PetsColumn).
 			Where(sql.InInts(user.PetsColumn, ids...)).
 			Where(sql.InInts(pet.FieldID, eids...)).
@@ -384,7 +390,7 @@ func (uuo *UserUpdateOne) sqlSave(ctx context.Context) (u *User, err error) {
 			for eid := range uuo.pets {
 				p.Or().EQ(pet.FieldID, eid)
 			}
-			query, args := sql.Update(user.PetsTable).
+			query, args := builder.Update(user.PetsTable).
 				Set(user.PetsColumn, id).
 				Where(sql.And(p, sql.IsNull(user.PetsColumn))).
 				Query()