diff --git a/drizzle-kit/src/serializer/pgSchema.ts b/drizzle-kit/src/serializer/pgSchema.ts index d7604d645c..5de349b613 100644 --- a/drizzle-kit/src/serializer/pgSchema.ts +++ b/drizzle-kit/src/serializer/pgSchema.ts @@ -183,7 +183,7 @@ const column = object({ uniqueName: string().optional(), nullsNotDistinct: boolean().optional(), generated: object({ - type: literal('stored'), + type: enumType(['virtual', 'stored']), as: string(), }).optional(), identity: sequenceSchema @@ -207,7 +207,7 @@ const columnSquashed = object({ uniqueName: string().optional(), nullsNotDistinct: boolean().optional(), generated: object({ - type: literal('stored'), + type: enumType(['virtual', 'stored']), as: string(), }).optional(), identity: string().optional(), diff --git a/drizzle-kit/src/serializer/pgSerializer.ts b/drizzle-kit/src/serializer/pgSerializer.ts index a53353dbc2..158d64ae23 100644 --- a/drizzle-kit/src/serializer/pgSerializer.ts +++ b/drizzle-kit/src/serializer/pgSerializer.ts @@ -192,7 +192,7 @@ export const generatePgSnapshot = ( : typeof generated.as === 'function' ? dialect.sqlToQuery(generated.as() as SQL).sql : (generated.as as any), - type: 'stored', + type: generated.mode ?? 'stored', } : undefined, identity: identity @@ -780,7 +780,7 @@ export const generatePgSnapshot = ( : typeof generated.as === 'function' ? dialect.sqlToQuery(generated.as() as SQL).sql : (generated.as as any), - type: 'stored', + type: generated.mode ?? 'stored', } : undefined, identity: identity diff --git a/drizzle-kit/src/sqlgenerator.ts b/drizzle-kit/src/sqlgenerator.ts index 64d3c4063c..315c77f8b9 100644 --- a/drizzle-kit/src/sqlgenerator.ts +++ b/drizzle-kit/src/sqlgenerator.ts @@ -413,7 +413,9 @@ class PgCreateTableConvertor extends Convertor { const type = parseType(schemaPrefix, column.type); const generated = column.generated; - const generatedStatement = generated ? ` GENERATED ALWAYS AS (${generated?.as}) STORED` : ''; + const generatedStatement = generated + ? ` GENERATED ALWAYS AS (${generated?.as}) ${generated?.type.toUpperCase()}` + : ''; const unsquashedIdentity = column.identity ? PgSquasher.unsquashIdentity(column.identity) @@ -1793,7 +1795,9 @@ class PgAlterTableAddColumnConvertor extends Convertor { })` : ''; - const generatedStatement = generated ? ` GENERATED ALWAYS AS (${generated?.as}) STORED` : ''; + const generatedStatement = generated + ? ` GENERATED ALWAYS AS (${generated?.as}) ${generated?.type.toUpperCase()}` + : ''; return `ALTER TABLE ${tableNameWithSchema} ADD COLUMN "${name}" ${fixedType}${primaryKeyStatement}${defaultStatement}${generatedStatement}${notNullStatement}${identityStatement};`; } diff --git a/drizzle-kit/tests/pg-generated.test.ts b/drizzle-kit/tests/pg-generated.test.ts index e9f294891f..f14486b2f5 100644 --- a/drizzle-kit/tests/pg-generated.test.ts +++ b/drizzle-kit/tests/pg-generated.test.ts @@ -527,3 +527,229 @@ test('generated as string: change generated constraint', async () => { 'ALTER TABLE "users" ADD COLUMN "gen_name" text GENERATED ALWAYS AS ("users"."name" || \'hello\') STORED;', ]); }); + +test('generated as callback: add column with virtual generated constraint', async () => { + const from = { + users: pgTable('users', { + id: integer('id'), + id2: integer('id2'), + name: text('name'), + }), + }; + const to = { + users: pgTable('users', { + id: integer('id'), + id2: integer('id2'), + name: text('name'), + generatedName: text('gen_name').generatedAlwaysAs( + (): SQL => sql`${to.users.name} || 'hello'`, + { mode: 'virtual' }, + ), + }), + }; + + const { statements, sqlStatements } = await diffTestSchemas(from, to, []); + + expect(statements).toStrictEqual([ + { + column: { + generated: { + as: '"users"."name" || \'hello\'', + type: 'virtual', + }, + name: 'gen_name', + notNull: false, + primaryKey: false, + type: 'text', + }, + schema: '', + tableName: 'users', + type: 'alter_table_add_column', + }, + ]); + expect(sqlStatements).toStrictEqual([ + `ALTER TABLE "users" ADD COLUMN "gen_name" text GENERATED ALWAYS AS (\"users\".\"name\" || 'hello') VIRTUAL;`, + ]); +}); + +test('generated as SQL: add column with virtual generated constraint', async () => { + const from = { + users: pgTable('users', { + id: integer('id'), + id2: integer('id2'), + name: text('name'), + }), + }; + const to = { + users: pgTable('users', { + id: integer('id'), + id2: integer('id2'), + name: text('name'), + generatedName: text('gen_name').generatedAlwaysAs( + sql`concat("users"."name", 'hello')`, + { mode: 'virtual' }, + ), + }), + }; + + const { statements, sqlStatements } = await diffTestSchemas(from, to, []); + + expect(statements).toStrictEqual([ + { + column: { + generated: { + as: 'concat("users"."name", \'hello\')', + type: 'virtual', + }, + name: 'gen_name', + notNull: false, + primaryKey: false, + type: 'text', + }, + schema: '', + tableName: 'users', + type: 'alter_table_add_column', + }, + ]); + expect(sqlStatements).toStrictEqual([ + `ALTER TABLE "users" ADD COLUMN "gen_name" text GENERATED ALWAYS AS (concat("users"."name", 'hello')) VIRTUAL;`, + ]); +}); + +test('generated as string: add column with virtual generated constraint', async () => { + const from = { + users: pgTable('users', { + id: integer('id'), + id2: integer('id2'), + name: text('name'), + }), + }; + const to = { + users: pgTable('users', { + id: integer('id'), + id2: integer('id2'), + name: text('name'), + generatedName: text('gen_name').generatedAlwaysAs( + `\"users\".\"name\" || 'hello'`, + { mode: 'virtual' }, + ), + }), + }; + + const { statements, sqlStatements } = await diffTestSchemas(from, to, []); + + expect(statements).toStrictEqual([ + { + column: { + generated: { + as: '"users"."name" || \'hello\'', + type: 'virtual', + }, + name: 'gen_name', + notNull: false, + primaryKey: false, + type: 'text', + }, + schema: '', + tableName: 'users', + type: 'alter_table_add_column', + }, + ]); + expect(sqlStatements).toStrictEqual([ + `ALTER TABLE "users" ADD COLUMN "gen_name" text GENERATED ALWAYS AS ("users"."name" || 'hello') VIRTUAL;`, + ]); +}); + +test('change from stored to virtual generated constraint', async () => { + const from = { + users: pgTable('users', { + id: integer('id'), + id2: integer('id2'), + name: text('name'), + generatedName: text('gen_name').generatedAlwaysAs( + sql`"users"."name" || 'hello'`, + { mode: 'stored' }, + ), + }), + }; + const to = { + users: pgTable('users', { + id: integer('id'), + id2: integer('id2'), + name: text('name'), + generatedName: text('gen_name').generatedAlwaysAs( + sql`"users"."name" || 'hello'`, + { mode: 'virtual' }, + ), + }), + }; + + const { statements, sqlStatements } = await diffTestSchemas(from, to, []); + + expect(statements).toStrictEqual([ + { + columnAutoIncrement: undefined, + columnDefault: undefined, + columnGenerated: { as: '"users"."name" || \'hello\'', type: 'virtual' }, + columnName: 'gen_name', + columnNotNull: false, + columnOnUpdate: undefined, + columnPk: false, + newDataType: 'text', + schema: '', + tableName: 'users', + type: 'alter_table_alter_column_alter_generated', + }, + ]); + expect(sqlStatements).toStrictEqual([ + 'ALTER TABLE "users" drop column "gen_name";', + 'ALTER TABLE "users" ADD COLUMN "gen_name" text GENERATED ALWAYS AS ("users"."name" || \'hello\') VIRTUAL;', + ]); +}); + +test('change from virtual to stored generated constraint', async () => { + const from = { + users: pgTable('users', { + id: integer('id'), + id2: integer('id2'), + name: text('name'), + generatedName: text('gen_name').generatedAlwaysAs( + sql`"users"."name" || 'hello'`, + { mode: 'virtual' }, + ), + }), + }; + const to = { + users: pgTable('users', { + id: integer('id'), + id2: integer('id2'), + name: text('name'), + generatedName: text('gen_name').generatedAlwaysAs( + sql`"users"."name" || 'hello'`, + { mode: 'stored' }, + ), + }), + }; + + const { statements, sqlStatements } = await diffTestSchemas(from, to, []); + + expect(statements).toStrictEqual([ + { + columnAutoIncrement: undefined, + columnDefault: undefined, + columnGenerated: { as: '"users"."name" || \'hello\'', type: 'stored' }, + columnName: 'gen_name', + columnNotNull: false, + columnOnUpdate: undefined, + columnPk: false, + newDataType: 'text', + schema: '', + tableName: 'users', + type: 'alter_table_alter_column_alter_generated', + }, + ]); + expect(sqlStatements).toStrictEqual([ + 'ALTER TABLE "users" drop column "gen_name";', + 'ALTER TABLE "users" ADD COLUMN "gen_name" text GENERATED ALWAYS AS ("users"."name" || \'hello\') STORED;', + ]); +}); diff --git a/drizzle-orm/src/pg-core/columns/common.ts b/drizzle-orm/src/pg-core/columns/common.ts index aa1d6c1f23..187b795e6f 100644 --- a/drizzle-orm/src/pg-core/columns/common.ts +++ b/drizzle-orm/src/pg-core/columns/common.ts @@ -30,6 +30,10 @@ export interface ReferenceConfig { }; } +export interface PgGeneratedColumnConfig { + mode?: 'virtual' | 'stored'; +} + export interface PgColumnBuilderBase< T extends ColumnBuilderBaseConfig = ColumnBuilderBaseConfig, TTypeConfig extends object = object, @@ -83,13 +87,13 @@ export abstract class PgColumnBuilder< return this; } - generatedAlwaysAs(as: SQL | T['data'] | (() => SQL)): HasGenerated SQL), config?: PgGeneratedColumnConfig): HasGenerated { this.config.generated = { as, type: 'always', - mode: 'stored', + mode: config?.mode ?? 'stored', }; return this as HasGenerated >(); } + +const usersWithVirtual = pgTable( + 'users', + { + id: serial('id').primaryKey(), + firstName: varchar('first_name', { length: 255 }), + lastName: varchar('last_name', { length: 255 }), + email: text('email').notNull(), + fullName: text('full_name').generatedAlwaysAs( + sql`concat_ws(first_name, ' ', last_name)`, + { mode: 'virtual' }, + ).notNull(), + upperName: text('upper_name').generatedAlwaysAs( + sql` case when first_name is null then null else upper(first_name) end `, + { mode: 'virtual' }, + ), + }, +); + +{ + type User = typeof usersWithVirtual.$inferSelect; + type NewUser = typeof usersWithVirtual.$inferInsert; + + Expect< + Equal< + { + id: number; + firstName: string | null; + lastName: string | null; + email: string; + fullName: string; + upperName: string | null; + }, + User + > + >(); + + Expect< + Equal< + { + email: string; + id?: number | undefined; + firstName?: string | null | undefined; + lastName?: string | null | undefined; + }, + NewUser + > + >(); +} + +{ + type User = InferSelectModel; + type NewUser = InferInsertModel; + + Expect< + Equal< + { + id: number; + firstName: string | null; + lastName: string | null; + email: string; + fullName: string; + upperName: string | null; + }, + User + > + >(); + + Expect< + Equal< + { + email: string; + id?: number | undefined; + firstName?: string | null | undefined; + lastName?: string | null | undefined; + }, + NewUser + > + >(); +} + +{ + const dbUsers = await db.select().from(usersWithVirtual); + + Expect< + Equal< + { + id: number; + firstName: string | null; + lastName: string | null; + email: string; + fullName: string; + upperName: string | null; + }[], + typeof dbUsers + > + >(); +} + +{ + const db = drizzle({} as any, { schema: { users: usersWithVirtual } }); + + const dbUser = await db.query.users.findFirst(); + + Expect< + Equal< + { + id: number; + firstName: string | null; + lastName: string | null; + email: string; + fullName: string; + upperName: string | null; + } | undefined, + typeof dbUser + > + >(); +} + +{ + const db = drizzle({} as any, { schema: { users: usersWithVirtual } }); + + const dbUser = await db.query.users.findMany(); + + Expect< + Equal< + { + id: number; + firstName: string | null; + lastName: string | null; + email: string; + fullName: string; + upperName: string | null; + }[], + typeof dbUser + > + >(); +} + +{ + // @ts-expect-error - Can't use the fullName because it's a generated column + await db.insert(usersWithVirtual).values({ + firstName: 'test', + lastName: 'test', + email: 'test', + fullName: 'test', + }); +} + +{ + await db.update(usersWithVirtual).set({ + firstName: 'test', + lastName: 'test', + email: 'test', + // @ts-expect-error - Can't use the fullName because it's a generated column + fullName: 'test', + }); +}