main.go

package main

import (
	"bytes"
	"context"
	"flag"
	"fmt"
	"os"
	"strings"
	"text/template"

	"entgo.io/ent/entc"
	"entgo.io/ent/entc/gen"

	"oss.terrastruct.com/d2/d2graph"
	"oss.terrastruct.com/d2/d2layouts/d2elklayout"
	"oss.terrastruct.com/d2/d2lib"
	"oss.terrastruct.com/d2/d2renderers/d2svg"
	"oss.terrastruct.com/d2/d2themes/d2themescatalog"
	"oss.terrastruct.com/d2/lib/textmeasure"
)

var (
	flagOutputSvg = flag.String("o", "", "output svg")
	flagOutputD2  = flag.String("d2", "", "output d2")
)

func main() {
	flag.Parse()
	path := "./ent/schema"
	if len(flag.Args()) > 0 {
		path = flag.Args()[0]
	}
	if err := run(path); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}

}

func run(path string) error {
	graph, err := entc.LoadGraph(path, &gen.Config{})
	if err != nil {
		return err
	}
	// capture to bytes buffer
	buf := new(bytes.Buffer)
	if err := tmpl.Execute(buf, graph); err != nil {
		return err
	}
	if *flagOutputSvg == "" {
		fmt.Println(buf.String())
	} else {
		if err := renderSvg(buf.String(), *flagOutputSvg); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
	}
	return nil
}

func ref[T any](v T) *T {
	return &v
}

func renderSvg(contents string, outFilePath string) error {
	ruler, _ := textmeasure.NewRuler()
	defaultLayout := func(engine string) (d2graph.LayoutGraph, error) {
		return d2elklayout.DefaultLayout, nil
	}
	diagram, _, err := d2lib.Compile(context.Background(), contents, &d2lib.CompileOptions{
		LayoutResolver: defaultLayout,
		Ruler:          ruler,
	}, &d2svg.RenderOpts{})
	if err != nil {
		return err
	}
	out, err := d2svg.Render(diagram, &d2svg.RenderOpts{
		ThemeID: ref(d2themescatalog.GrapeSoda.ID),
	})
	if err != nil {
		return err
	}
	return os.WriteFile(outFilePath, out, 0600)
}

var tmpl = template.Must(template.New("d2-diagram").
	Funcs(template.FuncMap{
		"fmtType": func(s string) string {
			// if has [ then put in single quotes:
			if strings.Contains(s, "[") {
				return "'" + s + "'"
			}
			return s
		},
	}).
	Parse(`
{{- with $.Nodes }}
# Generated by ent2d2 (https://github.com/tmc/ent2d2).

# Edges
{{- range $n := . }}
    {{- range $e := $n.Edges }}
	{{- if not $e.IsInverse }}
		{{- $rt := "<->" }}{{ if $e.O2M }}{{ $rt = "<->" }}{{ else if $e.M2O }}{{ $rt = "<->" }}{{ else if $e.M2M }}{{ $rt = "<->" }}{{ end }}
    	{{ $n.Name }} {{ $rt }} {{ $e.Type.Name }} : "{{ $e.Name }}{{ with $e.Ref }}/{{ .Name }}{{ end }}" {
		{{ if $e.O2O -}}
		# o2o {{ if $e.Optional }}optional{{ else }}required{{ end }}
		source-arrowhead: {
			shape: cf-one-required
		}
		target-arrowhead: {
			shape: {{ if $e.Optional }}cf-one{{ else }}cf-one-required{{ end }}
		}
		{{ else if $e.O2M -}}
		# o2m {{ if $e.Optional }}optional{{ else }}required{{ end }}
		source-arrowhead: {
			shape: cf-one-required
		}
		target-arrowhead: {
			shape: cf-many
		}
		{{ else if $e.M2O -}}
		# m2o {{ if $e.Optional }}optional{{ else }}required{{ end }}
		source-arrowhead: {
			shape: cf-many
		}
		target-arrowhead: {
			shape: cf-one
		}
		{{ else if $e.M2M }}
		source-arrowhead: {
			shape: cf-many
		}
		target-arrowhead: {
			shape: {{ if $e.Optional }}cf-many{{ else }}cf-many-required{{ end }}
		}
		{{ end }}
	}
	{{- end }}
	{{- end }}
{{- end }}

# Tables
{{- range $n := . }}
{{ $n.Name }}: {
	shape: sql_table
	{{- if $n.HasOneFieldID }}
	{{ $n.ID.Name }}: {{ fmtType $n.ID.Type.String }} {constraint: primary_key}
	{{- end }}
	{{- range $f := $n.Fields }}
	{{ $f.Name }}: {{ fmtType $f.Type.String }} {{ if $f.IsEdgeField }} {constraint: foreign_key}{{ end }}{{ if $f.Unique }}{constraint: unique}{{ end }}
	# this is a comment

	{{- end }}
}
{{- end }}


{{- end }}
`))