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 }}

entx/timestamps.go

@@ -0,0 +1,72 @@
+package entx
+
+import (
+	"time"
+
+	"entgo.io/contrib/entgql"
+	"entgo.io/ent"
+	"entgo.io/ent/schema"
+	"entgo.io/ent/schema/field"
+	"entgo.io/ent/schema/index"
+	"entgo.io/ent/schema/mixin"
+)
+
+// TimestampsMixin defines an interface of a Mixin that provides the created_at and updated_at timestamp fields
+type TimestampsMixin interface {
+	ent.Mixin
+	CreatedAtAnnotations(...schema.Annotation) TimestampsMixin
+	UpdatedAtAnnotations(...schema.Annotation) TimestampsMixin
+}
+
+type timestampsMixin struct {
+	mixin.Schema
+	createdAnnotations []schema.Annotation
+	updatedAnnotations []schema.Annotation
+}
+
+// NewTimestampMixin returns a Mixin that provides the created_at and updated_at timestamp fields
+func NewTimestampMixin() TimestampsMixin {
+	return timestampsMixin{
+		createdAnnotations: []schema.Annotation{
+			entgql.Skip(entgql.SkipMutationCreateInput, entgql.SkipMutationUpdateInput),
+			entgql.OrderField("CREATED_AT"),
+		},
+		updatedAnnotations: []schema.Annotation{
+			entgql.Skip(entgql.SkipMutationCreateInput, entgql.SkipMutationUpdateInput),
+			entgql.OrderField("UPDATED_AT"),
+		},
+	}
+}
+
+func (m timestampsMixin) CreatedAtAnnotations(ants ...schema.Annotation) TimestampsMixin {
+	m.createdAnnotations = ants
+	return m
+}
+
+func (m timestampsMixin) UpdatedAtAnnotations(ants ...schema.Annotation) TimestampsMixin {
+	m.updatedAnnotations = ants
+	return m
+}
+
+// Fields provides the created_at and updated_at fields
+func (m timestampsMixin) Fields() []ent.Field {
+	return []ent.Field{
+		field.Time("created_at").
+			Default(time.Now).
+			Immutable().
+			Annotations(m.createdAnnotations...),
+		field.Time("updated_at").
+			Default(time.Now).
+			UpdateDefault(time.Now).
+			Immutable().
+			Annotations(m.updatedAnnotations...),
+	}
+}
+
+// Indexes provides indexes on both created_at and updated_at fields
+func (timestampsMixin) Indexes() []ent.Index {
+	return []ent.Index{
+		index.Fields("created_at"),
+		index.Fields("updated_at"),
+	}
+}

go.mod

@@ -3,6 +3,9 @@ module go.infratographer.com/x
 go 1.20
 
 require (
+	entgo.io/contrib v0.4.5
+	entgo.io/ent v0.12.3
+	github.com/99designs/gqlgen v0.17.31
 	github.com/MicahParks/keyfunc v1.9.0
 	github.com/XSAM/otelsql v0.21.0
 	github.com/cockroachdb/cockroach-go/v2 v2.3.3
@@ -21,6 +24,7 @@ require (
 	github.com/spf13/cobra v1.7.0
 	github.com/spf13/pflag v1.0.5
 	github.com/spf13/viper v1.15.0
+	github.com/vektah/gqlparser/v2 v2.5.1
 	github.com/zsais/go-gin-prometheus v0.1.0
 	go.opentelemetry.io/contrib/instrumentation/github.com/gin-gonic/gin/otelgin v0.40.0
 	go.opentelemetry.io/contrib/instrumentation/github.com/labstack/echo/otelecho v0.40.0
@@ -38,18 +42,32 @@ require (
 )
 
 require (
+	ariga.io/atlas v0.10.2-0.20230427182402-87a07dfb83bf // indirect
+	github.com/agext/levenshtein v1.2.1 // indirect
+	github.com/agnivade/levenshtein v1.1.1 // indirect
+	github.com/apparentlymart/go-textseg/v13 v13.0.0 // indirect
 	github.com/bytedance/sonic v1.8.0 // indirect
 	github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/go-openapi/inflect v0.19.0 // indirect
 	github.com/golang-jwt/jwt v3.2.2+incompatible // indirect
+	github.com/google/go-cmp v0.5.9 // indirect
+	github.com/hashicorp/hcl/v2 v2.13.0 // indirect
 	github.com/klauspost/cpuid/v2 v2.0.9 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect
+	github.com/mitchellh/go-wordwrap v0.0.0-20150314170334-ad45545899c7 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/twitchyliquid64/golang-asm v0.15.1 // indirect
 	github.com/valyala/bytebufferpool v1.0.0 // indirect
 	github.com/valyala/fasttemplate v1.2.2 // indirect
+	github.com/vmihailenco/msgpack/v5 v5.0.0-beta.9 // indirect
+	github.com/vmihailenco/tagparser v0.1.2 // indirect
+	github.com/zclconf/go-cty v1.8.0 // indirect
 	golang.org/x/arch v0.0.0-20210923205945-b76863e36670 // indirect
+	golang.org/x/exp v0.0.0-20221230185412-738e83a70c30 // indirect
+	golang.org/x/mod v0.10.0 // indirect
 	golang.org/x/time v0.3.0 // indirect
+	golang.org/x/tools v0.8.1-0.20230428195545-5283a0178901 // indirect
 	google.golang.org/appengine v1.6.7 // indirect
 )