feat(entx): Add initial shared ent functions (#58)

This adds our shared ent functions that are used to customize the
resulting enter code.

---------

Signed-off-by: GitHub <noreply@github.com>

entx/annotation.go

@@ -0,0 +1,14 @@
+package entx
+
+// AnnotationName is the value of the annotation when read during ent compilation
+var AnnotationName = "I12R_ENTX"
+
+// Annotation provides a ent.Annotaion spec
+type Annotation struct {
+	IsNamespacedDataJSONField bool
+}
+
+// Name implements the ent Annotation interface.
+func (a Annotation) Name() string {
+	return AnnotationName
+}

entx/doc.go

@@ -0,0 +1,18 @@
+// Copyright 2023 The Infratographer Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package entx is a package of tools for interacting with ent. Providing tools
+// to make generating federated gql easier, working with JSON values in ent, and
+// provided helpers for using idx as your ID on types.
+package entx

entx/generator.go

@@ -0,0 +1,68 @@
+package entx
+
+import (
+	"entgo.io/contrib/entgql"
+	"entgo.io/ent/entc"
+	"entgo.io/ent/entc/gen"
+)
+
+// Extension is an implementation of entc.Extension that adds all the templates
+// that entx needs.
+type Extension struct {
+	entc.DefaultExtension
+
+	templates []*gen.Template
+
+	gqlSchemaHooks []entgql.SchemaHook
+}
+
+// ExtensionOption allow for control over the behavior of the generator
+type ExtensionOption func(*Extension) error
+
+// WithFederation adds support for graphql federation by adding the Entity interface
+// to all types, as well as removing the node() and nodes() query calls.
+func WithFederation() ExtensionOption {
+	return func(ex *Extension) error {
+		ex.templates = append(ex.templates, FederationTemplate)
+		ex.gqlSchemaHooks = append(ex.gqlSchemaHooks, removeNodeGoModel, removeNodeQueries)
+
+		return nil
+	}
+}
+
+// WithJSONScalar adds the JSON scalar definition
+func WithJSONScalar() ExtensionOption {
+	return func(ex *Extension) error {
+		ex.gqlSchemaHooks = append(ex.gqlSchemaHooks, addJSONScalar)
+		return nil
+	}
+}
+
+// NewExtension returns an entc Extension that allows the entx package to generate
+// the schema changes and templates needed to function
+func NewExtension(opts ...ExtensionOption) (*Extension, error) {
+	e := &Extension{
+		templates:      MixinTemplates,
+		gqlSchemaHooks: []entgql.SchemaHook{},
+	}
+
+	for _, opt := range opts {
+		if err := opt(e); err != nil {
+			return nil, err
+		}
+	}
+
+	return e, nil
+}
+
+// Templates of the extension
+func (e *Extension) Templates() []*gen.Template {
+	return e.templates
+}
+
+// GQLSchemaHooks of the extension to seamlessly edit the final gql interface.
+func (e *Extension) GQLSchemaHooks() []entgql.SchemaHook {
+	return e.gqlSchemaHooks
+}
+
+var _ entc.Extension = (*Extension)(nil)

entx/gql_hooks.go

@@ -0,0 +1,71 @@
+package entx
+
+import (
+	"errors"
+
+	"entgo.io/ent/entc"
+	"entgo.io/ent/entc/gen"
+	"github.com/vektah/gqlparser/v2/ast"
+)
+
+// Skipping err113 linting since these errors are returned during generation and not runtime
+//
+//nolint:goerr113
+var (
+	removeNodeGoModel = func(g *gen.Graph, s *ast.Schema) error {
+		n, ok := s.Types["Node"]
+		if !ok {
+			return errors.New("failed to find node interface in schema")
+		}
+
+		dirs := ast.DirectiveList{}
+
+		for _, d := range n.Directives {
+			switch d.Name {
+			case "goModel":
+				continue
+			default:
+				dirs = append(dirs, d)
+			}
+		}
+		n.Directives = dirs
+
+		return nil
+	}
+
+	removeNodeQueries = func(g *gen.Graph, s *ast.Schema) error {
+		q, ok := s.Types["Query"]
+		if !ok {
+			return errors.New("failed to find query definition in schema")
+		}
+
+		fields := ast.FieldList{}
+
+		for _, f := range q.Fields {
+			switch f.Name {
+			case "node":
+			case "nodes":
+				continue
+			default:
+				fields = append(fields, f)
+			}
+		}
+		q.Fields = fields
+
+		return nil
+	}
+
+	addJSONScalar = func(g *gen.Graph, s *ast.Schema) error {
+		s.Types["JSON"] = &ast.Definition{
+			Kind:        ast.Scalar,
+			Description: "A valid JSON string.",
+			Name:        "JSON",
+		}
+		return nil
+	}
+)
+
+// import string mutations from entc
+var (
+	_ entc.Extension = (*Extension)(nil)
+)

entx/json.go

@@ -0,0 +1,41 @@
+package entx
+
+import (
+	"encoding/json"
+	"io"
+
+	"github.com/99designs/gqlgen/graphql"
+)
+
+// MarshalRawMessage provides a graphql.Marshaler for json.RawMessage
+func MarshalRawMessage(t json.RawMessage) graphql.Marshaler {
+	return graphql.WriterFunc(func(w io.Writer) {
+		s, _ := t.MarshalJSON()
+		_, _ = io.WriteString(w, string(s))
+	})
+}
+
+// UnmarshalRawMessage provides a graphql.Unmarshaler for json.RawMessage
+func UnmarshalRawMessage(v interface{}) (json.RawMessage, error) {
+	switch j := v.(type) {
+	case string:
+		return UnmarshalRawMessage([]byte(j))
+	case []byte:
+		return json.RawMessage(j), nil
+	case map[string]interface{}:
+		js, err := json.Marshal(v)
+		if err != nil {
+			return nil, err
+		}
+
+		return json.RawMessage(js), nil
+	default:
+		// Attempt to cast it as a fall back but return an error if it fails
+		js, err := json.Marshal(v)
+		if err != nil {
+			return nil, err
+		}
+
+		return json.RawMessage(js), nil
+	}
+}

entx/key_directive.go

@@ -0,0 +1,27 @@
+package entx
+
+import (
+	"entgo.io/contrib/entgql"
+	"github.com/vektah/gqlparser/v2/ast"
+)
+
+// GraphKeyDirective returns an entgql.Directive for setting the @key field on
+// a graphql type
+func GraphKeyDirective(fields string) entgql.Annotation {
+	return entgql.Directives(keyDirective(fields))
+}
+
+func keyDirective(fields string) entgql.Directive {
+	var args []*ast.Argument
+	if fields != "" {
+		args = append(args, &ast.Argument{
+			Name: "fields",
+			Value: &ast.Value{
+				Raw:  fields,
+				Kind: ast.StringValue,
+			},
+		})
+	}
+
+	return entgql.NewDirective("key", args...)
+}

entx/namespaced_data.go

@@ -0,0 +1,38 @@
+package entx
+
+import (
+	"encoding/json"
+
+	"entgo.io/contrib/entgql"
+	"entgo.io/ent"
+	"entgo.io/ent/schema/field"
+	"entgo.io/ent/schema/mixin"
+)
+
+var (
+	namespaceMinLength = 5
+	namespaceMaxLength = 64
+)
+
+// NamespacedDataMixin defines an ent Mixin that captures raw json associated with a namespace.
+type NamespacedDataMixin struct {
+	mixin.Schema
+}
+
+// Fields provides the namespace and data fields used in this mixin.
+func (m NamespacedDataMixin) Fields() []ent.Field {
+	return []ent.Field{
+		field.Text("namespace").
+			NotEmpty().
+			MinLen(namespaceMinLength).
+			MaxLen(namespaceMaxLength).
+			Annotations(
+				entgql.OrderField("NAMESPACE"),
+			),
+		field.JSON("data", json.RawMessage{}).
+			Annotations(
+				entgql.Type("JSON"),
+				Annotation{IsNamespacedDataJSONField: true},
+			),
+	}
+}

entx/template.go

@@ -0,0 +1,33 @@
+package entx
+
+import (
+	"embed"
+	"text/template"
+
+	"entgo.io/ent/entc/gen"
+)
+
+var (
+	// FederationTemplate adds support for generating the required output to support gql federation
+	FederationTemplate = parseT("template/gql_federation.tmpl")
+
+	// NamespacedDataWhereFuncsTemplate adds support for generating <T>WhereInput filters for schema types using the NamespacedData mixin
+	NamespacedDataWhereFuncsTemplate = parseT("template/namespaceddata_where_funcs.tmpl")
+
+	// TemplateFuncs contains the extra template functions used by entx.
+	TemplateFuncs = template.FuncMap{}
+
+	// MixinTemplates includes all templates for extending ent to support entx mixins.
+	MixinTemplates = []*gen.Template{
+		NamespacedDataWhereFuncsTemplate,
+	}
+
+	//go:embed template/*
+	_templates embed.FS
+)
+
+func parseT(path string) *gen.Template {
+	return gen.MustParse(gen.NewTemplate(path).
+		Funcs(TemplateFuncs).
+		ParseFS(_templates, path))
+}

entx/template/gql_federation.tmpl

@@ -0,0 +1,4 @@
+{{ define "model/additional/gql_federation" }}
+	// IsEntity implement fedruntime.Entity
+	func ({{ $.Receiver }} {{ $.Name }}) IsEntity() {}
+{{ end }}

entx/template/namespaceddata_where_funcs.tmpl

@@ -0,0 +1,18 @@
+{{/* gotype: entgo.io/ent/entc/gen.Type */}}
+
+{{ define "where/additional/entx_namespaced_data" }}
+	{{- range $f := $.Fields }}
+		{{- if $annotation := $f.Annotations.I12R_ENTX }}
+			{{- if $annotation.IsNamespacedDataJSONField }}
+				{{ $arg := "v" }}
+				{{ $func := print $f.StructField "HasKey" }}
+				// {{ $func }} checks if {{$f.StructField}} contains given value
+				func {{ $func }}({{ $arg}} string) predicate.{{ $.Name }} {
+						return predicate.{{ $.Name }}(func(s *sql.Selector) {
+								s.Where(sqljson.HasKey(s.C({{ $f.Constant }}), sqljson.DotPath({{ $arg }})))
+						})
+				}
+			{{ end }}
+		{{ end }}
+	{{ end }}
+{{ end }}