feat(operator): gate Service traffic on operator-set ready label (#58)

Author: passcod

src/controllers/replica.rs

@@ -190,6 +190,14 @@ pub async fn reconcile(replica: Arc<PostgresPhysicalReplica>, ctx: Arc<Context>)
 			"performing blue-green switchover"
 		);
 
+		// Mark the new restore's pod ready for traffic BEFORE pointing the
+		// Service at it. The Service selector requires the
+		// `ready-for-traffic=true` label (see [[READY_FOR_TRAFFIC_LABEL]]),
+		// so until the operator sets the label, no external client can
+		// reach the restore via the Service — operator-side prep work
+		// (DROP SCHEMA, migration Job) runs without external interference.
+		switching.mark_pod_ready_for_traffic(client).await?;
+
 		// Update Service selector to point to the new restore
 		switching.update_service_selector(client, &name).await?;
 

src/controllers/replica/resources.rs

@@ -6,7 +6,7 @@ use k8s_openapi::{
 	api::{
 		batch::v1::{Job, JobSpec},
 		core::v1::{
-			Container, EnvVar, LocalObjectReference, PodSpec, PodTemplateSpec,
+			Container, EnvVar, LocalObjectReference, Pod, PodSpec, PodTemplateSpec,
 			ResourceRequirements, Secret, Service, ServicePort, ServiceSpec,
 		},
 	},
@@ -415,13 +415,60 @@ impl PostgresPhysicalReplica {
 	}
 }
 
+/// Label the operator sets on a restore's postgres pod once schema
+/// migration (and anything else that must run pre-handover) has completed
+/// and the pod is safe to receive external traffic. The per-replica
+/// Service selector requires this label, so a restore in `Switching`
+/// can't be reached via the Service — operator-side work (DROP SCHEMA,
+/// `pg_dump | psql` migration Job, etc.) runs without external clients
+/// racing to grab locks on the schemas being touched.
+pub const READY_FOR_TRAFFIC_LABEL: &str = "pgro.bes.au/ready-for-traffic";
+
 impl PostgresPhysicalRestore {
+	/// Patch this restore's postgres pod to add the [`READY_FOR_TRAFFIC_LABEL`].
+	/// Idempotent and resilient to the pod not existing yet (the restore's
+	/// deployment may be mid-rollout); callers that need the label to be
+	/// present should retry on the next reconcile pass.
+	pub async fn mark_pod_ready_for_traffic(&self, client: &Client) -> Result<()> {
+		let pods: Api<Pod> = Api::namespaced(client.clone(), &self.ns());
+		let selector = format!("pgro.bes.au/restore={}", self.name_any());
+		let list = pods.list(&ListParams::default().labels(&selector)).await?;
+		let pod = list
+			.items
+			.into_iter()
+			.find(|p| p.status.as_ref().and_then(|s| s.phase.as_deref()) == Some("Running"));
+		let Some(pod) = pod else {
+			warn!(
+				restore = self.name_any(),
+				"no running pod for restore yet; will retry next reconcile"
+			);
+			return Ok(());
+		};
+		let pod_name = pod.name_any();
+		let patch = serde_json::json!({
+			"metadata": {
+				"labels": {
+					READY_FOR_TRAFFIC_LABEL: "true",
+				}
+			}
+		});
+		pods.patch(&pod_name, &PatchParams::default(), &Patch::Merge(&patch))
+			.await?;
+		info!(
+			restore = self.name_any(),
+			pod = pod_name,
+			"marked restore pod ready for traffic"
+		);
+		Ok(())
+	}
+
 	pub async fn update_service_selector(&self, client: &Client, service_name: &str) -> Result<()> {
 		let services: Api<Service> = Api::namespaced(client.clone(), &self.ns());
 		let patch = serde_json::json!({
 			"spec": {
 				"selector": {
 					"pgro.bes.au/restore": self.name_any(),
+					READY_FOR_TRAFFIC_LABEL: "true",
 				}
 			}
 		});

src/controllers/restore.rs

@@ -617,6 +617,13 @@ async fn reconcile_active(
 		apply_restore_deployment(client, restore, &replica, name, namespace).await?;
 	}
 
+	// Defensively ensure the ready-for-traffic label is on the pod. If the
+	// pod restarted (OOM, eviction, node loss), the label is gone from the
+	// new pod and the Service stops routing to it; re-applying every pass
+	// closes that gap. Also handles the upgrade path where existing Active
+	// pods predate the label.
+	restore.mark_pod_ready_for_traffic(client).await?;
+
 	Ok(Action::requeue(Duration::from_secs(300)))
 }
 

tests/switchover.rs

@@ -102,6 +102,13 @@ async fn second_restore_and_switchover() {
 		Some(first_restore_name.as_str()),
 		"service selector should point to first restore"
 	);
+	assert_eq!(
+		svc_selector
+			.get("pgro.bes.au/ready-for-traffic")
+			.map(|s| s.as_str()),
+		Some("true"),
+		"service selector should gate on ready-for-traffic label"
+	);
 
 	// Manually create a second PostgresPhysicalRestore to trigger switchover
 	let second_restore_name = "switchover-replica-manual-second";
@@ -213,6 +220,36 @@ async fn second_restore_and_switchover() {
 		Some(second_restore_name),
 		"service selector should point to second restore after switchover"
 	);
+	assert_eq!(
+		svc_selector
+			.get("pgro.bes.au/ready-for-traffic")
+			.map(|s| s.as_str()),
+		Some("true"),
+		"service selector should still gate on ready-for-traffic label after switchover"
+	);
+
+	// The second restore's pod should be labeled ready-for-traffic.
+	let pods: kube::Api<k8s_openapi::api::core::v1::Pod> =
+		kube::Api::namespaced(client.clone(), &ns);
+	let pod_list = pods
+		.list(
+			&kube::api::ListParams::default()
+				.labels(&format!("pgro.bes.au/restore={second_restore_name}")),
+		)
+		.await
+		.expect("failed to list second restore pods");
+	let labeled = pod_list.items.iter().any(|p| {
+		p.metadata
+			.labels
+			.as_ref()
+			.and_then(|l| l.get("pgro.bes.au/ready-for-traffic"))
+			.map(|v| v.as_str())
+			== Some("true")
+	});
+	assert!(
+		labeled,
+		"second restore's pod should carry pgro.bes.au/ready-for-traffic=true after switchover"
+	);
 
 	// Second restore should have activated_at set
 	let second_restore = restores