Skip to content

Commit 3efa9b9

Browse files
authored
Fix: Support GENERATED ALWAYS AS columns to reduce migration failures (#232)
1 parent f40bbbe commit 3efa9b9

File tree

7 files changed

+275
-37
lines changed

7 files changed

+275
-37
lines changed

internal/migration_acceptance_tests/column_cases_test.go

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1184,6 +1184,149 @@ var columnAcceptanceTestCases = []acceptanceTestCase{
11841184
`,
11851185
},
11861186
},
1187+
{
1188+
name: "Add generated column",
1189+
oldSchemaDDL: []string{
1190+
`
1191+
CREATE TABLE tabs (
1192+
id SERIAL PRIMARY KEY,
1193+
title TEXT NOT NULL,
1194+
artist TEXT
1195+
);
1196+
`,
1197+
},
1198+
newSchemaDDL: []string{
1199+
`
1200+
CREATE TABLE tabs (
1201+
id SERIAL PRIMARY KEY,
1202+
title TEXT NOT NULL,
1203+
artist TEXT,
1204+
search_vector tsvector GENERATED ALWAYS AS (
1205+
to_tsvector('simple', title || ' ' || coalesce(artist, ''))
1206+
) STORED
1207+
);
1208+
`,
1209+
},
1210+
},
1211+
{
1212+
name: "Drop generated column",
1213+
oldSchemaDDL: []string{
1214+
`
1215+
CREATE TABLE tabs (
1216+
id SERIAL PRIMARY KEY,
1217+
title TEXT NOT NULL,
1218+
artist TEXT,
1219+
search_vector tsvector GENERATED ALWAYS AS (
1220+
to_tsvector('simple', title || ' ' || coalesce(artist, ''))
1221+
) STORED
1222+
);
1223+
`,
1224+
},
1225+
newSchemaDDL: []string{
1226+
`
1227+
CREATE TABLE tabs (
1228+
id SERIAL PRIMARY KEY,
1229+
title TEXT NOT NULL,
1230+
artist TEXT
1231+
);
1232+
`,
1233+
},
1234+
expectedHazardTypes: []diff.MigrationHazardType{
1235+
diff.MigrationHazardTypeDeletesData,
1236+
},
1237+
},
1238+
{
1239+
name: "Add multiple generated columns",
1240+
oldSchemaDDL: []string{
1241+
`
1242+
CREATE TABLE products (
1243+
id SERIAL PRIMARY KEY,
1244+
name TEXT NOT NULL,
1245+
price NUMERIC(10,2) NOT NULL
1246+
);
1247+
`,
1248+
},
1249+
newSchemaDDL: []string{
1250+
`
1251+
CREATE TABLE products (
1252+
id SERIAL PRIMARY KEY,
1253+
name TEXT NOT NULL,
1254+
price NUMERIC(10,2) NOT NULL,
1255+
price_with_tax NUMERIC(10,2) GENERATED ALWAYS AS (price * 1.1) STORED,
1256+
display_name TEXT GENERATED ALWAYS AS (upper(name)) STORED
1257+
);
1258+
`,
1259+
},
1260+
// expectedDBSchemaDDL reflects the actual column order after migration
1261+
// PostgreSQL adds columns in the order of execution, not declaration order
1262+
expectedDBSchemaDDL: []string{
1263+
`
1264+
CREATE TABLE products (
1265+
id SERIAL PRIMARY KEY,
1266+
name TEXT NOT NULL,
1267+
price NUMERIC(10,2) NOT NULL,
1268+
display_name TEXT GENERATED ALWAYS AS (upper(name)) STORED,
1269+
price_with_tax NUMERIC(10,2) GENERATED ALWAYS AS (price * 1.1) STORED
1270+
);
1271+
`,
1272+
},
1273+
},
1274+
{
1275+
name: "Generated column with index",
1276+
oldSchemaDDL: []string{
1277+
`
1278+
CREATE TABLE articles (
1279+
id SERIAL PRIMARY KEY,
1280+
title TEXT NOT NULL,
1281+
content TEXT
1282+
);
1283+
`,
1284+
},
1285+
newSchemaDDL: []string{
1286+
`
1287+
CREATE TABLE articles (
1288+
id SERIAL PRIMARY KEY,
1289+
title TEXT NOT NULL,
1290+
content TEXT,
1291+
search_vector tsvector GENERATED ALWAYS AS (
1292+
to_tsvector('english', title || ' ' || coalesce(content, ''))
1293+
) STORED
1294+
);
1295+
CREATE INDEX idx_articles_search_vector ON articles USING gin (search_vector);
1296+
`,
1297+
},
1298+
expectedHazardTypes: []diff.MigrationHazardType{
1299+
diff.MigrationHazardTypeIndexBuild,
1300+
},
1301+
},
1302+
{
1303+
name: "Generated column no-op",
1304+
oldSchemaDDL: []string{
1305+
`
1306+
CREATE TABLE tabs (
1307+
id SERIAL PRIMARY KEY,
1308+
title TEXT NOT NULL,
1309+
artist TEXT,
1310+
search_vector tsvector GENERATED ALWAYS AS (
1311+
to_tsvector('simple', title || ' ' || coalesce(artist, ''))
1312+
) STORED
1313+
);
1314+
`,
1315+
},
1316+
newSchemaDDL: []string{
1317+
`
1318+
CREATE TABLE tabs (
1319+
id SERIAL PRIMARY KEY,
1320+
title TEXT NOT NULL,
1321+
artist TEXT,
1322+
search_vector tsvector GENERATED ALWAYS AS (
1323+
to_tsvector('simple', title || ' ' || coalesce(artist, ''))
1324+
) STORED
1325+
);
1326+
`,
1327+
},
1328+
expectEmptyPlan: true,
1329+
},
11871330
}
11881331

11891332
func TestColumnTestCases(t *testing.T) {

internal/queries/queries.sql

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -89,11 +89,6 @@ WITH identity_col_seq AS (
8989

9090
SELECT
9191
a.attname::TEXT AS column_name,
92-
COALESCE(coll.collname, '')::TEXT AS collation_name,
93-
COALESCE(collation_namespace.nspname, '')::TEXT AS collation_schema_name,
94-
COALESCE(
95-
pg_catalog.pg_get_expr(d.adbin, d.adrelid), ''
96-
)::TEXT AS default_value,
9792
a.attnotnull AS is_not_null,
9893
a.attlen AS column_size,
9994
a.attidentity::TEXT AS identity_type,
@@ -103,6 +98,22 @@ SELECT
10398
identity_col_seq.seqmin AS min_value,
10499
identity_col_seq.seqcache AS cache_size,
105100
identity_col_seq.seqcycle AS is_cycle,
101+
COALESCE(coll.collname, '')::TEXT AS collation_name,
102+
COALESCE(collation_namespace.nspname, '')::TEXT AS collation_schema_name,
103+
COALESCE(
104+
CASE
105+
WHEN a.attgenerated = 's' THEN ''
106+
ELSE pg_catalog.pg_get_expr(d.adbin, d.adrelid)
107+
END, ''
108+
)::TEXT AS default_value,
109+
COALESCE(
110+
CASE
111+
WHEN a.attgenerated = 's'
112+
THEN pg_catalog.pg_get_expr(d.adbin, d.adrelid)
113+
ELSE ''
114+
END, ''
115+
)::TEXT AS generation_expression,
116+
(a.attgenerated = 's') AS is_generated,
106117
pg_catalog.format_type(a.atttypid, a.atttypmod) AS column_type
107118
FROM pg_catalog.pg_attribute AS a
108119
LEFT JOIN

internal/queries/queries.sql.go

Lines changed: 37 additions & 22 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

internal/schema/schema.go

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -264,8 +264,16 @@ type (
264264
// ''::text
265265
// CURRENT_TIMESTAMP
266266
// If empty, indicates that there is no default value.
267-
Default string
268-
IsNullable bool
267+
Default string
268+
// If the column is a generated column, this will be true.
269+
IsGenerated bool
270+
// If the column is a generated column, this will be the generation expression.
271+
// Examples:
272+
// to_tsvector('simple', title || ' ' || coalesce(artist, ''))
273+
// (price * 1.1)
274+
// Only populated if IsGenerated is true.
275+
GenerationExpression string
276+
IsNullable bool
269277
// Size is the number of bytes required to store the value.
270278
// It is used for data-packing purposes
271279
Size int
@@ -981,9 +989,11 @@ func (s *schemaFetcher) buildTable(
981989
// ''::text
982990
// CURRENT_TIMESTAMP
983991
// If empty, indicates that there is no default value.
984-
Default: column.DefaultValue,
985-
Size: int(column.ColumnSize),
986-
Identity: identity,
992+
Default: column.DefaultValue,
993+
IsGenerated: column.IsGenerated,
994+
GenerationExpression: column.GenerationExpression,
995+
Size: int(column.ColumnSize),
996+
Identity: identity,
987997
})
988998
}
989999

internal/schema/schema_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -229,7 +229,7 @@ var (
229229
SELECT id, author
230230
FROM schema_2.foo;
231231
`},
232-
expectedHash: "f0fb3f95f68ba482",
232+
expectedHash: "ff9ed400558572aa",
233233
expectedSchema: Schema{
234234
NamedSchemas: []NamedSchema{
235235
{Name: "public"},
@@ -571,7 +571,7 @@ var (
571571
ALTER TABLE foo_fk_1 ADD CONSTRAINT foo_fk_1_fk FOREIGN KEY (author, content) REFERENCES foo_1 (author, content)
572572
NOT VALID;
573573
`},
574-
expectedHash: "bcad7c978a081c30",
574+
expectedHash: "9647ef46a878d426",
575575
expectedSchema: Schema{
576576
NamedSchemas: []NamedSchema{
577577
{Name: "public"},

pkg/diff/sql_generator.go

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2637,12 +2637,14 @@ func buildColumnDefinition(column schema.Column) (string, error) {
26372637
if column.IsCollated() {
26382638
sb.WriteString(fmt.Sprintf(" COLLATE %s", column.Collation.GetFQEscapedName()))
26392639
}
2640+
if column.IsGenerated {
2641+
sb.WriteString(fmt.Sprintf(" GENERATED ALWAYS AS (%s) STORED", column.GenerationExpression))
2642+
} else if len(column.Default) > 0 {
2643+
sb.WriteString(fmt.Sprintf(" DEFAULT %s", column.Default))
2644+
}
26402645
if !column.IsNullable {
26412646
sb.WriteString(" NOT NULL")
26422647
}
2643-
if len(column.Default) > 0 {
2644-
sb.WriteString(fmt.Sprintf(" DEFAULT %s", column.Default))
2645-
}
26462648
if column.Identity != nil {
26472649
identityDef, err := buildColumnIdentityDefinition(*column.Identity)
26482650
if err != nil {

0 commit comments

Comments
 (0)