exportable_gen.go

package entx

import (
	"bytes"
	"fmt"
	"os"
	"sort"
	"text/template"

	"entgo.io/ent/entc"
	"entgo.io/ent/entc/gen"
	"github.com/stoewer/go-strcase"
	"golang.org/x/tools/imports"
)

// ExportableGenerator handles the generation of exportable schema validation.
//
// This generator scans ent schemas for the Exportable annotation and generates
// validation code that can be used to verify export requests.
//
// Usage:
//
//	// In your entc.go file:
//	generator := entx.NewExportableGenerator("./schema", "./internal/hooks")
//	err := generator.Generate()
//	if err != nil {
//		log.Fatal(err)
//	}
//
// The generator will create a file with functions like:
//
//	func IsSchemaExportable(schemaName string) bool {
//		return ExportableSchemas[strings.ToLower(schemaName)]
//	}
//
//	func ValidateExportType(exportType string) error {
//		// Returns error if schema is not exportable
//	}
//
// Example schema annotation:
//
//	func (MySchema) Annotations() []schema.Annotation {
//		return []schema.Annotation{
//			entx.NewExportable(), // Marks this schema as exportable
//		}
//	}
type ExportableGenerator struct {
	// SchemaPath is the path to the schema directory
	SchemaPath string
	// OutputDir is the directory where the generated file will be written
	OutputDir string
	// Package is the package name for the generated file
	Package string
	// ImportPath is the import path for the enums package
	ImportPath string
}

// NewExportableGenerator creates a new ExportableGenerator with default settings.
func NewExportableGenerator(schemaPath, outputDir string) *ExportableGenerator {
	return &ExportableGenerator{
		SchemaPath: schemaPath,
		OutputDir:  outputDir,
		Package:    "hooks",
		ImportPath: "github.com/theopenlane/core/common/enums",
	}
}

// WithPackage sets the package name for the generated file.
func (e *ExportableGenerator) WithPackage(pkg string) *ExportableGenerator {
	e.Package = pkg
	return e
}

// WithImportPath sets the import path for the enums package.
func (e *ExportableGenerator) WithImportPath(importPath string) *ExportableGenerator {
	e.ImportPath = importPath
	return e
}

const exportableGeneratedFileTemplate = `// Code generated by entx. DO NOT EDIT.

package {{.Package}}

import (
	"fmt"
)

type info struct {
	hasOwnerField       bool
	hasSystemOwnedField bool
}

// ExportableSchemas contains all schemas that have Exportable annotation
var ExportableSchemas = map[string]info{
{{- range .Schemas}}
{{- range $key, $val := .}}	"{{$key}}": info{
        hasOwnerField: {{ $val.HasOwnerField}},
        hasSystemOwnedField: {{ $val.HasSystemOwnedField}},
    },
{{- end}}
{{- end}}
}

// IsSchemaExportable checks if a schema name is exportable
func IsSchemaExportable(schemaName string) bool {
	_, ok := ExportableSchemas[schemaName]
	return ok
}

// HasOwnerField checks if a schema has an owner field
func HasOwnerField(schemaName string) bool {
	info, ok := ExportableSchemas[schemaName]
	if !ok {
		return false
	}

	return info.hasOwnerField
}

// HasSystemOwnedField checks if a schema has a system owned field
func HasSystemOwnedField(schemaName string) bool {
	info, ok := ExportableSchemas[schemaName]
	if !ok {
		return false
	}

	return info.hasSystemOwnedField
}

// ValidateExportType validates that an export type corresponds to an exportable schema
func ValidateExportType(exportType string) error {
	if !IsSchemaExportable(exportType) {
		return fmt.Errorf("invalid export type provided ( %s )", exportType)
	}
	return nil
}
`

// Generate generates the exportable schemas validation code and ExportType enum.
func (e *ExportableGenerator) Generate(flags ...string) error {
	graph, err := entc.LoadGraph(e.SchemaPath, &gen.Config{
		BuildFlags: flags,
	})
	if err != nil {
		return fmt.Errorf("loading graph: %w", err)
	}

	schemas := e.findExportableSchemas(graph)

	err = e.generateValidationFile(schemas)
	if err != nil {
		return err
	}

	return nil
}

type Info struct {
	HasOwnerField       bool
	HasSystemOwnedField bool
}

func (e *ExportableGenerator) findExportableSchemas(graph *gen.Graph) []map[string]Info {
	var exportableSchemas []map[string]Info

	for _, n := range graph.Nodes {
		if !HasAnnotation[Exportable](n) {
			continue
		}

		exportableSchemas = append(exportableSchemas, map[string]Info{
			strcase.UpperSnakeCase(n.Name): {
				HasOwnerField:       HasAnnotation[OrgOwnedSchema](n),
				HasSystemOwnedField: HasAnnotation[SystemOwnedSchema](n),
			},
		})
	}

	// sort schemas for consistent output by the schema name which is a string comparison
	return sortExportableSchemas(exportableSchemas)
}

func sortExportableSchemas(schemas []map[string]Info) []map[string]Info {
	sort.SliceStable(schemas, func(i, j int) bool {
		for keyI := range schemas[i] {
			for keyJ := range schemas[j] {
				return keyI < keyJ
			}
		}

		return false
	})

	return schemas
}

func (e *ExportableGenerator) generateValidationFile(schemas []map[string]Info) error {
	tmpl, err := template.New("generated").Parse(exportableGeneratedFileTemplate)
	if err != nil {
		return err
	}

	filePath := fmt.Sprintf("%s/exportable_generated.go", e.OutputDir)

	file, err := os.Create(filePath)
	if err != nil {
		return err
	}

	defer file.Close()

	data := struct {
		Package    string
		ImportPath string
		Schemas    []map[string]Info
	}{
		Package:    e.Package,
		ImportPath: e.ImportPath,
		Schemas:    schemas,
	}

	var buf bytes.Buffer
	if err := tmpl.ExecuteTemplate(&buf, "generated", data); err != nil {
		return err
	}

	// run gofmt and goimports on the file contents
	formatted, err := imports.Process(filePath, buf.Bytes(), nil)
	if err != nil {
		return fmt.Errorf("%w: failed to format file", err)
	}

	if _, err := file.Write(formatted); err != nil {
		return err
	}

	return nil
}