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

entc/integration/json/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))

entc/integration/json/ent/user_create.go

@@ -79,22 +79,20 @@ 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}
+		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.url; value != nil {
 		buf, err := json.Marshal(*value)
 		if err != nil {
 			return nil, err
 		}
-		builder.Set(user.FieldURL, buf)
+		insert.Set(user.FieldURL, buf)
 		u.URL = *value
 	}
 	if value := uc.raw; value != nil {
@@ -102,7 +100,7 @@ func (uc *UserCreate) sqlSave(ctx context.Context) (*User, error) {
 		if err != nil {
 			return nil, err
 		}
-		builder.Set(user.FieldRaw, buf)
+		insert.Set(user.FieldRaw, buf)
 		u.Raw = *value
 	}
 	if value := uc.dirs; value != nil {
@@ -110,7 +108,7 @@ func (uc *UserCreate) sqlSave(ctx context.Context) (*User, error) {
 		if err != nil {
 			return nil, err
 		}
-		builder.Set(user.FieldDirs, buf)
+		insert.Set(user.FieldDirs, buf)
 		u.Dirs = *value
 	}
 	if value := uc.ints; value != nil {
@@ -118,7 +116,7 @@ func (uc *UserCreate) sqlSave(ctx context.Context) (*User, error) {
 		if err != nil {
 			return nil, err
 		}
-		builder.Set(user.FieldInts, buf)
+		insert.Set(user.FieldInts, buf)
 		u.Ints = *value
 	}
 	if value := uc.floats; value != nil {
@@ -126,7 +124,7 @@ func (uc *UserCreate) sqlSave(ctx context.Context) (*User, error) {
 		if err != nil {
 			return nil, err
 		}
-		builder.Set(user.FieldFloats, buf)
+		insert.Set(user.FieldFloats, buf)
 		u.Floats = *value
 	}
 	if value := uc.strings; value != nil {
@@ -134,14 +132,10 @@ func (uc *UserCreate) sqlSave(ctx context.Context) (*User, error) {
 		if err != nil {
 			return nil, err
 		}
-		builder.Set(user.FieldStrings, buf)
+		insert.Set(user.FieldStrings, buf)
 		u.Strings = *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)
 	}

entc/integration/json/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
 	}

entc/integration/json/ent/user_query.go

@@ -314,8 +314,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...)...)
@@ -584,5 +585,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))
 }

entc/integration/json/ent/user_update.go

@@ -148,7 +148,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)
 	}
@@ -176,70 +179,70 @@ 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.url; value != nil {
 		buf, err := json.Marshal(*value)
 		if err != nil {
 			return 0, err
 		}
-		builder.Set(user.FieldURL, buf)
+		updater.Set(user.FieldURL, buf)
 	}
 	if uu.clearurl {
-		builder.SetNull(user.FieldURL)
+		updater.SetNull(user.FieldURL)
 	}
 	if value := uu.raw; value != nil {
 		buf, err := json.Marshal(*value)
 		if err != nil {
 			return 0, err
 		}
-		builder.Set(user.FieldRaw, buf)
+		updater.Set(user.FieldRaw, buf)
 	}
 	if uu.clearraw {
-		builder.SetNull(user.FieldRaw)
+		updater.SetNull(user.FieldRaw)
 	}
 	if value := uu.dirs; value != nil {
 		buf, err := json.Marshal(*value)
 		if err != nil {
 			return 0, err
 		}
-		builder.Set(user.FieldDirs, buf)
+		updater.Set(user.FieldDirs, buf)
 	}
 	if uu.cleardirs {
-		builder.SetNull(user.FieldDirs)
+		updater.SetNull(user.FieldDirs)
 	}
 	if value := uu.ints; value != nil {
 		buf, err := json.Marshal(*value)
 		if err != nil {
 			return 0, err
 		}
-		builder.Set(user.FieldInts, buf)
+		updater.Set(user.FieldInts, buf)
 	}
 	if uu.clearints {
-		builder.SetNull(user.FieldInts)
+		updater.SetNull(user.FieldInts)
 	}
 	if value := uu.floats; value != nil {
 		buf, err := json.Marshal(*value)
 		if err != nil {
 			return 0, err
 		}
-		builder.Set(user.FieldFloats, buf)
+		updater.Set(user.FieldFloats, buf)
 	}
 	if uu.clearfloats {
-		builder.SetNull(user.FieldFloats)
+		updater.SetNull(user.FieldFloats)
 	}
 	if value := uu.strings; value != nil {
 		buf, err := json.Marshal(*value)
 		if err != nil {
 			return 0, err
 		}
-		builder.Set(user.FieldStrings, buf)
+		updater.Set(user.FieldStrings, buf)
 	}
 	if uu.clearstrings {
-		builder.SetNull(user.FieldStrings)
+		updater.SetNull(user.FieldStrings)
 	}
-	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)
 		}
@@ -374,7 +377,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()
@@ -405,88 +411,88 @@ 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.url; value != nil {
 		buf, err := json.Marshal(*value)
 		if err != nil {
 			return nil, err
 		}
-		builder.Set(user.FieldURL, buf)
+		updater.Set(user.FieldURL, buf)
 		u.URL = *value
 	}
 	if uuo.clearurl {
 		var value *url.URL
 		u.URL = value
-		builder.SetNull(user.FieldURL)
+		updater.SetNull(user.FieldURL)
 	}
 	if value := uuo.raw; value != nil {
 		buf, err := json.Marshal(*value)
 		if err != nil {
 			return nil, err
 		}
-		builder.Set(user.FieldRaw, buf)
+		updater.Set(user.FieldRaw, buf)
 		u.Raw = *value
 	}
 	if uuo.clearraw {
 		var value json.RawMessage
 		u.Raw = value
-		builder.SetNull(user.FieldRaw)
+		updater.SetNull(user.FieldRaw)
 	}
 	if value := uuo.dirs; value != nil {
 		buf, err := json.Marshal(*value)
 		if err != nil {
 			return nil, err
 		}
-		builder.Set(user.FieldDirs, buf)
+		updater.Set(user.FieldDirs, buf)
 		u.Dirs = *value
 	}
 	if uuo.cleardirs {
 		var value []http.Dir
 		u.Dirs = value
-		builder.SetNull(user.FieldDirs)
+		updater.SetNull(user.FieldDirs)
 	}
 	if value := uuo.ints; value != nil {
 		buf, err := json.Marshal(*value)
 		if err != nil {
 			return nil, err
 		}
-		builder.Set(user.FieldInts, buf)
+		updater.Set(user.FieldInts, buf)
 		u.Ints = *value
 	}
 	if uuo.clearints {
 		var value []int
 		u.Ints = value
-		builder.SetNull(user.FieldInts)
+		updater.SetNull(user.FieldInts)
 	}
 	if value := uuo.floats; value != nil {
 		buf, err := json.Marshal(*value)
 		if err != nil {
 			return nil, err
 		}
-		builder.Set(user.FieldFloats, buf)
+		updater.Set(user.FieldFloats, buf)
 		u.Floats = *value
 	}
 	if uuo.clearfloats {
 		var value []float64
 		u.Floats = value
-		builder.SetNull(user.FieldFloats)
+		updater.SetNull(user.FieldFloats)
 	}
 	if value := uuo.strings; value != nil {
 		buf, err := json.Marshal(*value)
 		if err != nil {
 			return nil, err
 		}
-		builder.Set(user.FieldStrings, buf)
+		updater.Set(user.FieldStrings, buf)
 		u.Strings = *value
 	}
 	if uuo.clearstrings {
 		var value []string
 		u.Strings = value
-		builder.SetNull(user.FieldStrings)
+		updater.SetNull(user.FieldStrings)
 	}
-	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)
 		}