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/o2obidi/ent/client.go

@@ -164,7 +164,8 @@ func (c *UserClient) GetX(ctx context.Context, id int) *User {
 func (c *UserClient) QuerySpouse(u *User) *UserQuery {
 	query := &UserQuery{config: c.config}
 	id := u.ID
-	query.sql = sql.Select().From(sql.Table(user.Table)).
+	builder := sql.Dialect(u.driver.Dialect())
+	query.sql = builder.Select().From(builder.Table(user.Table)).
 		Where(sql.EQ(user.SpouseColumn, id))
 
 	return query

examples/o2obidi/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/o2obidi/ent/user/where.go

@@ -364,11 +364,12 @@ func HasSpouse() 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(SpouseColumn).
-						From(sql.Table(SpouseTable)).
+					builder.Select(SpouseColumn).
+						From(builder.Table(SpouseTable)).
 						Where(sql.NotNull(SpouseColumn)),
 				),
 			)
@@ -380,8 +381,9 @@ func HasSpouse() predicate.User {
 func HasSpouseWith(preds ...predicate.User) predicate.User {
 	return predicate.User(
 		func(s *sql.Selector) {
+			builder := sql.Dialect(s.Dialect())
 			t1 := s.Table()
-			t2 := sql.Select(SpouseColumn).From(sql.Table(SpouseTable))
+			t2 := builder.Select(SpouseColumn).From(builder.Table(SpouseTable))
 			for _, p := range preds {
 				p(t2)
 			}

examples/o2obidi/ent/user_create.go

@@ -82,42 +82,37 @@ 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)
 	}
 	u.ID = int(id)
 	if len(uc.spouse) > 0 {
 		for eid := range uc.spouse {
-			query, args := sql.Update(user.SpouseTable).
+			query, args := builder.Update(user.SpouseTable).
 				Set(user.SpouseColumn, eid).
 				Where(sql.EQ(user.FieldID, id)).Query()
 			if err := tx.Exec(ctx, query, args, &res); err != nil {
 				return nil, rollback(tx, err)
 			}
-			query, args = sql.Update(user.SpouseTable).
+			query, args = builder.Update(user.SpouseTable).
 				Set(user.SpouseColumn, id).
 				Where(sql.EQ(user.FieldID, eid).And().IsNull(user.SpouseColumn)).Query()
 			if err := tx.Exec(ctx, query, args, &res); err != nil {

examples/o2obidi/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/o2obidi/ent/user_query.go

@@ -56,10 +56,12 @@ func (uq *UserQuery) Order(o ...Order) *UserQuery {
 // QuerySpouse chains the current query on the spouse edge.
 func (uq *UserQuery) QuerySpouse() *UserQuery {
 	query := &UserQuery{config: uq.config}
-	t1 := sql.Table(user.Table)
+
+	builder := sql.Dialect(uq.driver.Dialect())
+	t1 := builder.Table(user.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.SpouseColumn), t2.C(user.FieldID))
@@ -327,8 +329,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...)...)
@@ -597,5 +600,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/o2obidi/ent/user_update.go

@@ -115,7 +115,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)
 	}
@@ -143,32 +146,32 @@ 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)
 		}
 	}
 	if uu.clearedSpouse {
-		query, args := sql.Update(user.SpouseTable).
+		query, args := builder.Update(user.SpouseTable).
 			SetNull(user.SpouseColumn).
 			Where(sql.InInts(user.FieldID, ids...)).
 			Query()
 		if err := tx.Exec(ctx, query, args, &res); err != nil {
 			return 0, rollback(tx, err)
 		}
-		query, args = sql.Update(user.SpouseTable).
+		query, args = builder.Update(user.SpouseTable).
 			SetNull(user.SpouseColumn).
 			Where(sql.InInts(user.SpouseColumn, ids...)).
 			Query()
@@ -181,13 +184,13 @@ func (uu *UserUpdate) sqlSave(ctx context.Context) (n int, err error) {
 			return 0, rollback(tx, fmt.Errorf("ent: can't link O2O edge \"spouse\" to %d vertices (> 1)", n))
 		}
 		for eid := range uu.spouse {
-			query, args := sql.Update(user.SpouseTable).
+			query, args := builder.Update(user.SpouseTable).
 				Set(user.SpouseColumn, eid).
 				Where(sql.EQ(user.FieldID, ids[0])).Query()
 			if err := tx.Exec(ctx, query, args, &res); err != nil {
 				return 0, rollback(tx, err)
 			}
-			query, args = sql.Update(user.SpouseTable).
+			query, args = builder.Update(user.SpouseTable).
 				Set(user.SpouseColumn, ids[0]).
 				Where(sql.EQ(user.FieldID, eid).And().IsNull(user.SpouseColumn)).Query()
 			if err := tx.Exec(ctx, query, args, &res); err != nil {
@@ -301,7 +304,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()
@@ -332,35 +338,35 @@ 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)
 		}
 	}
 	if uuo.clearedSpouse {
-		query, args := sql.Update(user.SpouseTable).
+		query, args := builder.Update(user.SpouseTable).
 			SetNull(user.SpouseColumn).
 			Where(sql.InInts(user.FieldID, ids...)).
 			Query()
 		if err := tx.Exec(ctx, query, args, &res); err != nil {
 			return nil, rollback(tx, err)
 		}
-		query, args = sql.Update(user.SpouseTable).
+		query, args = builder.Update(user.SpouseTable).
 			SetNull(user.SpouseColumn).
 			Where(sql.InInts(user.SpouseColumn, ids...)).
 			Query()
@@ -373,13 +379,13 @@ func (uuo *UserUpdateOne) sqlSave(ctx context.Context) (u *User, err error) {
 			return nil, rollback(tx, fmt.Errorf("ent: can't link O2O edge \"spouse\" to %d vertices (> 1)", n))
 		}
 		for eid := range uuo.spouse {
-			query, args := sql.Update(user.SpouseTable).
+			query, args := builder.Update(user.SpouseTable).
 				Set(user.SpouseColumn, eid).
 				Where(sql.EQ(user.FieldID, ids[0])).Query()
 			if err := tx.Exec(ctx, query, args, &res); err != nil {
 				return nil, rollback(tx, err)
 			}
-			query, args = sql.Update(user.SpouseTable).
+			query, args = builder.Update(user.SpouseTable).
 				Set(user.SpouseColumn, ids[0]).
 				Where(sql.EQ(user.FieldID, eid).And().IsNull(user.SpouseColumn)).Query()
 			if err := tx.Exec(ctx, query, args, &res); err != nil {