relspecgo/pkg/writers/pgsql/directives.go

package pgsql

import (
	"fmt"
	"strings"

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

// directiveNamespace is the dialect namespace this writer consumes. Directives
// for other namespaces (e.g. "sqlite") are ignored and never emitted as SQL.
const directiveNamespace = "postgres"

// pgHandledDirectives maps a directive location to the set of postgres keys this
// writer knows how to translate. In strict mode an unknown key for this
// namespace at a supported location is a hard error.
var pgHandledDirectives = map[string]map[string]bool{
	models.DirectiveLocationTable:  {"partition": true, "inherits": true, "with": true, "tablespace": true},
	models.DirectiveLocationColumn: {"storage": true, "compression": true, "identity": true},
	models.DirectiveLocationIndex:  {"with": true, "tablespace": true},
}

// checkDirectives validates postgres directives across a schema when strict mode
// is enabled. It returns an error for any postgres directive whose key this
// writer cannot translate. With strict mode off it is a no-op.
func (w *Writer) checkDirectives(schema *models.Schema) error {
	if w.options == nil || !w.options.StrictDirectives {
		return nil
	}
	for _, table := range schema.Tables {
		if err := checkObjectDirectives(table.Metadata, models.DirectiveLocationTable, table.Name); err != nil {
			return err
		}
		for _, col := range table.Columns {
			if err := checkObjectDirectives(col.Metadata, models.DirectiveLocationColumn, table.Name+"."+col.Name); err != nil {
				return err
			}
		}
		for _, idx := range table.Indexes {
			if err := checkObjectDirectives(idx.Metadata, models.DirectiveLocationIndex, idx.Name); err != nil {
				return err
			}
		}
	}
	return nil
}

func checkObjectDirectives(meta map[string]any, location, owner string) error {
	for _, d := range models.DirectivesForNamespace(meta, directiveNamespace) {
		if !pgHandledDirectives[location][d.Key] {
			return fmt.Errorf("pgsql: %s: unsupported @postgres directive %q at %s level (strict mode)", owner, d.Key, location)
		}
	}
	return nil
}

// upperLeadingClause upcases a known leading keyword phrase in a directive
// argument so the emitted SQL reads conventionally. Identifiers that follow are
// left untouched.
func upperLeadingClause(args, lowerPrefix, upperPrefix string) string {
	args = strings.TrimSpace(args)
	if strings.HasPrefix(strings.ToLower(args), lowerPrefix) {
		return upperPrefix + args[len(lowerPrefix):]
	}
	return args
}

// pgTableDirectiveSuffix returns the clause appended after the closing ")" of a
// CREATE TABLE statement, e.g. " PARTITION BY RANGE (created_at) TABLESPACE fast".
func pgTableDirectiveSuffix(table *models.Table) string {
	directives := models.DirectivesForNamespace(table.Metadata, directiveNamespace)
	if len(directives) == 0 {
		return ""
	}

	byKey := firstByKey(directives)

	var parts []string
	if d, ok := byKey["partition"]; ok {
		parts = append(parts, upperLeadingClause(d.Args, "partition by", "PARTITION BY"))
	}
	if d, ok := byKey["inherits"]; ok {
		parts = append(parts, upperLeadingClause(d.Args, "inherits", "INHERITS"))
	}
	if d, ok := byKey["with"]; ok {
		parts = append(parts, upperLeadingClause(d.Args, "with", "WITH"))
	}
	if d, ok := byKey["tablespace"]; ok {
		parts = append(parts, upperLeadingClause(d.Args, "tablespace", "TABLESPACE"))
	}
	if len(parts) == 0 {
		return ""
	}
	return " " + strings.Join(parts, " ")
}

// pgColumnDirectiveSuffix returns the clause appended to a column definition,
// e.g. " STORAGE PLAIN" or " GENERATED ALWAYS AS IDENTITY".
func pgColumnDirectiveSuffix(col *models.Column) string {
	directives := models.DirectivesForNamespace(col.Metadata, directiveNamespace)
	if len(directives) == 0 {
		return ""
	}
	byKey := firstByKey(directives)

	var parts []string
	if d, ok := byKey["storage"]; ok {
		parts = append(parts, upperLeadingClause(d.Args, "storage", "STORAGE"))
	}
	if d, ok := byKey["compression"]; ok {
		parts = append(parts, upperLeadingClause(d.Args, "compression", "COMPRESSION"))
	}
	if d, ok := byKey["identity"]; ok {
		parts = append(parts, identityClause(d.Args))
	}
	if len(parts) == 0 {
		return ""
	}
	return " " + strings.Join(parts, " ")
}

// identityClause maps the two documented identity forms to standard SQL,
// falling back to a verbatim (upcased-keyword) rendering.
func identityClause(args string) string {
	switch strings.ToLower(strings.Join(strings.Fields(args), " ")) {
	case "identity always":
		return "GENERATED ALWAYS AS IDENTITY"
	case "identity default", "identity by default":
		return "GENERATED BY DEFAULT AS IDENTITY"
	default:
		return upperLeadingClause(args, "identity", "IDENTITY")
	}
}

// pgIndexDirectiveWith returns the parenthesised storage-parameter list from an
// @postgres: with (...) index directive, e.g. "fillfactor=90", or "".
func pgIndexDirectiveWith(index *models.Index) string {
	for _, d := range models.DirectivesForNamespace(index.Metadata, directiveNamespace) {
		if d.Key != "with" {
			continue
		}
		inner := d.Args
		if i := strings.Index(inner, "("); i >= 0 {
			if j := strings.LastIndex(inner, ")"); j > i {
				return strings.TrimSpace(inner[i+1 : j])
			}
		}
		return strings.TrimSpace(strings.TrimPrefix(strings.ToLower(inner), "with"))
	}
	return ""
}

// pgIndexWithParams returns the storage-parameter list to use for an index,
// preferring an @postgres: with (...) directive over the given fallback (e.g.
// one derived from the index comment).
func pgIndexWithParams(index *models.Index, fallback string) string {
	if p := pgIndexDirectiveWith(index); p != "" {
		return p
	}
	return fallback
}

// pgIndexDirectiveTablespace returns the tablespace name from an
// @postgres: tablespace <name> index directive, or "".
func pgIndexDirectiveTablespace(index *models.Index) string {
	for _, d := range models.DirectivesForNamespace(index.Metadata, directiveNamespace) {
		if d.Key == "tablespace" {
			return strings.TrimSpace(strings.TrimPrefix(strings.ToLower(d.Args), "tablespace"))
		}
	}
	return ""
}

// firstByKey indexes directives by key, keeping the first occurrence (the
// documented postgres keys used here are all singletons).
func firstByKey(directives []models.Directive) map[string]models.Directive {
	byKey := make(map[string]models.Directive, len(directives))
	for _, d := range directives {
		if _, exists := byKey[d.Key]; !exists {
			byKey[d.Key] = d
		}
	}
	return byKey
}