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

dialect/sql/builder.go

@@ -486,10 +486,11 @@ func (d *DropIndexBuilder) Query() (string, []interface{}) {
 // InsertBuilder is a builder for `INSERT INTO` statement.
 type InsertBuilder struct {
 	Builder
-	table    string
-	columns  []string
-	defaults string
-	values   [][]interface{}
+	table     string
+	columns   []string
+	defaults  string
+	returning []string
+	values    [][]interface{}
 }
 
 // Insert creates a builder for the `INSERT INTO` statement.
@@ -536,23 +537,35 @@ func (i *InsertBuilder) Default() *InsertBuilder {
 	return i
 }
 
+// Returning adds the `RETURNING` clause to the insert statement. PostgreSQL only.
+func (i *InsertBuilder) Returning(columns ...string) *InsertBuilder {
+	i.returning = append(i.returning, columns...)
+	return i
+}
+
 // Query returns query representation of an `INSERT INTO` statement.
 func (i *InsertBuilder) Query() (string, []interface{}) {
 	i.WriteString("INSERT INTO ")
+	i.Ident(i.table).Pad()
 	if i.defaults != "" && len(i.columns) == 0 {
-		return i.Ident(i.table).Pad().String() + i.defaults, nil
-	}
-	i.Ident(i.table).Pad().Nested(func(b *Builder) {
-		b.IdentComma(i.columns...)
-	})
-	i.WriteString(" VALUES ")
-	for j, v := range i.values {
-		if j > 0 {
-			i.Comma()
-		}
+		i.WriteString(i.defaults)
+	} else {
 		i.Nested(func(b *Builder) {
-			b.Args(v...)
+			b.IdentComma(i.columns...)
 		})
+		i.WriteString(" VALUES ")
+		for j, v := range i.values {
+			if j > 0 {
+				i.Comma()
+			}
+			i.Nested(func(b *Builder) {
+				b.Args(v...)
+			})
+		}
+	}
+	if len(i.returning) > 0 && i.postgres() {
+		i.WriteString(" RETURNING ")
+		i.IdentComma(i.returning...)
 	}
 	return i.String(), i.args
 }
@@ -1457,7 +1470,7 @@ func (s *Selector) Having(p *Predicate) *Selector {
 }
 
 // Query returns query representation of a `SELECT` statement.
-func (s Selector) Query() (string, []interface{}) {
+func (s *Selector) Query() (string, []interface{}) {
 	b := s.Builder.clone()
 	b.WriteString("SELECT ")
 	if s.distinct {
@@ -1475,13 +1488,11 @@ func (s Selector) Query() (string, []interface{}) {
 		b.WriteString(t.ref())
 	case *Selector:
 		t.SetDialect(s.dialect)
-		query, args := t.Query()
 		b.Nested(func(b *Builder) {
-			b.WriteString(query)
+			b.Join(t)
 		})
 		b.WriteString(" AS ")
 		b.Ident(t.as)
-		b.args = append(b.args, args...)
 	}
 	for _, join := range s.joins {
 		b.WriteString(" " + join.kind + " ")
@@ -1491,13 +1502,11 @@ func (s Selector) Query() (string, []interface{}) {
 			b.WriteString(view.ref())
 		case *Selector:
 			view.SetDialect(s.dialect)
-			query, args := view.Query()
 			b.Nested(func(b *Builder) {
-				b.WriteString(query)
+				b.Join(view)
 			})
 			b.WriteString(" AS ")
 			b.Ident(view.as)
-			b.args = append(b.args, args...)
 		}
 		if join.on != "" {
 			b.WriteString(" ON ")
@@ -1528,6 +1537,7 @@ func (s Selector) Query() (string, []interface{}) {
 		b.WriteString(" OFFSET ")
 		b.Arg(*s.offset)
 	}
+	s.total = b.total
 	return b.String(), b.args
 }
 
@@ -1817,7 +1827,7 @@ func (b Builder) Query() (string, []interface{}) {
 
 // clone returns a shallow clone of a builder.
 func (b Builder) clone() Builder {
-	c := Builder{dialect: b.dialect}
+	c := Builder{dialect: b.dialect, total: b.total}
 	if len(b.args) > 0 {
 		c.args = append(c.args, b.args...)
 	}

dialect/sql/builder_test.go

@@ -9,7 +9,6 @@ import (
 	"testing"
 
 	"github.com/facebookincubator/ent/dialect"
-
 	"github.com/stretchr/testify/require"
 )
 
@@ -196,6 +195,11 @@ func TestBuilder(t *testing.T) {
 			wantQuery: `INSERT INTO "users" ("age") VALUES ($1)`,
 			wantArgs:  []interface{}{1},
 		},
+		{
+			input:     Dialect(dialect.Postgres).Insert("users").Columns("age").Values(1).Returning("id"),
+			wantQuery: `INSERT INTO "users" ("age") VALUES ($1) RETURNING "id"`,
+			wantArgs:  []interface{}{1},
+		},
 		{
 			input:     Insert("users").Columns("name", "age").Values("a8m", 10),
 			wantQuery: "INSERT INTO `users` (`name`, `age`) VALUES (?, ?)",
@@ -1004,6 +1008,25 @@ func TestBuilder(t *testing.T) {
 			wantQuery: "WITH groups AS (SELECT * FROM `groups` WHERE `name` = ?) SELECT `age` FROM `groups`",
 			wantArgs:  []interface{}{"bar"},
 		},
+		{
+			input: func() Querier {
+				builder := Dialect(dialect.Postgres)
+				t1 := builder.Table("groups")
+				t2 := builder.Table("users")
+				t3 := builder.Table("user_groups")
+				t4 := builder.Select(t3.C("id")).
+					From(t3).
+					Join(t2).
+					On(t3.C("id"), t2.C("id2")).
+					Where(EQ(t2.C("id"), "baz"))
+				return builder.Select().
+					From(t1).
+					Join(t4).
+					On(t1.C("id"), t4.C("id")).Limit(1)
+			}(),
+			wantQuery: `SELECT * FROM "groups" JOIN (SELECT "user_groups"."id" FROM "user_groups" JOIN "users" AS "t0" ON "user_groups"."id" = "t0"."id2" WHERE "t0"."id" = $1) AS "t1" ON "groups"."id" = "t1"."id" LIMIT $2`,
+			wantArgs:  []interface{}{"baz", 1},
+		},
 		{
 			input:     CreateIndex("name_index").Table("users").Column("name"),
 			wantQuery: "CREATE INDEX `name_index` ON `users`(`name`)",