From 478271fa5a77df2d206bb60dc4ce74de339039c6 Mon Sep 17 00:00:00 2001
From: Dmitry Chepurovskiy <dm3ch@dm3ch.net>
Date: Sun, 2 Aug 2020 14:18:15 +0300
Subject: [PATCH] Added labels support to create command of SQL driver

Signed-off-by: Dmitry Chepurovskiy <dm3ch@dm3ch.net>
---
 cmd/helm/install.go       |  1 +
 pkg/action/install.go     |  8 ++++++++
 pkg/storage/driver/sql.go | 33 +++++++++++++++++++++++++++++++++
 3 files changed, 42 insertions(+)

diff --git a/cmd/helm/install.go b/cmd/helm/install.go
index f7d0239fb..459e68911 100644
--- a/cmd/helm/install.go
+++ b/cmd/helm/install.go
@@ -149,6 +149,7 @@ func addInstallFlags(cmd *cobra.Command, f *pflag.FlagSet, client *action.Instal
 	f.BoolVar(&client.Atomic, "atomic", false, "if set, the installation process deletes the installation on failure. The --wait flag will be set automatically if --atomic is used")
 	f.BoolVar(&client.SkipCRDs, "skip-crds", false, "if set, no CRDs will be installed. By default, CRDs are installed if not already present")
 	f.BoolVar(&client.SubNotes, "render-subchart-notes", false, "if set, render subchart notes along with the parent")
+	f.StringVarP(&client.Labels, "labels", "l", "", "Labels that would be added to relese metadata. Should be divided by comma")
 	addValueOptionsFlags(f, valueOpts)
 	addChartPathOptionsFlags(f, &client.ChartPathOptions)
 
diff --git a/pkg/action/install.go b/pkg/action/install.go
index 00fb208b0..55e7db960 100644
--- a/pkg/action/install.go
+++ b/pkg/action/install.go
@@ -100,6 +100,7 @@ type Install struct {
 	// OutputDir/<ReleaseName>
 	UseReleaseName bool
 	PostRenderer   postrender.PostRenderer
+	Labels         string
 }
 
 // ChartPathOptions captures common options used for controlling chart paths
@@ -235,6 +236,13 @@ func (i *Install) Run(chrt *chart.Chart, vals map[string]interface{}) (*release.
 
 	rel := i.createRelease(chrt, vals)
 
+	// FIXME: Rework me
+	rel.Labels = make(map[string]string)
+	for _, label := range strings.Split(i.Labels, ",") {
+		parts := strings.Split(label, "=")
+		rel.Labels[parts[0]] = parts[1]
+	}
+
 	var manifestDoc *bytes.Buffer
 	rel.Hooks, manifestDoc, rel.Info.Notes, err = i.cfg.renderResources(chrt, valuesToRender, i.ReleaseName, i.OutputDir, i.SubNotes, i.UseReleaseName, i.IncludeCRDs, i.PostRenderer, i.DryRun)
 	// Even for errors, attach this if available
diff --git a/pkg/storage/driver/sql.go b/pkg/storage/driver/sql.go
index f68f50f54..6d2b90592 100644
--- a/pkg/storage/driver/sql.go
+++ b/pkg/storage/driver/sql.go
@@ -150,6 +150,11 @@ func (s *SQL) ensureDBSetup() error {
 					`, sqlReleaseTableName),
 				},
 			},
+			{
+				Id:   "labels",
+				Up:   []string{"CREATE TABLE labels (release_key VARCHAR(67), key VARCHAR(64), value VARCHAR(67));"},
+				Down: []string{"DELETE TABLE labels;"},
+			},
 		},
 	}
 
@@ -399,6 +404,34 @@ func (s *SQL) Create(key string, rls *rspb.Release) error {
 		s.Log("failed to store release %s in SQL database: %v", key, err)
 		return err
 	}
+
+	for lk, lv := range rls.Labels {
+		insertLabelsQuery, args, err := s.statementBuilder.
+			Insert("labels").
+			Columns(
+				"release_key",
+				"key",
+				"value",
+			).
+			Values(
+				key,
+				lk,
+				lv,
+			).ToSql()
+
+		if err != nil {
+			defer transaction.Rollback()
+			s.Log("failed to build insert query: %v", err)
+			return err
+		}
+
+		if _, err := transaction.Exec(insertLabelsQuery, args...); err != nil {
+			defer transaction.Rollback()
+			s.Log("failed to write labels: %v", err)
+			return err
+		}
+	}
+
 	defer transaction.Commit()
 
 	return nil