PITR: fix clones and restore

Author: boddumanohar

src/api/organization/project/branch/__init__.py

@@ -50,6 +50,7 @@
 from .....deployment.kubernetes.volume_clone import (
     clone_branch_database_volume,
     restore_branch_database_volume_from_snapshot,
+    restore_branch_wal_volume_from_snapshot,
 )
 from .....deployment.settings import get_settings as get_deployment_settings
 from .....exceptions import VelaDeploymentError, VelaError, VelaKubernetesError, VelaSimplyblockAPIError
@@ -359,6 +360,7 @@ class RestoreSnapshotContext(TypedDict):
     namespace: str
     name: str
     content_name: str | None
+    wal_snapshot_name: str | None
 
 
 def snapshot_pgbouncer_config(config: PgbouncerConfig | None) -> PgbouncerConfigSnapshot:
@@ -765,6 +767,7 @@ async def _clone_branch_environment_task(
                 pvc_timeout_seconds=_PVC_CLONE_TIMEOUT_SECONDS,
                 pvc_poll_interval_seconds=_PVC_POLL_INTERVAL_SECONDS,
                 database_size=parameters.database_size,
+                pitr_enabled=pitr_enabled,
             )
         except VelaError:
             await _persist_branch_status(branch_id, BranchServiceStatus.ERROR)
@@ -821,6 +824,7 @@ async def _restore_branch_environment_task(
     snapshot_namespace: str,
     snapshot_name: str,
     snapshot_content_name: str | None,
+    wal_snapshot_name: str | None,
     restore_database_size: int,
     pgbouncer_config: PgbouncerConfigSnapshot,
     pitr_enabled: bool,
@@ -844,6 +848,19 @@ async def _restore_branch_environment_task(
             pvc_poll_interval_seconds=_PVC_POLL_INTERVAL_SECONDS,
             database_size=restore_database_size,
         )
+        if wal_snapshot_name is not None:
+            await restore_branch_wal_volume_from_snapshot(
+                source_branch_id=source_branch_id,
+                target_branch_id=branch_id,
+                snapshot_namespace=snapshot_namespace,
+                snapshot_name=wal_snapshot_name,
+                snapshot_class=_VOLUME_SNAPSHOT_CLASS,
+                storage_class_name=storage_class_name,
+                snapshot_timeout_seconds=_SNAPSHOT_TIMEOUT_SECONDS,
+                snapshot_poll_interval_seconds=_SNAPSHOT_POLL_INTERVAL_SECONDS,
+                pvc_timeout_seconds=_PVC_TIMEOUT_SECONDS,
+                pvc_poll_interval_seconds=_PVC_POLL_INTERVAL_SECONDS,
+            )
     except VelaError:
         await _persist_branch_status(branch_id, BranchServiceStatus.ERROR)
         await _cleanup_failed_branch_deployment(branch_id)
@@ -1320,6 +1337,7 @@ def _schedule_branch_environment_tasks(
                 snapshot_namespace=restore_snapshot["namespace"],
                 snapshot_name=restore_snapshot["name"],
                 snapshot_content_name=restore_snapshot["content_name"],
+                wal_snapshot_name=restore_snapshot["wal_snapshot_name"],
                 restore_database_size=restore_database_size,
                 pgbouncer_config=pgbouncer_config,
                 pitr_enabled=branch.pitr_enabled,
@@ -1386,6 +1404,7 @@ async def create(  # noqa: C901
             namespace=cast("str", backup_entry.snapshot_namespace),
             name=cast("str", backup_entry.snapshot_name),
             content_name=backup_entry.snapshot_content_name,
+            wal_snapshot_name=backup_entry.wal_snapshot_name,
         )
     source_id: Identifier | None = getattr(source, "id", None)
     clone_parameters: DeploymentParameters | None = None

src/deployment/kubernetes/volume_clone.py

@@ -5,12 +5,15 @@
 from datetime import UTC, datetime
 from typing import Any
 
+from kubernetes.utils import parse_quantity
+
 from ..._util import Identifier
 from ...exceptions import VelaKubernetesError
 from .. import (
     _POD_SECURITY_LABELS,
     AUTOSCALER_PVC_SUFFIX,
     AUTOSCALER_WAL_PVC_SUFFIX,
+    PITR_WAL_PVC_SIZE,
     get_autoscaler_vm_identity,
     kube_service,
 )
@@ -36,6 +39,8 @@
 
 logger = logging.getLogger(__name__)
 
+_PITR_WAL_PVC_SIZE_BYTES: int = int(parse_quantity(PITR_WAL_PVC_SIZE))
+
 
 @dataclass(frozen=True)
 class CloneTimeouts:
@@ -319,6 +324,7 @@ class _SnapshotRestoreOperation:
     storage_class_name: str
     target_database_size: int
     timeouts: CloneTimeouts
+    pvc_suffix: str = AUTOSCALER_PVC_SUFFIX
     ids: CloneIdentifiers = field(init=False)
     created_target_snapshot: bool = field(default=False, init=False)
     created_content: bool = field(default=False, init=False)
@@ -327,8 +333,8 @@ def __post_init__(self) -> None:
         source_ns = self.snapshot_namespace
         _, source_vm_name = get_autoscaler_vm_identity(self.source_branch_id)
         target_ns, target_vm_name = get_autoscaler_vm_identity(self.target_branch_id)
-        pvc_name = f"{source_vm_name}{AUTOSCALER_PVC_SUFFIX}"
-        target_pvc_name = f"{target_vm_name}{AUTOSCALER_PVC_SUFFIX}"
+        pvc_name = f"{source_vm_name}{self.pvc_suffix}"
+        target_pvc_name = f"{target_vm_name}{self.pvc_suffix}"
         if target_pvc_name != pvc_name:
             raise VelaKubernetesError(
                 f"Autoscaler PVC name mismatch between source ({pvc_name}) and target ({target_pvc_name})"
@@ -520,7 +526,7 @@ async def clone_branch_database_volume(
             target_branch_id=target_branch_id,
             snapshot_class=snapshot_class,
             storage_class_name=storage_class_name,
-            target_database_size=database_size,
+            target_database_size=_PITR_WAL_PVC_SIZE_BYTES,
             timeouts=timeouts,
             volume_label="wal",
             pvc_suffix=AUTOSCALER_WAL_PVC_SUFFIX,
@@ -563,3 +569,40 @@ async def restore_branch_database_volume_from_snapshot(
         ),
     )
     await operation.run()
+
+
+async def restore_branch_wal_volume_from_snapshot(
+    *,
+    source_branch_id: Identifier,
+    target_branch_id: Identifier,
+    snapshot_namespace: str,
+    snapshot_name: str,
+    snapshot_class: str,
+    storage_class_name: str,
+    snapshot_timeout_seconds: float,
+    snapshot_poll_interval_seconds: float,
+    pvc_timeout_seconds: float,
+    pvc_poll_interval_seconds: float,
+) -> None:
+    """
+    Restore the WAL volume for a PITR-enabled branch from an existing VolumeSnapshot.
+    The WAL PVC is always restored to PITR_WAL_PVC_SIZE (100Gi).
+    """
+    operation = _SnapshotRestoreOperation(
+        source_branch_id=source_branch_id,
+        target_branch_id=target_branch_id,
+        snapshot_namespace=snapshot_namespace,
+        snapshot_name=snapshot_name,
+        snapshot_content_name=None,
+        snapshot_class=snapshot_class,
+        storage_class_name=storage_class_name,
+        target_database_size=_PITR_WAL_PVC_SIZE_BYTES,
+        pvc_suffix=AUTOSCALER_WAL_PVC_SUFFIX,
+        timeouts=CloneTimeouts(
+            snapshot_ready=snapshot_timeout_seconds,
+            snapshot_poll=snapshot_poll_interval_seconds,
+            pvc_ready=pvc_timeout_seconds,
+            pvc_poll=pvc_poll_interval_seconds,
+        ),
+    )
+    await operation.run()