Merge pull request #11 from psteinroe/chore/increase-days-ahead

chore: create 7 partitions in advance instead of 3

Author: psteinroe

README.md

@@ -329,7 +329,7 @@ When the replication slot is invalidated (WAL exceeded `max_slot_wal_keep_size`)
 Postgres Stream automatically manages daily partitions in the background:
 
 **Retention policy:**
-- **Creates partitions** 3 days in advance (today, tomorrow, day after)
+- **Creates partitions** 7 days in advance
 - **Drops partitions** older than 7 days
 - **Runs on startup** and then daily
 

src/maintenance.rs

@@ -12,6 +12,9 @@ use crate::{
     store::StreamStore,
 };
 
+const DAYS_AHEAD_TO_CREATE: u32 = 7;
+const RETENTION_DAYS: u32 = 7;
+
 /// Domain representation of a partition with its date-based metadata.
 #[derive(Debug, Clone)]
 pub struct PartitionInfo {
@@ -117,9 +120,13 @@ pub async fn run_maintenance<S: SchemaStore>(store: &StreamStore<S>) -> EtlResul
     let result = async {
         let existing_partitions = store.load_partitions(SCHEMA_NAME, EVENTS_TABLE).await?;
 
-        let policy = RetentionPolicy::new(7);
-        let (to_drop, to_create) =
-            policy.plan_maintenance(existing_partitions, EVENTS_TABLE, 3, start);
+        let policy = RetentionPolicy::new(RETENTION_DAYS);
+        let (to_drop, to_create) = policy.plan_maintenance(
+            existing_partitions,
+            EVENTS_TABLE,
+            DAYS_AHEAD_TO_CREATE,
+            start,
+        );
 
         for partition in to_drop {
             info!("Dropping partition: {}", partition.name);
@@ -196,7 +203,7 @@ mod tests {
 
     #[test]
     fn test_retention_policy_plan_maintenance_drop_old() {
-        let policy = RetentionPolicy::new(7);
+        let policy = RetentionPolicy::new(RETENTION_DAYS);
         let now = Utc.with_ymd_and_hms(2024, 3, 15, 12, 0, 0).unwrap();
 
         // Partitions: one old (should drop), one current, one future
@@ -206,7 +213,8 @@ mod tests {
             PartitionInfo::for_date("events", NaiveDate::from_ymd_opt(2024, 3, 16).unwrap()), // tomorrow
         ];
 
-        let (to_drop, _to_create) = policy.plan_maintenance(partitions, "events", 3, now);
+        let (to_drop, _to_create) =
+            policy.plan_maintenance(partitions, "events", DAYS_AHEAD_TO_CREATE, now);
 
         assert_eq!(to_drop.len(), 1);
         assert_eq!(
@@ -217,50 +225,57 @@ mod tests {
 
     #[test]
     fn test_retention_policy_plan_maintenance_create_future() {
-        let policy = RetentionPolicy::new(7);
+        let policy = RetentionPolicy::new(RETENTION_DAYS);
         let now = Utc.with_ymd_and_hms(2024, 3, 15, 12, 0, 0).unwrap();
 
-        // Only today exists, should create 2 more (tomorrow and day after)
+        // Only today exists, should create 6 more (days 1-6 ahead)
         let partitions = vec![PartitionInfo::for_date(
             "events",
             NaiveDate::from_ymd_opt(2024, 3, 15).unwrap(),
         )];
 
-        let (to_drop, to_create) = policy.plan_maintenance(partitions, "events", 3, now);
+        let (to_drop, to_create) =
+            policy.plan_maintenance(partitions, "events", DAYS_AHEAD_TO_CREATE, now);
 
         assert_eq!(to_drop.len(), 0);
-        assert_eq!(to_create.len(), 2);
+        assert_eq!(to_create.len(), 6);
+        // Verify first and last created partitions
         assert_eq!(
             to_create.first().expect("first element exists").date,
             NaiveDate::from_ymd_opt(2024, 3, 16).unwrap()
         );
         assert_eq!(
-            to_create.get(1).expect("second element exists").date,
-            NaiveDate::from_ymd_opt(2024, 3, 17).unwrap()
+            to_create.last().expect("last element exists").date,
+            NaiveDate::from_ymd_opt(2024, 3, 21).unwrap()
         );
     }
 
     #[test]
     fn test_retention_policy_plan_maintenance_no_changes_needed() {
-        let policy = RetentionPolicy::new(7);
+        let policy = RetentionPolicy::new(RETENTION_DAYS);
         let now = Utc.with_ymd_and_hms(2024, 3, 15, 12, 0, 0).unwrap();
 
-        // All partitions already exist
+        // All partitions already exist (today + 6 days ahead)
         let partitions = vec![
             PartitionInfo::for_date("events", NaiveDate::from_ymd_opt(2024, 3, 15).unwrap()),
             PartitionInfo::for_date("events", NaiveDate::from_ymd_opt(2024, 3, 16).unwrap()),
             PartitionInfo::for_date("events", NaiveDate::from_ymd_opt(2024, 3, 17).unwrap()),
+            PartitionInfo::for_date("events", NaiveDate::from_ymd_opt(2024, 3, 18).unwrap()),
+            PartitionInfo::for_date("events", NaiveDate::from_ymd_opt(2024, 3, 19).unwrap()),
+            PartitionInfo::for_date("events", NaiveDate::from_ymd_opt(2024, 3, 20).unwrap()),
+            PartitionInfo::for_date("events", NaiveDate::from_ymd_opt(2024, 3, 21).unwrap()),
         ];
 
-        let (to_drop, to_create) = policy.plan_maintenance(partitions, "events", 3, now);
+        let (to_drop, to_create) =
+            policy.plan_maintenance(partitions, "events", DAYS_AHEAD_TO_CREATE, now);
 
         assert_eq!(to_drop.len(), 0);
         assert_eq!(to_create.len(), 0);
     }
 
     #[test]
     fn test_retention_policy_plan_maintenance_at_retention_boundary() {
-        let policy = RetentionPolicy::new(7);
+        let policy = RetentionPolicy::new(RETENTION_DAYS);
         let now = Utc.with_ymd_and_hms(2024, 3, 15, 12, 0, 0).unwrap();
 
         // Partition exactly at cutoff (7 days old) should be kept
@@ -269,6 +284,7 @@ mod tests {
             PartitionInfo::for_date("events", NaiveDate::from_ymd_opt(2024, 3, 7).unwrap()), // 8 days old - should drop
         ];
 
+        // Using days_ahead=1 to focus on testing the retention boundary behavior
         let (to_drop, _to_create) = policy.plan_maintenance(partitions, "events", 1, now);
 
         assert_eq!(to_drop.len(), 1);
@@ -280,17 +296,18 @@ mod tests {
 
     #[test]
     fn test_retention_policy_combined_drop_and_create() {
-        let policy = RetentionPolicy::new(7);
+        let policy = RetentionPolicy::new(RETENTION_DAYS);
         let now = Utc.with_ymd_and_hms(2024, 3, 15, 12, 0, 0).unwrap();
 
         let partitions = vec![
             PartitionInfo::for_date("events", NaiveDate::from_ymd_opt(2024, 3, 1).unwrap()), // old - drop
             PartitionInfo::for_date("events", NaiveDate::from_ymd_opt(2024, 3, 15).unwrap()), // today - keep
         ];
 
-        let (to_drop, to_create) = policy.plan_maintenance(partitions, "events", 3, now);
+        let (to_drop, to_create) =
+            policy.plan_maintenance(partitions, "events", DAYS_AHEAD_TO_CREATE, now);
 
         assert_eq!(to_drop.len(), 1);
-        assert_eq!(to_create.len(), 2); // tomorrow and day after
+        assert_eq!(to_create.len(), 6); // days 1-6 ahead
     }
 }

tests/maintenance_tests.rs

@@ -17,7 +17,7 @@ async fn test_initial_partitions_created() {
     let _stream: PgStream<MemorySink, PostgresStore> =
         PgStream::create(config, sink, store).await.unwrap();
 
-    // Check that partitions exist (today + 2 days ahead = 3 partitions)
+    // Check that partitions exist (today + 6 days ahead = 7 partitions)
     let count: (i64,) = sqlx::query_as(
         "select count(*) from pg_tables
          where schemaname = 'pgstream'
@@ -27,7 +27,7 @@ async fn test_initial_partitions_created() {
     .await
     .unwrap();
 
-    assert_eq!(count.0, 3, "Should create 3 initial partitions");
+    assert_eq!(count.0, 7, "Should create 7 initial partitions");
 }
 
 #[tokio::test(flavor = "multi_thread")]
@@ -77,7 +77,7 @@ async fn test_maintenance_creates_future_partitions() {
         .unwrap();
     run_maintenance(&stream_store).await.unwrap();
 
-    // Should now have 3 partitions again
+    // Should now have 7 partitions again (today + 6 days ahead)
     let count_after: (i64,) = sqlx::query_as(
         "select count(*) from pg_tables
          where schemaname = 'pgstream'
@@ -86,7 +86,7 @@ async fn test_maintenance_creates_future_partitions() {
     .fetch_one(&db.pool)
     .await
     .unwrap();
-    assert_eq!(count_after.0, 3);
+    assert_eq!(count_after.0, 7);
 }
 
 #[tokio::test(flavor = "multi_thread")]
@@ -213,5 +213,5 @@ async fn test_maintenance_idempotent() {
     .await
     .unwrap();
 
-    assert_eq!(count.0, 3, "Running maintenance twice should be idempotent");
+    assert_eq!(count.0, 7, "Running maintenance twice should be idempotent");
 }