feat: track restore stage in _pgro.restore_info (#25)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: passcod

.rules

@@ -785,21 +785,33 @@ else
 ALTER ROLE ${{ANALYTICS_USERNAME}} WITH SUPERUSER;
 SQLEOF
 fi
-echo "Writing restore metadata..."
+if [ -f /pgdata/needs-reindex ]; then
+  PGRO_STAGE=restored
+else
+  PGRO_STAGE=ready
+fi
+
+echo "Writing restore metadata (stage=${{PGRO_STAGE}})..."
 psql -U postgres -d postgres << SQLEOF
 CREATE SCHEMA IF NOT EXISTS _pgro;
 CREATE TABLE IF NOT EXISTS _pgro.restore_info (
   id integer PRIMARY KEY DEFAULT 1,
   snapshot_id text NOT NULL,
   snapshot_time timestamptz,
-  restored_at timestamptz NOT NULL DEFAULT now()
+  restored_at timestamptz NOT NULL DEFAULT now(),
+  stage text NOT NULL DEFAULT 'restored',
+  last_transition_time timestamptz NOT NULL DEFAULT now()
 );
-INSERT INTO _pgro.restore_info (id, snapshot_id, snapshot_time)
-VALUES (1, '${{PGRO_SNAPSHOT_ID}}', CASE WHEN '${{PGRO_SNAPSHOT_TIME}}' = '' THEN NULL ELSE '${{PGRO_SNAPSHOT_TIME}}'::timestamptz END)
+ALTER TABLE _pgro.restore_info ADD COLUMN IF NOT EXISTS stage text NOT NULL DEFAULT 'restored';
+ALTER TABLE _pgro.restore_info ADD COLUMN IF NOT EXISTS last_transition_time timestamptz NOT NULL DEFAULT now();
+INSERT INTO _pgro.restore_info (id, snapshot_id, snapshot_time, stage, last_transition_time)
+VALUES (1, '${{PGRO_SNAPSHOT_ID}}', CASE WHEN '${{PGRO_SNAPSHOT_TIME}}' = '' THEN NULL ELSE '${{PGRO_SNAPSHOT_TIME}}'::timestamptz END, '${{PGRO_STAGE}}', now())
 ON CONFLICT (id) DO UPDATE
   SET snapshot_id = EXCLUDED.snapshot_id,
       snapshot_time = EXCLUDED.snapshot_time,
-      restored_at = now();
+      restored_at = now(),
+      stage = EXCLUDED.stage,
+      last_transition_time = now();
 SQLEOF
 
 echo "Stopping temporary postgres..."
@@ -951,6 +963,7 @@ if [ -f /pgdata/needs-reindex ]; then
   PG_MAJOR=$(cat /pgdata/pgdata/PG_VERSION)
   (
     while ! pg_isready -q -U postgres -d postgres; do sleep 2; done
+    psql -U postgres -d postgres -c "UPDATE _pgro.restore_info SET stage = 'reindexing', last_transition_time = now() WHERE id = 1;"
     for db in $(psql -U postgres -d postgres -At -c "SELECT datname FROM pg_database WHERE datallowconn AND datname <> 'template0'"); do
       INDEXES=$(psql -U postgres -d "$db" -At -c "
         SELECT DISTINCT indexrelid::regclass::text
@@ -973,6 +986,7 @@ if [ -f /pgdata/needs-reindex ]; then
       done
     done
     rm -f /pgdata/needs-reindex
+    psql -U postgres -d postgres -c "UPDATE _pgro.restore_info SET stage = 'ready', last_transition_time = now() WHERE id = 1;"
     echo "Background reindex complete"
   ) &
 fi

src/controllers/restore/builders.rs

@@ -307,6 +307,90 @@ fn deployment_init_script_sets_shared_buffers() {
 	);
 }
 
+#[test]
+fn init_script_sets_initial_stage_based_on_reindex_flag() {
+	let (mut restore, replica) = test_restore_and_replica();
+	restore.status = Some(PostgresPhysicalRestoreStatus {
+		postgres_version: Some("16".to_string()),
+		..Default::default()
+	});
+
+	let deploy = build_deployment(&restore, "test-restore", "default", &replica).unwrap();
+	let pod_spec = deploy.spec.unwrap().template.spec.unwrap();
+
+	let setup_auth = pod_spec
+		.init_containers
+		.as_ref()
+		.unwrap()
+		.iter()
+		.find(|c| c.name == "setup-auth")
+		.expect("setup-auth init container must exist");
+	let script = &setup_auth.args.as_ref().unwrap()[0];
+
+	assert!(
+		script.contains("stage text NOT NULL DEFAULT 'restored'"),
+		"table must include stage column"
+	);
+	assert!(
+		script.contains("last_transition_time timestamptz NOT NULL DEFAULT now()"),
+		"table must include last_transition_time column"
+	);
+	assert!(
+		script.contains("ADD COLUMN IF NOT EXISTS stage"),
+		"must add stage column for existing tables carried in snapshot"
+	);
+	assert!(
+		script.contains("ADD COLUMN IF NOT EXISTS last_transition_time"),
+		"must add last_transition_time column for existing tables carried in snapshot"
+	);
+	assert!(
+		script.contains("PGRO_STAGE=restored") && script.contains("PGRO_STAGE=ready"),
+		"init must pick stage based on needs-reindex flag"
+	);
+	assert!(
+		script.contains("'${PGRO_STAGE}'"),
+		"insert must use the chosen stage"
+	);
+}
+
+#[test]
+fn postgres_container_updates_stage_around_reindex() {
+	let (mut restore, replica) = test_restore_and_replica();
+	restore.status = Some(PostgresPhysicalRestoreStatus {
+		postgres_version: Some("16".to_string()),
+		..Default::default()
+	});
+
+	let deploy = build_deployment(&restore, "test-restore", "default", &replica).unwrap();
+	let pod_spec = deploy.spec.unwrap().template.spec.unwrap();
+
+	let postgres = pod_spec
+		.containers
+		.iter()
+		.find(|c| c.name == "postgres")
+		.expect("postgres container must exist");
+	let script = &postgres.args.as_ref().unwrap()[0];
+
+	let reindexing_pos = script
+		.find("stage = 'reindexing'")
+		.expect("must update stage to reindexing before the REINDEX loop");
+	let ready_pos = script
+		.find("stage = 'ready'")
+		.expect("must update stage to ready after the REINDEX loop");
+	assert!(
+		reindexing_pos < ready_pos,
+		"reindexing update must come before ready update"
+	);
+	assert!(
+		reindexing_pos < script.find("REINDEX INDEX").unwrap(),
+		"reindexing update must come before any REINDEX call"
+	);
+	assert!(
+		script[..ready_pos].contains("rm -f /pgdata/needs-reindex"),
+		"ready update must happen after the needs-reindex flag is cleared"
+	);
+}
+
 #[test]
 fn deployment_shared_buffers_with_custom_resources() {
 	let (mut restore, mut replica) = test_restore_and_replica();