ent/entc: configure storage driver in codegen

Summary: Pull Request resolved: https://github.com/facebookexternal/fbc/pull/1229

Reviewed By: alexsn

Differential Revision: D16539934

fbshipit-source-id: b3a8bf1f1be6f65ad3f649cd921ea20fc24182bf

entc/gen/template/dialect/sql/by.tmpl

@@ -0,0 +1,28 @@
+{{ define "dialect/sql/order" -}}
+	{{- $f := $.Scope.Func -}}
+	func(s *sql.Selector) {
+		for _, f := range fields {
+			s.OrderBy(sql.{{ $f }}(f))
+		}
+	}
+{{- end }}
+
+{{/* custom signature for group-by function */}}
+{{ define "dialect/sql/group/signature" -}}
+	// SQL the column wrapped with the aggregation function.
+	SQL     func(*sql.Selector) string
+{{- end }}
+
+{{ define "dialect/sql/group/as" -}}
+	func(s *sql.Selector) string {
+			return sql.As(fn.SQL(s), end)
+	}
+{{- end }}
+
+{{ define "dialect/sql/group/func" -}}
+	{{- $fn := $.Scope.Func }}
+	{{- $withField := $.Scope.WithField }}
+	func(s *sql.Selector) string {
+		return sql.{{ if eq $fn "Mean" }}Avg{{ else }}{{ $fn }}{{ end }}({{ if $withField }}s.C(field){{ else }}"*"{{ end }})
+	}
+{{- end }}

entc/gen/template/dialect/sql/decode.tmpl

@@ -0,0 +1,59 @@
+{{ define "dialect/sql/decode/one" }}
+{{ $receiver := $.Receiver }}
+
+// FromRows scans the sql response data into {{ $.Name }}.
+func ({{ $receiver }} *{{ $.Name }}) FromRows(rows *sql.Rows) error {
+	{{- $scan := print "v" $receiver }}
+	var {{ $scan }} struct {
+		ID   {{ if $.ID.IsString }}int{{ else }}{{ $.ID.Type }}{{ end }}
+		{{ range $_, $f := $.Fields }}
+			{{- pascal $f.Name }} {{ if or $f.Nullable $f.Optional }}{{ $f.NullType }}{{ else }}{{ $f.Type }}{{ end }}
+		{{ end }}
+	}
+	// the order here should be the same as in the `{{ $.Package }}.Columns`.
+	if err := rows.Scan(
+			&{{ $scan }}.ID,
+			{{- range $_, $f := $.Fields }}
+				&{{ $scan }}.{{- pascal $f.Name }},
+			{{- end }}
+	); err != nil {
+		return err
+	}
+	{{ $receiver }}.ID = {{ if $.ID.IsString }}strconv.Itoa({{ $scan }}.ID){{ else }}{{ $scan }}.ID{{ end }}
+	{{- range $_, $f := $.Fields }}
+		{{- if $f.Nullable }}
+			{{- if $f.IsTime }}
+				{{ $receiver }}.{{ pascal $f.Name }} = &{{ $scan }}.{{ pascal $f.Name }}
+			{{- else }}
+				if {{ $scan }}.{{- pascal $f.Name }}.Valid {
+					{{ $receiver }}.{{ pascal $f.Name }} = new({{ $f.Type }})
+					*{{ $receiver }}.{{ pascal $f.Name }} = {{ printf "%s.%s" $scan (pascal $f.Name) | $f.NullTypeField }}
+				}
+			{{- end }}
+		{{- else if $f.Optional }}
+			{{ $receiver }}.{{ pascal $f.Name }} = {{ printf "%s.%s" $scan (pascal $f.Name) | $f.NullTypeField }}
+		{{- else }}
+			{{ $receiver }}.{{ pascal $f.Name }} = {{ $scan }}.{{ pascal $f.Name }}
+		{{- end }}
+	{{- end }}
+	return nil
+}
+{{ end }}
+
+{{ define "dialect/sql/decode/many" }}
+{{ $receiver := $.Receiver }}
+{{ $slice := $.Scope.Slice }}
+
+// FromRows scans the sql response data into {{ $slice }}.
+func ({{ $receiver }} *{{ $slice }}) FromRows(rows *sql.Rows) error {
+	for rows.Next() {
+		{{- $scan := print "v" $receiver }}
+		{{ $scan }} := &{{ $.Name }}{}
+		if err := {{ $scan }}.FromRows(rows); err != nil {
+			return err
+		}
+		*{{ $receiver }} = append(*{{ $receiver }}, {{ $scan }})
+	}
+	return nil
+}
+{{ end }}

entc/gen/template/dialect/sql/errors.tmpl

@@ -0,0 +1,21 @@
+{{/* custom errors and errors handlers from sql dialects */}}
+{{ define "dialect/sql/errors" }}
+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
+	}
+	return nil, false
+}
+
+// rollback calls to tx.Rollback and wraps the given error with the rollback error if occurred.
+func rollback(tx dialect.Tx, err error) error {
+	if rerr := tx.Rollback(); rerr != nil {
+		err = fmt.Errorf("%s: %v", err.Error(), rerr)
+	}
+	if err, ok := isSQLConstraintError(err); ok {
+		return err
+	}
+	return err
+}
+{{ end }}

entc/gen/template/dialect/sql/meta.tmpl

@@ -0,0 +1,42 @@
+{{/* constants needed for sql dialects. */}}
+{{ define "dialect/sql/meta/constants" }}
+	// Table holds the table name of the {{ lower $.Name }} in the database.
+	Table = "{{ $.Table }}"
+	{{- range $_, $e := $.Edges }}
+		// {{ $e.TableConstant }} is the table the holds the {{ $e.Name }} relation/edge.
+		{{- if $e.M2M }} The primary key declared below.{{ end }}
+		{{ $e.TableConstant }} = "{{ $e.Rel.Table }}"
+		{{- if eq $.Table $e.Type.Table | not }}
+			// {{ $e.InverseTableConstant }} is the table name for the {{ $e.Type.Name }} entity.
+			// It exists in this package in order to avoid circular dependency with the "{{ $e.Type.Package }}" package.
+			{{ $e.InverseTableConstant }} = "{{ $e.Type.Table }}"
+		{{- end }}
+		{{- if not $e.M2M }}
+			// {{ $e.ColumnConstant }} is the table column denoting the {{ $e.Name }} relation/edge.
+			{{ $e.ColumnConstant }} = "{{ $e.Rel.Column }}"
+		{{- end }}
+	{{- end }}
+{{ end }}
+
+{{/* variables needed for sql dialects. */}}
+{{ define "dialect/sql/meta/variables" }}
+	// Columns holds all SQL columns are {{ lower $.Name }} fields.
+	var Columns = []string{
+		{{ $.ID.Constant }},
+		{{- range $_, $f := $.Fields }}
+			{{ $f.Constant }},
+		{{- end }}
+	}
+
+	{{ with $.NumM2M }}
+		var (
+			{{- range $_, $e := $.Edges }}
+				{{- if $e.M2M }}
+					// {{ $e.PKConstant }} and {{ $e.ColumnConstant }}2 are the table columns denoting the
+					// primary key for the {{ $e.Name }} relation (M2M).
+					{{ $e.PKConstant }} = []string{"{{ index $e.Rel.Columns 0 }}", "{{ index $e.Rel.Columns 1 }}"}
+				{{- end }}
+			{{- end }}
+		)
+	{{ end }}
+{{ end }}

entc/gen/template/dialect/sql/predicate.tmpl

@@ -0,0 +1,145 @@
+{{ define "dialect/sql/predicate/id" -}}
+	func(s *sql.Selector) {
+		{{- if $.ID.IsString }}id, _ := strconv.Atoi(id){{- end }}
+		s.Where(sql.EQ(s.C({{ $.ID.Constant }}), id))
+	}
+{{- end }}
+
+{{ define "dialect/sql/predicate/id/ops" -}}
+	{{- $op := $.Scope.Op -}}
+	{{- $arg := $.Scope.Arg -}}
+	func(s *sql.Selector) {
+		{{- if $op.Variadic }}
+			// if not arguments were provided, append the FALSE constants,
+			// since we can't apply "IN ()". This will make this predicate falsy.
+			if len({{ $arg }}) == 0 {
+				s.Where(sql.False())
+				return
+			}
+			v := make([]interface{}, len({{ $arg }}))
+			for i := range v {
+				{{ if $.ID.IsString }}v[i], _ = strconv.Atoi({{ $arg }}[i]){{ else }}v[i] = {{ $arg }}[i]{{ end }}
+			}
+		{{- else if $.ID.IsString }}
+			v, _ := strconv.Atoi({{ $arg }})
+		{{- end }}
+		s.Where(sql.{{ $op.Name }}(s.C({{ $.ID.Constant }}), v{{ if $op.Variadic }}...{{ end }}))
+	}
+{{- end }}
+
+{{ define "dialect/sql/predicate/field" -}}
+	{{- $f := $.Scope.Field -}}
+	func(s *sql.Selector) {
+		s.Where(sql.EQ(s.C({{ $f.Constant }}), v))
+	}
+{{- end }}
+
+{{ define "dialect/sql/predicate/field/ops" -}}
+	{{- $f := $.Scope.Field -}}
+	{{- $op := $.Scope.Op -}}
+	{{- $arg := $.Scope.Arg -}}
+	func(s *sql.Selector) {
+		{{- if $op.Variadic }}
+			// if not arguments were provided, append the FALSE constants,
+			// since we can't apply "IN ()". This will make this predicate falsy.
+			if len({{ $arg }}) == 0 {
+				s.Where(sql.False())
+				return
+			}
+		{{- end }}
+		s.Where(sql.{{ $op.Name }}(s.C({{ $f.Constant }}), v{{ if $op.Variadic }}...{{ end }}))
+	}
+{{- end }}
+
+{{ define "dialect/sql/predicate/edge/has" -}}
+	{{- $e := $.Scope.Edge -}}
+	func(s *sql.Selector) {
+		{{- if $e.M2M  }}
+			t1 := s.Table()
+			s.Where(
+				sql.In(
+					t1.C({{ $.ID.Constant }}),
+					sql.Select({{ $e.PKConstant }}[{{ if $e.IsInverse }}1{{ else }}0{{ end }}]).From(sql.Table({{ $e.TableConstant }})),
+				),
+			)
+		{{- else if or $e.M2O (and $e.O2O $e.IsInverse) }}{{/* M2O || (O2O with inverse edge) */}}
+			t1 := s.Table()
+			s.Where(sql.NotNull(t1.C({{ $e.ColumnConstant }})))
+		{{- else }}{{/* O2M || (O2O with assoc edge) */}}
+			t1 := s.Table()
+			s.Where(
+				sql.In(
+					t1.C({{ $.ID.Constant }}),
+					sql.Select({{ $e.ColumnConstant }}).
+						From(sql.Table({{ $e.TableConstant }})).
+						Where(sql.NotNull({{ $e.ColumnConstant }})),
+				),
+			)
+		{{- end }}
+	}
+{{- end }}
+
+{{ define "dialect/sql/predicate/edge/haswith" -}}
+	{{- $e := $.Scope.Edge -}}
+	func(s *sql.Selector) {
+		{{- if $e.M2M  }}
+			{{ $i := 1 }}{{ $j := 0 }}{{- if $e.IsInverse }}{{ $i = 0 }}{{ $j = 1 }}{{ end -}}
+			t1 := s.Table()
+			t2 := sql.Table(
+				{{- if ne $.Table $e.Type.Table -}}
+					{{ $e.InverseTableConstant }}
+				{{- else -}}
+					Table
+				{{- end -}}
+			)
+			t3 := sql.Table({{ $e.TableConstant }})
+			t4 := sql.Select(t3.C({{ $e.PKConstant }}[{{ $j }}])).
+				From(t3).
+				Join(t2).
+				On(t3.C({{ $e.PKConstant }}[{{ $i }}]), t2.C({{ $e.Type.ID.Constant }}))
+			t5 := sql.Select().From(t2)
+			for _, p := range preds {
+				p(t5)
+			}
+			t4.FromSelect(t5)
+			s.Where(sql.In(t1.C({{ $.ID.Constant }}), t4))
+		{{- else if or $e.M2O (and $e.O2O $e.IsInverse) }}{{/* M2O || (O2O with inverse edge) */}}
+			t1 := s.Table()
+			t2 := sql.Select({{ $e.Type.ID.Constant }}).From(sql.Table(
+				{{- if ne $.Table $e.Type.Table -}}
+					{{ $e.InverseTableConstant }}
+				{{- else -}}
+					{{ $e.TableConstant }}
+				{{- end -}}
+			))
+			for _, p := range preds {
+				p(t2)
+			}
+			s.Where(sql.In(t1.C({{ $e.ColumnConstant }}), t2))
+		{{- else }}{{/* O2M || (O2O with assoc edge) */}}
+			t1 := s.Table()
+			t2 := sql.Select({{ $e.ColumnConstant }}).From(sql.Table({{ $e.TableConstant }}))
+			for _, p := range preds {
+				p(t2)
+			}
+			s.Where(sql.In(t1.C({{ $.ID.Constant }}), t2))
+		{{- end }}
+	}
+{{- end }}
+
+{{ define "dialect/sql/predicate/or" -}}
+	func(s *sql.Selector) {
+		for i, p := range predicates {
+			if i > 0 {
+				s.Or()
+			}
+			p(s)
+		}
+	}
+{{- end }}
+
+{{ define "dialect/sql/predicate/not" -}}
+	func(s *sql.Selector) {
+		p(s.Not())
+	}
+{{- end }}

entc/gen/template/dialect/sql/query.tmpl

@@ -66,7 +66,6 @@ func ({{ $receiver }} *{{ $builder }}) sqlIDs(ctx context.Context) ([]{{ $.ID.Ty
 	return ids, nil
 }
 
-
 func ({{ $receiver }} *{{ $builder }}) sqlQuery() *sql.Selector {
 	t1 := sql.Table({{ $.Package }}.Table)
 	selector := sql.Select(t1.Columns({{ $.Package }}.Columns...)...).From(t1)
@@ -78,7 +77,7 @@ func ({{ $receiver }} *{{ $builder }}) sqlQuery() *sql.Selector {
 		p(selector)
 	}
 	for _, p := range {{ $receiver }}.order {
-		p.SQL(selector)
+		p(selector)
 	}
 	if offset := {{ $receiver }}.offset; offset != nil {
 		// limit is mandatory for offset clause. We start
@@ -90,5 +89,75 @@ func ({{ $receiver }} *{{ $builder }}) sqlQuery() *sql.Selector {
 	}
 	return selector
 }
+{{ end }}
 
+{{/* query/path defines the query generation for path of a given edge. */}}
+{{ define "dialect/sql/query/path" }}
+	{{- $e := $.Scope.Edge }} {{/* the edge we need to genegrate the path to. */}}
+	{{- $receiver := $.Scope.Receiver }}
+
+	{{- if $e.M2M }}
+		{{ $i := 1 }}{{ $j := 0 }}{{- if $e.IsInverse }}{{ $i = 0 }}{{ $j = 1 }}{{ end -}}
+		t1 := sql.Table({{ $e.Type.Package }}.Table)
+		t2 := {{ $receiver }}.sqlQuery()
+		t2.Select(t2.C({{ $.Package }}.{{ $.ID.Constant }}))
+		t3 := sql.Table({{ $.Package }}.{{ $e.TableConstant }})
+		t4 := sql.Select(t3.C({{ $.Package }}.{{ $e.PKConstant }}[{{ $i }}])).
+				From(t3).
+				Join(t2).
+				On(t3.C({{ $.Package }}.{{ $e.PKConstant }}[{{ $j }}]), t2.C({{ $.Package }}.{{ $.ID.Constant }}))
+		query.sql = sql.Select().
+			From(t1).
+			Join(t4).
+			On(t1.C({{ $e.Type.Package }}.{{ $e.Type.ID.Constant }}), t4.C({{ $.Package }}.{{ $e.PKConstant }}[{{ $i }}]))
+	{{- else if or $e.M2O (and $e.O2O $e.IsInverse) }}{{/* M2O || (O2O with inverse edge) */}}
+		t1 := sql.Table({{ $e.Type.Package }}.Table)
+		t2 := {{ $receiver }}.sqlQuery()
+		t2.Select(t2.C({{ $.Package }}.{{ $e.ColumnConstant }}))
+		query.sql = sql.Select(t1.Columns({{ $e.Type.Package }}.Columns...)...).
+			From(t1).
+			Join(t2).
+			On(t1.C({{ $e.Type.Package }}.{{ $e.Type.ID.Constant }}), t2.C({{ $.Package }}.{{ $e.ColumnConstant }}))
+	{{- else }}{{/* O2M || (O2O with assoc edge) */}}
+		t1 := sql.Table({{ $e.Type.Package }}.Table)
+		t2 := {{ $receiver }}.sqlQuery()
+		t2.Select(t2.C({{ $.Package }}.{{ $.ID.Constant }}))
+		query.sql = sql.Select().
+			From(t1).
+			Join(t2).
+			On(t1.C({{ $.Package }}.{{ $e.ColumnConstant }}), t2.C({{ $.Package }}.{{ $.ID.Constant }}))
+	{{- end }}
+{{ end }}
+
+{{/* query/from defines the query generation for an edge query from a given node. */}}
+{{ define "dialect/sql/query/from" }}
+	{{- $n := $ }} {{/* the node we start the query from. */}}
+	{{- $e := $.Scope.Edge }} {{/* the edge we need to genegrate the path to. */}}
+	{{- $receiver := $.Scope.Receiver -}}
+
+	id := {{ $receiver }}.{{- if $n.ID.IsString }}id(){{ else }}ID{{ end }}
+	{{- if $e.M2M  }}
+		{{ $i := 1 }}{{ $j := 0 }}{{- if $e.IsInverse }}{{ $i = 0 }}{{ $j = 1 }}{{ end -}}
+		t1 := sql.Table({{ $e.Type.Package }}.Table)
+		t2 := sql.Table({{ $n.Package }}.Table)
+		t3 := sql.Table({{ $n.Package }}.{{ $e.TableConstant }})
+		t4 := sql.Select(t3.C({{ $n.Package }}.{{ $e.PKConstant }}[{{ $i }}])).
+				From(t3).
+				Join(t2).
+				On(t3.C({{ $n.Package }}.{{ $e.PKConstant }}[{{ $j }}]), t2.C({{ $n.Package }}.{{ $n.ID.Constant }})).
+				Where(sql.EQ(t2.C({{ $n.Package }}.{{ $n.ID.Constant }}), id))
+		query.sql = sql.Select().
+			From(t1).
+			Join(t4).
+			On(t1.C({{ $e.Type.Package }}.{{ $e.Type.ID.Constant }}), t4.C({{ $n.Package }}.{{ $e.PKConstant }}[{{ $i }}]))
+	{{- else if or $e.M2O (and $e.O2O $e.IsInverse) }}{{/* M2O || (O2O with inverse edge) */}}
+		t1 := sql.Table({{ $e.Type.Package }}.Table)
+		t2 := sql.Select({{ $n.Package }}.{{ $e.ColumnConstant }}).
+			From(sql.Table({{ $n.Package }}.{{ $e.TableConstant }})).
+			Where(sql.EQ({{ $n.Package }}.{{ $n.ID.Constant }}, id))
+		query.sql = sql.Select().From(t1).Join(t2).On(t1.C({{ $e.Type.Package }}.{{ $e.Type.ID.Constant }}), t2.C({{ $n.Package }}.{{ $e.ColumnConstant }}))
+	{{- else }}{{/* O2M || (O2O with assoc edge) */}}
+		query.sql = sql.Select().From(sql.Table({{ $e.Type.Package }}.Table)).
+			Where(sql.EQ({{ $n.Package }}.{{ $e.ColumnConstant }}, id))
+	{{- end }}
 {{ end }}