From 21fab6a86bb1114480b33143eb208c5a436364ba Mon Sep 17 00:00:00 2001
From: Bryce Plunkett <130488190+bplunkett-stripe@users.noreply.github.com>
Date: Wed, 27 Nov 2024 12:29:41 -0800
Subject: [PATCH] Fix defaults with identity column (#183)

---
 .../column_cases_test.go                      | 17 ++++++++++++
 pkg/diff/sql_generator.go                     | 27 ++++++++++++-------
 2 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/internal/migration_acceptance_tests/column_cases_test.go b/internal/migration_acceptance_tests/column_cases_test.go
index 5cf0809..c117446 100644
--- a/internal/migration_acceptance_tests/column_cases_test.go
+++ b/internal/migration_acceptance_tests/column_cases_test.go
@@ -995,6 +995,23 @@ var columnAcceptanceTestCases = []acceptanceTestCase{
 			`,
 		},
 	},
+	{
+		name: "Add identity to column with existing default",
+		oldSchemaDDL: []string{
+			`
+            CREATE TABLE foobar(
+                identity_always BIGINT DEFAULT 5
+            );
+			`,
+		},
+		newSchemaDDL: []string{
+			`
+            CREATE TABLE foobar(
+                identity_always BIGINT GENERATED ALWAYS AS IDENTITY 
+            );
+			`,
+		},
+	},
 	{
 		name: "Alter identity type - to by default",
 		oldSchemaDDL: []string{
diff --git a/pkg/diff/sql_generator.go b/pkg/diff/sql_generator.go
index cf786ab..7919d7d 100644
--- a/pkg/diff/sql_generator.go
+++ b/pkg/diff/sql_generator.go
@@ -1181,6 +1181,23 @@ func (csg *columnSQLVertexGenerator) Alter(diff columnDiff) ([]Statement, error)
 		})
 	}
 
+	// Drop the default before type conversion. This will allow type conversions
+	// between incompatible types if the previous column has a default and the new column is dropping its default.
+	// It also must be dropped before an identity is added to the column, otherwise adding the identity errors
+	// with "a default already exists."
+	//
+	// To keep the code simpler, put dropping the default before updating the column identity AND dropping the default.
+	// There is an argument to drop the default after removing the not null constraint, since writes will continue to succeed
+	// on columns migrating from not-null and a default to nullable with no default; however, this would not work
+	// for the case where an identity is being added to the column.
+	if len(oldColumn.Default) > 0 && len(newColumn.Default) == 0 {
+		stmts = append(stmts, Statement{
+			DDL:         fmt.Sprintf("%s DROP DEFAULT", alterColumnPrefix),
+			Timeout:     statementTimeoutDefault,
+			LockTimeout: lockTimeoutDefault,
+		})
+	}
+
 	updateIdentityStmts, err := csg.buildUpdateIdentityStatements(oldColumn, newColumn)
 	if err != nil {
 		return nil, fmt.Errorf("building update identity statements: %w", err)
@@ -1197,16 +1214,6 @@ func (csg *columnSQLVertexGenerator) Alter(diff columnDiff) ([]Statement, error)
 		})
 	}
 
-	if len(oldColumn.Default) > 0 && len(newColumn.Default) == 0 {
-		// Drop the default before type conversion. This will allow type conversions
-		// between incompatible types if the previous column has a default and the new column is dropping its default
-		stmts = append(stmts, Statement{
-			DDL:         fmt.Sprintf("%s DROP DEFAULT", alterColumnPrefix),
-			Timeout:     statementTimeoutDefault,
-			LockTimeout: lockTimeoutDefault,
-		})
-	}
-
 	if !strings.EqualFold(oldColumn.Type, newColumn.Type) ||
 		!strings.EqualFold(oldColumn.Collation.GetFQEscapedName(), newColumn.Collation.GetFQEscapedName()) {
 		stmts = append(stmts,