fix: index slot recovery lookups by LSN

migrations/1787724000000_add_events_lsn_brin_index.sql

@@ -0,0 +1,73 @@
+-- Add the partitioned BRIN index used to look up events by WAL LSN.
+--
+-- A new database has no event partitions yet, so creating the parent index is
+-- safe and all partitions created later will inherit a child index. Existing
+-- databases must install the partitioned index online before this migration.
+DO $$
+DECLARE
+    parent_index_oid oid;
+    parent_index_is_expected boolean;
+    has_leaf_partitions boolean;
+BEGIN
+    SELECT
+        index_relation.oid,
+        coalesce(
+            index_relation.relkind = 'I'
+            AND access_method.amname = 'brin'
+            AND index_metadata.indrelid = 'pgstream.events'::regclass
+            AND index_metadata.indisready
+            AND index_metadata.indisvalid
+            AND index_metadata.indnatts = 1
+            AND index_metadata.indkey[0] = lsn_attribute.attnum
+            AND index_metadata.indpred IS NULL
+            AND index_metadata.indexprs IS NULL,
+            false
+        )
+    INTO parent_index_oid, parent_index_is_expected
+    FROM pg_class AS index_relation
+    JOIN pg_namespace AS index_namespace
+      ON index_namespace.oid = index_relation.relnamespace
+    LEFT JOIN pg_am AS access_method
+      ON access_method.oid = index_relation.relam
+    LEFT JOIN pg_index AS index_metadata
+      ON index_metadata.indexrelid = index_relation.oid
+    LEFT JOIN pg_attribute AS lsn_attribute
+      ON lsn_attribute.attrelid = 'pgstream.events'::regclass
+     AND lsn_attribute.attname = 'lsn'
+     AND NOT lsn_attribute.attisdropped
+    WHERE index_namespace.nspname = 'pgstream'
+      AND index_relation.relname = 'events_lsn_brin_idx';
+
+    IF parent_index_oid IS NOT NULL THEN
+        IF NOT parent_index_is_expected THEN
+            RAISE EXCEPTION
+                'pgstream.events_lsn_brin_idx exists but is not a ready, valid, non-partial partitioned BRIN index on pgstream.events(lsn); inspect or remove it before retrying the migration';
+        END IF;
+
+        -- Production databases that completed the online rollout already have
+        -- the expected index.
+        RETURN;
+    END IF;
+
+    -- Prevent maintenance from adding a partition between the safety check and
+    -- CREATE INDEX. Existing databases take the online concurrent rollout path.
+    LOCK TABLE pgstream.events IN SHARE MODE NOWAIT;
+
+    SELECT EXISTS (
+        SELECT 1
+        FROM pg_partition_tree('pgstream.events'::regclass) AS partition_tree
+        WHERE partition_tree.isleaf
+    )
+    INTO has_leaf_partitions;
+
+    IF has_leaf_partitions THEN
+        RAISE EXCEPTION
+            'pgstream.events has existing leaf partitions but pgstream.events_lsn_brin_idx is absent; install and attach the BRIN index online before retrying the migration';
+    END IF;
+
+    CREATE INDEX events_lsn_brin_idx
+        ON pgstream.events
+        USING brin (lsn)
+        WITH (pages_per_range = 64, autosummarize = on);
+END
+$$;

src/slot_recovery.rs

@@ -120,7 +120,9 @@ pub async fn handle_slot_recovery(pool: &PgPool, stream_id: u64) -> EtlResult<()
         "found confirmed_flush_lsn from invalidated slot"
     );
 
-    // 2. Find the first event after the confirmed LSN
+    // 2. Find the first event after the confirmed LSN.
+    // `id` intentionally resolves to the selected text output column so its
+    // ordering matches the persisted checkpoint representation.
     let lsn_checkpoint: Option<Checkpoint> = sqlx::query_as(
         "SELECT id::text, created_at FROM pgstream.events
          WHERE lsn > $1::pg_lsn AND stream_id = $2

tests/maintenance_tests.rs

@@ -5,6 +5,83 @@ use postgres_stream::sink::memory::MemorySink;
 use postgres_stream::stream::PgStream;
 use postgres_stream::test_utils::{TestDatabase, create_postgres_store, test_stream_config};
 
+async fn assert_events_lsn_brin_indexes(db: &TestDatabase) {
+    let (parent_index_valid,): (bool,) = sqlx::query_as(
+        "select exists(
+            select 1
+            from pg_class as index_relation
+            join pg_namespace as index_namespace
+              on index_namespace.oid = index_relation.relnamespace
+            join pg_am as access_method
+              on access_method.oid = index_relation.relam
+            join pg_index as index_metadata
+              on index_metadata.indexrelid = index_relation.oid
+            where index_namespace.nspname = 'pgstream'
+              and index_relation.relname = 'events_lsn_brin_idx'
+              and index_relation.relkind = 'I'
+              and access_method.amname = 'brin'
+              and index_metadata.indisready
+              and index_metadata.indisvalid
+        )",
+    )
+    .fetch_one(&db.pool)
+    .await
+    .unwrap();
+    assert!(parent_index_valid, "parent BRIN index should be valid");
+
+    let (partition_count, attached_index_count): (i64, i64) = sqlx::query_as(
+        "with leaf_partitions as (
+            select partition_tree.relid, partition_relation.relname
+            from pg_partition_tree('pgstream.events'::regclass) as partition_tree
+            join pg_class as partition_relation
+              on partition_relation.oid = partition_tree.relid
+            where partition_tree.isleaf
+        ), parent_index as (
+            select index_relation.oid
+            from pg_class as index_relation
+            join pg_namespace as index_namespace
+              on index_namespace.oid = index_relation.relnamespace
+            where index_namespace.nspname = 'pgstream'
+              and index_relation.relname = 'events_lsn_brin_idx'
+        )
+        select
+            (select count(*) from leaf_partitions),
+            (select count(*)
+             from leaf_partitions
+             where exists(
+                 select 1
+                 from pg_index as child_index
+                 join pg_class as child_index_relation
+                   on child_index_relation.oid = child_index.indexrelid
+                 join pg_am as access_method
+                   on access_method.oid = child_index_relation.relam
+                 join pg_inherits as index_attachment
+                   on index_attachment.inhrelid = child_index.indexrelid
+                 where child_index.indrelid = leaf_partitions.relid
+                   and child_index.indnatts = 1
+                   and child_index.indkey[0] = (
+                       select attnum
+                       from pg_attribute
+                       where attrelid = leaf_partitions.relid
+                         and attname = 'lsn'
+                         and not attisdropped
+                   )
+                   and access_method.amname = 'brin'
+                   and child_index.indisready
+                   and child_index.indisvalid
+                   and index_attachment.inhparent = (select oid from parent_index)
+             ))",
+    )
+    .fetch_one(&db.pool)
+    .await
+    .unwrap();
+
+    assert_eq!(
+        attached_index_count, partition_count,
+        "every events leaf partition should have a valid attached BRIN index"
+    );
+}
+
 #[tokio::test(flavor = "multi_thread")]
 async fn test_initial_partitions_created() {
     let db = TestDatabase::spawn().await;
@@ -27,6 +104,7 @@ async fn test_initial_partitions_created() {
     .unwrap();
 
     assert_eq!(count.0, 7, "Should create 7 initial partitions");
+    assert_events_lsn_brin_indexes(&db).await;
 }
 
 #[tokio::test(flavor = "multi_thread")]
@@ -86,6 +164,7 @@ async fn test_maintenance_creates_future_partitions() {
     .await
     .unwrap();
     assert_eq!(count_after.0, 7);
+    assert_events_lsn_brin_indexes(&db).await;
 }
 
 #[tokio::test(flavor = "multi_thread")]

tests/subscription_extension_tests.rs

@@ -361,7 +361,6 @@ async fn upgrade_moves_subscriptions_without_recreating_rows_or_triggers() {
 
     apply_legacy_core_migrations(&database).await;
 
-    database.ensure_today_partition().await;
     create_users_table(&database).await;
 
     let subscription_id: sqlx::types::Uuid = sqlx::query_scalar(
@@ -445,6 +444,8 @@ async fn upgrade_moves_subscriptions_without_recreating_rows_or_triggers() {
     .expect("Failed to inspect migration history");
     assert!(extraction_applied);
 
+    database.ensure_today_partition().await;
+
     sqlx::query("insert into public.users (email) values ('upgrade@example.com')")
         .execute(&database.pool)
         .await