Handle columns who use functions as defaults

stripe · Feb 6, 2025 · 6b33a65 · 6b33a65
1 parent 7e225f6
commit 6b33a65
Show file tree

Hide file tree

Showing 6 changed files with 61 additions and 13 deletions.
diff --git a/internal/queries/dml.sql.go b/internal/queries/dml.sql.go
diff --git a/internal/queries/models.sql.go b/internal/queries/models.sql.go
diff --git a/internal/queries/queries.sql b/internal/queries/queries.sql
@@ -94,7 +94,8 @@ SELECT
     identity_col_seq.seqmin AS min_value,
     identity_col_seq.seqcache AS cache_size,
     identity_col_seq.seqcycle AS is_cycle,
-    pg_catalog.format_type(a.atttypid, a.atttypmod) AS column_type
+    pg_catalog.format_type(a.atttypid, a.atttypmod) AS column_type,
+    default_value.oid AS default_value_oid
 FROM pg_catalog.pg_attribute AS a
 LEFT JOIN
     pg_catalog.pg_attrdef AS d
@@ -108,6 +109,11 @@ LEFT JOIN
     ON
         a.attrelid = identity_col_seq.owner_relid
         AND a.attnum = identity_col_seq.owner_attnum
+LEFT JOIN
+    pg_catalog.pg_attrdef AS default_value
+    ON
+        a.attrelid = default_value.adrelid
+        AND a.attnum = default_value.adnum
 WHERE
     a.attrelid = $1
     AND a.attnum > 0

diff --git a/internal/queries/queries.sql.go b/internal/queries/queries.sql.go
diff --git a/internal/schema/schema.go b/internal/schema/schema.go
@@ -234,6 +234,8 @@ type (
 		// It is used for data-packing purposes
 		Size     int
 		Identity *ColumnIdentity
+
+		DependsOnFunctions []SchemaQualifiedName
 	}
 )
 
@@ -878,6 +880,11 @@ func (s *schemaFetcher) buildTable(
 			}
 		}
 
+		dependsOnFunctions, _ := s.fetchDependsOnFunctions(ctx, "pg_attrdef", column.DefaultValueOid)
+		if err != nil {
+			return Table{}, fmt.Errorf("fetchDependsOnFunctions(%s): %w", column.DefaultValueOid, err)
+		}
+
 		columns = append(columns, Column{
 			Name:       column.ColumnName,
 			Type:       column.ColumnType,
@@ -888,9 +895,10 @@ func (s *schemaFetcher) buildTable(
 			//   ''::text
 			//   CURRENT_TIMESTAMP
 			// If empty, indicates that there is no default value.
-			Default:  column.DefaultValue,
-			Size:     int(column.ColumnSize),
-			Identity: identity,
+			Default:            column.DefaultValue,
+			Size:               int(column.ColumnSize),
+			Identity:           identity,
+			DependsOnFunctions: dependsOnFunctions,
 		})
 	}
 
@@ -905,6 +913,7 @@ func (s *schemaFetcher) buildTable(
 		SchemaName:  table.TableSchemaName,
 		EscapedName: EscapeIdentifier(table.TableName),
 	}
+
 	return Table{
 		SchemaQualifiedName: schemaQualifiedName,
 		Columns:             columns,

diff --git a/pkg/diff/sql_generator.go b/pkg/diff/sql_generator.go
@@ -1074,6 +1074,13 @@ func (t *tableSQLVertexGenerator) GetAddAlterDependencies(table, _ schema.Table)
 			mustRun(t.GetSQLVertexId(table, diffTypeAddAlter)).after(buildTableVertexId(*table.ParentTable, diffTypeAddAlter)),
 		)
 	}
+
+	for _, col := range table.Columns {
+		for _, depFunction := range col.DependsOnFunctions {
+			deps = append(deps, mustRun(t.GetSQLVertexId(table, diffTypeDelete)).after(buildFunctionVertexId(depFunction, diffTypeDelete)))
+		}
+	}
+
 	return deps, nil
 }
 
@@ -1142,6 +1149,13 @@ func (t *tableSQLVertexGenerator) GetDeleteDependencies(table schema.Table) ([]d
 			mustRun(t.GetSQLVertexId(table, diffTypeDelete)).after(buildTableVertexId(*table.ParentTable, diffTypeDelete)),
 		)
 	}
+
+	for _, col := range table.Columns {
+		for _, depFunction := range col.DependsOnFunctions {
+			deps = append(deps, mustRun(t.GetSQLVertexId(table, diffTypeDelete)).before(buildFunctionVertexId(depFunction, diffTypeDelete)))
+		}
+	}
+
 	return deps, nil
 }
 
@@ -1407,13 +1421,26 @@ func buildColumnVertexId(columnName string, diffType diffType) sqlVertexId {
 }
 
 func (csg *columnSQLVertexGenerator) GetAddAlterDependencies(col, _ schema.Column) ([]dependency, error) {
-	return []dependency{
+
+	var deps []dependency = []dependency{
 		mustRun(csg.GetSQLVertexId(col, diffTypeDelete)).before(csg.GetSQLVertexId(col, diffTypeAddAlter)),
-	}, nil
+	}
+
+	for _, depFunction := range col.DependsOnFunctions {
+		deps = append(deps, mustRun(csg.GetSQLVertexId(col, diffTypeDelete)).after(buildFunctionVertexId(depFunction, diffTypeDelete)))
+	}
+
+	return deps, nil
+
 }
 
-func (csg *columnSQLVertexGenerator) GetDeleteDependencies(_ schema.Column) ([]dependency, error) {
-	return nil, nil
+func (csg *columnSQLVertexGenerator) GetDeleteDependencies(col schema.Column) ([]dependency, error) {
+
+	var deps []dependency
+	for _, depFunction := range col.DependsOnFunctions {
+		deps = append(deps, mustRun(csg.GetSQLVertexId(col, diffTypeDelete)).before(buildFunctionVertexId(depFunction, diffTypeDelete)))
+	}
+	return deps, nil
 }
 
 type renameConflictingIndexSQLVertexGenerator struct {