relspecgo/pkg/writers/sqlexec/writer.go

package sqlexec

import (
	"context"
	"fmt"
	"sort"

	"github.com/jackc/pgx/v5"

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

// Writer implements the writers.Writer interface for executing SQL scripts
type Writer struct {
	options *writers.WriterOptions
}

// NewWriter creates a new SQL executor writer
func NewWriter(options *writers.WriterOptions) *Writer {
	return &Writer{
		options: options,
	}
}

// Options returns the writer options (useful for reading execution results)
func (w *Writer) Options() *writers.WriterOptions {
	return w.options
}

// WriteDatabase executes all scripts from all schemas in the database
func (w *Writer) WriteDatabase(db *models.Database) error {
	if db == nil {
		return fmt.Errorf("database is nil")
	}

	// Get connection string from metadata
	connString, ok := w.options.Metadata["connection_string"].(string)
	if !ok || connString == "" {
		return fmt.Errorf("connection_string is required in writer metadata")
	}

	// Preserve schema order and script ordering within each schema, but aggregate once.
	var scripts []*models.Script
	for _, schema := range db.Schemas {
		for _, script := range sortedScripts(schema.Scripts) {
			copy := *script
			copy.Name = schema.Name + "." + script.Name
			scripts = append(scripts, &copy)
		}
	}

	// Connect to database
	ctx := context.Background()
	conn, err := pgsql.Connect(ctx, connString, "writer-sqlexec")
	if err != nil {
		return writers.FailExecution(w.options, "PostgreSQL scripts", scriptOperations(scripts), fmt.Errorf("failed to connect to database: %w", err))
	}
	defer conn.Close(ctx)

	return w.runScripts(ctx, conn, scripts)
}

// WriteSchema executes all scripts from a single schema
func (w *Writer) WriteSchema(schema *models.Schema) error {
	if schema == nil {
		return fmt.Errorf("schema is nil")
	}

	// Get connection string from metadata
	connString, ok := w.options.Metadata["connection_string"].(string)
	if !ok || connString == "" {
		return fmt.Errorf("connection_string is required in writer metadata")
	}

	// Connect to database
	ctx := context.Background()
	conn, err := pgsql.Connect(ctx, connString, "writer-sqlexec")
	if err != nil {
		return writers.FailExecution(w.options, "PostgreSQL scripts", scriptOperations(sortedScripts(schema.Scripts)), fmt.Errorf("failed to connect to database: %w", err))
	}
	defer conn.Close(ctx)

	// Execute scripts
	if err := w.executeScripts(ctx, conn, schema.Scripts); err != nil {
		return fmt.Errorf("failed to execute scripts: %w", err)
	}

	return nil
}

// WriteTable is not applicable for SQL script execution
func (w *Writer) WriteTable(table *models.Table) error {
	return fmt.Errorf("WriteTable is not supported for SQL script execution")
}

// executeScripts executes scripts in Priority, Sequence, then Name order.
func (w *Writer) executeScripts(ctx context.Context, conn *pgx.Conn, scripts []*models.Script) error {
	return w.runScripts(ctx, conn, sortedScripts(scripts))
}

func sortedScripts(scripts []*models.Script) []*models.Script {
	ordered := append([]*models.Script(nil), scripts...)
	sort.SliceStable(ordered, func(i, j int) bool {
		if ordered[i].Priority != ordered[j].Priority {
			return ordered[i].Priority < ordered[j].Priority
		}
		if ordered[i].Sequence != ordered[j].Sequence {
			return ordered[i].Sequence < ordered[j].Sequence
		}
		return ordered[i].Name < ordered[j].Name
	})
	return ordered
}

func (w *Writer) runScripts(ctx context.Context, conn *pgx.Conn, scripts []*models.Script) error {
	operations := scriptOperations(scripts)
	backend := writers.ExecutionBackend{ScriptManaged: true, Exec: func(ctx context.Context, i int, _ string) error {
		script := scripts[i]
		fmt.Printf("Executing script: %s (Priority=%d, Sequence=%d)\n", script.Name, script.Priority, script.Sequence)
		assetRoot, _ := w.options.Metadata["asset_root"].(string)
		sql, err := processEmbedDirectives(script, assetRoot)
		if err != nil {
			return fmt.Errorf("preparing script %s: %w", script.Name, err)
		}
		if _, err := conn.Exec(ctx, sql); err != nil {
			return fmt.Errorf("script %s: %w", script.Name, err)
		}
		return nil
	}}
	// Supplied scripts can manage transactions themselves. Do not promise whole-script rollback.
	return writers.RunExecution(ctx, w.options, &writers.ExecutionReport{Dialect: "PostgreSQL scripts"}, operations, backend)
}

func scriptOperations(scripts []*models.Script) []writers.ExecutionOperation {
	operations := make([]writers.ExecutionOperation, len(scripts))
	for i, script := range scripts {
		operations[i] = writers.ExecutionOperation{Name: script.Name, Statement: script.SQL}
	}
	return operations
}

func processEmbedDirectives(script *models.Script, assetRoot string) (string, error) {
	if !assetloader.HasEmbedDirectives(script.SQL) {
		return script.SQL, nil
	}
	if assetRoot == "" {
		return "", fmt.Errorf("asset_root is required to resolve embed directives from a model; source_path metadata does not authorize file access")
	}
	sqlPath, _ := script.Metadata[assetloader.ScriptSourcePathMetadataKey].(string)
	if sqlPath == "" {
		return "", fmt.Errorf("source_path is required for model embed directives")
	}
	return assetloader.ProcessEmbedDirectivesWithOptions(sqlPath, script.SQL, assetloader.EmbedOptions{AssetRoot: assetRoot})
}