relspecgo/pkg/writers/sqlite/writer.go

package sqlite

import (
	"context"
	"database/sql"
	"fmt"
	"io"
	"os"
	"strings"

	_ "modernc.org/sqlite" // SQLite driver

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

// Writer implements the Writer interface for SQLite SQL output
type Writer struct {
	options  *writers.WriterOptions
	writer   io.Writer
	executor *TemplateExecutor
}

// NewWriter creates a new SQLite SQL writer
// SQLite doesn't support schemas, so FlattenSchema is automatically enabled
func NewWriter(options *writers.WriterOptions) *Writer {
	// Force schema flattening for SQLite
	options.FlattenSchema = true

	executor, _ := NewTemplateExecutor(options)
	return &Writer{
		options:  options,
		executor: executor,
	}
}

// WriteDatabase writes the entire database schema as SQLite SQL.
//
// If Metadata["connection_string"] is set (a path to a SQLite database file),
// the generated DDL is executed directly against that file instead of being
// written out as a .sql script.
func (w *Writer) WriteDatabase(db *models.Database) error {
	if dbPath, ok := w.options.Metadata["connection_string"].(string); ok && dbPath != "" {
		return w.executeDatabaseSQL(db, dbPath)
	}

	var writer io.Writer
	var file *os.File
	var err error

	// Use existing writer if already set (for testing)
	if w.writer != nil {
		writer = w.writer
	} else if w.options.OutputPath != "" {
		// Determine output destination
		file, err = os.Create(w.options.OutputPath)
		if err != nil {
			return fmt.Errorf("failed to create output file: %w", err)
		}
		defer file.Close()
		writer = file
	} else {
		writer = os.Stdout
	}

	w.writer = writer
	return w.writeContent(db)
}

// writeContent writes the header, pragma, and every schema's DDL to w.writer.
func (w *Writer) writeContent(db *models.Database) error {
	// Write header comment
	fmt.Fprintf(w.writer, "-- SQLite Database Schema\n")
	fmt.Fprintf(w.writer, "-- Database: %s\n", db.Name)
	fmt.Fprintf(w.writer, "-- Generated by RelSpec\n")
	fmt.Fprintf(w.writer, "-- Note: SQLite has no schema concept; non-default schema names are flattened into table name prefixes (e.g., auth.sessions -> auth_sessions)\n\n")

	// Enable foreign keys
	pragma, err := w.executor.ExecutePragmaForeignKeys()
	if err != nil {
		return fmt.Errorf("failed to generate pragma statement: %w", err)
	}
	fmt.Fprintf(w.writer, "%s\n", pragma)

	// Process each schema in the database
	for _, schema := range db.Schemas {
		if err := w.WriteSchema(schema); err != nil {
			return fmt.Errorf("failed to write schema %s: %w", schema.Name, err)
		}
	}

	return nil
}

// statementCollector captures each Write call as a single SQL statement (or
// comment line), matching the writer's convention of one Fprintf per statement.
type statementCollector struct {
	statements []string
}

func (c *statementCollector) Write(p []byte) (int, error) {
	if s := strings.TrimSpace(string(p)); s != "" {
		c.statements = append(c.statements, s)
	}
	return len(p), nil
}

// executeDatabaseSQL generates the DDL for db and executes it directly
// against the SQLite database file at dbPath.
func (w *Writer) executeDatabaseSQL(db *models.Database, dbPath string) error {
	collector := &statementCollector{}
	w.writer = collector
	if err := w.writeContent(db); err != nil {
		return fmt.Errorf("failed to generate SQL statements: %w", err)
	}

	conn, err := sql.Open("sqlite", dbPath)
	if err != nil {
		return fmt.Errorf("failed to open sqlite database %q: %w", dbPath, err)
	}
	defer conn.Close()

	ctx := context.Background()
	ignoreErrors := false
	if val, ok := w.options.Metadata["ignore_errors"].(bool); ok {
		ignoreErrors = val
	}

	total, executed := 0, 0
	var execErrors []string
	for _, stmt := range collector.statements {
		if strings.HasPrefix(stmt, "--") {
			continue
		}

		total++
		if _, err := conn.ExecContext(ctx, stmt); err != nil {
			execErrors = append(execErrors, fmt.Sprintf("statement %d (%s): %v", total, truncateStatement(stmt), err))
			if !ignoreErrors {
				break
			}
			continue
		}
		executed++
	}

	w.options.Metadata["execution_total"] = total
	w.options.Metadata["execution_success"] = executed
	w.options.Metadata["execution_failed"] = len(execErrors)

	if len(execErrors) > 0 {
		return fmt.Errorf("failed to execute %d/%d statement(s) against %q:\n%s", len(execErrors), total, dbPath, strings.Join(execErrors, "\n"))
	}

	return nil
}

// truncateStatement shortens a SQL statement for error messages.
func truncateStatement(stmt string) string {
	const maxLen = 80
	stmt = strings.Join(strings.Fields(stmt), " ")
	if len(stmt) > maxLen {
		return stmt[:maxLen] + "..."
	}
	return stmt
}

// defaultSchemaNames are treated as "no schema" for SQLite output: SQLite has
// no schema concept, and a lone default schema (e.g. DBML's implicit "public")
// should produce bare table names rather than a "public_" prefix.
var defaultSchemaNames = map[string]bool{
	"public": true,
	"main":   true,
}

// tableSchemaName returns the schema name to use for table/constraint naming,
// collapsing default schema names to "" so they aren't prefixed onto table names.
func tableSchemaName(schema string) string {
	if defaultSchemaNames[strings.ToLower(schema)] {
		return ""
	}
	return schema
}

// WriteSchema writes a single schema as SQLite SQL
func (w *Writer) WriteSchema(schema *models.Schema) error {
	tableSchema := tableSchemaName(schema.Name)

	// SQLite doesn't have schemas, so we just write a comment (skip for the
	// default schema, since its tables aren't actually being prefixed)
	if tableSchema != "" {
		fmt.Fprintf(w.writer, "-- Schema: %s (flattened into table names)\n\n", schema.Name)
	}

	// Phase 1: Create tables
	for _, table := range schema.Tables {
		if err := w.writeTable(tableSchema, table); err != nil {
			return fmt.Errorf("failed to write table %s: %w", table.Name, err)
		}
	}

	// Phase 2: Create indexes
	for _, table := range schema.Tables {
		if err := w.writeIndexes(tableSchema, table); err != nil {
			return fmt.Errorf("failed to write indexes for table %s: %w", table.Name, err)
		}
	}

	// Phase 3: Create unique constraints (as unique indexes)
	for _, table := range schema.Tables {
		if err := w.writeUniqueConstraints(tableSchema, table); err != nil {
			return fmt.Errorf("failed to write unique constraints for table %s: %w", table.Name, err)
		}
	}

	// Phase 4: Check constraints (as comments, since SQLite requires them in CREATE TABLE)
	for _, table := range schema.Tables {
		if err := w.writeCheckConstraints(tableSchema, table); err != nil {
			return fmt.Errorf("failed to write check constraints for table %s: %w", table.Name, err)
		}
	}

	return nil
}

// WriteTable writes a single table as SQLite SQL
func (w *Writer) WriteTable(table *models.Table) error {
	return w.writeTable("", table)
}

// writeTable is the internal implementation
func (w *Writer) writeTable(schema string, table *models.Table) error {
	// Build table template data
	data := BuildTableTemplateData(schema, table)

	// Execute template
	sql, err := w.executor.ExecuteCreateTable(data)
	if err != nil {
		return fmt.Errorf("failed to execute create table template: %w", err)
	}

	fmt.Fprintf(w.writer, "%s\n", sql)
	return nil
}

// writeIndexes writes indexes for a table
func (w *Writer) writeIndexes(schema string, table *models.Table) error {
	for _, index := range sortIndexes(table.Indexes) {
		// Skip primary key indexes
		if strings.HasSuffix(index.Name, "_pkey") {
			continue
		}

		// Skip unique indexes (handled separately as unique constraints)
		if index.Unique {
			continue
		}

		data := IndexTemplateData{
			Schema:  schema,
			Table:   table.Name,
			Name:    index.Name,
			Columns: index.Columns,
		}

		sql, err := w.executor.ExecuteCreateIndex(data)
		if err != nil {
			return fmt.Errorf("failed to execute create index template: %w", err)
		}

		fmt.Fprintf(w.writer, "%s\n", sql)
	}

	return nil
}

// writeUniqueConstraints writes unique constraints as unique indexes
func (w *Writer) writeUniqueConstraints(schema string, table *models.Table) error {
	for _, constraint := range sortConstraints(table.Constraints) {
		if constraint.Type != models.UniqueConstraint {
			continue
		}

		data := ConstraintTemplateData{
			Schema:  schema,
			Table:   table.Name,
			Name:    constraint.Name,
			Columns: constraint.Columns,
		}

		sql, err := w.executor.ExecuteCreateUniqueConstraint(data)
		if err != nil {
			return fmt.Errorf("failed to execute create unique constraint template: %w", err)
		}

		fmt.Fprintf(w.writer, "%s\n", sql)
	}

	// Also handle unique indexes from the Indexes map
	for _, index := range sortIndexes(table.Indexes) {
		if !index.Unique {
			continue
		}

		// Skip if already handled as a constraint
		alreadyHandled := false
		for _, constraint := range table.Constraints {
			if constraint.Type == models.UniqueConstraint && constraint.Name == index.Name {
				alreadyHandled = true
				break
			}
		}
		if alreadyHandled {
			continue
		}

		data := ConstraintTemplateData{
			Schema:  schema,
			Table:   table.Name,
			Name:    index.Name,
			Columns: index.Columns,
		}

		sql, err := w.executor.ExecuteCreateUniqueConstraint(data)
		if err != nil {
			return fmt.Errorf("failed to execute create unique index template: %w", err)
		}

		fmt.Fprintf(w.writer, "%s\n", sql)
	}

	return nil
}

// writeCheckConstraints writes check constraints as comments
func (w *Writer) writeCheckConstraints(schema string, table *models.Table) error {
	for _, constraint := range sortConstraints(table.Constraints) {
		if constraint.Type != models.CheckConstraint {
			continue
		}

		data := ConstraintTemplateData{
			Schema:     schema,
			Table:      table.Name,
			Name:       constraint.Name,
			Expression: constraint.Expression,
		}

		sql, err := w.executor.ExecuteCreateCheckConstraint(data)
		if err != nil {
			return fmt.Errorf("failed to execute create check constraint template: %w", err)
		}

		fmt.Fprintf(w.writer, "%s\n", sql)
	}

	return nil
}