entc/gen: add option to check feature-flags in templates (#826)

entc/gen/template/privacy/privacy.tmpl

@@ -0,0 +1,239 @@
+{{/*
+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 "privacy" }}
+
+{{- with extend $ "Package" "privacy" -}}
+	{{ template "header" . }}
+{{ end }}
+
+import "{{ $.Config.Package }}"
+
+{{ $pkg := base $.Config.Package }}
+
+var (
+	// Allow may be returned by rules to indicate that the policy
+	// evaluation should terminate with an allow decision.
+	Allow = errors.New("ent/privacy: allow rule")
+
+	// Deny may be returned by rules to indicate that the policy
+	// evaluation should terminate with an deny decision.
+	Deny = errors.New("ent/privacy: deny rule")
+
+	// Skip may be returned by rules to indicate that the policy
+	// evaluation should continue to the next rule.
+	Skip = errors.New("ent/privacy: skip rule")
+)
+
+{{- range $decision := list "Allow" "Deny" "Skip" }}
+	// {{ $decision }}f returns an formatted wrapped {{ $decision }} decision.
+	func {{ $decision }}f(format string, a ...interface{}) error {
+		return fmt.Errorf(format+": %w", append(a, {{ $decision }})...)
+	}
+{{- end }}
+
+type decisionCtxKey struct {}
+
+// DecisionContext creates a decision context.
+func DecisionContext(parent context.Context, decision error) context.Context {
+	if decision == nil || errors.Is(decision, Skip) {
+		return parent
+	}
+	return context.WithValue(parent, decisionCtxKey{}, decision)
+}
+
+func decisionFromContext(ctx context.Context) (error, bool) {
+	decision, ok := ctx.Value(decisionCtxKey{}).(error)
+	if ok && errors.Is(decision, Allow) {
+		decision = nil
+	}
+	return decision, ok
+}
+
+type (
+	// QueryPolicy combines multiple query rules into a single policy.
+	QueryPolicy []QueryRule
+
+	// QueryRule defines the interface deciding whether a
+	// query is allowed and optionally modify it.
+	QueryRule interface {
+		EvalQuery(context.Context, {{ $pkg }}.Query) error
+	}
+)
+
+// EvalQuery evaluates a query against a query policy.
+func (policy QueryPolicy) EvalQuery(ctx context.Context, q {{ $pkg }}.Query) error {
+	if decision, ok := decisionFromContext(ctx); ok {
+		return decision
+	}
+	for _, rule := range policy {
+		switch decision := rule.EvalQuery(ctx, q); {
+		case decision == nil || errors.Is(decision, Skip):
+		case errors.Is(decision, Allow):
+			return nil
+		default:
+			return decision
+		}
+	}
+	return nil
+}
+
+// QueryRuleFunc type is an adapter to allow the use of
+// ordinary functions as query rules.
+type QueryRuleFunc func(context.Context, {{ $pkg }}.Query) error
+
+// Eval returns f(ctx, q).
+func (f QueryRuleFunc) EvalQuery(ctx context.Context, q {{ $pkg }}.Query) error {
+	return f(ctx, q)
+}
+
+type (
+	// MutationPolicy combines multiple mutation rules into a single policy.
+	MutationPolicy []MutationRule
+
+	// MutationRule defines the interface deciding whether a
+	// mutation is allowed and optionally modify it.
+	MutationRule interface {
+		EvalMutation(context.Context, {{ $pkg }}.Mutation) error
+	}
+)
+
+// EvalMutation evaluates a mutation against a mutation policy.
+func (policy MutationPolicy) EvalMutation(ctx context.Context, m {{ $pkg }}.Mutation) error {
+	if decision, ok := decisionFromContext(ctx); ok {
+		return decision
+	}
+	for _, rule := range policy {
+		switch decision := rule.EvalMutation(ctx, m); {
+		case decision == nil || errors.Is(decision, Skip):
+		case errors.Is(decision, Allow):
+			return nil
+		default:
+			return decision
+		}
+	}
+	return nil
+}
+
+// MutationRuleFunc type is an adapter to allow the use of
+// ordinary functions as mutation rules.
+type MutationRuleFunc func(context.Context, {{ $pkg }}.Mutation) error
+
+// EvalMutation returns f(ctx, m).
+func (f MutationRuleFunc) EvalMutation(ctx context.Context, m {{ $pkg }}.Mutation) error {
+	return f(ctx, m)
+}
+
+// Policy groups query and mutation policies.
+type Policy struct {
+	Query QueryPolicy
+	Mutation MutationPolicy
+}
+
+// EvalQuery forwards evaluation to query policy.
+func (policy Policy) EvalQuery(ctx context.Context, q {{ $pkg }}.Query) error {
+	return policy.Query.EvalQuery(ctx, q)
+}
+
+// EvalMutation forwards evaluation to mutation policy.
+func (policy Policy) EvalMutation(ctx context.Context, m {{ $pkg }}.Mutation) error {
+	return policy.Mutation.EvalMutation(ctx, m)
+}
+
+// QueryMutationRule is the interface that groups query and mutation rules.
+type QueryMutationRule interface {
+	QueryRule
+	MutationRule
+}
+
+// AlwaysAllowRule returns a rule that returns an allow decision.
+func AlwaysAllowRule() QueryMutationRule {
+	return fixedDecision{Allow}
+}
+
+// AlwaysDenyRule returns a rule that returns a deny decision.
+func AlwaysDenyRule() QueryMutationRule {
+	return fixedDecision{Deny}
+}
+
+type fixedDecision struct {
+	decision error
+}
+
+func (f fixedDecision) EvalQuery(context.Context, {{ $pkg }}.Query) error {
+	return f.decision
+}
+
+func (f fixedDecision) EvalMutation(context.Context, {{ $pkg }}.Mutation) error {
+	return f.decision
+}
+
+type contextDecision struct {
+	eval func(context.Context) error
+}
+
+// ContextQueryMutationRule creates a query/mutation rule from a context eval func.
+func ContextQueryMutationRule(eval func(context.Context) error) QueryMutationRule {
+	return contextDecision{eval}
+}
+
+func (c contextDecision) EvalQuery(ctx context.Context, _ {{ $pkg }}.Query) error {
+	return c.eval(ctx)
+}
+
+func (c contextDecision) EvalMutation(ctx context.Context, _ {{ $pkg }}.Mutation) error {
+	return c.eval(ctx)
+}
+
+// OnMutationOperation evaluates the given rule only on a given mutation operation.
+func OnMutationOperation(rule MutationRule, op {{ $pkg }}.Op) MutationRule {
+	return MutationRuleFunc(func(ctx context.Context, m {{ $pkg }}.Mutation) error {
+		if m.Op().Is(op) {
+			return rule.EvalMutation(ctx, m)
+		}
+		return Skip
+	})
+}
+
+// DenyMutationOperationRule returns a rule denying specified mutation operation.
+func DenyMutationOperationRule(op {{ $pkg }}.Op) MutationRule {
+	rule := MutationRuleFunc(func(_ context.Context, m {{ $pkg }}.Mutation) error {
+		return Denyf("ent/privacy: operation %s is not allowed", m.Op())
+	})
+	return OnMutationOperation(rule, op)
+}
+
+{{- range $n := $.Nodes }}
+	{{ $name := print $n.Name "QueryRuleFunc" }}
+	{{ $type := printf "*%s.%s" $pkg $n.QueryName }}
+	// The {{ $name }} type is an adapter to allow the use of ordinary
+	// functions as a query rule.
+	type {{ $name }} func(context.Context, {{ $type }}) error
+
+	// EvalQuery return f(ctx, q).
+	func (f {{ $name }}) EvalQuery(ctx context.Context, q {{ $pkg }}.Query) error {
+		if q, ok := q.({{ $type }}); ok {
+			return f(ctx, q)
+		}
+		return Denyf("ent/privacy: unexpected query type %T, expect {{ $type }}", q)
+	}
+
+	{{ $name = print $n.Name "MutationRuleFunc" }}
+	{{ $type = printf "*%s.%s" $pkg $n.MutationName }}
+	// The {{ $name }} type is an adapter to allow the use of ordinary
+	// functions as a mutation rule.
+	type {{ $name }} func(context.Context, {{ $type }}) error
+
+	// EvalMutation calls f(ctx, m).
+	func (f {{ $name }}) EvalMutation(ctx context.Context, m {{ $pkg }}.Mutation) error {
+		if m, ok := m.({{ $type }}); ok {
+			return f(ctx, m)
+		}
+		return Denyf("ent/privacy: unexpected mutation type %T, expect {{ $type }}", m)
+	}
+{{- end }}
+
+{{ end }}