relspecgo/pkg/writers/dbml/writer.go

package dbml

import (
	"fmt"
	"os"
	"strings"

	"git.warky.dev/wdevs/relspecgo/pkg/models"
	"git.warky.dev/wdevs/relspecgo/pkg/writers"
)

// Writer implements the writers.Writer interface for DBML format
type Writer struct {
	options *writers.WriterOptions
}

// NewWriter creates a new DBML writer with the given options
func NewWriter(options *writers.WriterOptions) *Writer {
	return &Writer{
		options: options,
	}
}

// WriteDatabase writes a Database model to DBML format
func (w *Writer) WriteDatabase(db *models.Database) error {
	content := w.databaseToDBML(db)

	if w.options.OutputPath != "" {
		return os.WriteFile(w.options.OutputPath, []byte(content), 0644)
	}

	fmt.Print(content)
	return nil
}

// WriteSchema writes a Schema model to DBML format
func (w *Writer) WriteSchema(schema *models.Schema) error {
	content := w.schemaToDBML(schema)

	if w.options.OutputPath != "" {
		return os.WriteFile(w.options.OutputPath, []byte(content), 0644)
	}

	fmt.Print(content)
	return nil
}

// WriteTable writes a Table model to DBML format
func (w *Writer) WriteTable(table *models.Table) error {
	content := w.tableToDBML(table)

	if w.options.OutputPath != "" {
		return os.WriteFile(w.options.OutputPath, []byte(content), 0644)
	}

	fmt.Print(content)
	return nil
}

// databaseToDBML converts a Database to DBML format string
func (w *Writer) databaseToDBML(d *models.Database) string {
	var sb strings.Builder

	if d.Description != "" {
		sb.WriteString(fmt.Sprintf("// %s\n", d.Description))
	}
	if d.Comment != "" {
		sb.WriteString(fmt.Sprintf("// %s\n", d.Comment))
	}
	if d.Description != "" || d.Comment != "" {
		sb.WriteString("\n")
	}

	for _, schema := range d.Schemas {
		sb.WriteString(w.schemaToDBML(schema))
	}

	sb.WriteString("\n// Relationships\n")
	for _, schema := range d.Schemas {
		for _, table := range schema.Tables {
			for _, constraint := range table.Constraints {
				if constraint.Type == models.ForeignKeyConstraint {
					sb.WriteString(w.constraintToDBML(constraint, table))
				}
			}
		}
	}

	return sb.String()
}

// schemaToDBML converts a Schema to DBML format string
func (w *Writer) schemaToDBML(schema *models.Schema) string {
	var sb strings.Builder

	if schema.Description != "" {
		sb.WriteString(fmt.Sprintf("// Schema: %s - %s\n", schema.Name, schema.Description))
	}

	for _, table := range schema.Tables {
		sb.WriteString(w.tableToDBML(table))
		sb.WriteString("\n")
	}

	return sb.String()
}

// tableToDBML converts a Table to DBML format string
func (w *Writer) tableToDBML(t *models.Table) string {
	var sb strings.Builder

	tableName := fmt.Sprintf("%s.%s", t.Schema, t.Name)
	sb.WriteString(fmt.Sprintf("Table %s {\n", tableName))

	for _, column := range t.Columns {
		sb.WriteString(fmt.Sprintf("  %s %s", column.Name, column.Type))

		var attrs []string
		if column.IsPrimaryKey {
			attrs = append(attrs, "pk")
		}
		if column.NotNull && !column.IsPrimaryKey {
			attrs = append(attrs, "not null")
		}
		if column.AutoIncrement {
			attrs = append(attrs, "increment")
		}
		if column.Default != nil {
			attrs = append(attrs, fmt.Sprintf("default: `%v`", column.Default))
		}

		if len(attrs) > 0 {
			sb.WriteString(fmt.Sprintf(" [%s]", strings.Join(attrs, ", ")))
		}

		if column.Comment != "" {
			sb.WriteString(fmt.Sprintf(" // %s", column.Comment))
		}
		sb.WriteString("\n")
	}

	if len(t.Indexes) > 0 {
		sb.WriteString("\n  indexes {\n")
		for _, index := range t.Indexes {
			var indexAttrs []string
			if index.Unique {
				indexAttrs = append(indexAttrs, "unique")
			}
			if index.Name != "" {
				indexAttrs = append(indexAttrs, fmt.Sprintf("name: '%s'", index.Name))
			}
			if index.Type != "" {
				indexAttrs = append(indexAttrs, fmt.Sprintf("type: %s", index.Type))
			}

			sb.WriteString(fmt.Sprintf("    (%s)", strings.Join(index.Columns, ", ")))
			if len(indexAttrs) > 0 {
				sb.WriteString(fmt.Sprintf(" [%s]", strings.Join(indexAttrs, ", ")))
			}
			sb.WriteString("\n")
		}
		sb.WriteString("  }\n")
	}

	note := strings.TrimSpace(t.Description + " " + t.Comment)
	if note != "" {
		sb.WriteString(fmt.Sprintf("\n  Note: '%s'\n", note))
	}

	sb.WriteString("}\n")
	return sb.String()
}

// constraintToDBML converts a Constraint to DBML format string
func (w *Writer) constraintToDBML(c *models.Constraint, t *models.Table) string {
	if c.Type != models.ForeignKeyConstraint || c.ReferencedTable == "" {
		return ""
	}

	fromTable := fmt.Sprintf("%s.%s", c.Schema, c.Table)
	toTable := fmt.Sprintf("%s.%s", c.ReferencedSchema, c.ReferencedTable)

	relationship := ">" // Default to many-to-one
	for _, index := range t.Indexes {
		if index.Unique && strings.Join(index.Columns, ",") == strings.Join(c.Columns, ",") {
			relationship = "-" // one-to-one
			break
		}
	}
	for _, column := range c.Columns {
		if t.Columns[column].IsPrimaryKey {
			relationship = "-" // one-to-one
			break
		}
	}

	var fromRef, toRef string
	if len(c.Columns) == 1 {
		fromRef = fmt.Sprintf("%s.%s", fromTable, c.Columns[0])
	} else {
		fromRef = fmt.Sprintf("%s.(%s)", fromTable, strings.Join(c.Columns, ", "))
	}

	if len(c.ReferencedColumns) == 1 {
		toRef = fmt.Sprintf("%s.%s", toTable, c.ReferencedColumns[0])
	} else {
		toRef = fmt.Sprintf("%s.(%s)", toTable, strings.Join(c.ReferencedColumns, ", "))
	}

	var actions []string
	if c.OnDelete != "" {
		actions = append(actions, fmt.Sprintf("delete: %s", c.OnDelete))
	}
	if c.OnUpdate != "" {
		actions = append(actions, fmt.Sprintf("update: %s", c.OnUpdate))
	}

	refLine := fmt.Sprintf("Ref: %s %s %s", fromRef, relationship, toRef)
	if len(actions) > 0 {
		refLine += fmt.Sprintf(" [%s]", strings.Join(actions, ", "))
	}

	return refLine + "\n"
}