entc/gen: add feature-flag to set unique backref when their inverse i…

…s loaded (ent#3953)

entc/gen/feature.go

@@ -51,6 +51,16 @@ var (
 		Description: "NamedEdges provides an API for eager-loading edges with dynamic names",
 	}
 
+	// FeatureBidiEdgeRefs provides a feature-flag for sql dialect to set two-way
+	// references when loading (unique) edges. Note, users that use the standard
+	// encoding/json.MarshalJSON should detach the circular references before marshaling.
+	FeatureBidiEdgeRefs = Feature{
+		Name:        "bidiedges",
+		Stage:       Experimental,
+		Default:     false,
+		Description: "This features guides Ent to set two-way references when loading (O2M/O2O) edges",
+	}
+
 	// FeatureSnapshot stores a snapshot of ent/schema and auto-solve merge-conflict (issue #852).
 	FeatureSnapshot = Feature{
 		Name:        "schema/snapshot",
@@ -140,6 +150,7 @@ var (
 		FeatureIntercept,
 		FeatureEntQL,
 		FeatureNamedEdges,
+		FeatureBidiEdgeRefs,
 		FeatureSnapshot,
 		FeatureSchemaConfig,
 		FeatureLock,

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

@@ -80,7 +80,21 @@ func ({{ $receiver }} *{{ $builder }}) sqlAll(ctx context.Context, hooks ...quer
 		if query := {{ $receiver }}.{{ $e.EagerLoadField }}; query != nil {
 			if err := {{ $receiver }}.load{{ $e.StructField }}(ctx, query, nodes, {{ if $e.Unique }}nil{{ else }}
 				func(n *{{ $.Name }}){ n.Edges.{{ $e.StructField }} = []*{{ $e.Type.Name }}{} }{{ end }},
-				func(n *{{ $.Name }}, e *{{ $e.Type.Name }}){ n.Edges.{{ $e.StructField }} = {{ if $e.Unique }}e{{ else }}append(n.Edges.{{ $e.StructField }}, e){{ end }} }); err != nil {
+				{{- $lhs := printf "n.Edges.%s" $e.StructField }}
+				{{- $rhs := print "e" }}{{- if not $e.Unique }}{{ $rhs = printf "append(%s, e)" $lhs }}{{ end }}
+				{{- if and ($.FeatureEnabled "bidiedges") $e.Ref $e.Ref.Unique }}
+					func(n *{{ $.Name }}, e *{{ $e.Type.Name }}){
+						{{ printf "%s = %s" $lhs $rhs }}
+						{{- $idx := $e.Ref.Index }}
+						{{- /* Set only in case this type was not loaded explicitly (without custom options). */}}
+						if !e.Edges.loadedTypes[{{ $idx }}] {
+							e.Edges.{{ $e.Ref.StructField }} = n
+						}
+					}); err != nil {
+				{{- else }}
+					{{- /* Keep it one-liner if there is not inverse-condition. */}}
+					func(n *{{ $.Name }}, e *{{ $e.Type.Name }}){ {{ printf "%s = %s" $lhs $rhs }} }); err != nil {
+				{{- end }}
 				return nil, err
 			}
 		}

entc/gen/template/ent.tmpl

@@ -52,8 +52,9 @@ type {{ $.Name }} struct {
 }
 
 {{- with $.Edges }}
+{{- $edgesType := print $.Name "Edges"}}
 // {{ $.Name }}Edges holds the relations/edges for other nodes in the graph.
-type {{ $.Name }}Edges struct {
+type {{ $edgesType }} struct {
 	{{- range $e := . }}
 		{{- template "model/edgecomment" $e }}
 		{{ $e.StructField }} {{ if not $e.Unique }}[]{{ end }}*{{ $e.Type.Name }} {{ with $e.StructTag }}`{{ . }}`{{ end }}
@@ -68,16 +69,19 @@ type {{ $.Name }}Edges struct {
 {{- range $i, $e := . }}
 	// {{ $e.StructField }}OrErr returns the {{ $e.StructField }} value or an error if the edge
 	// was not loaded in eager-loading{{ if $e.Unique }}, or loaded but was not found{{ end }}.
-	func (e {{ $.Name }}Edges) {{ $e.StructField }}OrErr() ({{ if not $e.Unique }}[]{{ end }}*{{ $e.Type.Name }}, error) {
-		if e.loadedTypes[{{ $i }}] {
-			{{- if $e.Unique }}
-				if e.{{ $e.StructField }} == nil {
-					// Edge was loaded but was not found.
-					return nil, &NotFoundError{label: {{ $e.Type.Package }}.Label}
-				}
-			{{- end }}
-			return e.{{ $e.StructField }}, nil
-		}
+	func (e {{ $edgesType }}) {{ $e.StructField }}OrErr() ({{ if not $e.Unique }}[]{{ end }}*{{ $e.Type.Name }}, error) {
+		{{- if $e.Unique }}
+			if e.{{ $e.StructField }} != nil {
+				return e.{{ $e.StructField }}, nil
+			} else if e.loadedTypes[{{ $i }}] {
+				{{- /* Edge was loaded but was not found. */}}
+				return nil, &NotFoundError{label: {{ $e.Type.Package }}.Label}
+			}
+		{{- else }}
+			if e.loadedTypes[{{ $i }}] {
+				return e.{{ $e.StructField }}, nil
+			}
+		{{- end }}
 		return nil, &NotLoadedError{edge: "{{ $e.Name }}"}
 	}
 {{- end }}

entc/gen/type.go

@@ -2122,6 +2122,22 @@ func (e Edge) EntSQL() *entsql.Annotation {
 	return sqlAnnotate(e.Annotations)
 }
 
+// Index returns the index of the edge in the schema.
+// Used mainly to extract its position in the "loadedTypes" array.
+func (e Edge) Index() (int, error) {
+	// "owner" is the type that holds the edge.
+	owner := e.Owner
+	if e.IsInverse() {
+		owner = e.Ref.Type
+	}
+	for i, e1 := range owner.Edges {
+		if e1.Name == e.Name {
+			return i, nil
+		}
+	}
+	return 0, fmt.Errorf("edge %q was not found in its owner schema %q", e.Name, e.Owner.Name)
+}
+
 // Column returns the first element from the columns slice.
 func (r Relation) Column() string {
 	if len(r.Columns) == 0 {