ent/entc/gen/template/intercept.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.Graph */}}

{{ define "intercept" }}

{{ with extend $ "Package" "intercept" }}
	{{ template "header" . }}
{{ end }}

import (
	"fmt"
	"context"

	"entgo.io/ent/dialect/sql"
	"{{ $.Config.Package }}"
	"{{ $.Config.Package }}/predicate"
	{{- range $n := $.Nodes }}
		{{ $n.PackageAlias }} "{{ $n.Config.Package }}/{{ $n.PackageDir }}"
	{{- end }}
)

{{ $pkg := base $.Config.Package }}

// The Query interface represents an operation that queries a graph.
// By using this interface, users can write generic code that manipulates
// query builders of different types.
type Query interface {
	// Type returns the string representation of the query type.
	Type() string
	// Limit the number of records to be returned by this query.
	Limit(int)
	// Offset to start from.
	Offset(int)
	// Unique configures the query builder to filter duplicate records.
	Unique(bool)
	// Order specifies how the records should be ordered.
	Order(...func({{ $.Storage.Builder }}))
	// WhereP appends storage-level predicates to the query builder. Using this method, users
	// can use type-assertion to append predicates that do not depend on any generated package.
	WhereP(...func({{ $.Storage.Builder }}))
}

// The Func type is an adapter that allows ordinary functions to be used as interceptors.
// Unlike traversal functions, interceptors are skipped during graph traversals. Note that the
// implementation of Func is different from the one defined in entgo.io/ent.InterceptFunc.
type Func func(context.Context, Query) error

// Intercept calls f(ctx, q) and then applied the next Querier.
func (f Func) Intercept(next {{ $pkg }}.Querier) {{ $pkg }}.Querier {
	return {{ $pkg }}.QuerierFunc(func(ctx context.Context, q {{ $pkg }}.Query) ({{ $pkg }}.Value, error) {
		query, err := NewQuery(q)
		if err != nil {
			return nil, err
		}
		if err := f(ctx, query); err != nil {
			return nil, err
		}
		return next.Query(ctx, q)
	})
}

// The TraverseFunc type is an adapter to allow the use of ordinary function as Traverser.
// If f is a function with the appropriate signature, TraverseFunc(f) is a Traverser that calls f.
type TraverseFunc func(context.Context, Query) error

// Intercept is a dummy implementation of Intercept that returns the next Querier in the pipeline.
func (f TraverseFunc) Intercept(next {{ $pkg }}.Querier) {{ $pkg }}.Querier {
	return next
}

// Traverse calls f(ctx, q).
func (f TraverseFunc) Traverse(ctx context.Context, q {{ $pkg }}.Query) error {
	query, err := NewQuery(q)
	if err != nil {
		return err
	}
	return f(ctx, query)
}

{{- range $n := $.Nodes }}
	{{ $name := print $n.Name "Func" }}
	{{ $type := printf "*%s.%s" $pkg $n.QueryName }}
	// The {{ $name }} type is an adapter to allow the use of ordinary function as a Querier.
	type {{ $name }} func(context.Context, {{ $type }}) ({{ $pkg }}.Value, error)

	// Query calls f(ctx, q).
	func (f {{ $name }}) Query(ctx context.Context, q {{ $pkg }}.Query) ({{ $pkg }}.Value, error) {
		if q, ok := q.({{ $type }}); ok {
			return f(ctx, q)
		}
		return nil, fmt.Errorf("unexpected query type %T. expect {{ $type }}", q)
	}

	{{ $name = print "Traverse" $n.Name }}
	// The {{ $name }} type is an adapter to allow the use of ordinary function as Traverser.
	type {{ $name }} func(context.Context, {{ $type }}) error

	// Intercept is a dummy implementation of Intercept that returns the next Querier in the pipeline.
	func (f {{ $name }}) Intercept(next {{ $pkg }}.Querier) {{ $pkg }}.Querier {
		return next
	}

	// Traverse calls f(ctx, q).
	func (f {{ $name }}) Traverse(ctx context.Context, q {{ $pkg }}.Query) error {
		if q, ok := q.({{ $type }}); ok {
			return f(ctx, q)
		}
		return fmt.Errorf("unexpected query type %T. expect {{ $type }}", q)
	}
{{- end }}

// NewQuery returns the generic Query interface for the given typed query.
func NewQuery(q {{ $pkg }}.Query) (Query, error) {
	switch q := q.(type) {
	{{- range $n := $.Nodes }}
		case *{{ $pkg }}.{{ $n.QueryName }}:
		return &query[*{{ $pkg }}.{{ $n.QueryName }}, predicate.{{ $n.Name }}, {{ $n.Package }}.OrderOption]{typ: {{ $pkg }}.{{ $n.TypeName }}, tq: q}, nil
	{{- end }}
	default:
		return nil, fmt.Errorf("unknown query type %T", q)
	}
}

type query[T any, P ~func({{ $.Storage.Builder }}), R ~func({{ $.Storage.Builder }})] struct {
	typ string
	tq  interface {
		Limit(int) T
		Offset(int) T
		Unique(bool) T
		Order(...R) T
		Where(...P) T
	}
}

func (q query[T, P, R]) Type() string {
	return q.typ
}

func (q query[T, P, R]) Limit(limit int) {
	q.tq.Limit(limit)
}

func (q query[T, P, R]) Offset(offset int) {
	q.tq.Offset(offset)
}

func (q query[T, P, R]) Unique(unique bool) {
	q.tq.Unique(unique)
}

func (q query[T, P, R]) Order(orders ...func({{ $.Storage.Builder }})) {
	rs := make([]R, len(orders))
	for i := range orders {
		rs[i] = orders[i]
	}
	q.tq.Order(rs...)
}

func (q query[T, P, R]) WhereP(ps ...func({{ $.Storage.Builder }})) {
	p := make([]P, len(ps))
	for i := range ps {
		p[i] = ps[i]
	}
	q.tq.Where(p...)
}
{{ end }}