diff --git a/ocp_resources/data_source.py b/ocp_resources/data_source.py index 247bca42da..d796b9ff65 100644 --- a/ocp_resources/data_source.py +++ b/ocp_resources/data_source.py @@ -3,6 +3,7 @@ from ocp_resources.constants import TIMEOUT_4MINUTES from ocp_resources.persistent_volume_claim import PersistentVolumeClaim from ocp_resources.resource import NamespacedResource +from ocp_resources.volume_snapshot import VolumeSnapshot class DataSource(NamespacedResource): @@ -39,19 +40,29 @@ def to_dict(self): }, }) - @property - def pvc(self): - data_source_pvc = self.instance.spec.source.pvc - pvc_name = data_source_pvc.name - pvc_namespace = data_source_pvc.namespace + def _get_boot_source(self, boot_source_type): + boot_source = self.instance.spec.source.get(boot_source_type) + if not boot_source: + return None + boot_source_name = boot_source.name + boot_source_namespace = boot_source.namespace try: - return PersistentVolumeClaim( + boot_source_object = PersistentVolumeClaim if boot_source_type == "pvc" else VolumeSnapshot + return boot_source_object( client=self.client, - name=pvc_name, - namespace=pvc_namespace, + name=boot_source_name, + namespace=boot_source_namespace, ) except ResourceNotFoundError: self.logger.warning( - f"dataSource {self.name} is pointing to a non-existing PVC, name:" - f" {pvc_name}, namespace: {pvc_namespace}" + f"dataSource {self.name} is pointing to a non-existing {boot_source_type}, name:" + f" {boot_source_name}, namespace: {boot_source_namespace}" ) + + @property + def pvc(self): + return self._get_boot_source(boot_source_type="pvc") + + @property + def snapshot(self): + return self._get_boot_source(boot_source_type="snapshot")