From 165623bb1dfeb268087858f2574d501a0542910d Mon Sep 17 00:00:00 2001 From: Hein Date: Sat, 31 Jan 2026 21:04:43 +0200 Subject: [PATCH] =?UTF-8?q?feat(pgsql):=20=E2=9C=A8=20Add=20templates=20fo?= =?UTF-8?q?r=20constraints=20and=20sequences?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Introduce new templates for creating unique, check, and foreign key constraints with existence checks. * Add templates for setting sequence values and creating sequences. * Refactor existing SQL generation logic to utilize new templates for better maintainability and readability. * Ensure identifiers are properly quoted to handle special characters and reserved keywords. --- pkg/writers/pgsql/migration_writer_test.go | 67 +++ pkg/writers/pgsql/template_functions.go | 38 ++ pkg/writers/pgsql/template_functions_test.go | 27 +- pkg/writers/pgsql/templates.go | 136 ++++++ pkg/writers/pgsql/templates/add_column.tmpl | 4 +- .../templates/add_column_with_check.tmpl | 12 + .../pgsql/templates/alter_column_default.tmpl | 8 +- .../pgsql/templates/alter_column_type.tmpl | 4 +- .../pgsql/templates/comment_column.tmpl | 2 +- .../pgsql/templates/comment_table.tmpl | 2 +- .../templates/create_check_constraint.tmpl | 12 + .../pgsql/templates/create_foreign_key.tmpl | 10 +- .../create_foreign_key_with_check.tmpl | 18 + pkg/writers/pgsql/templates/create_index.tmpl | 4 +- .../pgsql/templates/create_primary_key.tmpl | 4 +- ...create_primary_key_with_autogen_check.tmpl | 27 ++ .../pgsql/templates/create_sequence.tmpl | 6 + pkg/writers/pgsql/templates/create_table.tmpl | 4 +- .../templates/create_unique_constraint.tmpl | 12 + .../pgsql/templates/drop_constraint.tmpl | 2 +- pkg/writers/pgsql/templates/drop_index.tmpl | 2 +- .../pgsql/templates/set_sequence_value.tmpl | 19 + pkg/writers/pgsql/writer.go | 414 ++++++++---------- 23 files changed, 588 insertions(+), 246 deletions(-) create mode 100644 pkg/writers/pgsql/templates/add_column_with_check.tmpl create mode 100644 pkg/writers/pgsql/templates/create_check_constraint.tmpl create mode 100644 pkg/writers/pgsql/templates/create_foreign_key_with_check.tmpl create mode 100644 pkg/writers/pgsql/templates/create_primary_key_with_autogen_check.tmpl create mode 100644 pkg/writers/pgsql/templates/create_sequence.tmpl create mode 100644 pkg/writers/pgsql/templates/create_unique_constraint.tmpl create mode 100644 pkg/writers/pgsql/templates/set_sequence_value.tmpl diff --git a/pkg/writers/pgsql/migration_writer_test.go b/pkg/writers/pgsql/migration_writer_test.go index 093628e..51d8f76 100644 --- a/pkg/writers/pgsql/migration_writer_test.go +++ b/pkg/writers/pgsql/migration_writer_test.go @@ -215,3 +215,70 @@ func TestTemplateExecutor_AuditFunction(t *testing.T) { t.Error("SQL missing DELETE handling") } } + +func TestWriteMigration_NumericConstraintNames(t *testing.T) { + // Current database (empty) + current := models.InitDatabase("testdb") + currentSchema := models.InitSchema("entity") + current.Schemas = append(current.Schemas, currentSchema) + + // Model database (with constraint starting with number) + model := models.InitDatabase("testdb") + modelSchema := models.InitSchema("entity") + + // Create individual_actor_relationship table + table := models.InitTable("individual_actor_relationship", "entity") + idCol := models.InitColumn("id", "individual_actor_relationship", "entity") + idCol.Type = "integer" + idCol.IsPrimaryKey = true + table.Columns["id"] = idCol + + actorIDCol := models.InitColumn("actor_id", "individual_actor_relationship", "entity") + actorIDCol.Type = "integer" + table.Columns["actor_id"] = actorIDCol + + // Add constraint with name starting with number + constraint := &models.Constraint{ + Name: "215162_fk_actor", + Type: models.ForeignKeyConstraint, + Columns: []string{"actor_id"}, + ReferencedSchema: "entity", + ReferencedTable: "actor", + ReferencedColumns: []string{"id"}, + OnDelete: "CASCADE", + OnUpdate: "NO ACTION", + } + table.Constraints["215162_fk_actor"] = constraint + + modelSchema.Tables = append(modelSchema.Tables, table) + model.Schemas = append(model.Schemas, modelSchema) + + // Generate migration + var buf bytes.Buffer + writer, err := NewMigrationWriter(&writers.WriterOptions{}) + if err != nil { + t.Fatalf("Failed to create writer: %v", err) + } + writer.writer = &buf + + err = writer.WriteMigration(model, current) + if err != nil { + t.Fatalf("WriteMigration failed: %v", err) + } + + output := buf.String() + t.Logf("Generated migration:\n%s", output) + + // Verify constraint name is properly quoted + if !strings.Contains(output, `"215162_fk_actor"`) { + t.Error("Constraint name starting with number should be quoted") + } + + // Verify the SQL is syntactically correct (contains required keywords) + if !strings.Contains(output, "ADD CONSTRAINT") { + t.Error("Migration missing ADD CONSTRAINT") + } + if !strings.Contains(output, "FOREIGN KEY") { + t.Error("Migration missing FOREIGN KEY") + } +} diff --git a/pkg/writers/pgsql/template_functions.go b/pkg/writers/pgsql/template_functions.go index 419ecf1..a6b0638 100644 --- a/pkg/writers/pgsql/template_functions.go +++ b/pkg/writers/pgsql/template_functions.go @@ -21,6 +21,7 @@ func TemplateFunctions() map[string]interface{} { "quote": quote, "escape": escape, "safe_identifier": safeIdentifier, + "quote_ident": quoteIdent, // Type conversion "goTypeToSQL": goTypeToSQL, @@ -122,6 +123,43 @@ func safeIdentifier(s string) string { return strings.ToLower(safe) } +// quoteIdent quotes a PostgreSQL identifier if necessary +// Identifiers need quoting if they: +// - Start with a digit +// - Contain special characters +// - Are reserved keywords +// - Contain uppercase letters (to preserve case) +func quoteIdent(s string) string { + if s == "" { + return `""` + } + + // Check if quoting is needed + needsQuoting := unicode.IsDigit(rune(s[0])) + + // Starts with digit + + // Contains uppercase letters or special characters + for _, r := range s { + if unicode.IsUpper(r) { + needsQuoting = true + break + } + if !unicode.IsLetter(r) && !unicode.IsDigit(r) && r != '_' { + needsQuoting = true + break + } + } + + if needsQuoting { + // Escape double quotes by doubling them + escaped := strings.ReplaceAll(s, `"`, `""`) + return `"` + escaped + `"` + } + + return s +} + // Type conversion functions // goTypeToSQL converts Go type to PostgreSQL type diff --git a/pkg/writers/pgsql/template_functions_test.go b/pkg/writers/pgsql/template_functions_test.go index 73d5983..d1034ad 100644 --- a/pkg/writers/pgsql/template_functions_test.go +++ b/pkg/writers/pgsql/template_functions_test.go @@ -101,6 +101,31 @@ func TestSafeIdentifier(t *testing.T) { } } +func TestQuoteIdent(t *testing.T) { + tests := []struct { + input string + expected string + }{ + {"valid_name", "valid_name"}, + {"ValidName", `"ValidName"`}, + {"123column", `"123column"`}, + {"215162_fk_constraint", `"215162_fk_constraint"`}, + {"user-id", `"user-id"`}, + {"user@domain", `"user@domain"`}, + {`"quoted"`, `"""quoted"""`}, + {"", `""`}, + {"lowercase", "lowercase"}, + {"with_underscore", "with_underscore"}, + } + + for _, tt := range tests { + result := quoteIdent(tt.input) + if result != tt.expected { + t.Errorf("quoteIdent(%q) = %q, want %q", tt.input, result, tt.expected) + } + } +} + func TestGoTypeToSQL(t *testing.T) { tests := []struct { input string @@ -243,7 +268,7 @@ func TestTemplateFunctions(t *testing.T) { // Check that all expected functions are registered expectedFuncs := []string{ "upper", "lower", "snake_case", "camelCase", - "indent", "quote", "escape", "safe_identifier", + "indent", "quote", "escape", "safe_identifier", "quote_ident", "goTypeToSQL", "sqlTypeToGo", "isNumeric", "isText", "first", "last", "filter", "mapFunc", "join_with", "join", diff --git a/pkg/writers/pgsql/templates.go b/pkg/writers/pgsql/templates.go index 63793bd..eb73f3a 100644 --- a/pkg/writers/pgsql/templates.go +++ b/pkg/writers/pgsql/templates.go @@ -177,6 +177,72 @@ type AuditTriggerData struct { Events string } +// CreateUniqueConstraintData contains data for create unique constraint template +type CreateUniqueConstraintData struct { + SchemaName string + TableName string + ConstraintName string + Columns string +} + +// CreateCheckConstraintData contains data for create check constraint template +type CreateCheckConstraintData struct { + SchemaName string + TableName string + ConstraintName string + Expression string +} + +// CreateForeignKeyWithCheckData contains data for create foreign key with existence check template +type CreateForeignKeyWithCheckData struct { + SchemaName string + TableName string + ConstraintName string + SourceColumns string + TargetSchema string + TargetTable string + TargetColumns string + OnDelete string + OnUpdate string + Deferrable bool +} + +// SetSequenceValueData contains data for set sequence value template +type SetSequenceValueData struct { + SchemaName string + TableName string + SequenceName string + ColumnName string +} + +// CreateSequenceData contains data for create sequence template +type CreateSequenceData struct { + SchemaName string + SequenceName string + Increment int + MinValue int64 + MaxValue int64 + StartValue int64 + CacheSize int +} + +// AddColumnWithCheckData contains data for add column with existence check template +type AddColumnWithCheckData struct { + SchemaName string + TableName string + ColumnName string + ColumnDefinition string +} + +// CreatePrimaryKeyWithAutoGenCheckData contains data for primary key with auto-generated key check template +type CreatePrimaryKeyWithAutoGenCheckData struct { + SchemaName string + TableName string + ConstraintName string + AutoGenNames string // Comma-separated list of names like "'name1', 'name2'" + Columns string +} + // Execute methods for each template // ExecuteCreateTable executes the create table template @@ -319,6 +385,76 @@ func (te *TemplateExecutor) ExecuteAuditTrigger(data AuditTriggerData) (string, return buf.String(), nil } +// ExecuteCreateUniqueConstraint executes the create unique constraint template +func (te *TemplateExecutor) ExecuteCreateUniqueConstraint(data CreateUniqueConstraintData) (string, error) { + var buf bytes.Buffer + err := te.templates.ExecuteTemplate(&buf, "create_unique_constraint.tmpl", data) + if err != nil { + return "", fmt.Errorf("failed to execute create_unique_constraint template: %w", err) + } + return buf.String(), nil +} + +// ExecuteCreateCheckConstraint executes the create check constraint template +func (te *TemplateExecutor) ExecuteCreateCheckConstraint(data CreateCheckConstraintData) (string, error) { + var buf bytes.Buffer + err := te.templates.ExecuteTemplate(&buf, "create_check_constraint.tmpl", data) + if err != nil { + return "", fmt.Errorf("failed to execute create_check_constraint template: %w", err) + } + return buf.String(), nil +} + +// ExecuteCreateForeignKeyWithCheck executes the create foreign key with check template +func (te *TemplateExecutor) ExecuteCreateForeignKeyWithCheck(data CreateForeignKeyWithCheckData) (string, error) { + var buf bytes.Buffer + err := te.templates.ExecuteTemplate(&buf, "create_foreign_key_with_check.tmpl", data) + if err != nil { + return "", fmt.Errorf("failed to execute create_foreign_key_with_check template: %w", err) + } + return buf.String(), nil +} + +// ExecuteSetSequenceValue executes the set sequence value template +func (te *TemplateExecutor) ExecuteSetSequenceValue(data SetSequenceValueData) (string, error) { + var buf bytes.Buffer + err := te.templates.ExecuteTemplate(&buf, "set_sequence_value.tmpl", data) + if err != nil { + return "", fmt.Errorf("failed to execute set_sequence_value template: %w", err) + } + return buf.String(), nil +} + +// ExecuteCreateSequence executes the create sequence template +func (te *TemplateExecutor) ExecuteCreateSequence(data CreateSequenceData) (string, error) { + var buf bytes.Buffer + err := te.templates.ExecuteTemplate(&buf, "create_sequence.tmpl", data) + if err != nil { + return "", fmt.Errorf("failed to execute create_sequence template: %w", err) + } + return buf.String(), nil +} + +// ExecuteAddColumnWithCheck executes the add column with check template +func (te *TemplateExecutor) ExecuteAddColumnWithCheck(data AddColumnWithCheckData) (string, error) { + var buf bytes.Buffer + err := te.templates.ExecuteTemplate(&buf, "add_column_with_check.tmpl", data) + if err != nil { + return "", fmt.Errorf("failed to execute add_column_with_check template: %w", err) + } + return buf.String(), nil +} + +// ExecuteCreatePrimaryKeyWithAutoGenCheck executes the create primary key with auto-generated key check template +func (te *TemplateExecutor) ExecuteCreatePrimaryKeyWithAutoGenCheck(data CreatePrimaryKeyWithAutoGenCheckData) (string, error) { + var buf bytes.Buffer + err := te.templates.ExecuteTemplate(&buf, "create_primary_key_with_autogen_check.tmpl", data) + if err != nil { + return "", fmt.Errorf("failed to execute create_primary_key_with_autogen_check template: %w", err) + } + return buf.String(), nil +} + // Helper functions to build template data from models // BuildCreateTableData builds CreateTableData from a models.Table diff --git a/pkg/writers/pgsql/templates/add_column.tmpl b/pkg/writers/pgsql/templates/add_column.tmpl index dd9a3ad..f52f803 100644 --- a/pkg/writers/pgsql/templates/add_column.tmpl +++ b/pkg/writers/pgsql/templates/add_column.tmpl @@ -1,4 +1,4 @@ -ALTER TABLE {{.SchemaName}}.{{.TableName}} - ADD COLUMN IF NOT EXISTS {{.ColumnName}} {{.ColumnType}} +ALTER TABLE {{quote_ident .SchemaName}}.{{quote_ident .TableName}} + ADD COLUMN IF NOT EXISTS {{quote_ident .ColumnName}} {{.ColumnType}} {{- if .Default}} DEFAULT {{.Default}}{{end}} {{- if .NotNull}} NOT NULL{{end}}; \ No newline at end of file diff --git a/pkg/writers/pgsql/templates/add_column_with_check.tmpl b/pkg/writers/pgsql/templates/add_column_with_check.tmpl new file mode 100644 index 0000000..0d86d07 --- /dev/null +++ b/pkg/writers/pgsql/templates/add_column_with_check.tmpl @@ -0,0 +1,12 @@ +DO $$ +BEGIN + IF NOT EXISTS ( + SELECT 1 FROM information_schema.columns + WHERE table_schema = '{{.SchemaName}}' + AND table_name = '{{.TableName}}' + AND column_name = '{{.ColumnName}}' + ) THEN + ALTER TABLE {{quote_ident .SchemaName}}.{{quote_ident .TableName}} ADD COLUMN {{.ColumnDefinition}}; + END IF; +END; +$$; diff --git a/pkg/writers/pgsql/templates/alter_column_default.tmpl b/pkg/writers/pgsql/templates/alter_column_default.tmpl index c38093a..0543a06 100644 --- a/pkg/writers/pgsql/templates/alter_column_default.tmpl +++ b/pkg/writers/pgsql/templates/alter_column_default.tmpl @@ -1,7 +1,7 @@ {{- if .SetDefault -}} -ALTER TABLE {{.SchemaName}}.{{.TableName}} - ALTER COLUMN {{.ColumnName}} SET DEFAULT {{.DefaultValue}}; +ALTER TABLE {{quote_ident .SchemaName}}.{{quote_ident .TableName}} + ALTER COLUMN {{quote_ident .ColumnName}} SET DEFAULT {{.DefaultValue}}; {{- else -}} -ALTER TABLE {{.SchemaName}}.{{.TableName}} - ALTER COLUMN {{.ColumnName}} DROP DEFAULT; +ALTER TABLE {{quote_ident .SchemaName}}.{{quote_ident .TableName}} + ALTER COLUMN {{quote_ident .ColumnName}} DROP DEFAULT; {{- end -}} \ No newline at end of file diff --git a/pkg/writers/pgsql/templates/alter_column_type.tmpl b/pkg/writers/pgsql/templates/alter_column_type.tmpl index d6387a7..02d9168 100644 --- a/pkg/writers/pgsql/templates/alter_column_type.tmpl +++ b/pkg/writers/pgsql/templates/alter_column_type.tmpl @@ -1,2 +1,2 @@ -ALTER TABLE {{.SchemaName}}.{{.TableName}} - ALTER COLUMN {{.ColumnName}} TYPE {{.NewType}}; \ No newline at end of file +ALTER TABLE {{quote_ident .SchemaName}}.{{quote_ident .TableName}} + ALTER COLUMN {{quote_ident .ColumnName}} TYPE {{.NewType}}; \ No newline at end of file diff --git a/pkg/writers/pgsql/templates/comment_column.tmpl b/pkg/writers/pgsql/templates/comment_column.tmpl index be7b5c4..9bfcef9 100644 --- a/pkg/writers/pgsql/templates/comment_column.tmpl +++ b/pkg/writers/pgsql/templates/comment_column.tmpl @@ -1 +1 @@ -COMMENT ON COLUMN {{.SchemaName}}.{{.TableName}}.{{.ColumnName}} IS '{{.Comment}}'; \ No newline at end of file +COMMENT ON COLUMN {{quote_ident .SchemaName}}.{{quote_ident .TableName}}.{{quote_ident .ColumnName}} IS '{{.Comment}}'; \ No newline at end of file diff --git a/pkg/writers/pgsql/templates/comment_table.tmpl b/pkg/writers/pgsql/templates/comment_table.tmpl index 2bc8dc9..43a5b8b 100644 --- a/pkg/writers/pgsql/templates/comment_table.tmpl +++ b/pkg/writers/pgsql/templates/comment_table.tmpl @@ -1 +1 @@ -COMMENT ON TABLE {{.SchemaName}}.{{.TableName}} IS '{{.Comment}}'; \ No newline at end of file +COMMENT ON TABLE {{quote_ident .SchemaName}}.{{quote_ident .TableName}} IS '{{.Comment}}'; \ No newline at end of file diff --git a/pkg/writers/pgsql/templates/create_check_constraint.tmpl b/pkg/writers/pgsql/templates/create_check_constraint.tmpl new file mode 100644 index 0000000..abc15d3 --- /dev/null +++ b/pkg/writers/pgsql/templates/create_check_constraint.tmpl @@ -0,0 +1,12 @@ +DO $$ +BEGIN + IF NOT EXISTS ( + SELECT 1 FROM information_schema.table_constraints + WHERE table_schema = '{{.SchemaName}}' + AND table_name = '{{.TableName}}' + AND constraint_name = '{{.ConstraintName}}' + ) THEN + ALTER TABLE {{quote_ident .SchemaName}}.{{quote_ident .TableName}} ADD CONSTRAINT {{quote_ident .ConstraintName}} CHECK ({{.Expression}}); + END IF; +END; +$$; \ No newline at end of file diff --git a/pkg/writers/pgsql/templates/create_foreign_key.tmpl b/pkg/writers/pgsql/templates/create_foreign_key.tmpl index e6583d2..06bc250 100644 --- a/pkg/writers/pgsql/templates/create_foreign_key.tmpl +++ b/pkg/writers/pgsql/templates/create_foreign_key.tmpl @@ -1,10 +1,10 @@ -ALTER TABLE {{.SchemaName}}.{{.TableName}} - DROP CONSTRAINT IF EXISTS {{.ConstraintName}}; +ALTER TABLE {{quote_ident .SchemaName}}.{{quote_ident .TableName}} + DROP CONSTRAINT IF EXISTS {{quote_ident .ConstraintName}}; -ALTER TABLE {{.SchemaName}}.{{.TableName}} - ADD CONSTRAINT {{.ConstraintName}} +ALTER TABLE {{quote_ident .SchemaName}}.{{quote_ident .TableName}} + ADD CONSTRAINT {{quote_ident .ConstraintName}} FOREIGN KEY ({{.SourceColumns}}) - REFERENCES {{.TargetSchema}}.{{.TargetTable}} ({{.TargetColumns}}) + REFERENCES {{quote_ident .TargetSchema}}.{{quote_ident .TargetTable}} ({{.TargetColumns}}) ON DELETE {{.OnDelete}} ON UPDATE {{.OnUpdate}} DEFERRABLE; \ No newline at end of file diff --git a/pkg/writers/pgsql/templates/create_foreign_key_with_check.tmpl b/pkg/writers/pgsql/templates/create_foreign_key_with_check.tmpl new file mode 100644 index 0000000..d3fde82 --- /dev/null +++ b/pkg/writers/pgsql/templates/create_foreign_key_with_check.tmpl @@ -0,0 +1,18 @@ +DO $$ +BEGIN + IF NOT EXISTS ( + SELECT 1 FROM information_schema.table_constraints + WHERE table_schema = '{{.SchemaName}}' + AND table_name = '{{.TableName}}' + AND constraint_name = '{{.ConstraintName}}' + ) THEN + ALTER TABLE {{quote_ident .SchemaName}}.{{quote_ident .TableName}} + ADD CONSTRAINT {{quote_ident .ConstraintName}} + FOREIGN KEY ({{.SourceColumns}}) + REFERENCES {{quote_ident .TargetSchema}}.{{quote_ident .TargetTable}} ({{.TargetColumns}}) + ON DELETE {{.OnDelete}} + ON UPDATE {{.OnUpdate}}{{if .Deferrable}} + DEFERRABLE{{end}}; + END IF; +END; +$$; \ No newline at end of file diff --git a/pkg/writers/pgsql/templates/create_index.tmpl b/pkg/writers/pgsql/templates/create_index.tmpl index 44937eb..8cb55fb 100644 --- a/pkg/writers/pgsql/templates/create_index.tmpl +++ b/pkg/writers/pgsql/templates/create_index.tmpl @@ -1,2 +1,2 @@ -CREATE {{if .Unique}}UNIQUE {{end}}INDEX IF NOT EXISTS {{.IndexName}} - ON {{.SchemaName}}.{{.TableName}} USING {{.IndexType}} ({{.Columns}}); \ No newline at end of file +CREATE {{if .Unique}}UNIQUE {{end}}INDEX IF NOT EXISTS {{quote_ident .IndexName}} + ON {{quote_ident .SchemaName}}.{{quote_ident .TableName}} USING {{.IndexType}} ({{.Columns}}); \ No newline at end of file diff --git a/pkg/writers/pgsql/templates/create_primary_key.tmpl b/pkg/writers/pgsql/templates/create_primary_key.tmpl index 615165e..bc5fec0 100644 --- a/pkg/writers/pgsql/templates/create_primary_key.tmpl +++ b/pkg/writers/pgsql/templates/create_primary_key.tmpl @@ -6,8 +6,8 @@ BEGIN AND table_name = '{{.TableName}}' AND constraint_name = '{{.ConstraintName}}' ) THEN - ALTER TABLE {{.SchemaName}}.{{.TableName}} - ADD CONSTRAINT {{.ConstraintName}} PRIMARY KEY ({{.Columns}}); + ALTER TABLE {{quote_ident .SchemaName}}.{{quote_ident .TableName}} + ADD CONSTRAINT {{quote_ident .ConstraintName}} PRIMARY KEY ({{.Columns}}); END IF; END; $$; \ No newline at end of file diff --git a/pkg/writers/pgsql/templates/create_primary_key_with_autogen_check.tmpl b/pkg/writers/pgsql/templates/create_primary_key_with_autogen_check.tmpl new file mode 100644 index 0000000..6789a80 --- /dev/null +++ b/pkg/writers/pgsql/templates/create_primary_key_with_autogen_check.tmpl @@ -0,0 +1,27 @@ +DO $$ +DECLARE + auto_pk_name text; +BEGIN + -- Drop auto-generated primary key if it exists + SELECT constraint_name INTO auto_pk_name + FROM information_schema.table_constraints + WHERE table_schema = '{{.SchemaName}}' + AND table_name = '{{.TableName}}' + AND constraint_type = 'PRIMARY KEY' + AND constraint_name IN ({{.AutoGenNames}}); + + IF auto_pk_name IS NOT NULL THEN + EXECUTE 'ALTER TABLE {{quote_ident .SchemaName}}.{{quote_ident .TableName}} DROP CONSTRAINT ' || quote_ident(auto_pk_name); + END IF; + + -- Add named primary key if it doesn't exist + IF NOT EXISTS ( + SELECT 1 FROM information_schema.table_constraints + WHERE table_schema = '{{.SchemaName}}' + AND table_name = '{{.TableName}}' + AND constraint_name = '{{.ConstraintName}}' + ) THEN + ALTER TABLE {{quote_ident .SchemaName}}.{{quote_ident .TableName}} ADD CONSTRAINT {{quote_ident .ConstraintName}} PRIMARY KEY ({{.Columns}}); + END IF; +END; +$$; diff --git a/pkg/writers/pgsql/templates/create_sequence.tmpl b/pkg/writers/pgsql/templates/create_sequence.tmpl new file mode 100644 index 0000000..0b74e90 --- /dev/null +++ b/pkg/writers/pgsql/templates/create_sequence.tmpl @@ -0,0 +1,6 @@ +CREATE SEQUENCE IF NOT EXISTS {{quote_ident .SchemaName}}.{{quote_ident .SequenceName}} + INCREMENT {{.Increment}} + MINVALUE {{.MinValue}} + MAXVALUE {{.MaxValue}} + START {{.StartValue}} + CACHE {{.CacheSize}}; diff --git a/pkg/writers/pgsql/templates/create_table.tmpl b/pkg/writers/pgsql/templates/create_table.tmpl index 88f5be2..3f471ba 100644 --- a/pkg/writers/pgsql/templates/create_table.tmpl +++ b/pkg/writers/pgsql/templates/create_table.tmpl @@ -1,7 +1,7 @@ -CREATE TABLE IF NOT EXISTS {{.SchemaName}}.{{.TableName}} ( +CREATE TABLE IF NOT EXISTS {{quote_ident .SchemaName}}.{{quote_ident .TableName}} ( {{- range $i, $col := .Columns}} {{- if $i}},{{end}} - {{$col.Name}} {{$col.Type}} + {{quote_ident $col.Name}} {{$col.Type}} {{- if $col.Default}} DEFAULT {{$col.Default}}{{end}} {{- if $col.NotNull}} NOT NULL{{end}} {{- end}} diff --git a/pkg/writers/pgsql/templates/create_unique_constraint.tmpl b/pkg/writers/pgsql/templates/create_unique_constraint.tmpl new file mode 100644 index 0000000..478edbb --- /dev/null +++ b/pkg/writers/pgsql/templates/create_unique_constraint.tmpl @@ -0,0 +1,12 @@ +DO $$ +BEGIN + IF NOT EXISTS ( + SELECT 1 FROM information_schema.table_constraints + WHERE table_schema = '{{.SchemaName}}' + AND table_name = '{{.TableName}}' + AND constraint_name = '{{.ConstraintName}}' + ) THEN + ALTER TABLE {{quote_ident .SchemaName}}.{{quote_ident .TableName}} ADD CONSTRAINT {{quote_ident .ConstraintName}} UNIQUE ({{.Columns}}); + END IF; +END; +$$; \ No newline at end of file diff --git a/pkg/writers/pgsql/templates/drop_constraint.tmpl b/pkg/writers/pgsql/templates/drop_constraint.tmpl index f5f916e..439103f 100644 --- a/pkg/writers/pgsql/templates/drop_constraint.tmpl +++ b/pkg/writers/pgsql/templates/drop_constraint.tmpl @@ -1 +1 @@ -ALTER TABLE {{.SchemaName}}.{{.TableName}} DROP CONSTRAINT IF EXISTS {{.ConstraintName}}; \ No newline at end of file +ALTER TABLE {{quote_ident .SchemaName}}.{{quote_ident .TableName}} DROP CONSTRAINT IF EXISTS {{quote_ident .ConstraintName}}; \ No newline at end of file diff --git a/pkg/writers/pgsql/templates/drop_index.tmpl b/pkg/writers/pgsql/templates/drop_index.tmpl index 5f8b3bf..938af10 100644 --- a/pkg/writers/pgsql/templates/drop_index.tmpl +++ b/pkg/writers/pgsql/templates/drop_index.tmpl @@ -1 +1 @@ -DROP INDEX IF EXISTS {{.SchemaName}}.{{.IndexName}} CASCADE; \ No newline at end of file +DROP INDEX IF EXISTS {{quote_ident .SchemaName}}.{{quote_ident .IndexName}} CASCADE; \ No newline at end of file diff --git a/pkg/writers/pgsql/templates/set_sequence_value.tmpl b/pkg/writers/pgsql/templates/set_sequence_value.tmpl new file mode 100644 index 0000000..e4f8cd2 --- /dev/null +++ b/pkg/writers/pgsql/templates/set_sequence_value.tmpl @@ -0,0 +1,19 @@ +DO $$ +DECLARE + m_cnt bigint; +BEGIN + IF EXISTS ( + SELECT 1 FROM pg_class c + INNER JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE c.relname = '{{.SequenceName}}' + AND n.nspname = '{{.SchemaName}}' + AND c.relkind = 'S' + ) THEN + SELECT COALESCE(MAX({{quote_ident .ColumnName}}), 0) + 1 + FROM {{quote_ident .SchemaName}}.{{quote_ident .TableName}} + INTO m_cnt; + + PERFORM setval('{{quote_ident .SchemaName}}.{{quote_ident .SequenceName}}'::regclass, m_cnt); + END IF; +END; +$$; \ No newline at end of file diff --git a/pkg/writers/pgsql/writer.go b/pkg/writers/pgsql/writer.go index d143007..5f8e1b7 100644 --- a/pkg/writers/pgsql/writer.go +++ b/pkg/writers/pgsql/writer.go @@ -22,6 +22,7 @@ type Writer struct { options *writers.WriterOptions writer io.Writer executionReport *ExecutionReport + executor *TemplateExecutor } // ExecutionReport tracks the execution status of SQL statements @@ -57,8 +58,10 @@ type ExecutionError struct { // NewWriter creates a new PostgreSQL SQL writer func NewWriter(options *writers.WriterOptions) *Writer { + executor, _ := NewTemplateExecutor() return &Writer{ - options: options, + options: options, + executor: executor, } } @@ -215,36 +218,19 @@ func (w *Writer) GenerateSchemaStatements(schema *models.Schema) ([]string, erro fmt.Sprintf("%s_%s_pkey", schema.Name, table.Name), } - // Wrap in DO block to drop auto-generated PK and add our named PK - stmt := fmt.Sprintf("DO $$\nDECLARE\n"+ - " auto_pk_name text;\n"+ - "BEGIN\n"+ - " -- Drop auto-generated primary key if it exists\n"+ - " SELECT constraint_name INTO auto_pk_name\n"+ - " FROM information_schema.table_constraints\n"+ - " WHERE table_schema = '%s'\n"+ - " AND table_name = '%s'\n"+ - " AND constraint_type = 'PRIMARY KEY'\n"+ - " AND constraint_name IN (%s);\n"+ - "\n"+ - " IF auto_pk_name IS NOT NULL THEN\n"+ - " EXECUTE 'ALTER TABLE %s.%s DROP CONSTRAINT ' || quote_ident(auto_pk_name);\n"+ - " END IF;\n"+ - "\n"+ - " -- Add named primary key if it doesn't exist\n"+ - " IF NOT EXISTS (\n"+ - " SELECT 1 FROM information_schema.table_constraints\n"+ - " WHERE table_schema = '%s'\n"+ - " AND table_name = '%s'\n"+ - " AND constraint_name = '%s'\n"+ - " ) THEN\n"+ - " ALTER TABLE %s.%s ADD CONSTRAINT %s PRIMARY KEY (%s);\n"+ - " END IF;\n"+ - "END;\n$$", - schema.Name, table.Name, formatStringList(autoGenPKNames), - schema.SQLName(), table.SQLName(), - schema.Name, table.Name, pkName, - schema.SQLName(), table.SQLName(), pkName, strings.Join(pkColumns, ", ")) + // Use template to generate primary key statement + data := CreatePrimaryKeyWithAutoGenCheckData{ + SchemaName: schema.Name, + TableName: table.Name, + ConstraintName: pkName, + AutoGenNames: formatStringList(autoGenPKNames), + Columns: strings.Join(pkColumns, ", "), + } + + stmt, err := w.executor.ExecuteCreatePrimaryKeyWithAutoGenCheck(data) + if err != nil { + return nil, fmt.Errorf("failed to generate primary key for %s.%s: %w", schema.Name, table.Name, err) + } statements = append(statements, stmt) } } @@ -290,7 +276,7 @@ func (w *Writer) GenerateSchemaStatements(schema *models.Schema) ([]string, erro } stmt := fmt.Sprintf("CREATE %sINDEX IF NOT EXISTS %s ON %s.%s USING %s (%s)%s", - uniqueStr, index.Name, schema.SQLName(), table.SQLName(), indexType, strings.Join(columnExprs, ", "), whereClause) + uniqueStr, quoteIdentifier(index.Name), schema.SQLName(), table.SQLName(), indexType, strings.Join(columnExprs, ", "), whereClause) statements = append(statements, stmt) } } @@ -302,20 +288,18 @@ func (w *Writer) GenerateSchemaStatements(schema *models.Schema) ([]string, erro continue } - // Wrap in DO block to check for existing constraint - stmt := fmt.Sprintf("DO $$\nBEGIN\n"+ - " IF NOT EXISTS (\n"+ - " SELECT 1 FROM information_schema.table_constraints\n"+ - " WHERE table_schema = '%s'\n"+ - " AND table_name = '%s'\n"+ - " AND constraint_name = '%s'\n"+ - " ) THEN\n"+ - " ALTER TABLE %s.%s ADD CONSTRAINT %s UNIQUE (%s);\n"+ - " END IF;\n"+ - "END;\n$$", - schema.Name, table.Name, constraint.Name, - schema.SQLName(), table.SQLName(), constraint.Name, - strings.Join(constraint.Columns, ", ")) + // Use template to generate unique constraint statement + data := CreateUniqueConstraintData{ + SchemaName: schema.Name, + TableName: table.Name, + ConstraintName: constraint.Name, + Columns: strings.Join(constraint.Columns, ", "), + } + + stmt, err := w.executor.ExecuteCreateUniqueConstraint(data) + if err != nil { + return nil, fmt.Errorf("failed to generate unique constraint for %s.%s: %w", schema.Name, table.Name, err) + } statements = append(statements, stmt) } } @@ -327,20 +311,18 @@ func (w *Writer) GenerateSchemaStatements(schema *models.Schema) ([]string, erro continue } - // Wrap in DO block to check for existing constraint - stmt := fmt.Sprintf("DO $$\nBEGIN\n"+ - " IF NOT EXISTS (\n"+ - " SELECT 1 FROM information_schema.table_constraints\n"+ - " WHERE table_schema = '%s'\n"+ - " AND table_name = '%s'\n"+ - " AND constraint_name = '%s'\n"+ - " ) THEN\n"+ - " ALTER TABLE %s.%s ADD CONSTRAINT %s CHECK (%s);\n"+ - " END IF;\n"+ - "END;\n$$", - schema.Name, table.Name, constraint.Name, - schema.SQLName(), table.SQLName(), constraint.Name, - constraint.Expression) + // Use template to generate check constraint statement + data := CreateCheckConstraintData{ + SchemaName: schema.Name, + TableName: table.Name, + ConstraintName: constraint.Name, + Expression: constraint.Expression, + } + + stmt, err := w.executor.ExecuteCreateCheckConstraint(data) + if err != nil { + return nil, fmt.Errorf("failed to generate check constraint for %s.%s: %w", schema.Name, table.Name, err) + } statements = append(statements, stmt) } } @@ -367,23 +349,24 @@ func (w *Writer) GenerateSchemaStatements(schema *models.Schema) ([]string, erro onUpdate = "NO ACTION" } - // Wrap in DO block to check for existing constraint - stmt := fmt.Sprintf("DO $$\nBEGIN\n"+ - " IF NOT EXISTS (\n"+ - " SELECT 1 FROM information_schema.table_constraints\n"+ - " WHERE table_schema = '%s'\n"+ - " AND table_name = '%s'\n"+ - " AND constraint_name = '%s'\n"+ - " ) THEN\n"+ - " ALTER TABLE %s.%s ADD CONSTRAINT %s FOREIGN KEY (%s) REFERENCES %s.%s(%s) ON DELETE %s ON UPDATE %s;\n"+ - " END IF;\n"+ - "END;\n$$", - schema.Name, table.Name, constraint.Name, - schema.SQLName(), table.SQLName(), constraint.Name, - strings.Join(constraint.Columns, ", "), - strings.ToLower(refSchema), strings.ToLower(constraint.ReferencedTable), - strings.Join(constraint.ReferencedColumns, ", "), - onDelete, onUpdate) + // Use template to generate foreign key statement + data := CreateForeignKeyWithCheckData{ + SchemaName: schema.Name, + TableName: table.Name, + ConstraintName: constraint.Name, + SourceColumns: strings.Join(constraint.Columns, ", "), + TargetSchema: refSchema, + TargetTable: constraint.ReferencedTable, + TargetColumns: strings.Join(constraint.ReferencedColumns, ", "), + OnDelete: onDelete, + OnUpdate: onUpdate, + Deferrable: false, + } + + stmt, err := w.executor.ExecuteCreateForeignKeyWithCheck(data) + if err != nil { + return nil, fmt.Errorf("failed to generate foreign key for %s.%s: %w", schema.Name, table.Name, err) + } statements = append(statements, stmt) } } @@ -431,19 +414,18 @@ func (w *Writer) GenerateAddColumnStatements(schema *models.Schema) ([]string, e for _, col := range columns { colDef := w.generateColumnDefinition(col) - // Generate DO block that checks if column exists before adding - stmt := fmt.Sprintf("DO $$\nBEGIN\n"+ - " IF NOT EXISTS (\n"+ - " SELECT 1 FROM information_schema.columns\n"+ - " WHERE table_schema = '%s'\n"+ - " AND table_name = '%s'\n"+ - " AND column_name = '%s'\n"+ - " ) THEN\n"+ - " ALTER TABLE %s.%s ADD COLUMN %s;\n"+ - " END IF;\n"+ - "END;\n$$", - schema.Name, table.Name, col.Name, - schema.SQLName(), table.SQLName(), colDef) + // Use template to generate add column statement + data := AddColumnWithCheckData{ + SchemaName: schema.Name, + TableName: table.Name, + ColumnName: col.Name, + ColumnDefinition: colDef, + } + + stmt, err := w.executor.ExecuteAddColumnWithCheck(data) + if err != nil { + return nil, fmt.Errorf("failed to generate add column for %s.%s.%s: %w", schema.Name, table.Name, col.Name, err) + } statements = append(statements, stmt) } } @@ -699,13 +681,23 @@ func (w *Writer) writeSequences(schema *models.Schema) error { } seqName := fmt.Sprintf("identity_%s_%s", table.SQLName(), pk.SQLName()) - fmt.Fprintf(w.writer, "CREATE SEQUENCE IF NOT EXISTS %s.%s\n", - schema.SQLName(), seqName) - fmt.Fprintf(w.writer, " INCREMENT 1\n") - fmt.Fprintf(w.writer, " MINVALUE 1\n") - fmt.Fprintf(w.writer, " MAXVALUE 9223372036854775807\n") - fmt.Fprintf(w.writer, " START 1\n") - fmt.Fprintf(w.writer, " CACHE 1;\n\n") + + data := CreateSequenceData{ + SchemaName: schema.Name, + SequenceName: seqName, + Increment: 1, + MinValue: 1, + MaxValue: 9223372036854775807, + StartValue: 1, + CacheSize: 1, + } + + sql, err := w.executor.ExecuteCreateSequence(data) + if err != nil { + return fmt.Errorf("failed to generate create sequence for %s.%s: %w", schema.Name, seqName, err) + } + fmt.Fprint(w.writer, sql) + fmt.Fprint(w.writer, "\n") } return nil @@ -747,18 +739,19 @@ func (w *Writer) writeAddColumns(schema *models.Schema) error { for _, col := range columns { colDef := w.generateColumnDefinition(col) - // Generate DO block that checks if column exists before adding - fmt.Fprintf(w.writer, "DO $$\nBEGIN\n") - fmt.Fprintf(w.writer, " IF NOT EXISTS (\n") - fmt.Fprintf(w.writer, " SELECT 1 FROM information_schema.columns\n") - fmt.Fprintf(w.writer, " WHERE table_schema = '%s'\n", schema.Name) - fmt.Fprintf(w.writer, " AND table_name = '%s'\n", table.Name) - fmt.Fprintf(w.writer, " AND column_name = '%s'\n", col.Name) - fmt.Fprintf(w.writer, " ) THEN\n") - fmt.Fprintf(w.writer, " ALTER TABLE %s.%s ADD COLUMN %s;\n", - schema.SQLName(), table.SQLName(), colDef) - fmt.Fprintf(w.writer, " END IF;\n") - fmt.Fprintf(w.writer, "END;\n$$;\n\n") + data := AddColumnWithCheckData{ + SchemaName: schema.Name, + TableName: table.Name, + ColumnName: col.Name, + ColumnDefinition: colDef, + } + + sql, err := w.executor.ExecuteAddColumnWithCheck(data) + if err != nil { + return fmt.Errorf("failed to generate add column for %s.%s.%s: %w", schema.Name, table.Name, col.Name, err) + } + fmt.Fprint(w.writer, sql) + fmt.Fprint(w.writer, "\n") } } @@ -812,37 +805,20 @@ func (w *Writer) writePrimaryKeys(schema *models.Schema) error { fmt.Sprintf("%s_%s_pkey", schema.Name, table.Name), } - fmt.Fprintf(w.writer, "DO $$\nDECLARE\n") - fmt.Fprintf(w.writer, " auto_pk_name text;\nBEGIN\n") + data := CreatePrimaryKeyWithAutoGenCheckData{ + SchemaName: schema.Name, + TableName: table.Name, + ConstraintName: pkName, + AutoGenNames: formatStringList(autoGenPKNames), + Columns: strings.Join(columnNames, ", "), + } - // Check for and drop auto-generated primary keys - fmt.Fprintf(w.writer, " -- Drop auto-generated primary key if it exists\n") - fmt.Fprintf(w.writer, " SELECT constraint_name INTO auto_pk_name\n") - fmt.Fprintf(w.writer, " FROM information_schema.table_constraints\n") - fmt.Fprintf(w.writer, " WHERE table_schema = '%s'\n", schema.Name) - fmt.Fprintf(w.writer, " AND table_name = '%s'\n", table.Name) - fmt.Fprintf(w.writer, " AND constraint_type = 'PRIMARY KEY'\n") - fmt.Fprintf(w.writer, " AND constraint_name IN (%s);\n", formatStringList(autoGenPKNames)) - fmt.Fprintf(w.writer, "\n") - fmt.Fprintf(w.writer, " IF auto_pk_name IS NOT NULL THEN\n") - fmt.Fprintf(w.writer, " EXECUTE 'ALTER TABLE %s.%s DROP CONSTRAINT ' || quote_ident(auto_pk_name);\n", - schema.SQLName(), table.SQLName()) - fmt.Fprintf(w.writer, " END IF;\n") - fmt.Fprintf(w.writer, "\n") - - // Add our named primary key if it doesn't exist - fmt.Fprintf(w.writer, " -- Add named primary key if it doesn't exist\n") - fmt.Fprintf(w.writer, " IF NOT EXISTS (\n") - fmt.Fprintf(w.writer, " SELECT 1 FROM information_schema.table_constraints\n") - fmt.Fprintf(w.writer, " WHERE table_schema = '%s'\n", schema.Name) - fmt.Fprintf(w.writer, " AND table_name = '%s'\n", table.Name) - fmt.Fprintf(w.writer, " AND constraint_name = '%s'\n", pkName) - fmt.Fprintf(w.writer, " ) THEN\n") - fmt.Fprintf(w.writer, " ALTER TABLE %s.%s\n", schema.SQLName(), table.SQLName()) - fmt.Fprintf(w.writer, " ADD CONSTRAINT %s PRIMARY KEY (%s);\n", - pkName, strings.Join(columnNames, ", ")) - fmt.Fprintf(w.writer, " END IF;\n") - fmt.Fprintf(w.writer, "END;\n$$;\n\n") + sql, err := w.executor.ExecuteCreatePrimaryKeyWithAutoGenCheck(data) + if err != nil { + return fmt.Errorf("failed to generate primary key for %s.%s: %w", schema.Name, table.Name, err) + } + fmt.Fprint(w.writer, sql) + fmt.Fprint(w.writer, "\n") } return nil @@ -954,20 +930,17 @@ func (w *Writer) writeUniqueConstraints(schema *models.Schema) error { continue } - // Wrap in DO block to check for existing constraint - fmt.Fprintf(w.writer, "DO $$\n") - fmt.Fprintf(w.writer, "BEGIN\n") - fmt.Fprintf(w.writer, " IF NOT EXISTS (\n") - fmt.Fprintf(w.writer, " SELECT 1 FROM information_schema.table_constraints\n") - fmt.Fprintf(w.writer, " WHERE table_schema = '%s'\n", schema.Name) - fmt.Fprintf(w.writer, " AND table_name = '%s'\n", table.Name) - fmt.Fprintf(w.writer, " AND constraint_name = '%s'\n", constraint.Name) - fmt.Fprintf(w.writer, " ) THEN\n") - fmt.Fprintf(w.writer, " ALTER TABLE %s.%s ADD CONSTRAINT %s UNIQUE (%s);\n", - schema.SQLName(), table.SQLName(), constraint.Name, strings.Join(columnExprs, ", ")) - fmt.Fprintf(w.writer, " END IF;\n") - fmt.Fprintf(w.writer, "END;\n") - fmt.Fprintf(w.writer, "$$;\n\n") + sql, err := w.executor.ExecuteCreateUniqueConstraint(CreateUniqueConstraintData{ + SchemaName: schema.Name, + TableName: table.Name, + ConstraintName: constraint.Name, + Columns: strings.Join(columnExprs, ", "), + }) + if err != nil { + return fmt.Errorf("failed to generate unique constraint: %w", err) + } + + fmt.Fprintf(w.writer, "%s\n\n", sql) } } @@ -996,20 +969,17 @@ func (w *Writer) writeCheckConstraints(schema *models.Schema) error { continue } - // Wrap in DO block to check for existing constraint - fmt.Fprintf(w.writer, "DO $$\n") - fmt.Fprintf(w.writer, "BEGIN\n") - fmt.Fprintf(w.writer, " IF NOT EXISTS (\n") - fmt.Fprintf(w.writer, " SELECT 1 FROM information_schema.table_constraints\n") - fmt.Fprintf(w.writer, " WHERE table_schema = '%s'\n", schema.Name) - fmt.Fprintf(w.writer, " AND table_name = '%s'\n", table.Name) - fmt.Fprintf(w.writer, " AND constraint_name = '%s'\n", constraint.Name) - fmt.Fprintf(w.writer, " ) THEN\n") - fmt.Fprintf(w.writer, " ALTER TABLE %s.%s ADD CONSTRAINT %s CHECK (%s);\n", - schema.SQLName(), table.SQLName(), constraint.Name, constraint.Expression) - fmt.Fprintf(w.writer, " END IF;\n") - fmt.Fprintf(w.writer, "END;\n") - fmt.Fprintf(w.writer, "$$;\n\n") + sql, err := w.executor.ExecuteCreateCheckConstraint(CreateCheckConstraintData{ + SchemaName: schema.Name, + TableName: table.Name, + ConstraintName: constraint.Name, + Expression: constraint.Expression, + }) + if err != nil { + return fmt.Errorf("failed to generate check constraint: %w", err) + } + + fmt.Fprintf(w.writer, "%s\n\n", sql) } } @@ -1093,24 +1063,24 @@ func (w *Writer) writeForeignKeys(schema *models.Schema) error { refTable = rel.ToTable } - // Use DO block to check if constraint exists before adding - fmt.Fprintf(w.writer, "DO $$\nBEGIN\n") - fmt.Fprintf(w.writer, " IF NOT EXISTS (\n") - fmt.Fprintf(w.writer, " SELECT 1 FROM information_schema.table_constraints\n") - fmt.Fprintf(w.writer, " WHERE table_schema = '%s'\n", schema.Name) - fmt.Fprintf(w.writer, " AND table_name = '%s'\n", table.Name) - fmt.Fprintf(w.writer, " AND constraint_name = '%s'\n", fkName) - fmt.Fprintf(w.writer, " ) THEN\n") - fmt.Fprintf(w.writer, " ALTER TABLE %s.%s\n", schema.SQLName(), table.SQLName()) - fmt.Fprintf(w.writer, " ADD CONSTRAINT %s\n", fkName) - fmt.Fprintf(w.writer, " FOREIGN KEY (%s)\n", strings.Join(sourceColumns, ", ")) - fmt.Fprintf(w.writer, " REFERENCES %s.%s (%s)\n", - refSchema, refTable, strings.Join(targetColumns, ", ")) - fmt.Fprintf(w.writer, " ON DELETE %s\n", onDelete) - fmt.Fprintf(w.writer, " ON UPDATE %s\n", onUpdate) - fmt.Fprintf(w.writer, " DEFERRABLE;\n") - fmt.Fprintf(w.writer, " END IF;\n") - fmt.Fprintf(w.writer, "END;\n$$;\n\n") + // Use template executor to generate foreign key with existence check + data := CreateForeignKeyWithCheckData{ + SchemaName: schema.Name, + TableName: table.Name, + ConstraintName: fkName, + SourceColumns: strings.Join(sourceColumns, ", "), + TargetSchema: refSchema, + TargetTable: refTable, + TargetColumns: strings.Join(targetColumns, ", "), + OnDelete: onDelete, + OnUpdate: onUpdate, + Deferrable: true, + } + sql, err := w.executor.ExecuteCreateForeignKeyWithCheck(data) + if err != nil { + return fmt.Errorf("failed to generate foreign key for %s.%s: %w", schema.Name, table.Name, err) + } + fmt.Fprint(w.writer, sql) } // Also process any foreign key constraints that don't have a relationship @@ -1172,23 +1142,24 @@ func (w *Writer) writeForeignKeys(schema *models.Schema) error { } refTable := constraint.ReferencedTable - // Use DO block to check if constraint exists before adding - fmt.Fprintf(w.writer, "DO $$\nBEGIN\n") - fmt.Fprintf(w.writer, " IF NOT EXISTS (\n") - fmt.Fprintf(w.writer, " SELECT 1 FROM information_schema.table_constraints\n") - fmt.Fprintf(w.writer, " WHERE table_schema = '%s'\n", schema.Name) - fmt.Fprintf(w.writer, " AND table_name = '%s'\n", table.Name) - fmt.Fprintf(w.writer, " AND constraint_name = '%s'\n", constraint.Name) - fmt.Fprintf(w.writer, " ) THEN\n") - fmt.Fprintf(w.writer, " ALTER TABLE %s.%s\n", schema.SQLName(), table.SQLName()) - fmt.Fprintf(w.writer, " ADD CONSTRAINT %s\n", constraint.Name) - fmt.Fprintf(w.writer, " FOREIGN KEY (%s)\n", strings.Join(sourceColumns, ", ")) - fmt.Fprintf(w.writer, " REFERENCES %s.%s (%s)\n", - refSchema, refTable, strings.Join(targetColumns, ", ")) - fmt.Fprintf(w.writer, " ON DELETE %s\n", onDelete) - fmt.Fprintf(w.writer, " ON UPDATE %s;\n", onUpdate) - fmt.Fprintf(w.writer, " END IF;\n") - fmt.Fprintf(w.writer, "END;\n$$;\n\n") + // Use template executor to generate foreign key with existence check + data := CreateForeignKeyWithCheckData{ + SchemaName: schema.Name, + TableName: table.Name, + ConstraintName: constraint.Name, + SourceColumns: strings.Join(sourceColumns, ", "), + TargetSchema: refSchema, + TargetTable: refTable, + TargetColumns: strings.Join(targetColumns, ", "), + OnDelete: onDelete, + OnUpdate: onUpdate, + Deferrable: false, + } + sql, err := w.executor.ExecuteCreateForeignKeyWithCheck(data) + if err != nil { + return fmt.Errorf("failed to generate foreign key for %s.%s: %w", schema.Name, table.Name, err) + } + fmt.Fprint(w.writer, sql) } } @@ -1207,26 +1178,19 @@ func (w *Writer) writeSetSequenceValues(schema *models.Schema) error { seqName := fmt.Sprintf("identity_%s_%s", table.SQLName(), pk.SQLName()) - fmt.Fprintf(w.writer, "DO $$\n") - fmt.Fprintf(w.writer, "DECLARE\n") - fmt.Fprintf(w.writer, " m_cnt bigint;\n") - fmt.Fprintf(w.writer, "BEGIN\n") - fmt.Fprintf(w.writer, " IF EXISTS (\n") - fmt.Fprintf(w.writer, " SELECT 1 FROM pg_class c\n") - fmt.Fprintf(w.writer, " INNER JOIN pg_namespace n ON n.oid = c.relnamespace\n") - fmt.Fprintf(w.writer, " WHERE c.relname = '%s'\n", seqName) - fmt.Fprintf(w.writer, " AND n.nspname = '%s'\n", schema.Name) - fmt.Fprintf(w.writer, " AND c.relkind = 'S'\n") - fmt.Fprintf(w.writer, " ) THEN\n") - fmt.Fprintf(w.writer, " SELECT COALESCE(MAX(%s), 0) + 1\n", pk.SQLName()) - fmt.Fprintf(w.writer, " FROM %s.%s\n", schema.SQLName(), table.SQLName()) - fmt.Fprintf(w.writer, " INTO m_cnt;\n") - fmt.Fprintf(w.writer, " \n") - fmt.Fprintf(w.writer, " PERFORM setval('%s.%s'::regclass, m_cnt);\n", - schema.SQLName(), seqName) - fmt.Fprintf(w.writer, " END IF;\n") - fmt.Fprintf(w.writer, "END;\n") - fmt.Fprintf(w.writer, "$$;\n\n") + // Use template executor to generate set sequence value statement + data := SetSequenceValueData{ + SchemaName: schema.Name, + TableName: table.Name, + SequenceName: seqName, + ColumnName: pk.Name, + } + sql, err := w.executor.ExecuteSetSequenceValue(data) + if err != nil { + return fmt.Errorf("failed to generate set sequence value for %s.%s: %w", schema.Name, table.Name, err) + } + fmt.Fprint(w.writer, sql) + fmt.Fprint(w.writer, "\n") } return nil @@ -1580,3 +1544,9 @@ func truncateStatement(stmt string) string { func getCurrentTimestamp() string { return time.Now().Format("2006-01-02 15:04:05") } + +// quoteIdentifier wraps an identifier in double quotes if necessary +// This is needed for identifiers that start with numbers or contain special characters +func quoteIdentifier(s string) string { + return quoteIdent(s) +}