From ef95e13e2ba97b8cfe4c7980f5c4740895360299 Mon Sep 17 00:00:00 2001 From: nemanull Date: Thu, 2 Jul 2026 15:30:22 -0700 Subject: [PATCH] fix/drop-replication-slots-on-remove --- src/commands/remove.spec.ts | 137 +++++++++++++++++++++++++++++++++++- src/commands/remove.ts | 9 +-- src/core/database.spec.ts | 68 ++++++++++++++++++ src/core/database.ts | 18 +++++ 4 files changed, 226 insertions(+), 6 deletions(-) create mode 100644 src/core/database.spec.ts diff --git a/src/commands/remove.spec.ts b/src/commands/remove.spec.ts index bf07f4f..aea1153 100644 --- a/src/commands/remove.spec.ts +++ b/src/commands/remove.spec.ts @@ -1,5 +1,53 @@ -import { describe, expect, it } from '@jest/globals'; -import { parseRemoveTargets } from './remove'; +import { afterEach, beforeEach, describe, expect, it, jest } from '@jest/globals'; +import * as fs from 'node:fs'; +import * as os from 'node:os'; +import * as path from 'node:path'; + +jest.mock('../core/registry', () => ({ + readRegistry: jest.fn(), + writeRegistry: jest.fn(), + removeAllocation: jest.fn(), + findByPath: jest.fn(), +})); + +jest.mock('../core/database', () => ({ + dropDatabase: jest.fn(), +})); + +jest.mock('../core/docker-services', () => ({ + removeDockerServices: jest.fn(), + usesDockerServices: jest.fn(), +})); + +jest.mock('../core/git', () => ({ + getMainWorktreePath: jest.fn(), + removeWorktree: jest.fn(), + getUncommittedChanges: jest.fn(), + getUnsyncedStatus: jest.fn(), +})); + +jest.mock('./setup', () => ({ + loadConfig: jest.fn(), +})); + +import { readRegistry, writeRegistry, removeAllocation } from '../core/registry'; +import { dropDatabase } from '../core/database'; +import { removeDockerServices, usesDockerServices } from '../core/docker-services'; +import { getMainWorktreePath } from '../core/git'; +import { loadConfig } from './setup'; +import { parseRemoveTargets, removeCommand } from './remove'; +import type { Allocation, Registry, WtConfig } from '../types'; + +const mockReadRegistry = readRegistry as jest.MockedFunction; +const mockWriteRegistry = writeRegistry as jest.MockedFunction; +const mockRemoveAllocation = removeAllocation as jest.MockedFunction; +const mockDropDatabase = dropDatabase as jest.MockedFunction; +const mockRemoveDockerServices = + removeDockerServices as jest.MockedFunction; +const mockUsesDockerServices = usesDockerServices as jest.MockedFunction; +const mockGetMainWorktreePath = + getMainWorktreePath as jest.MockedFunction; +const mockLoadConfig = loadConfig as jest.MockedFunction; describe('remove command target parsing', () => { it('parses comma-separated slots', () => { @@ -22,3 +70,88 @@ describe('remove command target parsing', () => { expect(parseRemoveTargets(['1,,2', ' , '])).toEqual(['1', '2']); }); }); + +describe('remove command teardown order', () => { + let tmpDir: string; + let stderrSpy: jest.SpiedFunction; + let consoleLogSpy: jest.SpiedFunction; + + const config: WtConfig = { + baseDatabaseName: 'myapp', + baseWorktreePath: '.worktrees', + portStride: 100, + maxSlots: 50, + services: [{ name: 'web', defaultPort: 3000 }], + dockerServices: [], + envFiles: [], + postSetup: [], + autoInstall: true, + }; + + const allocation: Allocation = { + worktreePath: '/repo/.worktrees/feat-x', + branchName: 'feat/x', + dbName: 'myapp_wt1', + docker: { projectName: 'wt-1-myapp-abcd1234', services: ['electric'] }, + ports: { web: 3100 }, + createdAt: '2026-03-24T00:00:00.000Z', + }; + + beforeEach(() => { + tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'wt-remove-test-')); + fs.writeFileSync( + path.join(tmpDir, '.env'), + 'DATABASE_URL=postgresql://user:pw@localhost:5432/myapp\n', + 'utf-8', + ); + + stderrSpy = jest.spyOn(process.stderr, 'write').mockImplementation(() => true); + consoleLogSpy = jest.spyOn(console, 'log').mockImplementation(() => {}); + + mockGetMainWorktreePath.mockReturnValue(tmpDir); + mockLoadConfig.mockReturnValue(config); + mockReadRegistry.mockReturnValue({ + version: 1, + allocations: { '1': allocation }, + } satisfies Registry); + mockUsesDockerServices.mockReturnValue(true); + mockRemoveDockerServices.mockReturnValue(true); + mockDropDatabase.mockResolvedValue(undefined); + mockRemoveAllocation.mockReturnValue({ version: 1, allocations: {} } satisfies Registry); + process.exitCode = 0; + }); + + afterEach(() => { + stderrSpy.mockRestore(); + consoleLogSpy.mockRestore(); + fs.rmSync(tmpDir, { recursive: true, force: true }); + jest.clearAllMocks(); + process.exitCode = 0; + }); + + it('removes Docker services before dropping the database', async () => { + // Act + await removeCommand(['1'], { json: true, keepDb: false, all: false, force: true }); + + // Assert: Electric (a Docker service) is torn down before DROP DATABASE, so + // the logical replication slot is released before Postgres drops the DB. + expect(mockRemoveDockerServices).toHaveBeenCalledTimes(1); + expect(mockDropDatabase).toHaveBeenCalledTimes(1); + const dockerOrder = mockRemoveDockerServices.mock.invocationCallOrder[0]; + const dropOrder = mockDropDatabase.mock.invocationCallOrder[0]; + if (dockerOrder === undefined || dropOrder === undefined) { + throw new Error('expected both teardown steps to run'); + } + expect(dockerOrder).toBeLessThan(dropOrder); + expect(mockWriteRegistry).toHaveBeenCalledTimes(1); + }); + + it('skips the database drop under --keep-db but still removes Docker services', async () => { + // Act + await removeCommand(['1'], { json: true, keepDb: true, all: false, force: true }); + + // Assert + expect(mockRemoveDockerServices).toHaveBeenCalledTimes(1); + expect(mockDropDatabase).not.toHaveBeenCalled(); + }); +}); diff --git a/src/commands/remove.ts b/src/commands/remove.ts index 8cfb660..045196b 100644 --- a/src/commands/remove.ts +++ b/src/commands/remove.ts @@ -218,6 +218,11 @@ export async function removeCommand( try { log(`Removing slot ${resolved.slot} (${resolved.worktreePath})`); + // Stop Docker (e.g. Electric) before the DB drop so its replication slot is inactive. + const dockerRemoved = dockerServicesEnabled || resolved.dockerProjectName !== undefined + ? removeDockerServices(mainRoot, resolved.slot, log) + : false; + if (dbContext !== null) { await dropDatabase( dbContext.databaseUrl, @@ -229,10 +234,6 @@ export async function removeCommand( log(`Skipping database drop for '${resolved.dbName}' (--keep-db).`); } - const dockerRemoved = dockerServicesEnabled || resolved.dockerProjectName !== undefined - ? removeDockerServices(mainRoot, resolved.slot, log) - : false; - if (fs.existsSync(resolved.worktreePath)) { removeWorktree( resolved.worktreePath, diff --git a/src/core/database.spec.ts b/src/core/database.spec.ts new file mode 100644 index 0000000..d5aeaf9 --- /dev/null +++ b/src/core/database.spec.ts @@ -0,0 +1,68 @@ +import { afterEach, describe, expect, it, jest } from '@jest/globals'; + +type QueryResult = { rows: Array> }; + +const mockConnect = jest.fn<() => Promise>(); +const mockEnd = jest.fn<() => Promise>(); +const mockQuery = jest.fn<(sql: string, params?: readonly unknown[]) => Promise>(); + +jest.mock('pg', () => ({ + Client: jest.fn(() => ({ + connect: mockConnect, + query: mockQuery, + end: mockEnd, + })), +})); + +import { dropDatabase } from './database'; + +const DATABASE_URL = 'postgresql://user:pw@localhost:5432/myapp'; + +describe('dropDatabase replication slot cleanup', () => { + afterEach(() => { + jest.clearAllMocks(); + }); + + it('drops logical replication slots bound to the database before DROP DATABASE', async () => { + // Arrange: the target database still has an Electric replication slot. + mockConnect.mockResolvedValue(undefined); + mockEnd.mockResolvedValue(undefined); + mockQuery.mockImplementation((sql) => { + if (sql.includes('pg_replication_slots')) { + return Promise.resolve({ rows: [{ slot_name: 'electric_slot_myapp_wt1' }] }); + } + return Promise.resolve({ rows: [] }); + }); + + // Act + await dropDatabase(DATABASE_URL, 'myapp_wt1', 'myapp'); + + // Assert: the slot is dropped, and strictly before the database itself. + const statements = mockQuery.mock.calls.map((call) => call[0]); + const dropSlotIndex = statements.findIndex((sql) => sql.includes('pg_drop_replication_slot')); + const dropDbIndex = statements.findIndex((sql) => sql.includes('DROP DATABASE')); + expect(dropSlotIndex).toBeGreaterThanOrEqual(0); + expect(dropDbIndex).toBeGreaterThanOrEqual(0); + expect(dropSlotIndex).toBeLessThan(dropDbIndex); + + const dropSlotCall = mockQuery.mock.calls.find((call) => + call[0].includes('pg_drop_replication_slot'), + ); + expect(dropSlotCall?.[1]).toEqual(['electric_slot_myapp_wt1']); + }); + + it('drops the database when no replication slots remain', async () => { + // Arrange: no slots reference the target database. + mockConnect.mockResolvedValue(undefined); + mockEnd.mockResolvedValue(undefined); + mockQuery.mockResolvedValue({ rows: [] }); + + // Act + await dropDatabase(DATABASE_URL, 'myapp_wt1', 'myapp'); + + // Assert + const statements = mockQuery.mock.calls.map((call) => call[0]); + expect(statements.some((sql) => sql.includes('DROP DATABASE'))).toBe(true); + expect(statements.some((sql) => sql.includes('pg_drop_replication_slot'))).toBe(false); + }); +}); diff --git a/src/core/database.ts b/src/core/database.ts index f195882..7c99c51 100644 --- a/src/core/database.ts +++ b/src/core/database.ts @@ -68,6 +68,22 @@ export async function createDatabase( }); } +/** Drops a database's logical replication slots; caller must make them inactive first. */ +async function dropReplicationSlots( + client: Client, + dbName: string, + logSql?: SqlLogger, +): Promise { + const selectSql = 'SELECT slot_name FROM pg_replication_slots WHERE database = $1'; + logSql?.(formatQueryLog(selectSql, [dbName])); + const result = await client.query<{ slot_name: string }>(selectSql, [dbName]); + for (const { slot_name } of result.rows) { + const dropSlotSql = 'SELECT pg_drop_replication_slot($1)'; + logSql?.(formatQueryLog(dropSlotSql, [slot_name])); + await client.query(dropSlotSql, [slot_name]); + } +} + /** Drop a database if it exists. Refuses to drop the template database. */ export async function dropDatabase( databaseUrl: string, @@ -85,6 +101,8 @@ export async function dropDatabase( logSql?.(formatQueryLog(terminateSql, [dbName])); await client.query(terminateSql, [dbName]); + await dropReplicationSlots(client, dbName, logSql); + const dropSql = `DROP DATABASE IF EXISTS ${quoteIdentifier(dbName)}`; logSql?.(formatQueryLog(dropSql)); await client.query(dropSql);