rake-db 2.0.1 → 2.0.3

package/src/commands/generate.ts

@@ -41,18 +41,18 @@ const makeContent = (name: string, args: string[]): string => {
   const [first, rest] = getFirstWordAndRest(name);
   if (rest) {
     if (first === 'create' || first === 'drop') {
-      content += `\n  db.${
+      content += `\n  await db.${
         first === 'create' ? 'createTable' : 'dropTable'
       }('${rest}', (t) => ({`;
       content += makeColumnsContent(args);
       content += '\n  }));';
     } else if (first === 'change') {
-      content += `\n  db.changeTable('${rest}', (t) => ({`;
+      content += `\n  await db.changeTable('${rest}', (t) => ({`;
       content += '\n  }));';
     } else if (first === 'add' || first === 'remove') {
       const table =
         first === 'add' ? getTextAfterTo(rest) : getTextAfterFrom(rest);
-      content += `\n  db.changeTable(${
+      content += `\n  await db.changeTable(${
         table ? `'${table}'` : 'tableName'
       }, (t) => ({`;
       content += makeColumnsContent(args, first);

package/src/commands/migrateOrRollback.test.ts

@@ -1,7 +1,7 @@
 import { migrate, rollback } from './migrateOrRollback';
 import { createSchemaMigrations, migrationConfigDefaults } from '../common';
 import { getMigrationFiles } from '../common';
-import { Adapter, TransactionAdapter } from 'pqb';
+import { Adapter, noop, TransactionAdapter } from 'pqb';
 import { Migration } from '../migration/migration';
 
 jest.mock('../common', () => ({
@@ -35,6 +35,11 @@ const requireTsMock = jest.fn();
 const config = {
   ...migrationConfigDefaults,
   requireTs: requireTsMock,
+  log: false,
+  logger: {
+    log: jest.fn(),
+    error: noop,
+  },
 };
 
 describe('migrateOrRollback', () => {
@@ -49,7 +54,7 @@ describe('migrateOrRollback', () => {
       queryMock.mockReturnValueOnce(undefined);
       requireTsMock.mockResolvedValue(undefined);
 
-      await migrate(options, config);
+      await migrate(options, config, []);
 
       expect(getMigrationFiles).toBeCalledWith(config, true);
 
@@ -62,6 +67,9 @@ describe('migrateOrRollback', () => {
       expect(transactionQueryMock).toBeCalledWith(
         `INSERT INTO "schemaMigrations" VALUES ('3')`,
       );
+
+      expect(config.logger.log).toBeCalledWith('file2 migrated');
+      expect(config.logger.log).toBeCalledWith('file3 migrated');
     });
 
     it('should create migrations table if it not exist', async () => {
@@ -69,12 +77,13 @@ describe('migrateOrRollback', () => {
       getMigratedVersionsArrayMock.mockRejectedValueOnce({ code: '42P01' });
       (createSchemaMigrations as jest.Mock).mockResolvedValueOnce(undefined);
 
-      await migrate(options, config);
+      await migrate(options, config, []);
 
       expect(getMigrationFiles).toBeCalledWith(config, true);
       expect(createSchemaMigrations).toBeCalled();
       expect(requireTsMock).not.toBeCalled();
       expect(transactionQueryMock).not.toBeCalled();
+      expect(config.logger.log).not.toBeCalled();
     });
   });
 
@@ -87,19 +96,20 @@ describe('migrateOrRollback', () => {
       queryMock.mockReturnValueOnce(undefined);
       requireTsMock.mockResolvedValue(undefined);
 
-      await rollback(options, config);
+      await rollback(options, config, []);
 
       expect(getMigrationFiles).toBeCalledWith(config, false);
 
+      expect(requireTsMock).toBeCalledTimes(1);
       expect(requireTsMock).toBeCalledWith('file2');
-      expect(requireTsMock).toBeCalledWith('file1');
 
+      expect(transactionQueryMock).toBeCalledTimes(1);
       expect(transactionQueryMock).toBeCalledWith(
         `DELETE FROM "schemaMigrations" WHERE version = '2'`,
       );
-      expect(transactionQueryMock).toBeCalledWith(
-        `DELETE FROM "schemaMigrations" WHERE version = '1'`,
-      );
+
+      expect(config.logger.log).toBeCalledTimes(1);
+      expect(config.logger.log).toBeCalledWith('file2 rolled back');
     });
 
     it('should create migrations table if it not exist', async () => {
@@ -107,12 +117,13 @@ describe('migrateOrRollback', () => {
       getMigratedVersionsArrayMock.mockRejectedValueOnce({ code: '42P01' });
       (createSchemaMigrations as jest.Mock).mockResolvedValueOnce(undefined);
 
-      await rollback(options, config);
+      await rollback(options, config, []);
 
       expect(getMigrationFiles).toBeCalledWith(config, false);
       expect(createSchemaMigrations).toBeCalled();
       expect(requireTsMock).not.toBeCalled();
       expect(transactionQueryMock).not.toBeCalled();
+      expect(config.logger.log).not.toBeCalled();
     });
   });
 });

package/src/commands/migrateOrRollback.ts

@@ -4,6 +4,7 @@ import {
   getMigrationFiles,
   MigrationConfig,
   MigrationFile,
+  quoteTable,
 } from '../common';
 import {
   getCurrentPromise,
@@ -15,10 +16,14 @@ import { Migration } from '../migration/migration';
 const migrateOrRollback = async (
   options: MaybeArray<AdapterOptions>,
   config: MigrationConfig,
+  args: string[],
   up: boolean,
 ) => {
   const files = await getMigrationFiles(config, up);
 
+  const argCount = args[0] === 'all' ? Infinity : parseInt(args[0]);
+  let count = isNaN(argCount) ? (up ? Infinity : 1) : argCount;
+
   for (const opts of toArray(options)) {
     const db = new Adapter(opts);
     const migratedVersions = await getMigratedVersionsMap(db, config);
@@ -31,7 +36,10 @@ const migrateOrRollback = async (
           continue;
         }
 
+        if (count-- <= 0) break;
+
         await processMigration(db, up, file, config);
+        config.logger?.log(`${file.path} ${up ? 'migrated' : 'rolled back'}`);
       }
     } finally {
       await db.destroy();
@@ -46,7 +54,7 @@ const processMigration = async (
   config: MigrationConfig,
 ) => {
   await db.transaction(async (tx) => {
-    const db = new Migration(tx, up);
+    const db = new Migration(tx, up, config);
     setCurrentMigration(db);
     setCurrentMigrationUp(up);
     config.requireTs(file.path);
@@ -65,7 +73,7 @@ const saveMigratedVersion = async (
   config: MigrationConfig,
 ) => {
   await db.query(
-    `INSERT INTO "${config.migrationsTable}" VALUES ('${version}')`,
+    `INSERT INTO ${quoteTable(config.migrationsTable)} VALUES ('${version}')`,
   );
 };
 
@@ -75,7 +83,9 @@ const removeMigratedVersion = async (
   config: MigrationConfig,
 ) => {
   await db.query(
-    `DELETE FROM "${config.migrationsTable}" WHERE version = '${version}'`,
+    `DELETE FROM ${quoteTable(
+      config.migrationsTable,
+    )} WHERE version = '${version}'`,
   );
 };
 
@@ -85,7 +95,7 @@ const getMigratedVersionsMap = async (
 ): Promise<Record<string, boolean>> => {
   try {
     const result = await db.arrays<[string]>(
-      `SELECT * FROM "${config.migrationsTable}"`,
+      `SELECT * FROM ${quoteTable(config.migrationsTable)}`,
     );
     return Object.fromEntries(result.rows.map((row) => [row[0], true]));
   } catch (err) {
@@ -100,9 +110,11 @@ const getMigratedVersionsMap = async (
 export const migrate = (
   options: MaybeArray<AdapterOptions>,
   config: MigrationConfig,
-) => migrateOrRollback(options, config, true);
+  args: string[],
+) => migrateOrRollback(options, config, args, true);
 
 export const rollback = (
   options: MaybeArray<AdapterOptions>,
   config: MigrationConfig,
-) => migrateOrRollback(options, config, false);
+  args: string[],
+) => migrateOrRollback(options, config, args, false);

package/src/common.test.ts

@@ -8,6 +8,7 @@ import {
   joinColumns,
   joinWords,
   migrationConfigDefaults,
+  quoteTable,
   setAdapterOptions,
   setAdminCredentialsToOptions,
   sortAsc,
@@ -45,6 +46,8 @@ describe('common', () => {
         migrationsPath: 'custom-path',
         migrationsTable: 'schemaMigrations',
         requireTs: expect.any(Function),
+        log: true,
+        logger: console,
       });
     });
   });
@@ -278,4 +281,14 @@ describe('common', () => {
       expect(joinColumns(['a', 'b', 'c'])).toBe('"a", "b", "c"');
     });
   });
+
+  describe('quoteTable', () => {
+    it('should quote a table', () => {
+      expect(quoteTable('table')).toBe('"table"');
+    });
+
+    it('should quote a table with schema', () => {
+      expect(quoteTable('schema.table')).toBe('"schema"."table"');
+    });
+  });
 });

package/src/common.ts

@@ -1,4 +1,4 @@
-import { Adapter, AdapterOptions } from 'pqb';
+import { Adapter, AdapterOptions, QueryLogOptions } from 'pqb';
 import Enquirer from 'enquirer';
 import path from 'path';
 import { readdir } from 'fs/promises';
@@ -7,7 +7,7 @@ export type MigrationConfig = {
   migrationsPath: string;
   migrationsTable: string;
   requireTs(path: string): void;
-};
+} & QueryLogOptions;
 
 const registered = false;
 
@@ -21,6 +21,8 @@ export const migrationConfigDefaults = {
     }
     require(path);
   },
+  log: true,
+  logger: console,
 };
 
 export const getMigrationConfigWithDefaults = (
@@ -116,7 +118,9 @@ export const createSchemaMigrations = async (
 ) => {
   try {
     await db.query(
-      `CREATE TABLE "${config.migrationsTable}" ( version TEXT NOT NULL )`,
+      `CREATE TABLE ${quoteTable(
+        config.migrationsTable,
+      )} ( version TEXT NOT NULL )`,
     );
     console.log('Created versions table');
   } catch (err) {
@@ -222,3 +226,12 @@ export const joinWords = (...words: string[]) => {
 export const joinColumns = (columns: string[]) => {
   return columns.map((column) => `"${column}"`).join(', ');
 };
+
+export const quoteTable = (table: string) => {
+  const index = table.indexOf('.');
+  if (index !== -1) {
+    return `"${table.slice(0, index)}"."${table.slice(index + 1)}"`;
+  } else {
+    return `"${table}"`;
+  }
+};

package/src/index.ts

@@ -1,5 +1,6 @@
 export * from './commands/createOrDrop';
 export * from './commands/generate';
-export * from './commands/migrateOrRollback.test';
+export * from './commands/migrateOrRollback';
 export { change } from './migration/change';
 export * from './migration/migration';
+export { rakeDb } from './rakeDb';

package/src/migration/changeTable.ts

@@ -29,6 +29,7 @@ import {
   migrateIndexes,
   primaryKeyToSql,
 } from './migrationUtils';
+import { quoteTable } from '../common';
 
 const newChangeTableData = () => ({
   add: [],
@@ -203,7 +204,8 @@ export const changeTable = async (
 
   if (state.alterTable.length) {
     await migration.query(
-      `ALTER TABLE "${tableName}"` + `\n  ${state.alterTable.join(',\n  ')}`,
+      `ALTER TABLE ${quoteTable(tableName)}` +
+        `\n  ${state.alterTable.join(',\n  ')}`,
     );
   }
 
@@ -227,7 +229,7 @@ const changeActions = {
       value = Array.isArray(comment) ? comment[0] : null;
     }
     return migration.query(
-      `COMMENT ON TABLE "${tableName}" IS ${quote(value)}`,
+      `COMMENT ON TABLE ${quoteTable(tableName)} IS ${quote(value)}`,
     );
   },
 

package/src/migration/createTable.ts

@@ -24,7 +24,7 @@ import {
   migrateIndexes,
   primaryKeyToSql,
 } from './migrationUtils';
-import { joinWords } from '../common';
+import { joinWords, quoteTable } from '../common';
 import { singular } from 'pluralize';
 
 class UnknownColumn extends ColumnType {
@@ -49,18 +49,23 @@ export const createJoinTable = async (
   }
 
   const tablesWithPrimaryKeys = await Promise.all(
-    tables.map(
-      async (table) =>
-        [
-          table,
-          await getPrimaryKeysOfTable(migration, table).then((items) =>
-            items.map((item) => ({
-              ...item,
-              joinedName: joinWords(singular(table), item.name),
-            })),
-          ),
-        ] as const,
-    ),
+    tables.map(async (table) => {
+      const primaryKeys = await getPrimaryKeysOfTable(migration, table).then(
+        (items) =>
+          items.map((item) => ({
+            ...item,
+            joinedName: joinWords(singular(table), item.name),
+          })),
+      );
+
+      if (!primaryKeys.length) {
+        throw new Error(
+          `Primary key for table ${quoteTable(table)} is not defined`,
+        );
+      }
+
+      return [table, primaryKeys] as const;
+    }),
   );
 
   return createTable(migration, up, tableName, options, (t) => {
@@ -114,7 +119,7 @@ export const createTable = async (
   if (!up) {
     const { dropMode } = options;
     await migration.query(
-      `DROP TABLE "${tableName}"${dropMode ? ` ${dropMode}` : ''}`,
+      `DROP TABLE ${quoteTable(tableName)}${dropMode ? ` ${dropMode}` : ''}`,
     );
     return;
   }
@@ -152,7 +157,7 @@ export const createTable = async (
   });
 
   await migration.query({
-    text: `CREATE TABLE "${tableName}" (${lines.join(',')}\n)`,
+    text: `CREATE TABLE ${quoteTable(tableName)} (${lines.join(',')}\n)`,
     values: state.values,
   });
 
@@ -163,7 +168,7 @@ export const createTable = async (
 
   if (options.comment) {
     await migration.query(
-      `COMMENT ON TABLE "${tableName}" IS ${quote(options.comment)}`,
+      `COMMENT ON TABLE ${quoteTable(tableName)} IS ${quote(options.comment)}`,
     );
   }
 };

package/src/migration/migration.test.ts

@@ -306,6 +306,99 @@ describe('migration', () => {
           ? expectDropTable
           : expectCreateTable)();
       });
+
+      it('should throw error if table has no primary key', async () => {
+        if (action === 'dropJoinTable') {
+          db.up = false;
+        }
+
+        (getPrimaryKeysOfTable as jest.Mock)
+          .mockResolvedValueOnce([
+            {
+              name: 'id',
+              type: 'integer',
+            },
+          ])
+          .mockResolvedValueOnce([]);
+
+        await expect(db[action](['posts', 'comments'])).rejects.toThrow(
+          'Primary key for table "comments" is not defined',
+        );
+      });
+    });
+  });
+
+  (['createSchema', 'dropSchema'] as const).forEach((action) => {
+    describe(action, () => {
+      it(`should ${
+        action === 'createSchema' ? 'add' : 'drop'
+      } a schema`, async () => {
+        const fn = () => {
+          return db[action]('schemaName');
+        };
+
+        const expectCreateSchema = () => {
+          expectSql(`
+            CREATE SCHEMA "schemaName"
+          `);
+        };
+
+        const expectDropSchema = () => {
+          expectSql(`
+            DROP SCHEMA "schemaName"
+          `);
+        };
+
+        await fn();
+        (action === 'createSchema' ? expectCreateSchema : expectDropSchema)();
+
+        db.up = false;
+        queryMock.mockClear();
+        await fn();
+        (action === 'createSchema' ? expectDropSchema : expectCreateSchema)();
+      });
+    });
+  });
+
+  (['createExtension', 'dropExtension'] as const).forEach((action) => {
+    describe(action, () => {
+      it(`should ${
+        action === 'createExtension' ? 'add' : 'drop'
+      } an extension`, async () => {
+        const fn = () => {
+          return db[action]('extensionName', {
+            ifExists: true,
+            ifNotExists: true,
+            schema: 'schemaName',
+            version: '123',
+            cascade: true,
+          });
+        };
+
+        const expectCreateExtension = () => {
+          expectSql(`
+            CREATE EXTENSION IF NOT EXISTS "extensionName" SCHEMA "schemaName" VERSION '123' CASCADE
+          `);
+        };
+
+        const expectDropExtension = () => {
+          expectSql(`
+            DROP EXTENSION IF EXISTS "extensionName" CASCADE
+          `);
+        };
+
+        await fn();
+        (action === 'createExtension'
+          ? expectCreateExtension
+          : expectDropExtension)();
+
+        db.up = false;
+        queryMock.mockClear();
+        await fn();
+        (action === 'createExtension'
+          ? expectDropExtension
+          : expectCreateExtension)();
+      });
     });
   });
 

package/src/migration/migration.ts

@@ -5,11 +5,21 @@ import {
   ForeignKeyOptions,
   IndexColumnOptions,
   IndexOptions,
+  logParamToLogObject,
   MaybeArray,
+  QueryArraysResult,
+  QueryInput,
+  QueryLogObject,
+  QueryLogOptions,
+  QueryResult,
+  QueryResultRow,
+  Sql,
   TransactionAdapter,
+  TypeParsers,
 } from 'pqb';
 import { createJoinTable, createTable } from './createTable';
 import { changeTable, TableChangeData, TableChanger } from './changeTable';
+import { quoteTable } from '../common';
 
 export type DropMode = 'CASCADE' | 'RESTRICT';
 
@@ -31,9 +41,40 @@ export type JoinTableOptions = {
   dropMode?: DropMode;
 };
 
+export type ExtensionOptions = {
+  schema?: string;
+  version?: string;
+  cascade?: boolean;
+};
+
 export class Migration extends TransactionAdapter {
-  constructor(tx: TransactionAdapter, public up: boolean) {
+  public log?: QueryLogObject;
+
+  constructor(
+    tx: TransactionAdapter,
+    public up: boolean,
+    options: QueryLogOptions,
+  ) {
     super(tx.pool, tx.client, tx.types);
+    this.log = logParamToLogObject(options.logger || console, options.log);
+  }
+
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  async query<T extends QueryResultRow = any>(
+    query: QueryInput,
+    types: TypeParsers = this.types,
+    log = this.log,
+  ): Promise<QueryResult<T>> {
+    return wrapWithLog(log, query, () => super.query(query, types));
+  }
+
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  async arrays<R extends any[] = any[]>(
+    query: QueryInput,
+    types: TypeParsers = this.types,
+    log = this.log,
+  ): Promise<QueryArraysResult<R>> {
+    return wrapWithLog(log, query, () => super.arrays(query, types));
   }
 
   createTable(
@@ -123,7 +164,7 @@ export class Migration extends TransactionAdapter {
 
   async renameTable(from: string, to: string): Promise<void> {
     const [table, newName] = this.up ? [from, to] : [to, from];
-    await this.query(`ALTER TABLE "${table}" RENAME TO "${newName}"`);
+    await this.query(`ALTER TABLE ${quoteTable(table)} RENAME TO "${newName}"`);
   }
 
   addColumn(
@@ -216,6 +257,34 @@ export class Migration extends TransactionAdapter {
     }));
   }
 
+  createSchema(schemaName: string) {
+    return createSchema(this, this.up, schemaName);
+  }
+
+  dropSchema(schemaName: string) {
+    return createSchema(this, !this.up, schemaName);
+  }
+
+  createExtension(
+    name: string,
+    options: ExtensionOptions & { ifNotExists?: boolean } = {},
+  ) {
+    return createExtension(this, this.up, name, {
+      ...options,
+      checkExists: options.ifNotExists,
+    });
+  }
+
+  dropExtension(
+    name: string,
+    options: { ifExists?: boolean; cascade?: boolean } = {},
+  ) {
+    return createExtension(this, !this.up, name, {
+      ...options,
+      checkExists: options.ifExists,
+    });
+  }
+
   async tableExists(tableName: string) {
     return queryExists(this, {
       text: `SELECT 1 FROM "information_schema"."tables" WHERE "table_name" = $1`,
@@ -238,6 +307,35 @@ export class Migration extends TransactionAdapter {
   }
 }
 
+const wrapWithLog = async <Result>(
+  log: QueryLogObject | undefined,
+  query: QueryInput,
+  fn: () => Promise<Result>,
+): Promise<Result> => {
+  if (!log) {
+    return fn();
+  } else {
+    const sql = (
+      typeof query === 'string'
+        ? { text: query, values: [] }
+        : query.values
+        ? query
+        : { ...query, values: [] }
+    ) as Sql;
+
+    const logData = log.beforeQuery(sql);
+
+    try {
+      const result = await fn();
+      log.afterQuery(sql, logData);
+      return result;
+    } catch (err) {
+      log.onError(err as Error, sql, logData);
+      throw err;
+    }
+  }
+};
+
 const addColumn = (
   migration: Migration,
   up: boolean,
@@ -288,6 +386,43 @@ const addPrimaryKey = (
   }));
 };
 
+const createSchema = (
+  migration: Migration,
+  up: boolean,
+  schemaName: string,
+) => {
+  if (up) {
+    return migration.query(`CREATE SCHEMA "${schemaName}"`);
+  } else {
+    return migration.query(`DROP SCHEMA "${schemaName}"`);
+  }
+};
+
+const createExtension = (
+  migration: Migration,
+  up: boolean,
+  name: string,
+  options: ExtensionOptions & {
+    checkExists?: boolean;
+  },
+) => {
+  if (!up) {
+    return migration.query(
+      `DROP EXTENSION${options.checkExists ? ' IF EXISTS' : ''} "${name}"${
+        options.cascade ? ' CASCADE' : ''
+      }`,
+    );
+  }
+
+  return migration.query(
+    `CREATE EXTENSION${options.checkExists ? ' IF NOT EXISTS' : ''} "${name}"${
+      options.schema ? ` SCHEMA "${options.schema}"` : ''
+    }${options.version ? ` VERSION '${options.version}'` : ''}${
+      options.cascade ? ' CASCADE' : ''
+    }`,
+  );
+};
+
 const queryExists = (
   db: Migration,
   sql: { text: string; values: unknown[] },