feat: add support for trigger UPDATE OF columns (#342) (#344)

* feat: add support for trigger UPDATE OF columns (#342)

Triggers with column-specific UPDATE events (e.g., UPDATE OF email) were
losing the column specification during inspection, causing incorrect
migration plans that would fire triggers on all updates instead of only
on specified column changes.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* fix: guard UPDATE OF column extraction with tgtype bitmask check

Only extract UPDATE OF columns when the trigger actually has an UPDATE
event, preventing false positives if the substring appears elsewhere.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: tianzhou

internal/diff/trigger.go

@@ -46,6 +46,16 @@ func triggersEqual(old, new *ir.Trigger) bool {
 		}
 	}
 
+	// Compare update columns
+	if len(old.UpdateColumns) != len(new.UpdateColumns) {
+		return false
+	}
+	for i, col := range old.UpdateColumns {
+		if col != new.UpdateColumns[i] {
+			return false
+		}
+	}
+
 	// Compare constraint trigger properties
 	if old.IsConstraint != new.IsConstraint {
 		return false
@@ -215,7 +225,11 @@ func generateTriggerSQLWithMode(trigger *ir.Trigger, targetSchema string) string
 	for _, orderEvent := range eventOrder {
 		for _, triggerEvent := range trigger.Events {
 			if triggerEvent == orderEvent {
-				events = append(events, string(triggerEvent))
+				if triggerEvent == ir.TriggerEventUpdate && len(trigger.UpdateColumns) > 0 {
+					events = append(events, "UPDATE OF "+strings.Join(trigger.UpdateColumns, ", "))
+				} else {
+					events = append(events, string(triggerEvent))
+				}
 				break
 			}
 		}

ir/inspector.go

@@ -1452,6 +1452,65 @@ func extractWhenClauseFromTriggerDef(triggerDef string) string {
 	return strings.TrimSpace(triggerDef[start:end])
 }
 
+// extractUpdateColumnsFromTriggerDef extracts column names from UPDATE OF col1, col2 in a trigger definition
+// returned by pg_get_triggerdef(). The format is:
+// "CREATE TRIGGER name BEFORE INSERT OR UPDATE OF col1, col2 ON table ..."
+func extractUpdateColumnsFromTriggerDef(triggerDef string) []string {
+	upper := strings.ToUpper(triggerDef)
+	updateOfIdx := strings.Index(upper, "UPDATE OF ")
+	if updateOfIdx == -1 {
+		return nil
+	}
+
+	// Start after "UPDATE OF "
+	start := updateOfIdx + len("UPDATE OF ")
+
+	// Find " ON " which terminates the event/column list
+	onIdx := strings.Index(upper[start:], " ON ")
+	if onIdx == -1 {
+		return nil
+	}
+
+	// Extract the column list substring
+	colListStr := strings.TrimSpace(triggerDef[start : start+onIdx])
+
+	// Handle potential " OR " separating additional events after the columns
+	// e.g., "UPDATE OF col1, col2 OR DELETE ON ..."
+	// We need to check if there's an OR followed by another event keyword
+	eventKeywords := []string{"INSERT", "UPDATE", "DELETE", "TRUNCATE"}
+	parts := strings.Split(colListStr, " OR ")
+	// Only keep the first part (column list); the rest would be other events
+	if len(parts) > 1 {
+		// Check if subsequent parts are event keywords
+		for i := 1; i < len(parts); i++ {
+			trimmed := strings.TrimSpace(strings.ToUpper(parts[i]))
+			isEvent := false
+			for _, kw := range eventKeywords {
+				if trimmed == kw || strings.HasPrefix(trimmed, kw+" ") {
+					isEvent = true
+					break
+				}
+			}
+			if isEvent {
+				colListStr = strings.TrimSpace(parts[0])
+				break
+			}
+		}
+	}
+
+	// Split by comma and trim each column name
+	rawCols := strings.Split(colListStr, ",")
+	var columns []string
+	for _, col := range rawCols {
+		col = strings.TrimSpace(col)
+		if col != "" {
+			columns = append(columns, col)
+		}
+	}
+
+	return columns
+}
+
 // extractFunctionCallFromTriggerDef extracts the function call (with arguments) from a trigger definition
 // returned by pg_get_triggerdef(). The format is:
 // "... EXECUTE FUNCTION function_name(arg1, arg2)"
@@ -1548,6 +1607,12 @@ func (i *Inspector) buildTriggers(ctx context.Context, schema *IR, targetSchema
 			condition = extractWhenClauseFromTriggerDef(triggerRow.TriggerDefinition.String)
 		}
 
+		// Extract UPDATE OF columns from trigger definition (only for triggers with UPDATE events)
+		var updateColumns []string
+		if triggerRow.TriggerDefinition.Valid && triggerRow.TriggerType&triggerTypeUpdate != 0 {
+			updateColumns = extractUpdateColumnsFromTriggerDef(triggerRow.TriggerDefinition.String)
+		}
+
 		// Extract transition table names
 		oldTable := ""
 		if triggerRow.OldTable.Valid {
@@ -1577,6 +1642,7 @@ func (i *Inspector) buildTriggers(ctx context.Context, schema *IR, targetSchema
 			Table:             tableName,
 			Timing:            timing,
 			Events:            events,
+			UpdateColumns:     updateColumns,
 			Level:             level,
 			Function:          functionCall,
 			Condition:         condition,

ir/ir.go

@@ -278,8 +278,9 @@ type Trigger struct {
 	Table             string         `json:"table"`
 	Name              string         `json:"name"`
 	Timing            TriggerTiming  `json:"timing"` // BEFORE, AFTER, INSTEAD OF
-	Events            []TriggerEvent `json:"events"` // INSERT, UPDATE, DELETE
-	Level             TriggerLevel   `json:"level"`  // ROW, STATEMENT
+	Events            []TriggerEvent `json:"events"`                        // INSERT, UPDATE, DELETE
+	UpdateColumns     []string       `json:"update_columns,omitempty"`      // Column names for UPDATE OF
+	Level             TriggerLevel   `json:"level"`                         // ROW, STATEMENT
 	Function          string         `json:"function"`
 	Condition         string         `json:"condition,omitempty"` // WHEN condition
 	Comment           string         `json:"comment,omitempty"`

testdata/diff/create_trigger/add_trigger/diff.sql

@@ -8,6 +8,11 @@ CREATE OR REPLACE TRIGGER employees_last_modified_trigger
     FOR EACH ROW
     EXECUTE FUNCTION update_last_modified();
 
+CREATE OR REPLACE TRIGGER employees_salary_update_trigger
+    BEFORE UPDATE OF salary ON employees
+    FOR EACH ROW
+    EXECUTE FUNCTION update_last_modified();
+
 CREATE OR REPLACE TRIGGER employees_truncate_log_trigger
     AFTER TRUNCATE ON employees
     FOR EACH STATEMENT

testdata/diff/create_trigger/add_trigger/new.sql

@@ -23,6 +23,11 @@ CREATE TRIGGER employees_insert_timestamp_trigger
     FOR EACH ROW
     EXECUTE FUNCTION public.update_last_modified();
 
+CREATE TRIGGER employees_salary_update_trigger
+    BEFORE UPDATE OF salary ON public.employees
+    FOR EACH ROW
+    EXECUTE FUNCTION public.update_last_modified();
+
 CREATE TRIGGER employees_truncate_log_trigger
     AFTER TRUNCATE ON public.employees
     FOR EACH STATEMENT

testdata/diff/create_trigger/add_trigger/plan.json

@@ -20,6 +20,12 @@
           "operation": "create",
           "path": "public.employees.employees_last_modified_trigger"
         },
+        {
+          "sql": "CREATE OR REPLACE TRIGGER employees_salary_update_trigger\n    BEFORE UPDATE OF salary ON employees\n    FOR EACH ROW\n    EXECUTE FUNCTION update_last_modified();",
+          "type": "table.trigger",
+          "operation": "create",
+          "path": "public.employees.employees_salary_update_trigger"
+        },
         {
           "sql": "CREATE OR REPLACE TRIGGER employees_truncate_log_trigger\n    AFTER TRUNCATE ON employees\n    FOR EACH STATEMENT\n    EXECUTE FUNCTION update_last_modified();",
           "type": "table.trigger",

testdata/diff/create_trigger/add_trigger/plan.sql

@@ -8,6 +8,11 @@ CREATE OR REPLACE TRIGGER employees_last_modified_trigger
     FOR EACH ROW
     EXECUTE FUNCTION update_last_modified();
 
+CREATE OR REPLACE TRIGGER employees_salary_update_trigger
+    BEFORE UPDATE OF salary ON employees
+    FOR EACH ROW
+    EXECUTE FUNCTION update_last_modified();
+
 CREATE OR REPLACE TRIGGER employees_truncate_log_trigger
     AFTER TRUNCATE ON employees
     FOR EACH STATEMENT

testdata/diff/create_trigger/add_trigger/plan.txt

@@ -8,6 +8,7 @@ Tables:
   ~ employees
     + employees_insert_timestamp_trigger (trigger)
     + employees_last_modified_trigger (trigger)
+    + employees_salary_update_trigger (trigger)
     + employees_truncate_log_trigger (trigger)
 
 Views:
@@ -27,6 +28,11 @@ CREATE OR REPLACE TRIGGER employees_last_modified_trigger
     FOR EACH ROW
     EXECUTE FUNCTION update_last_modified();
 
+CREATE OR REPLACE TRIGGER employees_salary_update_trigger
+    BEFORE UPDATE OF salary ON employees
+    FOR EACH ROW
+    EXECUTE FUNCTION update_last_modified();
+
 CREATE OR REPLACE TRIGGER employees_truncate_log_trigger
     AFTER TRUNCATE ON employees
     FOR EACH STATEMENT

testdata/diff/create_trigger/alter_trigger/diff.sql

@@ -1,5 +1,5 @@
 CREATE OR REPLACE TRIGGER employees_last_modified_trigger
-    BEFORE INSERT OR UPDATE ON employees
+    BEFORE INSERT OR UPDATE OF salary ON employees
     FOR EACH ROW
     WHEN (((NEW.salary IS NOT NULL)))
     EXECUTE FUNCTION update_last_modified();

testdata/diff/create_trigger/alter_trigger/new.sql

@@ -14,7 +14,7 @@ END;
 $$ LANGUAGE plpgsql;
 
 CREATE TRIGGER employees_last_modified_trigger
-    BEFORE INSERT OR UPDATE ON public.employees
+    BEFORE INSERT OR UPDATE OF salary ON public.employees
     FOR EACH ROW
     WHEN (NEW.salary IS NOT NULL)
     EXECUTE FUNCTION public.update_last_modified();