feat(writer): 🎉 Improve primary key handling by checking for explicit constraints and columns

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 {