snapshot of https://github.com/ent/ent.git (master, 2026-09-15; blobs > 131072B omitted)

entc/gen/template/intercept.tmpl

@@ -0,0 +1,171 @@
+{{/*
+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 }}