From a3df2c688cec490f8615ac5c882b25defe375ec4 Mon Sep 17 00:00:00 2001
From: Victor Morales <v.morales@samsung.com>
Date: Thu, 19 Dec 2024 20:42:39 -0800
Subject: [PATCH] Implement OwnerReference in new DataVolumes

---
 go.mod             |  2 +-
 internal/runner.go | 57 +++++++++++++++++++++++++++++++---------------
 2 files changed, 40 insertions(+), 19 deletions(-)

diff --git a/go.mod b/go.mod
index b32fba5..6bcf938 100644
--- a/go.mod
+++ b/go.mod
@@ -11,6 +11,7 @@ require (
 	github.com/spf13/viper v1.19.0
 	k8s.io/api v0.31.3
 	k8s.io/apimachinery v0.31.3
+	k8s.io/utils v0.0.0-20240711033017-18e509b52bc8
 	kubevirt.io/api v1.4.0
 	kubevirt.io/client-go v1.4.0
 	kubevirt.io/containerized-data-importer-api v1.57.0-alpha1
@@ -79,7 +80,6 @@ require (
 	k8s.io/client-go v0.31.3 // indirect
 	k8s.io/klog/v2 v2.130.1 // indirect
 	k8s.io/kube-openapi v0.31.0 // indirect
-	k8s.io/utils v0.0.0-20240711033017-18e509b52bc8 // indirect
 	kubevirt.io/controller-lifecycle-operator-sdk/api v0.0.0-20220329064328-f3cc58c6ed90 // indirect
 	sigs.k8s.io/json v0.0.0-20221116044647-bc3834ca7abd // indirect
 	sigs.k8s.io/structured-merge-diff/v4 v4.4.1 // indirect
diff --git a/internal/runner.go b/internal/runner.go
index 7b3fbbb..a195972 100644
--- a/internal/runner.go
+++ b/internal/runner.go
@@ -25,6 +25,7 @@ import (
 	"github.com/spf13/pflag"
 	k8scorev1 "k8s.io/api/core/v1"
 	k8smetav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/utils/ptr"
 	v1 "kubevirt.io/api/core/v1"
 	cdiclient "kubevirt.io/client-go/containerizeddataimporter"
 	"kubevirt.io/client-go/kubecli"
@@ -77,9 +78,12 @@ func (rc *Runner) getResources(ctx context.Context, vmTemplate, runnerName, jitC
 	for _, dvt := range virtualMachine.Spec.DataVolumeTemplates {
 		for _, volume := range virtualMachineInstance.Spec.Volumes {
 			if volume.DataVolume != nil && volume.DataVolume.Name == dvt.Name {
-				dataVolume = &v1beta1.DataVolume{}
-				dataVolume.Name = fmt.Sprintf("%s-%s", dvt.Name, runnerName)
-				dataVolume.Spec = dvt.Spec
+				dataVolume = &v1beta1.DataVolume{
+					ObjectMeta: k8smetav1.ObjectMeta{
+						Name: fmt.Sprintf("%s-%s", dvt.Name, runnerName),
+					},
+					Spec: dvt.Spec,
+				}
 
 				volume.DataVolume.Name = dataVolume.Name
 
@@ -88,7 +92,13 @@ func (rc *Runner) getResources(ctx context.Context, vmTemplate, runnerName, jitC
 		}
 	}
 
-	virtualMachineInstance.Spec.Volumes = append(virtualMachineInstance.Spec.Volumes, v1.Volume{
+	virtualMachineInstance.Spec.Volumes = append(virtualMachineInstance.Spec.Volumes, generateRunnerInfoVolume())
+
+	return virtualMachineInstance, dataVolume
+}
+
+func generateRunnerInfoVolume() v1.Volume {
+	return v1.Volume{
 		Name: runnerInfoVolume,
 		VolumeSource: v1.VolumeSource{
 			DownwardAPI: &v1.DownwardAPIVolumeSource{
@@ -102,9 +112,7 @@ func (rc *Runner) getResources(ctx context.Context, vmTemplate, runnerName, jitC
 				},
 			},
 		},
-	})
-
-	return virtualMachineInstance, dataVolume
+	}
 }
 
 func (rc *Runner) CreateResources(ctx context.Context,
@@ -112,23 +120,36 @@ func (rc *Runner) CreateResources(ctx context.Context,
 ) {
 	virtualMachineInstance, dataVolume := rc.getResources(ctx, vmTemplate, runnerName, jitConfig)
 
-	log.Printf("Creating %s Data Volume\n", dataVolume.Name)
-
-	if _, err := rc.cdiClient.CdiV1beta1().DataVolumes(
-		rc.namespace).Create(ctx, dataVolume, k8smetav1.CreateOptions{}); err != nil {
-		log.Fatalf("cannot create data volume: %v\n", err)
-	}
-
-	rc.dataVolume = dataVolume.Name
-
 	log.Printf("Creating %s Virtual Machine Instance\n", virtualMachineInstance.Name)
 
-	if _, err := rc.virtClient.VirtualMachineInstance(rc.namespace).Create(ctx,
-		virtualMachineInstance, k8smetav1.CreateOptions{}); err != nil {
+	vmi, err := rc.virtClient.VirtualMachineInstance(rc.namespace).Create(ctx,
+		virtualMachineInstance, k8smetav1.CreateOptions{})
+	if err != nil {
 		log.Fatal(err.Error())
 	}
 
 	rc.virtualMachineInstance = virtualMachineInstance.Name
+
+	if dataVolume != nil {
+		log.Printf("Creating %s Data Volume\n", dataVolume.Name)
+
+		dataVolume.OwnerReferences = []k8smetav1.OwnerReference{
+			{
+				APIVersion: "kubevirt.io/v1",
+				Kind:       "VirtualMachineInstance",
+				Name:       vmi.Name,
+				UID:        vmi.UID,
+				Controller: ptr.To(false),
+			},
+		}
+
+		if _, err := rc.cdiClient.CdiV1beta1().DataVolumes(
+			rc.namespace).Create(ctx, dataVolume, k8smetav1.CreateOptions{}); err != nil {
+			log.Fatalf("cannot create data volume: %v\n", err)
+		}
+
+		rc.dataVolume = dataVolume.Name
+	}
 }
 
 func (rc *Runner) WaitForVirtualMachineInstance(ctx context.Context) {