ent/entc/gen/template/dialect/gremlin/query.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/query" }}
{{ $builder := pascal $.Scope.Builder }}
{{ $receiver := receiver $builder }}

func ({{ $receiver }} *{{ $builder }}) gremlinIDs(ctx context.Context) ([]{{ $.ID.Type }}, error) {
	res := &gremlin.Response{}
	query, bindings := {{ $receiver }}.gremlinQuery().Query()
	if err := {{ $receiver }}.driver.Exec(ctx, query, bindings, res); err != nil {
		return nil, err
	}
	vertices, err := res.ReadVertices()
	if err != nil {
		return nil, err
	}
	ids := make([]{{ $.ID.Type }}, 0, len(vertices))
	for _, vertex := range vertices {
		ids = append(ids, vertex.ID.({{ $.ID.Type }}))
	}
	return ids, nil
}

func ({{ $receiver }} *{{ $builder }}) gremlinAll(ctx context.Context) ([]*{{ $.Name }}, error) {
	res := &gremlin.Response{}
	query, bindings := {{ $receiver }}.gremlinQuery().ValueMap(true).Query()
	if err := {{ $receiver }}.driver.Exec(ctx, query, bindings, res); err != nil {
		return nil, err
	}
	var {{ plural $.Receiver }} {{ plural $.Name  }}
	if err := {{ plural $.Receiver }}.FromResponse(res); err != nil {
		return nil, err
	}
	{{ plural $.Receiver }}.config({{ $receiver }}.config)
	return {{ plural $.Receiver }}, nil
}

func ({{ $receiver }} *{{ $builder }}) gremlinCount(ctx context.Context) (int, error) {
	res := &gremlin.Response{}
	query, bindings := {{ $receiver }}.gremlinQuery().Count().Query()
	if err := {{ $receiver }}.driver.Exec(ctx, query, bindings, res); err != nil {
		return 0, err
	}
	return res.ReadInt()
}

func ({{ $receiver }} *{{ $builder }}) gremlinExist(ctx context.Context) (bool, error) {
	res := &gremlin.Response{}
	query, bindings := {{ $receiver }}.gremlinQuery().HasNext().Query()
	if err := {{ $receiver }}.driver.Exec(ctx, query, bindings, res); err != nil {
		return false, err
	}
	return res.ReadBool()
}

func ({{ $receiver }} *{{ $builder }}) gremlinQuery() *dsl.Traversal {
	v := g.V().HasLabel({{ $.Package }}.Label)
	if {{ $receiver }}.gremlin != nil {
		v = {{ $receiver }}.gremlin.Clone()
	}
	for _, p := range {{ $receiver }}.predicates {
		p(v)
	}
	if len({{ $receiver }}.order) > 0 {
		v.Order()
		for _, p := range {{ $receiver }}.order {
			p(v)
		}
	}
	switch limit, offset := {{ $receiver }}.limit, {{ $receiver }}.offset; {
	case limit != nil && offset != nil:
		v.Range(*offset, *offset + *limit)
	case offset != nil:
		v.Range(*offset, math.MaxInt32)
	case limit != nil:
		v.Limit(*limit)
	}
	if unique := {{ $receiver }}.unique; len(unique) == 0 {
		v.Dedup()
	}
	return v
}
{{ end }}

{{/* query/path defines the query generation for path of a given edge. */}}
{{ define "dialect/gremlin/query/path" }}
	{{- $e := $.Scope.Edge }} {{/* the edge we need to genegrate the path to. */}}
	{{- $receiver := $.Scope.Receiver }}
	gremlin := {{ $receiver }}.gremlinQuery()
	{{- if $e.SelfRef }}
		query.gremlin = gremlin.Both({{ $.Package }}.{{ $e.Constant }})
	{{- else if $e.IsInverse }}
		query.gremlin = gremlin.InE({{ $e.Type.Package }}.{{ $e.Constant }}).OutV()
	{{- else }}
		query.gremlin = gremlin.OutE({{ $.Package }}.{{ $e.Constant }}).InV()
	{{- end }}
{{ end }}


{{/* query/from defines the query generation for an edge query from a given node. */}}
{{ define "dialect/gremlin/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 }}

	{{- if $e.SelfRef }}
		query.gremlin = g.V({{ $receiver }}.ID).Both({{ $n.Package }}.{{ $e.Constant }})
	{{- else if $e.IsInverse }}
		query.gremlin = g.V({{ $receiver }}.ID).InE({{ $e.Type.Package }}.{{ $e.Constant }}).OutV()
	{{- else }}
		query.gremlin = g.V({{ $receiver }}.ID).OutE({{ $n.Package }}.{{ $e.Constant }}).InV()
	{{- end }}
{{ end }}