feat: add support for partitioning (#1295)

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) {

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>
+    | 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(

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(', ');
+}

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';

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) {

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',
+      },
+    }
+  );
+};

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) => {