ent/entc/gen/template/dialect/gremlin/update.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.
*/}}

{{ define "dialect/gremlin/update" }}
{{ $builder := pascal $.Scope.Builder }}
{{ $receiver := receiver $builder }}
{{ $one := hasSuffix $builder "One" }}
{{ $zero := 0 }}{{ if $one }}{{ $zero = "nil" }}{{ end }}

func ({{ $receiver }} *{{ $builder }}) gremlinSave(ctx context.Context) ({{- if $one }}*{{ $.Name }}{{ else }}int{{ end }}, error) {
	res := &gremlin.Response{}
	query, bindings := {{ $receiver }}.gremlin({{- if $one }}{{ $receiver }}.id{{ end }}).Query()
	if err := {{ $receiver }}.driver.Exec(ctx, query, bindings, res); err != nil {
		return {{ $zero }}, err
	}
	if err, ok := isConstantError(res); ok {
		return {{ $zero }}, err
	}
	{{- if $one }}
		{{- $r := $.Receiver }}
		{{ $r }} := &{{ $.Name }}{config: {{ $receiver }}.config}
		if err := {{ $r }}.FromResponse(res); err != nil {
			return nil, err
		}
		return {{ $r }}, nil
	{{- else }}
		return res.ReadInt()
	{{- end }}
}

func ({{ $receiver }} *{{ $builder }}) gremlin({{ if $one }}id {{ $.ID.Type }}{{ end }}) *dsl.Traversal {
	{{- with .NumConstraint }}
		type constraint struct {
			pred *dsl.Traversal // constraint predicate.
			test *dsl.Traversal // test matches and its constant.
		}
		constraints := make([]*constraint, 0, {{ . }})
	{{- end }}
	{{- /* case of update specific vertex */}}
	{{- if $one }}
		v := g.V(id)
	{{- /* general update for N vertices */}}
	{{- else }}
		v := g.V().HasLabel({{ $.Package }}.Label)
		for _, p := range {{ $receiver }}.predicates {
			p(v)
		}
	{{- end }}
	var (
		{{ if or .NumConstraint (len $.Edges) }}
			rv = v.Clone()
			_ = rv
		{{ end }}
		trs []*dsl.Traversal
	)
	{{- range $_, $f := $.Fields }}
		{{- if or (not $f.Immutable) $f.UpdateDefault }}
			if value := {{ $receiver }}.{{ $f.BuilderField }}; value != nil {
				{{- if $f.Unique }}
					constraints = append(constraints, &constraint{
						pred: g.V().Has({{ $.Package }}.Label, {{ $.Package }}.{{ $f.Constant }}, *value).Count(),
						test: __.Is(p.NEQ(0)).Constant(NewErrUniqueField({{ $.Package }}.Label, {{ $.Package }}.{{ $f.Constant }}, *value)),
					})
				{{- end }}
				v.Property(dsl.Single, {{ $.Package }}.{{ $f.Constant }}, *value)
			}
			{{- if $f.Type.Numeric }}
				if value := {{ $receiver }}.add{{ $f.BuilderField }}; value != nil {
					{{- if $f.Unique }}
						addValue := rv.Clone().Union(__.Values({{ $.Package }}.{{ $f.Constant }}), __.Constant(*value)).Sum().Next()
						constraints = append(constraints, &constraint{
							pred: g.V().Has({{ $.Package }}.Label, {{ $.Package }}.{{ $f.Constant }}, addValue).Count(),
							test: __.Is(p.NEQ(0)).Constant(NewErrUniqueField({{ $.Package }}.Label, {{ $.Package }}.{{ $f.Constant }}, fmt.Sprintf("+= %v", *value))),
						})
					{{- end }}
					v.Property(dsl.Single, {{ $.Package }}.{{ $f.Constant }}, __.Union(__.Values({{ $.Package }}.{{ $f.Constant }}), __.Constant(*value)).Sum())
				}
			{{- end }}
		{{- end }}
	{{- end }}
	{{- /* clear optional fields. */}}
	{{- with $.HasOptional }}
		var properties []interface{}
		{{- range $_, $f := $.Fields }}
			{{- if $f.Optional }}
				if {{ $receiver }}.clear{{ $f.BuilderField }} {
					properties = append(properties, {{ $.Package }}.{{ $f.Constant }})
				}
			{{- end }}
		{{- end }}
		if len(properties) > 0 {
			v.SideEffect(__.Properties(properties...).Drop())
		}
	{{- end }}
	{{- range $_, $e := $.Edges }}
		{{- $direction := "In" }}
		{{- $name := printf "%s.%s" $.Package $e.Constant }}
		{{- if $e.IsInverse }}
			{{- $direction = "Out" }}
			{{- $name = printf "%s.%s" $e.Type.Package $e.Constant }}
		{{- end }}
		{{- /* remove edges */}}
		{{- if $e.Unique }}
		if {{ $receiver }}.cleared{{ $e.StructField }} {
		{{- else }}
		for id := range {{ $receiver }}.removed{{ $e.StructField }} {
		{{- end }}
		{{- if $e.SelfRef }}
			tr := rv.Clone().BothE({{ $name }}){{ if not $e.Unique }}.Where(__.Or(__.InV().HasID(id), __.OutV().HasID(id))){{ end }}.Drop().Iterate()
		{{- else if $e.IsInverse }}
			tr := rv.Clone().InE({{ $name }}){{ if not $e.Unique }}.Where(__.OtherV().HasID(id)){{ end }}.Drop().Iterate()
		{{- else }}
			tr := rv.Clone().OutE({{ $name }}){{ if not $e.Unique }}.Where(__.OtherV().HasID(id)){{ end }}.Drop().Iterate()
		{{- end }}
			trs = append(trs, tr)
		}
		{{- /* update edges */}}
		for id := range {{ $receiver }}.{{ $e.BuilderField }} {
		{{- if $e.IsInverse }}
			v.AddE({{ $name }}).From(g.V(id)).InV()
		{{- else }}
			v.AddE({{ $name }}).To(g.V(id)).OutV()
		{{- end }}
		{{- if $e.HasConstraint }}
			{{- if $e.SelfRef }}
				constraints = append(constraints, &constraint{
					pred: rv.Clone().Both({{ $name }}).Count(),
					test: __.Is(p.NEQ(0)).Constant(NewErrUniqueEdge({{ $.Package }}.Label, {{ $name }}, id)),
				})
				constraints = append(constraints, &constraint{
					pred: g.E().HasLabel({{ $name }}).Where(__.Or(__.InV().HasID(id), __.OutV().HasID(id))).Count(),
					test: __.Is(p.NEQ(0)).Constant(NewErrUniqueEdge({{ $.Package }}.Label, {{ $name }}, id)),
				})
			{{- else }}
				constraints = append(constraints, &constraint{
					pred: g.E().HasLabel({{ $name }}).{{ $direction }}V().HasID(id).Count(),
					test: __.Is(p.NEQ(0)).Constant(NewErrUniqueEdge({{ $.Package }}.Label, {{ $name }}, id)),
				})
			{{- end }}
		{{- end }}
		}
	{{- end }}
	{{ if $one }}v.ValueMap(true){{ else }}v.Count(){{ end }}
	{{- with .NumConstraint }}
		if len(constraints) > 0 {
			{{- /* make sure the traversal does not contain more than one vertex if we have constraint */}}
			{{- if not $one }}
				constraints = append(constraints, &constraint{
					pred: rv.Count(),
					test: __.Is(p.GT(1)).Constant(&ErrConstraintFailed{msg: "update traversal contains more than one vertex"}),
				})
			{{- end }}
			v = constraints[0].pred.Coalesce(constraints[0].test, v)
			for _, cr := range constraints[1:] {
				v = cr.pred.Coalesce(cr.test, v)
			}
		}
	{{- end }}
	trs = append(trs, v)
	return dsl.Join(trs...)
}
{{ end }}