feat: add a ReplicationSlot custom resource

cdk-postgresql/lib/handler.ts

@@ -4,6 +4,7 @@ import { VError } from "verror";
 import { handler as dbHandler } from "./database.handler";
 import { handler as roleHandler } from "./role.handler";
 import { handler as roleMembershipHandler } from "./role-membership.handler";
+import { handler as replicationSlotHandler } from "./replication-slot.handler";
 
 export const handler = async (event: CloudFormationCustomResourceEvent) => {
   switch (event.ResourceType) {
@@ -13,6 +14,8 @@ export const handler = async (event: CloudFormationCustomResourceEvent) => {
       return dbHandler(event);
     case "Custom::Postgresql-RoleMembership":
       return roleMembershipHandler(event);
+    case "Custom::Postgresql-ReplicationSlot":
+      return replicationSlotHandler(event);
     default:
       throw new VError(`unexpected ResourceType: ${event.ResourceType}`);
   }

cdk-postgresql/lib/index.ts

@@ -1,4 +1,5 @@
 export * from "./database";
 export * from "./role";
 export * from "./role-membership";
+export * from "./replication-slot";
 export * from "./provider";

cdk-postgresql/lib/postgres.ts

@@ -105,3 +105,34 @@ export const revokeRoleMembership = async (props: {
     console.warn(thrown.message);
   }
 };
+
+export const createReplicationSlot = async (props: {
+  client: Client;
+  name: string;
+  plugin: string;
+}) => {
+  const { client, name, plugin } = props;
+
+  await client.query("SELECT pg_create_logical_replication_slot($1, $2)", [
+    name,
+    plugin,
+  ]);
+};
+
+/**
+ * Dropping is scoped to the connected database. A slot name is unique across
+ * the whole cluster, so an unscoped drop would destroy a slot decoding another
+ * database, and a dropped slot loses its replication position permanently.
+ * A slot that is already gone is left alone.
+ */
+export const dropReplicationSlot = async (props: {
+  client: Client;
+  name: string;
+}) => {
+  const { client, name } = props;
+
+  await client.query(
+    "SELECT pg_drop_replication_slot(slot_name) FROM pg_replication_slots WHERE slot_name = $1 AND database = current_database()",
+    [name]
+  );
+};

cdk-postgresql/lib/replication-slot.handler.ts

@@ -0,0 +1,216 @@
+import {
+  CloudFormationCustomResourceEvent,
+  CloudFormationCustomResourceCreateEvent,
+  CloudFormationCustomResourceUpdateEvent,
+  CloudFormationCustomResourceDeleteEvent,
+} from "aws-lambda/trigger/cloudformation-custom-resource";
+
+import { validateConnection, hashCode, getConnectedClient } from "./util";
+import { Connection } from "./lambda.types";
+import * as postgres from "./postgres";
+
+interface Props {
+  ServiceToken: string;
+  Connection: Connection;
+  Name: string;
+  Plugin: string;
+  RequiredPublication?: string;
+}
+
+const SLOT_PROVENANCES = ["created", "adopted"] as const;
+
+/**
+ * Whether this resource created the slot or adopted one that already existed.
+ * It is recorded in the physical id, which CloudFormation hands back verbatim
+ * on delete, so only a slot this resource created is ever dropped.
+ */
+type SlotProvenance = (typeof SLOT_PROVENANCES)[number];
+
+export const handler = async (event: CloudFormationCustomResourceEvent) => {
+  switch (event.RequestType) {
+    case "Create":
+      return handleCreate(event);
+    case "Update":
+      return handleUpdate(event);
+    case "Delete":
+      return handleDelete(event);
+  }
+};
+
+const handleCreate = async (event: CloudFormationCustomResourceCreateEvent) => {
+  const props = event.ResourceProperties as Props;
+  validateProps(props);
+  const provenance = await createSlotIfAbsent(props);
+  return {
+    PhysicalResourceId: buildPhysicalId({
+      identity: deriveSlotIdentity(props),
+      provenance,
+    }),
+  };
+};
+
+/**
+ * A slot cannot be altered in place: the database it decodes and its name are
+ * its identity, and its plugin is fixed at creation. A change to either half of
+ * that identity reports a new physical id, so CloudFormation creates the new
+ * slot and then deletes the old resource. A plugin change under an unchanged
+ * identity is refused, because the existing slot would silently keep its old
+ * plugin. An unchanged identity keeps the physical id it was given, so a slot
+ * adopted at creation is never recorded as created.
+ */
+const handleUpdate = async (event: CloudFormationCustomResourceUpdateEvent) => {
+  const props = event.ResourceProperties as Props;
+  validateProps(props);
+
+  const oldProps = event.OldResourceProperties as Props;
+  const identity = deriveSlotIdentity(props);
+  const { identity: currentIdentity } = parsePhysicalId(
+    event.PhysicalResourceId
+  );
+
+  if (identity != currentIdentity) {
+    const provenance = await createSlotIfAbsent(props);
+    return { PhysicalResourceId: buildPhysicalId({ identity, provenance }) };
+  }
+
+  if (props.Plugin != oldProps.Plugin) {
+    throw new Error(
+      `The plugin of replication slot "${props.Name}" cannot change in place; drop and recreate the slot deliberately instead`
+    );
+  }
+
+  return { PhysicalResourceId: event.PhysicalResourceId };
+};
+
+/**
+ * Only a slot the physical id records this resource as having created is
+ * dropped. A slot that was adopted, or one whose physical id carries no
+ * provenance, may belong to someone else, and a dropped slot loses its
+ * replication position permanently.
+ */
+const handleDelete = async (event: CloudFormationCustomResourceDeleteEvent) => {
+  const props = event.ResourceProperties as Props;
+  validateProps(props);
+
+  const { provenance } = parsePhysicalId(event.PhysicalResourceId);
+  if (provenance != "created") {
+    console.log(
+      "Not dropping replication slot, this resource did not create it",
+      props.Name,
+      provenance
+    );
+    return {};
+  }
+
+  console.log("Dropping replication slot", props.Name);
+  const client = await getConnectedClient(props.Connection);
+
+  try {
+    await postgres.dropReplicationSlot({ client, name: props.Name });
+  } finally {
+    await client.end();
+  }
+
+  return {};
+};
+
+const validateProps = (props: Props) => {
+  if (!("Connection" in props)) {
+    throw "Connection property is required";
+  }
+  validateConnection(props.Connection);
+
+  if (!("Name" in props)) {
+    throw "Name property is required";
+  }
+  if (!("Plugin" in props)) {
+    throw "Plugin property is required";
+  }
+};
+
+const deriveSlotIdentity = (props: Props): string => {
+  const { Host, Port, Database } = props.Connection;
+  const suffix = Math.abs(
+    hashCode(`${Host}-${Port}-${Database}-${props.Name}`)
+  );
+  return `replication-slot-${suffix}`;
+};
+
+const buildPhysicalId = (params: {
+  identity: string;
+  provenance: SlotProvenance;
+}): string => `${params.identity}-${params.provenance}`;
+
+/**
+ * The one place that knows the shape of a physical id. An id carrying no
+ * provenance suffix comes from the rollback of a failed create, and the guards
+ * in the create path throw exactly when a slot this resource never created
+ * already exists, so such an id reports "unknown" rather than "created".
+ */
+const parsePhysicalId = (
+  physicalResourceId: string
+): { identity: string; provenance: SlotProvenance | "unknown" } => {
+  for (const provenance of SLOT_PROVENANCES) {
+    const suffix = `-${provenance}`;
+    if (physicalResourceId.endsWith(suffix)) {
+      return {
+        identity: physicalResourceId.slice(0, -suffix.length),
+        provenance,
+      };
+    }
+  }
+
+  return { identity: physicalResourceId, provenance: "unknown" };
+};
+
+const createSlotIfAbsent = async (props: Props): Promise<SlotProvenance> => {
+  const client = await getConnectedClient(props.Connection);
+
+  try {
+    if (props.RequiredPublication) {
+      const { rows } = await client.query(
+        "SELECT 1 FROM pg_publication WHERE pubname = $1",
+        [props.RequiredPublication]
+      );
+      if (rows.length === 0) {
+        // A slot created before its publication decodes from a catalog snapshot
+        // in which the publication does not exist, and replication then fails
+        // continuously, so refusing here turns a subtle runtime failure into a
+        // loud deploy failure:
+        throw new Error(
+          `publication "${props.RequiredPublication}" does not exist in database "${props.Connection.Database}"; deploy the migration that creates it before this slot`
+        );
+      }
+    }
+
+    // `pg_replication_slots` lists the whole cluster, so a slot of the same
+    // name decoding another database must not be adopted as this one; letting
+    // Postgres raise "replication slot already exists" is the loud failure:
+    const { rows: existingSlots } = await client.query(
+      "SELECT plugin FROM pg_replication_slots WHERE slot_name = $1 AND database = current_database()",
+      [props.Name]
+    );
+    const [existingSlot] = existingSlots;
+
+    if (!existingSlot) {
+      console.log("Creating replication slot", props.Name);
+      await postgres.createReplicationSlot({
+        client,
+        name: props.Name,
+        plugin: props.Plugin,
+      });
+      return "created";
+    }
+
+    if (existingSlot.plugin !== props.Plugin) {
+      throw new Error(
+        `replication slot "${props.Name}" already exists in database "${props.Connection.Database}" with plugin "${existingSlot.plugin}" instead of "${props.Plugin}"; a slot's plugin is fixed at creation, so drop and recreate the slot deliberately instead`
+      );
+    }
+
+    console.log("Replication slot already exists", props.Name);
+    return "adopted";
+  } finally {
+    await client.end();
+  }
+};

cdk-postgresql/lib/replication-slot.ts

@@ -0,0 +1,75 @@
+import { Construct } from "constructs";
+import * as cdk from "aws-cdk-lib";
+import { RemovalPolicy } from "aws-cdk-lib";
+import { Provider } from "./provider";
+
+export interface ReplicationSlotProps {
+  /**
+   * Provider required to connect to the Postgresql server. Logical slots decode
+   * one database, the one the provider's `database` prop names.
+   */
+  provider: Provider;
+
+  /**
+   * The name of the slot. Must be unique on the PostgreSQL server instance
+   * where it is configured.
+   */
+  name: string;
+
+  /**
+   * The logical decoding plugin the slot is created with. It cannot change
+   * after creation.
+   *
+   * @default - "pgoutput"
+   */
+  plugin?: string;
+
+  /**
+   * A publication that has to exist before the slot is created. A slot created
+   * earlier decodes from a catalog snapshot in which the publication does not
+   * exist, and replication fails continuously, so creation is refused until
+   * the publication is there.
+   */
+  requiredPublication?: string;
+
+  /**
+   * Policy to apply when the slot is removed from this stack. A dropped slot
+   * loses its position permanently, so stacks that replicate production data
+   * should retain it.
+   *
+   * @default - The slot will be dropped.
+   */
+  removalPolicy?: RemovalPolicy;
+}
+
+/**
+ * A PostgreSQL logical replication slot. The database retains write-ahead-log
+ * data until the slot's consumer confirms having read it, which is what makes
+ * replication resumable after the consumer restarts.
+ *
+ * A slot that already exists on the same database with the same plugin is
+ * adopted and managed in place, and removing the resource only ever drops a
+ * slot this resource itself created.
+ */
+export class ReplicationSlot extends Construct {
+  constructor(scope: Construct, id: string, props: ReplicationSlotProps) {
+    super(scope, id);
+
+    const { provider, name, plugin, requiredPublication, removalPolicy } = props;
+
+    const cr = new cdk.CustomResource(this, "CustomResource", {
+      serviceToken: provider.serviceToken,
+      resourceType: "Custom::Postgresql-ReplicationSlot",
+      properties: {
+        connection: provider.buildConnectionProperty(),
+        name,
+        plugin: plugin ?? "pgoutput",
+        requiredPublication,
+      },
+      pascalCaseProperties: true,
+    });
+
+    cr.applyRemovalPolicy(removalPolicy || cdk.RemovalPolicy.DESTROY);
+    cr.node.addDependency(provider);
+  }
+}

cdk-postgresql/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@botpress/cdk-postgresql",
-  "version": "3.0.0",
+  "version": "3.1.0",
   "description": "Postgresql constructs for AWS CDK",
   "main": "./dist/index.cjs",
   "types": "./dist/index.d.cts",

cdk-postgresql/test/helpers.ts

@@ -81,3 +81,14 @@ export const isMemberOf = async (props: {
   );
   return rows[0].is_member;
 };
+
+export const replicationSlotExists = async (
+  client: Client,
+  name: string
+): Promise<boolean> => {
+  const { rows } = await client.query(
+    "SELECT 1 FROM pg_replication_slots WHERE slot_name = $1 AND database = current_database()",
+    [name]
+  );
+  return rows.length > 0;
+};