relspecgo/pkg/readers/dbml/commented_refs.go

package dbml

import (
	"fmt"
	"regexp"
	"strings"

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

// CommentedRefsMetadataKey is the Database.Metadata key holding commented
// `// Ref:` lines not yet resolved against the model ([]string).
const CommentedRefsMetadataKey = "dbml.commented_refs"

// commentedRefRegex matches `// Ref: ...` and `// ref: ...`.
var commentedRefRegex = regexp.MustCompile(`^//\s*[Rr]ef\s*:\s*(.+)$`)

// commentedRef returns the ref body of a trimmed `// Ref:` comment line.
func commentedRef(line string) (string, bool) {
	m := commentedRefRegex.FindStringSubmatch(line)
	if m == nil {
		return "", false
	}
	ref := strings.TrimSpace(m[1])
	return ref, ref != ""
}

// PendingCommentedRefs returns the commented refs not yet resolved.
func PendingCommentedRefs(db *models.Database) []string {
	if db == nil || db.Metadata == nil {
		return nil
	}
	refs, _ := db.Metadata[CommentedRefsMetadataKey].([]string)
	return refs
}

func setPendingCommentedRefs(db *models.Database, refs []string) {
	if len(refs) == 0 {
		if db.Metadata != nil {
			delete(db.Metadata, CommentedRefsMetadataKey)
		}
		return
	}
	if db.Metadata == nil {
		db.Metadata = make(map[string]any)
	}
	db.Metadata[CommentedRefsMetadataKey] = refs
}

// addPendingCommentedRef queues a ref, skipping exact repeats.
func addPendingCommentedRef(db *models.Database, ref string) {
	refs := PendingCommentedRefs(db)
	for _, existing := range refs {
		if existing == ref {
			return
		}
	}
	setPendingCommentedRefs(db, append(refs, ref))
}

// ResolveCommentedRefs turns pending commented refs into foreign keys and
// relationships when both ends (schema.table.column) exist in db. A ref that
// matches an existing FK on the same columns is dropped as a duplicate.
// Unmatched refs stay pending unless final is set, in which case they are
// dropped with a warning. Returns human-readable warnings.
func ResolveCommentedRefs(db *models.Database, final bool) []string {
	refs := PendingCommentedRefs(db)
	if len(refs) == 0 {
		return nil
	}

	var warnings []string
	var pending []string
	parser := &Reader{}

	for _, ref := range refs {
		fk := parser.parseRef(ref)
		if fk == nil || len(fk.Columns) == 0 || len(fk.Columns) != len(fk.ReferencedColumns) {
			warnings = append(warnings, fmt.Sprintf("skipping commented ref %q: cannot parse", ref))
			continue
		}

		srcTable, srcCols, srcMissing := lookupColumns(db, fk.Schema, fk.Table, fk.Columns)
		dstTable, dstCols, dstMissing := lookupColumns(db, fk.ReferencedSchema, fk.ReferencedTable, fk.ReferencedColumns)
		if srcMissing != "" || dstMissing != "" {
			if final {
				missing := srcMissing
				if missing == "" {
					missing = dstMissing
				}
				warnings = append(warnings, fmt.Sprintf("skipping commented ref %q: %s not found", ref, missing))
			} else {
				pending = append(pending, ref)
			}
			continue
		}

		fk.Schema, fk.Table, fk.Columns = srcTable.Schema, srcTable.Name, columnNames(srcCols)
		fk.ReferencedSchema, fk.ReferencedTable, fk.ReferencedColumns = dstTable.Schema, dstTable.Name, columnNames(dstCols)

		if hasFKOnColumns(srcTable, fk.Columns) {
			continue // already declared by an uncommented Ref or inline ref
		}
		if _, taken := srcTable.Constraints[fk.Name]; taken {
			warnings = append(warnings, fmt.Sprintf("skipping commented ref %q: constraint %s already exists", ref, fk.Name))
			continue
		}

		for i := range srcCols {
			if !compatibleFKTypes(srcCols[i].Type, dstCols[i].Type) {
				warnings = append(warnings, fmt.Sprintf("commented ref %q: type mismatch %s.%s.%s (%s) -> %s.%s.%s (%s)",
					ref, srcTable.Schema, srcTable.Name, srcCols[i].Name, srcCols[i].Type,
					dstTable.Schema, dstTable.Name, dstCols[i].Name, dstCols[i].Type))
			}
		}

		if srcTable.Constraints == nil {
			srcTable.Constraints = make(map[string]*models.Constraint)
		}
		srcTable.Constraints[fk.Name] = fk
		addFKRelationship(srcTable, fk)
	}

	setPendingCommentedRefs(db, pending)
	return warnings
}

// lookupColumns finds a table and its columns, case-insensitively. missing
// names the first object not found, or is empty.
func lookupColumns(db *models.Database, schemaName, tableName string, cols []string) (*models.Table, []*models.Column, string) {
	qualified := schemaName + "." + tableName
	var table *models.Table
	for _, schema := range db.Schemas {
		if !strings.EqualFold(schema.Name, schemaName) {
			continue
		}
		for _, t := range schema.Tables {
			if strings.EqualFold(t.Name, tableName) {
				table = t
				break
			}
		}
	}
	if table == nil {
		return nil, nil, "table " + qualified
	}

	found := make([]*models.Column, 0, len(cols))
	for _, name := range cols {
		col := table.Columns[name]
		if col == nil {
			for _, c := range table.Columns {
				if strings.EqualFold(c.Name, name) {
					col = c
					break
				}
			}
		}
		if col == nil {
			return nil, nil, "column " + qualified + "." + name
		}
		found = append(found, col)
	}
	return table, found, ""
}

func columnNames(cols []*models.Column) []string {
	names := make([]string, len(cols))
	for i, c := range cols {
		names[i] = c.Name
	}
	return names
}

// hasFKOnColumns reports whether table already has a foreign key over cols.
func hasFKOnColumns(table *models.Table, cols []string) bool {
	for _, c := range table.Constraints {
		if c.Type != models.ForeignKeyConstraint || len(c.Columns) != len(cols) {
			continue
		}
		same := true
		for i := range cols {
			if !strings.EqualFold(c.Columns[i], cols[i]) {
				same = false
				break
			}
		}
		if same {
			return true
		}
	}
	return false
}

// fkTypeAliases maps serial and alias spellings to their storage type.
var fkTypeAliases = map[string]string{
	"smallserial": "smallint", "serial2": "smallint", "int2": "smallint",
	"serial": "integer", "serial4": "integer", "int": "integer", "int4": "integer",
	"bigserial": "bigint", "serial8": "bigint", "int8": "bigint",
}

// compatibleFKTypes compares column types ignoring case, length and serial
// vs. integer spelling. Unknown (empty) types are treated as compatible.
func compatibleFKTypes(a, b string) bool {
	na, nb := normalizeFKType(a), normalizeFKType(b)
	return na == "" || nb == "" || na == nb
}

func normalizeFKType(t string) string {
	t = strings.ToLower(strings.TrimSpace(t))
	if i := strings.Index(t, "("); i >= 0 {
		t = strings.TrimSpace(t[:i])
	}
	if alias, ok := fkTypeAliases[t]; ok {
		return alias
	}
	return t
}