From cc4850397c1a53230bdb74ffb33ec2c28cc5f2e5 Mon Sep 17 00:00:00 2001
From: houkunpeng <houkunpeng@lixiang.com>
Date: Tue, 21 Feb 2023 10:47:27 +0800
Subject: [PATCH] feat:support custom folder expression and by annotation
 setting folder name

---
 .../pvc-with-custom-folder/kustomization.yaml |  4 ++++
 examples/pvc-with-custom-folder/pvc.yaml      | 13 +++++++++++
 provisioner.go                                | 22 +++++++++++++++----
 test/pod_test.go                              |  9 +++++++-
 .../pod-with-custom-floder/kustomization.yaml | 10 +++++++++
 5 files changed, 53 insertions(+), 5 deletions(-)
 create mode 100644 examples/pvc-with-custom-folder/kustomization.yaml
 create mode 100644 examples/pvc-with-custom-folder/pvc.yaml
 create mode 100644 test/testdata/pod-with-custom-floder/kustomization.yaml

diff --git a/examples/pvc-with-custom-folder/kustomization.yaml b/examples/pvc-with-custom-folder/kustomization.yaml
new file mode 100644
index 000000000..7bfd4518a
--- /dev/null
+++ b/examples/pvc-with-custom-folder/kustomization.yaml
@@ -0,0 +1,4 @@
+apiVersion: kustomize.config.k8s.io/v1beta1
+kind: Kustomization
+resources:
+- pvc.yaml
diff --git a/examples/pvc-with-custom-folder/pvc.yaml b/examples/pvc-with-custom-folder/pvc.yaml
new file mode 100644
index 000000000..65c0e7921
--- /dev/null
+++ b/examples/pvc-with-custom-folder/pvc.yaml
@@ -0,0 +1,13 @@
+apiVersion: v1
+kind: PersistentVolumeClaim
+metadata:
+  name: custom-folder-pvc
+  annotations:
+    "rancher.io/customFolderName": "demo1"
+spec:
+  accessModes:
+    - ReadWriteOnce
+  storageClassName: local-path
+  resources:
+    requests:
+      storage: 128Mi
diff --git a/provisioner.go b/provisioner.go
index 795a423d4..bdfe96ef6 100644
--- a/provisioner.go
+++ b/provisioner.go
@@ -30,7 +30,8 @@ const (
 )
 
 const (
-	KeyNode = "kubernetes.io/hostname"
+	KeyNode                    = "kubernetes.io/hostname"
+	customFolderNameAnnotation = "rancher.io/customFolderName"
 
 	NodeDefaultNonListedNodes = "DEFAULT_PATH_FOR_NON_LISTED_NODES"
 
@@ -45,6 +46,7 @@ const (
 
 const (
 	defaultCmdTimeoutSeconds = 120
+	defaultFolderExpression  = "{{.pvName}}-{{.namespace}}-{{.pvcName}}"
 )
 
 var (
@@ -77,6 +79,7 @@ type ConfigData struct {
 	NodePathMap          []*NodePathMapData `json:"nodePathMap,omitempty"`
 	CmdTimeoutSeconds    int                `json:"cmdTimeoutSeconds,omitempty"`
 	SharedFileSystemPath string             `json:"sharedFileSystemPath,omitempty"`
+	FolderExpression     string             `json:"folderExpression,omitempty"`
 }
 
 type NodePathMap struct {
@@ -87,6 +90,7 @@ type Config struct {
 	NodePathMap          map[string]*NodePathMap
 	CmdTimeoutSeconds    int
 	SharedFileSystemPath string
+	FolderExpression     string
 }
 
 func NewProvisioner(ctx context.Context, kubeClient *clientset.Clientset,
@@ -258,15 +262,20 @@ func (p *LocalPathProvisioner) Provision(ctx context.Context, opts pvController.
 	}
 
 	name := opts.PVName
-	folderName := strings.Join([]string{name, opts.PVC.Namespace, opts.PVC.Name}, "_")
-
+	folderName := ""
+	if pvc.GetAnnotations()[customFolderNameAnnotation] != "" {
+		folderName = pvc.GetAnnotations()[customFolderNameAnnotation]
+	} else {
+		folderName := strings.Replace(p.config.FolderExpression, "{{.namespace}}", opts.PVC.Namespace, -1)
+		folderName = strings.Replace(folderName, "{{.pvName}}", name, -1)
+		folderName = strings.Replace(folderName, "{{.pvcName}}", opts.PVC.Name, -1)
+	}
 	path := filepath.Join(basePath, folderName)
 	if nodeName == "" {
 		logrus.Infof("Creating volume %v at %v", name, path)
 	} else {
 		logrus.Infof("Creating volume %v at %v:%v", name, nodeName, path)
 	}
-
 	storage := pvc.Spec.Resources.Requests[v1.ResourceName(v1.ResourceStorage)]
 	provisionCmd := []string{"/bin/sh", "/script/setup"}
 	if err := p.createHelperPod(ActionTypeCreate, provisionCmd, volumeOptions{
@@ -653,5 +662,10 @@ func canonicalizeConfig(data *ConfigData) (cfg *Config, err error) {
 	} else {
 		cfg.CmdTimeoutSeconds = defaultCmdTimeoutSeconds
 	}
+	if data.FolderExpression != "" {
+		cfg.FolderExpression = data.FolderExpression
+	} else {
+		cfg.FolderExpression = defaultFolderExpression
+	}
 	return cfg, nil
 }
diff --git a/test/pod_test.go b/test/pod_test.go
index f7be8106a..8061b7553 100644
--- a/test/pod_test.go
+++ b/test/pod_test.go
@@ -1,3 +1,4 @@
+//go:build e2e
 // +build e2e
 
 package test
@@ -6,9 +7,9 @@ import (
 	"fmt"
 	"github.com/kelseyhightower/envconfig"
 	"github.com/stretchr/testify/suite"
+	"strings"
 	"testing"
 	"time"
-	"strings"
 )
 
 const (
@@ -131,6 +132,12 @@ func (p *PodTestSuite) TestPodWithSubpath() {
 	runTest(p, []string{p.config.IMAGE}, "ready", hostPathVolumeType)
 }
 
+func (p *PodTestSuite) TestPodWithCustomFolder() {
+	p.kustomizeDir = "pod-with-custom-folder"
+
+	runTest(p, []string{p.config.IMAGE}, "ready", hostPathVolumeType)
+}
+
 func runTest(p *PodTestSuite, images []string, waitCondition, volumeType string) {
 	kustomizeDir := testdataFile(p.kustomizeDir)
 
diff --git a/test/testdata/pod-with-custom-floder/kustomization.yaml b/test/testdata/pod-with-custom-floder/kustomization.yaml
new file mode 100644
index 000000000..c43cb220d
--- /dev/null
+++ b/test/testdata/pod-with-custom-floder/kustomization.yaml
@@ -0,0 +1,10 @@
+apiVersion: kustomize.config.k8s.io/v1beta1
+kind: Kustomization
+resources:
+- ../../../deploy
+- ../../../examples/pvc-with-custom-folder
+commonLabels:
+  app: custom-folder-provisioner
+images:
+- name: rancher/local-path-provisioner
+  newTag: dev