Add ability to drop the replication slot completely on teardown

Author: ericallam

internal-packages/replication/src/client.test.ts

@@ -86,4 +86,99 @@ describe("Replication Client", () => {
       await client.stop();
     }
   );
+
+  postgresAndRedisTest(
+    "should be able to teardown",
+    async ({ postgresContainer, prisma, redisOptions }) => {
+      await prisma.$executeRawUnsafe(`ALTER TABLE public."TaskRun" REPLICA IDENTITY FULL;`);
+
+      const client = new LogicalReplicationClient({
+        name: "test",
+        slotName: "test_slot",
+        publicationName: "test_publication",
+        redisOptions,
+        table: "TaskRun",
+        pgConfig: {
+          connectionString: postgresContainer.getConnectionUri(),
+        },
+      });
+
+      const logs: Array<{
+        lsn: string;
+        log: unknown;
+      }> = [];
+
+      client.events.on("data", (data) => {
+        console.log(data);
+        logs.push(data);
+      });
+
+      client.events.on("error", (error) => {
+        console.error(error);
+      });
+
+      await client.subscribe();
+
+      const organization = await prisma.organization.create({
+        data: {
+          title: "test",
+          slug: "test",
+        },
+      });
+
+      const project = await prisma.project.create({
+        data: {
+          name: "test",
+          slug: "test",
+          organizationId: organization.id,
+          externalRef: "test",
+        },
+      });
+
+      const runtimeEnvironment = await prisma.runtimeEnvironment.create({
+        data: {
+          slug: "test",
+          type: "DEVELOPMENT",
+          projectId: project.id,
+          organizationId: organization.id,
+          apiKey: "test",
+          pkApiKey: "test",
+          shortcode: "test",
+        },
+      });
+
+      // Now we insert a row into the table
+      await prisma.taskRun.create({
+        data: {
+          friendlyId: "run_1234",
+          taskIdentifier: "my-task",
+          payload: JSON.stringify({ foo: "bar" }),
+          traceId: "1234",
+          spanId: "1234",
+          queue: "test",
+          runtimeEnvironmentId: runtimeEnvironment.id,
+          projectId: project.id,
+        },
+      });
+
+      // Wait for a bit of time
+      await setTimeout(50);
+
+      // Now we should see the row in the logs
+      expect(logs.length).toBeGreaterThan(0);
+
+      const slotDropped = await client.teardown();
+
+      expect(slotDropped).toBe(true);
+
+      // Now the replication slot should be gone
+      const slotExists = await prisma.$queryRaw<
+        { exists: boolean }[]
+      >`SELECT EXISTS (SELECT 1 FROM pg_replication_slots WHERE slot_name = 'test_slot');`;
+
+      console.log(slotExists);
+
+      expect(slotExists[0].exists).toBe(false);
+    }
+  );
 });

internal-packages/replication/src/client.ts

@@ -190,9 +190,20 @@ export class LogicalReplicationClient {
       return false;
     }
 
+    this.client = new Client({
+      ...this.options.pgConfig,
+      // @ts-expect-error
+      replication: "database",
+      application_name: this.options.name,
+    });
+    await this.client.connect();
+
     // Drop the slot
     const slotDropped = await this.#dropSlot();
 
+    await this.client.end();
+    this.client = null;
+
     await this.#releaseLeaderLock();
 
     return slotDropped;