feat: add root option to afterReset hooks of type command (#137)

Co-authored-by: Chandler Gonzales <cgonzales@logixboard.com>
Co-authored-by: Benjie Gillam <benjie@jemjie.com>

README.md

@@ -601,17 +601,21 @@ When the command is invoked it will have access to the following envvars:
 
 - `GM_DBURL` - the relevant database URL (e.g. the one that was just
   reset/migrated)
-- `GM_DBNAME` - the database name in `GM_DBURL`; you might use this if you need
-  to use separate superuser credentials to install extensions against the
-  database
-- `GM_DBUSER` - the database user in `GM_DBURL`
+- `GM_DBNAME` - the database name in `GM_DBURL`
+- `GM_DBUSER` - the database user in `GM_DBURL` if `root` is `false`.
 - `GM_SHADOW` - set to `1` if we're dealing with the shadow DB, unset otherwise
 
 **IMPORTANT NOTE** the `DATABASE_URL` envvar will be set to the nonsense value
 `postgres://PLEASE:USE@GM_DBURL/INSTEAD` to avoid ambiguity - you almost
 certainly mean to use `GM_DBURL` in your scripts since they will want to change
 whichever database was just reset/migrated/etc (which could be the shadow DB).
 
+The `root` property applies to `command` actions with the similar effects as
+`sql` actions (see above). When `true`, the command will be run with GM_DBURL
+envvar set using the superuser role (i.e. the one defined in
+`rootConnectionString`) but with the database name from `connectionString`. When
+`root` is true, `GM_DBUSER` is not set.
+
 ## Collaboration
 
 The intention is that developers can work on different migrations in parallel,

__tests__/actions.test.ts

@@ -6,11 +6,18 @@ import "./helpers"; // Has side-effects; must come first
 
 import { exec } from "child_process";
 import * as mockFs from "mock-fs";
+import { parse } from "pg-connection-string";
 
 import { executeActions } from "../src/actions";
 import { _migrate } from "../src/commands/migrate";
+import { withClient } from "../src/pg";
 import { parseSettings } from "../src/settings";
-import { mockPgClient, TEST_DATABASE_URL } from "./helpers";
+import {
+  mockPgClient,
+  TEST_DATABASE_NAME,
+  TEST_DATABASE_URL,
+  TEST_ROOT_DATABASE_URL,
+} from "./helpers";
 
 beforeAll(() => {
   // eslint-disable-next-line no-console
@@ -74,6 +81,51 @@ it("runs command actions", async () => {
   expect(typeof mockedExec.mock.calls[0][1].env.GM_DBURL).toBe("string");
 });
 
+it("runs sql afterReset action with correct connection string when root", async () => {
+  mockFs({
+    "migrations/sqlfile1.sql": `[CONTENT:migrations/sqlfile1.sql]`,
+  });
+  const parsedSettings = await parseSettings({
+    connectionString: TEST_DATABASE_URL,
+    afterReset: [{ _: "sql", file: "sqlfile1.sql", root: true }],
+  });
+  const mockedWithClient: jest.Mock<typeof withClient> = withClient as any;
+  mockedWithClient.mockClear();
+  await executeActions(parsedSettings, false, parsedSettings.afterReset);
+  expect(mockedWithClient).toHaveBeenCalledTimes(1);
+  expect(mockedWithClient.mock.calls[0][0]).toBe(TEST_DATABASE_NAME);
+});
+
+it("runs command afterReset action with correct env vars when root", async () => {
+  const parsedSettings = await parseSettings({
+    connectionString: TEST_DATABASE_URL,
+    rootConnectionString: TEST_ROOT_DATABASE_URL,
+    afterReset: [
+      { _: "command", command: "touch testCommandAction", root: true },
+    ],
+  });
+  const mockedExec: jest.Mock<typeof exec> = exec as any;
+  mockedExec.mockClear();
+  mockedExec.mockImplementationOnce((_cmd, _options, callback) =>
+    callback(null, { stdout: "", stderr: "" }),
+  );
+
+  await executeActions(parsedSettings, false, parsedSettings.afterReset);
+  // When `root: true`, GM_DBUSER may be perceived as ambiguous, so we must not set it.
+  expect(mockedExec.mock.calls[0][1].env.GM_DBUSER).toBe(undefined);
+  const connectionStringParts = parse(TEST_DATABASE_URL);
+  const rootConnectionStringParts = parse(TEST_ROOT_DATABASE_URL);
+  expect(rootConnectionStringParts.database).not.toBe(
+    connectionStringParts.database,
+  );
+  const execUrlParts = parse(mockedExec.mock.calls[0][1].env.GM_DBURL);
+  expect(execUrlParts.host).toBe(rootConnectionStringParts.host);
+  expect(execUrlParts.port).toBe(rootConnectionStringParts.port);
+  expect(execUrlParts.user).toBe(rootConnectionStringParts.user);
+  expect(execUrlParts.password).toBe(rootConnectionStringParts.password);
+  expect(execUrlParts.database).toBe(connectionStringParts.database);
+});
+
 it("run normal and non-shadow actions in non-shadow mode", async () => {
   mockFs({
     "migrations/non-shadow-only.sql": `[CONTENT:migrations/non-shadow-only.sql]`,

__tests__/helpers.ts

@@ -27,7 +27,7 @@ if (!/^[a-zA-Z0-9_-]+$/.test(TEST_DATABASE_NAME)) {
   throw new Error("Invalid database name " + TEST_DATABASE_NAME);
 }
 
-const TEST_ROOT_DATABASE_URL: string =
+export const TEST_ROOT_DATABASE_URL: string =
   process.env.TEST_ROOT_DATABASE_URL || "postgres";
 
 export const settings: Settings = {

__tests__/settings.test.ts

@@ -179,6 +179,19 @@ describe("makeRootDatabaseConnectionString", () => {
       "postgres://root:pass@localhost:5432/modified?ssl=true&sslrootcert=./__tests__/data/amazon-rds-ca-cert.pem",
     );
   });
+
+  it("handles a standalone DB name that is not a valid URL", async () => {
+    mockFs.restore();
+    const parsedSettings = await parseSettings({
+      connectionString: exampleConnectionString,
+      rootConnectionString: "just_a_db",
+    });
+    const connectionString = makeRootDatabaseConnectionString(
+      parsedSettings,
+      "modified",
+    );
+    expect(connectionString).toBe("modified");
+  });
 });
 
 describe("actions", () => {

src/actions.ts

@@ -18,13 +18,6 @@ import {
 interface ActionSpecBase {
   _: string;
   shadow?: boolean;
-}
-
-export const DO_NOT_USE_DATABASE_URL = "postgres://PLEASE:USE@GM_DBURL/INSTEAD";
-
-export interface SqlActionSpec extends ActionSpecBase {
-  _: "sql";
-  file: string;
 
   /**
    * USE THIS WITH CARE! Currently only supported by the afterReset hook, all
@@ -35,6 +28,13 @@ export interface SqlActionSpec extends ActionSpecBase {
   root?: boolean;
 }
 
+export const DO_NOT_USE_DATABASE_URL = "postgres://PLEASE:USE@GM_DBURL/INSTEAD";
+
+export interface SqlActionSpec extends ActionSpecBase {
+  _: "sql";
+  file: string;
+}
+
 export interface CommandActionSpec extends ActionSpecBase {
   _: "command";
   command: string;
@@ -70,14 +70,14 @@ export async function executeActions(
     if (actionSpec.shadow !== undefined && actionSpec.shadow !== shadow) {
       continue;
     }
+    const hookConnectionString = actionSpec.root
+      ? makeRootDatabaseConnectionString(parsedSettings, databaseName)
+      : connectionString;
     if (actionSpec._ === "sql") {
       const body = await fsp.readFile(
         `${parsedSettings.migrationsFolder}/${actionSpec.file}`,
         "utf8",
       );
-      const hookConnectionString = actionSpec.root
-        ? makeRootDatabaseConnectionString(parsedSettings, databaseName)
-        : connectionString;
       await withClient(
         hookConnectionString,
         parsedSettings,
@@ -101,8 +101,13 @@ export async function executeActions(
           },
           {
             GM_DBNAME: databaseName,
-            GM_DBUSER: databaseUser,
-            GM_DBURL: connectionString,
+            // When `root: true`, GM_DBUSER may be perceived as ambiguous, so we must not set it.
+            ...(actionSpec.root
+              ? null
+              : {
+                  GM_DBUSER: databaseUser,
+                }),
+            GM_DBURL: hookConnectionString,
             ...(shadow
               ? {
                   GM_SHADOW: "1",

src/settings.ts

@@ -407,6 +407,7 @@ export function makeRootDatabaseConnectionString(
     );
   }
   const parsed = parseURL(rootConnectionString, true);
+  const isJustADatabaseName = !parsed.protocol;
   if (parsed.protocol === "socket:") {
     parsed.query.db = databaseName;
     const query = querystring.stringify(parsed.query);
@@ -416,6 +417,8 @@ export function makeRootDatabaseConnectionString(
     } else {
       return `socket:${parsed.pathname}?${query}`;
     }
+  } else if (isJustADatabaseName) {
+    return databaseName;
   } else {
     parsed.pathname = `/${databaseName}`;
     return formatURL(parsed);