Implement list and delete snapshot functionality in Python SDK

clients/python/agentic-sandbox-client/k8s_agent_sandbox/gke_extensions/podsnapshot.md

@@ -24,7 +24,15 @@ A specialized Sandbox client for interacting with the GKE Pod Snapshot Controlle
 *   **`is_restored_from_snapshot(self, snapshot_uid: str) -> RestoreCheckResult`**:
     *   Checks if the sandbox pod was restored from the specified snapshot.
     *   Verifies restoration by checking the 'PodRestored' condition in the pod status and confirming the message contains the expected snapshot UID.
-    *   Returns RestoreResult object(success, error_code, error_reason).
+    *   Returns RestoreCheckResult object(success, error_code, error_reason).
+*   **`list_snapshots(self, grouping_labels: dict[str, str] | None = None, ready_only: bool = True) -> ListSnapshotResult`**:
+    *   Checks for existing snapshots matching the provided grouping labels associated with the sandbox.
+    *   If `ready_only` is True, only returns snapshots that are in the 'Ready' state.
+    *   Returns a `ListSnapshotResult` object (success, error_code, error_reason, snapshots) containing the mapped snapshots sorted by creation timestamp (newest first).
+*   **`delete_snapshots(self, grouping_labels: dict[str, str] | None = None, snapshot_uid: str | None = None) -> DeleteSnapshotResult`**:
+    *   Deletes snapshots based on provided parameters. **Note:** `snapshot_uid` and `grouping_labels` are mutually exclusive—provide only one.
+    *   Targets a specific snapshot if `snapshot_uid` is provided or all matching `grouping_labels`, or all snapshots for the pod if neither is provided.
+    *   Returns a `DeleteSnapshotResult` object (success, error_code, error_reason, deleted_snapshots) containing the list of successfully deleted snapshot UIDs.
 *   **`__exit__(self)`**:
     *   Cleans up the `PodSnapshotManualTrigger` resources.
     *   Cleans up the `SandboxClaim` resources.

clients/python/agentic-sandbox-client/k8s_agent_sandbox/gke_extensions/podsnapshot_client.py

@@ -26,6 +26,7 @@
     PODSNAPSHOT_API_KIND,
     PODSNAPSHOTMANUALTRIGGER_API_KIND,
     PODSNAPSHOTMANUALTRIGGER_PLURAL,
+    PODSNAPSHOT_PLURAL,
 )
 
 SNAPSHOT_SUCCESS_CODE = 0
@@ -63,6 +64,36 @@ class RestoreCheckResult:
     error_code: int
 
 
+@dataclass
+class SnapshotDetail:
+    """Detailed information about a snapshot."""
+
+    snapshot_uid: str
+    source_pod: str
+    creation_timestamp: str
+    status: str
+
+
+@dataclass
+class ListSnapshotResult:
+    """Result of a list snapshots operation."""
+
+    success: bool
+    snapshots: list[SnapshotDetail]
+    error_reason: str
+    error_code: int
+
+
+@dataclass
+class DeleteSnapshotResult:
+    """Result of a delete snapshot operation."""
+
+    success: bool
+    deleted_snapshots: list[str]
+    error_reason: str
+    error_code: int
+
+
 class PodSnapshotSandboxClient(SandboxClient):
     """
     A specialized Sandbox client for interacting with the GKE Pod Snapshot Controller.
@@ -147,10 +178,11 @@ def _parse_created_snapshot_info(self, obj: dict[str, Any]) -> SnapshotResult:
                     snapshot_uid=snapshot_uid,
                     snapshot_timestamp=snapshot_timestamp,
                 )
-            elif condition.get("status") == "False" and condition.get("reason") in [
-                "Failed",
-                "Error",
-            ]:
+            elif (
+                condition.get("type") == "Triggered"
+                and condition.get("status") == "False"
+                and condition.get("reason") in ["Failed", "Error"]
+            ):
                 raise RuntimeError(
                     f"Snapshot failed. Condition: {condition.get('message', 'Unknown error')}"
                 )
@@ -374,8 +406,8 @@ def is_restored_from_snapshot(self, snapshot_uid: str) -> RestoreCheckResult:
                                 error_code=SNAPSHOT_ERROR_CODE,
                             )
                     else:
-                        reason_val = condition.get("reason") or ""
-                        msg_val = condition.get("message") or ""
+                        reason_val = condition.reason or ""
+                        msg_val = condition.message or ""
                         reason = f" reason: '{reason_val}'"
                         msg = f" message: '{msg_val}'"
                         return RestoreCheckResult(
@@ -407,6 +439,222 @@ def is_restored_from_snapshot(self, snapshot_uid: str) -> RestoreCheckResult:
                 error_code=SNAPSHOT_ERROR_CODE,
             )
 
+    def list_snapshots(
+        self, grouping_labels: dict[str, str] | None = None, ready_only: bool = True
+    ) -> ListSnapshotResult:
+        """
+        Checks for existing snapshots matching the grouping labels associated with the sandbox.
+        Returns a ListSnapshotResult containing valid snapshots sorted by creation timestamp (newest first).
+
+        grouping_labels: Filters snapshots by their metadata labels.
+        ready_only: If True, only returns snapshots that are in the 'Ready' state.
+        """
+
+        valid_snapshots = []
+
+        selectors = []
+        if not self.pod_name:
+            logger.warning("Pod name not found. Ensure sandbox is created.")
+            return ListSnapshotResult(
+                success=False,
+                snapshots=[],
+                error_reason="Pod name not found. Ensure sandbox is created.",
+                error_code=SNAPSHOT_ERROR_CODE,
+            )
+
+        selectors.append(f"podsnapshot.gke.io/pod-name={self.pod_name}")
+
+        if grouping_labels:
+            for k, v in grouping_labels.items():
+                selectors.append(f"{k}={v}")
+
+        label_selector = ",".join(selectors)
+
+        logger.info(f"Listing snapshots with label selector: {label_selector}")
+        try:
+            # Fetch the PodSnapshots using label selector directly
+            response = self.custom_objects_api.list_namespaced_custom_object(
+                group=PODSNAPSHOT_API_GROUP,
+                version=PODSNAPSHOT_API_VERSION,
+                namespace=self.namespace,
+                plural=PODSNAPSHOT_PLURAL,
+                label_selector=label_selector,
+            )
+        except ApiException as e:
+            logger.error(f"Failed to list PodSnapshots: {e}")
+            return ListSnapshotResult(
+                success=False,
+                snapshots=[],
+                error_reason=f"Failed to list PodSnapshots: {e}",
+                error_code=SNAPSHOT_ERROR_CODE,
+            )
+        except Exception as e:
+            logger.exception(
+                f"Unexpected error during list snapshots for grouping labels '{grouping_labels}': {e}"
+            )
+            return ListSnapshotResult(
+                success=False,
+                snapshots=[],
+                error_reason=f"Unexpected error: {e}",
+                error_code=SNAPSHOT_ERROR_CODE,
+            )
+
+        for snapshot in response.get("items") or []:
+            status = snapshot.get("status", {})
+            conditions = status.get("conditions") or []
+            metadata = snapshot.get("metadata", {})
+
+            # Check for Ready=True
+            is_ready = False
+            for cond in conditions:
+                if cond.get("type") == "Ready" and cond.get("status") == "True":
+                    is_ready = True
+                    break
+
+            # Skip if only ready snapshots are requested
+            if ready_only and not is_ready:
+                continue
+
+            valid_snapshots.append(
+                SnapshotDetail(
+                    snapshot_uid=metadata.get("name"),
+                    source_pod=metadata.get("labels", {}).get(
+                        "podsnapshot.gke.io/pod-name", "Unknown"
+                    ),
+                    creation_timestamp=metadata.get("creationTimestamp", ""),
+                    status="Ready" if is_ready else "NotReady",
+                )
+            )
+
+        if not valid_snapshots:
+            logger.info("No snapshots found matching criteria.")
+            return ListSnapshotResult(
+                success=True,
+                snapshots=[],
+                error_reason="",
+                error_code=SNAPSHOT_SUCCESS_CODE,
+            )
+
+        # Sort snapshots by creation timestamp descending
+        valid_snapshots.sort(key=lambda x: x.creation_timestamp or "", reverse=True)
+        logger.info(f"Found {len(valid_snapshots)} snapshots.")
+        return ListSnapshotResult(
+            success=True,
+            snapshots=valid_snapshots,
+            error_reason="",
+            error_code=SNAPSHOT_SUCCESS_CODE,
+        )
+
+    def delete_snapshots(
+        self,
+        grouping_labels: dict[str, str] | None = None,
+        snapshot_uid: str | None = None,
+    ) -> DeleteSnapshotResult:
+        """
+        Deletes snapshots.
+        - If snapshot_uid is provided, deletes that specific snapshot.
+        - If grouping_labels is provided, deletes all snapshots matching the grouping labels.
+        - If not provided, deletes ALL snapshots for this pod.
+
+        Note: snapshot_uid and grouping_labels are mutually exclusive.
+
+        Returns a DeleteSnapshotResult containing the list of successfully deleted snapshots.
+        """
+        if snapshot_uid and grouping_labels:
+            raise ValueError(
+                "snapshot_uid and grouping_labels are mutually exclusive. "
+                "Provide only one of them."
+            )
+
+        snapshots_to_delete = []
+
+        if snapshot_uid:
+            snapshots_to_delete.append(snapshot_uid)
+        else:
+            if grouping_labels:
+                logger.info(
+                    f"No snapshot_uid provided. Deleting snapshots based on pod name and grouping_labels: {grouping_labels}"
+                )
+            else:
+                logger.info("No filters provided. Deleting ALL snapshots for this pod.")
+
+            # Fetch all snapshots using list_snapshots without filtering by ready status
+            snapshots_result = self.list_snapshots(
+                grouping_labels=grouping_labels, ready_only=False
+            )
+            if not snapshots_result.success:
+                return DeleteSnapshotResult(
+                    success=False,
+                    deleted_snapshots=[],
+                    error_reason=f"Failed to list snapshots before deletion: {snapshots_result.error_reason}",
+                    error_code=SNAPSHOT_ERROR_CODE,
+                )
+            if snapshots_result.snapshots:
+                snapshots_to_delete = [
+                    s.snapshot_uid for s in snapshots_result.snapshots
+                ]
+        logger.info(f"Snapshots to delete: {snapshots_to_delete}")
+
+        if not snapshots_to_delete:
+            logger.info("No snapshots found matching criteria to delete.")
+            return DeleteSnapshotResult(
+                success=True,
+                deleted_snapshots=[],
+                error_reason="",
+                error_code=SNAPSHOT_SUCCESS_CODE,
+            )
+
+        deleted_snapshots = []
+        errors = []
+        for uid in snapshots_to_delete:
+            # Delete PodSnapshot
+            try:
+                logger.info(f"Deleting PodSnapshot '{uid}'...")
+                self.custom_objects_api.delete_namespaced_custom_object(
+                    group=PODSNAPSHOT_API_GROUP,
+                    version=PODSNAPSHOT_API_VERSION,
+                    namespace=self.namespace,
+                    plural=PODSNAPSHOT_PLURAL,
+                    name=uid,
+                )
+                logger.info(f"PodSnapshot '{uid}' deleted.")
+                deleted_snapshots.append(uid)
+            except ApiException as e:
+                if e.status == 404:
+                    logger.info(
+                        f"PodSnapshot '{uid}' not found in K8s (already deleted?)."
+                    )
+                else:
+                    msg = f"Failed to delete PodSnapshot '{uid}': {e}"
+                    logger.error(msg)
+                    errors.append(msg)
+            except Exception as e:
+                msg = f"Unexpected error deleting PodSnapshot '{uid}': {e}"
+                logger.exception(msg)
+                errors.append(msg)
+
+        logger.info(
+            f"Snapshot deletion process completed. Deleted {len(deleted_snapshots)} snapshots."
+        )
+
+        if errors:
+            error_msg = "; ".join(errors)
+            if deleted_snapshots:
+                error_msg = f"Partial failure: deleted {len(deleted_snapshots)}/{len(snapshots_to_delete)} snapshots. Errors: {error_msg}"
+            return DeleteSnapshotResult(
+                success=False,
+                deleted_snapshots=deleted_snapshots,
+                error_reason=error_msg,
+                error_code=SNAPSHOT_ERROR_CODE,
+            )
+
+        return DeleteSnapshotResult(
+            success=True,
+            deleted_snapshots=deleted_snapshots,
+            error_reason="",
+            error_code=SNAPSHOT_SUCCESS_CODE,
+        )
+
     def __exit__(self, exc_type, exc_val, exc_tb):
         """
         Cleans up the PodSnapshotManualTrigger Resources.