From 122743ee439cb76d8d1a5f51d865201e97c59a7a Mon Sep 17 00:00:00 2001
From: "Hein (Warky)" <hein.puth@gmail.com>
Date: Wed, 28 Jan 2026 22:08:27 +0200
Subject: [PATCH] =?UTF-8?q?feat(writer):=20=F0=9F=8E=89=20Improve=20primar?=
 =?UTF-8?q?y=20key=20handling=20by=20checking=20for=20explicit=20constrain?=
 =?UTF-8?q?ts=20and=20columns?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pkg/writers/pgsql/migration_writer.go | 51 ++++++++++++++++++++++++-
 pkg/writers/pgsql/writer.go           | 55 ++++++++++++++++++++-------
 2 files changed, 92 insertions(+), 14 deletions(-)

diff --git a/pkg/writers/pgsql/migration_writer.go b/pkg/writers/pgsql/migration_writer.go
index 82489f6..671bb93 100644
--- a/pkg/writers/pgsql/migration_writer.go
+++ b/pkg/writers/pgsql/migration_writer.go
@@ -427,9 +427,11 @@ func (w *MigrationWriter) generateIndexScripts(model *models.Schema, current *mo
 	for _, modelTable := range model.Tables {
 		currentTable := currentTables[strings.ToLower(modelTable.Name)]
 
-		// Process primary keys first
+		// Process primary keys first - check explicit constraints
+		foundExplicitPK := false
 		for constraintName, constraint := range modelTable.Constraints {
 			if constraint.Type == models.PrimaryKeyConstraint {
+				foundExplicitPK = true
 				shouldCreate := true
 
 				if currentTable != nil {
@@ -464,6 +466,53 @@ func (w *MigrationWriter) generateIndexScripts(model *models.Schema, current *mo
 			}
 		}
 
+		// If no explicit PK constraint, check for columns with IsPrimaryKey = true
+		if !foundExplicitPK {
+			pkColumns := []string{}
+			for _, col := range modelTable.Columns {
+				if col.IsPrimaryKey {
+					pkColumns = append(pkColumns, col.SQLName())
+				}
+			}
+			if len(pkColumns) > 0 {
+				sort.Strings(pkColumns)
+				constraintName := fmt.Sprintf("pk_%s_%s", strings.ToLower(model.Name), strings.ToLower(modelTable.Name))
+				shouldCreate := true
+
+				if currentTable != nil {
+					// Check if a PK constraint already exists (by any name)
+					for _, constraint := range currentTable.Constraints {
+						if constraint.Type == models.PrimaryKeyConstraint {
+							shouldCreate = false
+							break
+						}
+					}
+				}
+
+				if shouldCreate {
+					sql, err := w.executor.ExecuteCreatePrimaryKey(CreatePrimaryKeyData{
+						SchemaName:     model.Name,
+						TableName:      modelTable.Name,
+						ConstraintName: constraintName,
+						Columns:        strings.Join(pkColumns, ", "),
+					})
+					if err != nil {
+						return nil, err
+					}
+
+					script := MigrationScript{
+						ObjectName: fmt.Sprintf("%s.%s.%s", model.Name, modelTable.Name, constraintName),
+						ObjectType: "create primary key",
+						Schema:     model.Name,
+						Priority:   160,
+						Sequence:   len(scripts),
+						Body:       sql,
+					}
+					scripts = append(scripts, script)
+				}
+			}
+		}
+
 		// Process indexes
 		for indexName, modelIndex := range modelTable.Indexes {
 			// Skip primary key indexes
diff --git a/pkg/writers/pgsql/writer.go b/pkg/writers/pgsql/writer.go
index 87397c7..f4071b5 100644
--- a/pkg/writers/pgsql/writer.go
+++ b/pkg/writers/pgsql/writer.go
@@ -127,13 +127,35 @@ func (w *Writer) GenerateSchemaStatements(schema *models.Schema) ([]string, erro
 
 	// Phase 4: Primary keys
 	for _, table := range schema.Tables {
+		// First check for explicit PrimaryKeyConstraint
+		var pkConstraint *models.Constraint
 		for _, constraint := range table.Constraints {
-			if constraint.Type != models.PrimaryKeyConstraint {
-				continue
+			if constraint.Type == models.PrimaryKeyConstraint {
+				pkConstraint = constraint
+				break
 			}
+		}
+
+		if pkConstraint != nil {
 			stmt := fmt.Sprintf("ALTER TABLE %s.%s ADD CONSTRAINT %s PRIMARY KEY (%s)",
-				schema.SQLName(), table.SQLName(), constraint.Name, strings.Join(constraint.Columns, ", "))
+				schema.SQLName(), table.SQLName(), pkConstraint.Name, strings.Join(pkConstraint.Columns, ", "))
 			statements = append(statements, stmt)
+		} else {
+			// No explicit constraint, check for columns with IsPrimaryKey = true
+			pkColumns := []string{}
+			for _, col := range table.Columns {
+				if col.IsPrimaryKey {
+					pkColumns = append(pkColumns, col.SQLName())
+				}
+			}
+			if len(pkColumns) > 0 {
+				// Sort for consistent output
+				sort.Strings(pkColumns)
+				pkName := fmt.Sprintf("pk_%s_%s", schema.SQLName(), table.SQLName())
+				stmt := fmt.Sprintf("ALTER TABLE %s.%s ADD CONSTRAINT %s PRIMARY KEY (%s)",
+					schema.SQLName(), table.SQLName(), pkName, strings.Join(pkColumns, ", "))
+				statements = append(statements, stmt)
+			}
 		}
 	}
 
@@ -458,19 +480,26 @@ func (w *Writer) writePrimaryKeys(schema *models.Schema) error {
 			}
 		}
 
-		if pkConstraint == nil {
-			// No explicit PK constraint, skip
-			continue
-		}
-
+		var columnNames []string
 		pkName := fmt.Sprintf("pk_%s_%s", schema.SQLName(), table.SQLName())
 
-		// Build column list
-		columnNames := make([]string, 0, len(pkConstraint.Columns))
-		for _, colName := range pkConstraint.Columns {
-			if col, ok := table.Columns[colName]; ok {
-				columnNames = append(columnNames, col.SQLName())
+		if pkConstraint != nil {
+			// Build column list from explicit constraint
+			columnNames = make([]string, 0, len(pkConstraint.Columns))
+			for _, colName := range pkConstraint.Columns {
+				if col, ok := table.Columns[colName]; ok {
+					columnNames = append(columnNames, col.SQLName())
+				}
 			}
+		} else {
+			// No explicit PK constraint, check for columns with IsPrimaryKey = true
+			for _, col := range table.Columns {
+				if col.IsPrimaryKey {
+					columnNames = append(columnNames, col.SQLName())
+				}
+			}
+			// Sort for consistent output
+			sort.Strings(columnNames)
 		}
 
 		if len(columnNames) == 0 {