From f3013e0e6ee7e000a35b62d4f17ff52454e0296d Mon Sep 17 00:00:00 2001 From: Eric Date: Thu, 14 Nov 2024 09:58:21 -0500 Subject: [PATCH] feat: add support for partitioning (#1295) --- src/operations/tables/createTable.ts | 15 ++- src/operations/tables/shared.ts | 18 +++ src/utils/formatPartitionColumns.ts | 35 ++++++ src/utils/index.ts | 1 + test/migration.spec.ts | 8 +- test/migrations/092_partition.js | 23 ++++ test/operations/tables/createTable.spec.ts | 140 +++++++++++++++++++++ 7 files changed, 234 insertions(+), 6 deletions(-) create mode 100644 src/utils/formatPartitionColumns.ts create mode 100644 test/migrations/092_partition.js diff --git a/src/operations/tables/createTable.ts b/src/operations/tables/createTable.ts index cbd06c1e..e65d7577 100644 --- a/src/operations/tables/createTable.ts +++ b/src/operations/tables/createTable.ts @@ -1,5 +1,10 @@ import type { MigrationOptions } from '../../types'; -import { formatLines, intersection, makeComment } from '../../utils'; +import { + formatLines, + formatPartitionColumns, + intersection, + makeComment, +} from '../../utils'; import type { Name, Reversible } from '../generalTypes'; import type { DropTableOptions } from './dropTable'; import { dropTable } from './dropTable'; @@ -27,6 +32,7 @@ export function createTable(mOptions: MigrationOptions): CreateTable { like, constraints: optionsConstraints = {}, comment, + partition, } = options; const { @@ -64,11 +70,16 @@ export function createTable(mOptions: MigrationOptions): CreateTable { const inheritsStr = inherits ? ` INHERITS (${mOptions.literal(inherits)})` : ''; + + const partitionStr = partition + ? ` PARTITION BY ${partition.strategy} (${formatPartitionColumns(partition, mOptions.literal)})` + : ''; + const tableNameStr = mOptions.literal(tableName); const createTableQuery = `CREATE${temporaryStr} TABLE${ifNotExistsStr} ${tableNameStr} ( ${formatLines(tableDefinition)} -)${inheritsStr};`; +)${inheritsStr}${partitionStr};`; const comments = [...columnComments, ...constraintComments]; if (comment !== undefined) { diff --git a/src/operations/tables/shared.ts b/src/operations/tables/shared.ts index b2b3c033..cf2c39f8 100644 --- a/src/operations/tables/shared.ts +++ b/src/operations/tables/shared.ts @@ -102,6 +102,22 @@ export interface ConstraintOptions { comment?: string; } +export type PartitionStrategy = 'RANGE' | 'LIST' | 'HASH'; + +export interface PartitionColumnOptions { + name: string; + collate?: string; + opclass?: string; +} + +export interface PartitionOptions { + strategy: PartitionStrategy; + columns: + | Array + | string + | PartitionColumnOptions; +} + export interface TableOptions extends IfNotExistsOption { temporary?: boolean; @@ -112,6 +128,8 @@ export interface TableOptions extends IfNotExistsOption { constraints?: ConstraintOptions; comment?: string | null; + + partition?: PartitionOptions; } export function parseReferences( diff --git a/src/utils/formatPartitionColumns.ts b/src/utils/formatPartitionColumns.ts new file mode 100644 index 00000000..e5defc0d --- /dev/null +++ b/src/utils/formatPartitionColumns.ts @@ -0,0 +1,35 @@ +import type { + PartitionColumnOptions, + PartitionOptions, +} from '../operations/tables'; +import { toArray } from './toArray'; + +function formatPartitionColumn( + column: string | PartitionColumnOptions, + literal: (value: string) => string +): string { + if (typeof column === 'string') { + return literal(column); + } + + let formatted = literal(column.name); + + if (column.collate) { + formatted += ` COLLATE ${column.collate}`; + } + + if (column.opclass) { + formatted += ` ${column.opclass}`; + } + + return formatted; +} + +// Helper function to format all partition columns +export function formatPartitionColumns( + partition: PartitionOptions, + literal: (value: string) => string +): string { + const columns = toArray(partition.columns); + return columns.map((col) => formatPartitionColumn(col, literal)).join(', '); +} diff --git a/src/utils/index.ts b/src/utils/index.ts index 83f8bcb9..45b0ad72 100644 --- a/src/utils/index.ts +++ b/src/utils/index.ts @@ -4,6 +4,7 @@ export { decamelize } from './decamelize'; export { escapeValue } from './escapeValue'; export { formatLines } from './formatLines'; export { formatParams } from './formatParams'; +export { formatPartitionColumns } from './formatPartitionColumns'; export { getMigrationTableSchema } from './getMigrationTableSchema'; export { getSchemas } from './getSchemas'; export { identity } from './identity'; diff --git a/test/migration.spec.ts b/test/migration.spec.ts index b4c33be4..20d3c7d6 100644 --- a/test/migration.spec.ts +++ b/test/migration.spec.ts @@ -57,7 +57,7 @@ describe('migration', () => { const filePaths = await getMigrationFilePaths(dir, { logger }); expect(Array.isArray(filePaths)).toBeTruthy(); - expect(filePaths).toHaveLength(91); + expect(filePaths).toHaveLength(92); expect(filePaths).not.toContainEqual(expect.stringContaining('nested')); for (const filePath of filePaths) { @@ -78,7 +78,7 @@ describe('migration', () => { }); expect(Array.isArray(filePaths)).toBeTruthy(); - expect(filePaths).toHaveLength(66); + expect(filePaths).toHaveLength(67); for (const filePath of filePaths) { expect(isAbsolute(filePath)).toBeTruthy(); @@ -94,7 +94,7 @@ describe('migration', () => { }); expect(Array.isArray(filePaths)).toBeTruthy(); - expect(filePaths).toHaveLength(104); + expect(filePaths).toHaveLength(105); expect(filePaths).toContainEqual(expect.stringContaining('nested')); for (const filePath of filePaths) { @@ -114,7 +114,7 @@ describe('migration', () => { }); expect(Array.isArray(filePaths)).toBeTruthy(); - expect(filePaths).toHaveLength(103); + expect(filePaths).toHaveLength(104); expect(filePaths).toContainEqual(expect.stringContaining('nested')); for (const filePath of filePaths) { diff --git a/test/migrations/092_partition.js b/test/migrations/092_partition.js new file mode 100644 index 00000000..c74350e7 --- /dev/null +++ b/test/migrations/092_partition.js @@ -0,0 +1,23 @@ +exports.up = (pgm) => { + pgm.createTable( + 't_partition', + { + id: 'serial', + string: { type: 'text', notNull: true }, + created: { + type: 'timestamp', + notNull: true, + default: pgm.func('current_timestamp'), + }, + }, + { + constraints: { + primaryKey: ['id', 'created'], + }, + partition: { + strategy: 'RANGE', + columns: 'created', + }, + } + ); +}; diff --git a/test/operations/tables/createTable.spec.ts b/test/operations/tables/createTable.spec.ts index a15e17bc..e73abb06 100644 --- a/test/operations/tables/createTable.spec.ts +++ b/test/operations/tables/createTable.spec.ts @@ -530,6 +530,146 @@ COMMENT ON CONSTRAINT "fkColB" ON "myTableName" IS $pga$fk b comment$pga$;`, COMMENT ON CONSTRAINT "my_table_name_fk_col_a" ON "my_table_name" IS $pga$fk a comment$pga$; COMMENT ON CONSTRAINT "fk_col_b" ON "my_table_name" IS $pga$fk b comment$pga$;`, ], + // Add to the existing it.each array: + [ + 'should support RANGE partitioning', + options1, + [ + 'events', + { + id: 'serial', + created_at: 'timestamp', + data: 'jsonb', + }, + { + partition: { + strategy: 'RANGE', + columns: 'created_at', + }, + }, + ], + `CREATE TABLE "events" ( + "id" serial, + "created_at" timestamp, + "data" jsonb +) PARTITION BY RANGE ("created_at");`, + ], + [ + 'should support LIST partitioning with multiple columns', + options1, + [ + 'metrics', + { + id: 'serial', + region: 'text', + category: 'text', + value: 'numeric', + }, + { + partition: { + strategy: 'LIST', + columns: ['region', 'category'], + }, + }, + ], + `CREATE TABLE "metrics" ( + "id" serial, + "region" text, + "category" text, + "value" numeric +) PARTITION BY LIST ("region", "category");`, + ], + [ + 'should support HASH partitioning with operator class', + options1, + [ + 'users', + { + id: 'uuid', + email: 'text', + name: 'text', + }, + { + partition: { + strategy: 'HASH', + columns: { name: 'id', opclass: 'hash_extension.uuid_ops' }, + }, + }, + ], + `CREATE TABLE "users" ( + "id" uuid, + "email" text, + "name" text +) PARTITION BY HASH ("id" hash_extension.uuid_ops);`, + ], + [ + 'should support partitioning with INHERITS', + options1, + [ + 'child_events', + { + id: 'serial', + created_at: 'timestamp', + }, + { + inherits: 'parent_events', + partition: { + strategy: 'RANGE', + columns: 'created_at', + }, + }, + ], + `CREATE TABLE "child_events" ( + "id" serial, + "created_at" timestamp +) INHERITS ("parent_events") PARTITION BY RANGE ("created_at");`, + ], + [ + 'should handle snake case naming with partitioning', + options2, + [ + 'userMetrics', + { + userId: 'uuid', + eventType: 'text', + createdAt: 'timestamp', + }, + { + partition: { + strategy: 'LIST', + columns: 'event_type', + }, + }, + ], + `CREATE TABLE "user_metrics" ( + "user_id" uuid, + "event_type" text, + "created_at" timestamp +) PARTITION BY LIST ("event_type");`, + ], + [ + 'should support partitioning with collation', + options1, + [ + 'posts', + { + id: 'serial', + title: 'text', + language: 'text', + }, + { + partition: { + strategy: 'LIST', + columns: { name: 'language', collate: 'en_US' }, + }, + }, + ], + `CREATE TABLE "posts" ( + "id" serial, + "title" text, + "language" text +) PARTITION BY LIST ("language" COLLATE en_US);`, + ], ] as const)( '%s', (_, optionPreset, [tableName, columns, options], expected) => {