From 2b86defba61907b8fcd9d80ff92ccc42366484ab Mon Sep 17 00:00:00 2001
From: Vivek Singh <vivek.singh@veeam.com>
Date: Fri, 8 Mar 2024 21:33:40 +0100
Subject: [PATCH] Add `provisioned-by` annotations to PVs created by
 `CreateVolumeFromSnapshot` kanister function (#2705)

* Add `provisioned-by` annotations to PVs created by `CreateVolumeFromSnapshot` fun

* Address review comment, fix go check import

* Add `provisioned-by` annotations to PVs created by `CreateVolumeFromSnapshot` fun

* Address review comment, fix go check import

---------

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
---
 pkg/consts/consts.go                          |  7 +++
 pkg/function/create_volume_from_snapshot.go   | 21 +++++++-
 .../create_volume_from_snapshot_test.go       | 51 +++++++++++++++++++
 3 files changed, 78 insertions(+), 1 deletion(-)

diff --git a/pkg/consts/consts.go b/pkg/consts/consts.go
index 948648672a..1130db2ebe 100644
--- a/pkg/consts/consts.go
+++ b/pkg/consts/consts.go
@@ -38,6 +38,13 @@ const (
 	ProfileResourceNamePlural   = "profiles"
 )
 
+const (
+	PVProvisionedByAnnotation = "pv.kubernetes.io/provisioned-by"
+
+	AWSEBSProvisionerInTree = "kubernetes.io/aws-ebs"
+	GCEPDProvisionerInTree  = "kubernetes.io/gce-pd"
+)
+
 // These consts are used to query Repository server API objects
 const RepositoryServerResourceName = "repositoryserver"
 const RepositoryServerResourceNamePlural = "repositoryservers"
diff --git a/pkg/function/create_volume_from_snapshot.go b/pkg/function/create_volume_from_snapshot.go
index 17c54ff3b9..30b14e9d96 100644
--- a/pkg/function/create_volume_from_snapshot.go
+++ b/pkg/function/create_volume_from_snapshot.go
@@ -28,6 +28,7 @@ import (
 	awsconfig "github.com/kanisterio/kanister/pkg/aws"
 	"github.com/kanisterio/kanister/pkg/blockstorage"
 	"github.com/kanisterio/kanister/pkg/blockstorage/getter"
+	"github.com/kanisterio/kanister/pkg/consts"
 	"github.com/kanisterio/kanister/pkg/field"
 	"github.com/kanisterio/kanister/pkg/kube"
 	kubevolume "github.com/kanisterio/kanister/pkg/kube/volume"
@@ -115,7 +116,9 @@ func createVolumeFromSnapshot(ctx context.Context, cli kubernetes.Interface, nam
 		if err != nil {
 			return nil, errors.Wrapf(err, "Unable to create PVC for volume %v", *vol)
 		}
-		pv, err := kubevolume.CreatePV(ctx, cli, vol, vol.Type, annotations, nil, nil)
+
+		pvAnnotations := addPVProvisionedByAnnotation(nil, provider)
+		pv, err := kubevolume.CreatePV(ctx, cli, vol, vol.Type, pvAnnotations, nil, nil)
 		if err != nil {
 			return nil, errors.Wrapf(err, "Unable to create PV for volume %v", *vol)
 		}
@@ -125,6 +128,22 @@ func createVolumeFromSnapshot(ctx context.Context, cli kubernetes.Interface, nam
 	return providerList, nil
 }
 
+func addPVProvisionedByAnnotation(annotations map[string]string, provider blockstorage.Provider) map[string]string {
+	if annotations == nil {
+		annotations = make(map[string]string)
+	}
+
+	storageType := provider.Type()
+	switch storageType {
+	case blockstorage.TypeGPD:
+		annotations[consts.PVProvisionedByAnnotation] = consts.GCEPDProvisionerInTree
+	case blockstorage.TypeEBS:
+		annotations[consts.PVProvisionedByAnnotation] = consts.AWSEBSProvisionerInTree
+	}
+
+	return annotations
+}
+
 func (c *createVolumeFromSnapshotFunc) Exec(ctx context.Context, tp param.TemplateParams, args map[string]interface{}) (map[string]interface{}, error) {
 	// Set progress percent
 	c.progressPercent = progress.StartedPercent
diff --git a/pkg/function/create_volume_from_snapshot_test.go b/pkg/function/create_volume_from_snapshot_test.go
index d2dd4a9f95..9bd184e34d 100644
--- a/pkg/function/create_volume_from_snapshot_test.go
+++ b/pkg/function/create_volume_from_snapshot_test.go
@@ -28,6 +28,9 @@ import (
 
 	crv1alpha1 "github.com/kanisterio/kanister/pkg/apis/cr/v1alpha1"
 	"github.com/kanisterio/kanister/pkg/blockstorage"
+	"github.com/kanisterio/kanister/pkg/blockstorage/awsebs"
+	"github.com/kanisterio/kanister/pkg/blockstorage/gcepd"
+	"github.com/kanisterio/kanister/pkg/consts"
 	"github.com/kanisterio/kanister/pkg/param"
 	"github.com/kanisterio/kanister/pkg/testutil/mockblockstorage"
 )
@@ -133,3 +136,51 @@ func (s *CreateVolumeFromSnapshotTestSuite) TestCreateVolumeFromSnapshot(c *C) {
 		}
 	}
 }
+
+func (s *CreateVolumeFromSnapshotTestSuite) TestAddPVProvisionedByAnnotation(c *C) {
+	for _, tc := range []struct {
+		st                  blockstorage.Provider
+		annotations         map[string]string
+		expectedAnnotations map[string]string
+	}{
+		{
+			st:          &gcepd.GpdStorage{},
+			annotations: nil,
+			expectedAnnotations: map[string]string{
+				consts.PVProvisionedByAnnotation: consts.GCEPDProvisionerInTree,
+			},
+		},
+		{
+			st: &gcepd.GpdStorage{},
+			annotations: map[string]string{
+				"key": "value",
+			},
+			expectedAnnotations: map[string]string{
+				"key":                            "value",
+				consts.PVProvisionedByAnnotation: consts.GCEPDProvisionerInTree,
+			},
+		},
+		{
+			st:          &gcepd.GpdStorage{},
+			annotations: map[string]string{},
+			expectedAnnotations: map[string]string{
+				consts.PVProvisionedByAnnotation: consts.GCEPDProvisionerInTree,
+			},
+		},
+		{
+			st: &awsebs.EbsStorage{},
+			annotations: map[string]string{
+				"keyone": "valueone",
+				"keytwo": "valuetwo",
+			},
+			expectedAnnotations: map[string]string{
+				"keyone":                         "valueone",
+				"keytwo":                         "valuetwo",
+				consts.PVProvisionedByAnnotation: consts.AWSEBSProvisionerInTree,
+			},
+		},
+	} {
+		op := addPVProvisionedByAnnotation(tc.annotations, tc.st)
+		c.Assert(op, DeepEquals, tc.expectedAnnotations)
+	}
+}