ent/entc/gen/template/dialect/sql/feature/namedges.tmpl

{{/*
Copyright 2019-present Facebook Inc. All rights reserved.
This source code is licensed under the Apache 2.0 license found
in the LICENSE file in the root directory of this source tree.
*/}}

{{/* gotype: entgo.io/ent/entc/gen.Type */}}

{{/* Templates used by the "namedges" feature-flag to allow eager-loading edges with dynamic names. */}}

{{ define "dialect/sql/model/edges/fields/additional/namedges" }}
    {{- if $.FeatureEnabled "namedges" }}
        {{- range $e := $.Edges }}
            {{- if not $e.Unique }}
                named{{ $e.StructField }} map[string][]*{{ $e.Type.Name }}
            {{- end }}
        {{- end }}
    {{- end }}
{{ end }}

{{ define "dialect/sql/model/additional/namedges" }}
    {{- if $.FeatureEnabled "namedges" }}
        {{ $receiver := $.Receiver }}
        {{- range $e := $.Edges }}
            {{- if not $e.Unique }}
                {{ $func := print "Named" $e.StructField }}
                // {{ $func }} returns the {{ $e.StructField }} named value or an error if the edge was not
                // loaded in eager-loading with this name.
                func ({{ $receiver }} *{{ $.Name }}) Named{{ $e.StructField }}(name string) ([]*{{ $e.Type.Name }}, error) {
                    if {{ $receiver }}.Edges.named{{ $e.StructField }} == nil {
                        return nil, &NotLoadedError{edge: name}
                    }
                    nodes, ok := {{ $receiver }}.Edges.named{{ $e.StructField }}[name]
                    if !ok {
                        return nil, &NotLoadedError{edge: name}
                    }
                    return nodes, nil
                }

                func ({{ $receiver }} *{{ $.Name }}) appendNamed{{ $e.StructField }}(name string, edges ...*{{ $e.Type.Name }}) {
                    if {{ $receiver }}.Edges.named{{ $e.StructField }} == nil {
                        {{ $receiver }}.Edges.named{{ $e.StructField }} = make(map[string][]*{{ $e.Type.Name }})
                    }
                    if len(edges) == 0 {
                        {{- /* Prefer empty array over nil to stay consistent with the standard eager-loading API. */}}
                        {{ $receiver }}.Edges.named{{ $e.StructField }}[name] = []*{{ $e.Type.Name }}{}
                    } else {
                        {{ $receiver }}.Edges.named{{ $e.StructField }}[name] = append({{ $receiver }}.Edges.named{{ $e.StructField }}[name], edges...)
                    }
                }
            {{- end }}
        {{- end }}
    {{- end }}
{{ end }}

{{- define "dialect/sql/query/fields/additional/namedges" }}
    {{- if $.FeatureEnabled "namedges" }}
        {{- range $e := $.Edges }}
            {{- if not $e.Unique }}
                {{ $e.EagerLoadNamedField }} map[string]*{{ $e.Type.QueryName }}
            {{- end }}
        {{- end }}
    {{- end }}
{{- end }}

{{ define "dialect/sql/query/additional/namedges" }}
    {{- if $.FeatureEnabled "namedges" }}
        {{ $builder := $.QueryName }}
        {{ $receiver := receiver $builder }}
        {{- range $e := $.Edges }}
            {{- if not $e.Unique }}
                {{ $ebuilder := $e.Type.QueryName }}
                {{ $func := print "WithNamed" $e.StructField }}
                // {{ $func }} tells the query-builder to eager-load the nodes that are connected to the "{{ $e.Name }}"
                // edge with the given name. The optional arguments are used to configure the query builder of the edge.
                func ({{ $receiver }} *{{ $builder }}) {{ $func }}(name string, opts ...func(*{{ $ebuilder }})) *{{ $builder }} {
                    query := (&{{ $e.Type.ClientName }}{config: {{ $receiver }}.config}).Query()
                    for _, opt := range opts {
                        opt(query)
                    }
                    if {{ $receiver }}.{{ $e.EagerLoadNamedField }} == nil {
                        {{ $receiver }}.{{ $e.EagerLoadNamedField }} = make(map[string]*{{ $e.Type.QueryName }})
                    }
                    {{ $receiver }}.{{ $e.EagerLoadNamedField }}[name] = query
                    return {{ $receiver }}
                }
            {{- end }}
        {{- end }}
    {{- end }}
{{ end }}

{{/* Process nodes before they are returned and resolve named-edges. */}}
{{- define "dialect/sql/query/all/nodes/namedges" }}
    {{- if $.FeatureEnabled "namedges" }}
        {{- $builder := pascal $.Scope.Builder }}
        {{- $receiver := receiver $builder }}
        {{- range $e := $.Edges }}
            {{- if not $e.Unique }}
                for name, query := range {{ $receiver }}.{{ $e.EagerLoadNamedField }} {
                    if err := {{ $receiver }}.load{{ $e.StructField }}(ctx, query, nodes,
                        func(n *{{ $.Name }}) { n.appendNamed{{ $e.StructField }}(name) },
                        {{- if and ($.FeatureEnabled "bidiedges") $e.Ref $e.Ref.Unique }}
                            func(n *{{ $.Name }}, e *{{ $e.Type.Name }}){
                                n.appendNamed{{ $e.StructField }}(name, e)
                                {{- $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 }}){ n.appendNamed{{ $e.StructField }}(name, e) }); err != nil {
                        {{- end }}
                        return nil, err
                    }
                }
            {{- end }}
        {{- end }}
    {{- end }}
{{- end }}