Added support for storing custom labels in SQL storage driver

Fix list command for SQL storage driver

Fix SQL storage drivers tests after adding custom labels support

Remove notes that SQL driver not supported for storing labels in install and upgrade

Signed-off-by: Dmitry Chepurovskiy <dm3ch@dm3ch.net>
Signed-off-by: Dmitry Chepurovskiy <me@dm3ch.net>

cmd/helm/install.go

@@ -155,7 +155,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.StringToStringVarP(&client.Labels, "labels", "l", nil, "Labels that would be added to relese metadata. Should be divided by comma. (Currently works only with configmap and secret storage drivers).")
+	f.StringToStringVarP(&client.Labels, "labels", "l", nil, "Labels that would be added to relese metadata. Should be divided by comma.")
 	addValueOptionsFlags(f, valueOpts)
 	addChartPathOptionsFlags(f, &client.ChartPathOptions)
 

cmd/helm/upgrade.go

@@ -229,7 +229,7 @@ func newUpgradeCmd(cfg *action.Configuration, out io.Writer) *cobra.Command {
 	f.IntVar(&client.MaxHistory, "history-max", settings.MaxHistory, "limit the maximum number of revisions saved per release. Use 0 for no limit")
 	f.BoolVar(&client.CleanupOnFail, "cleanup-on-fail", false, "allow deletion of new resources created in this upgrade when upgrade fails")
 	f.BoolVar(&client.SubNotes, "render-subchart-notes", false, "if set, render subchart notes along with the parent")
-	f.StringToStringVarP(&client.Labels, "labels", "l", nil, "Labels that would be added to relese metadata. Should be divided by comma. (Currently works only with configmap and secret storage drivers). Original release labels would be merged with upgrade labels. You can unset label using null.")
+	f.StringToStringVarP(&client.Labels, "labels", "l", nil, "Labels that would be added to relese metadata. Should be divided by comma. Original release labels would be merged with upgrade labels. You can unset label using null.")
 	f.StringVar(&client.Description, "description", "", "add a custom description")
 	f.BoolVar(&client.DependencyUpdate, "dependency-update", false, "update dependencies if they are missing before installing the chart")
 	addChartPathOptionsFlags(f, &client.ChartPathOptions)

pkg/storage/driver/sql.go

@@ -49,6 +49,7 @@ const postgreSQLDialect = "postgres"
 const SQLDriverName = "SQL"
 
 const sqlReleaseTableName = "releases_v1"
+const sqlCustomLabelsTableName = "custom_labels_v1"
 
 const (
 	sqlReleaseTableKeyColumn        = "key"
@@ -61,6 +62,17 @@ const (
 	sqlReleaseTableOwnerColumn      = "owner"
 	sqlReleaseTableCreatedAtColumn  = "createdAt"
 	sqlReleaseTableModifiedAtColumn = "modifiedAt"
+
+	sqlCustomLabelsTableReleaseKeyColumn       = "releaseKey"
+	sqlCustomLabelsTableReleaseNamespaceColumn = "releaseNamespace"
+	sqlCustomLabelsTableKeyColumn              = "key"
+	sqlCustomLabelsTableValueColumn            = "value"
+)
+
+// Following limits based on k8s labels limits - https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/#syntax-and-character-set
+const (
+	sqlCustomLabelsTableKeyMaxLenght   = 253 + 1 + 63
+	sqlCustomLabelsTableValueMaxLenght = 63
 )
 
 const (
@@ -150,6 +162,41 @@ func (s *SQL) ensureDBSetup() error {
 					`, sqlReleaseTableName),
 				},
 			},
+			{
+				Id: "custom_labels",
+				Up: []string{
+					fmt.Sprintf(`
+						CREATE TABLE %s (
+							%s VARCHAR(64),
+							%s VARCHAR(67),
+							%s VARCHAR(%d), 
+							%s VARCHAR(%d)
+						);
+						CREATE INDEX ON %s (%s, %s);
+						
+						GRANT ALL ON %s TO PUBLIC;
+						ALTER TABLE %s ENABLE ROW LEVEL SECURITY;
+					`,
+						sqlCustomLabelsTableName,
+						sqlCustomLabelsTableReleaseKeyColumn,
+						sqlCustomLabelsTableReleaseNamespaceColumn,
+						sqlCustomLabelsTableKeyColumn,
+						sqlCustomLabelsTableKeyMaxLenght,
+						sqlCustomLabelsTableValueColumn,
+						sqlCustomLabelsTableValueMaxLenght,
+						sqlCustomLabelsTableName,
+						sqlCustomLabelsTableReleaseKeyColumn,
+						sqlCustomLabelsTableReleaseNamespaceColumn,
+						sqlCustomLabelsTableName,
+						sqlCustomLabelsTableName,
+					),
+				},
+				Down: []string{
+					fmt.Sprintf(`
+						DELETE TABLE %s;
+					`, sqlCustomLabelsTableName),
+				},
+			},
 		},
 	}
 
@@ -180,6 +227,13 @@ type SQLReleaseWrapper struct {
 	ModifiedAt int    `db:"modifiedAt"`
 }
 
+type SQLReleaseCustomLabelWrapper struct {
+	ReleaseKey       string `db:"release_key"`
+	ReleaseNamespace string `db:"release_namespace"`
+	Key              string `db:"key"`
+	Value            string `db:"value"`
+}
+
 // NewSQL initializes a new sql driver.
 func NewSQL(connectionString string, logger func(string, ...interface{}), namespace string) (*SQL, error) {
 	db, err := sqlx.Connect(postgreSQLDialect, connectionString)
@@ -230,13 +284,18 @@ func (s *SQL) Get(key string) (*rspb.Release, error) {
 		return nil, err
 	}
 
+	if release.Labels, err = s.getReleaseCustomLabels(key, s.namespace); err != nil {
+		s.Log("failed to get release %s/%s custom labels: %v", s.namespace, key, err)
+		return nil, err
+	}
+
 	return release, nil
 }
 
 // List returns the list of all releases such that filter(release) == true
 func (s *SQL) List(filter func(*rspb.Release) bool) ([]*rspb.Release, error) {
 	sb := s.statementBuilder.
-		Select(sqlReleaseTableBodyColumn).
+		Select(sqlReleaseTableKeyColumn, sqlReleaseTableNamespaceColumn, sqlReleaseTableBodyColumn).
 		From(sqlReleaseTableName).
 		Where(sq.Eq{sqlReleaseTableOwnerColumn: sqlReleaseDefaultOwner})
 
@@ -264,6 +323,12 @@ func (s *SQL) List(filter func(*rspb.Release) bool) ([]*rspb.Release, error) {
 			s.Log("list: failed to decode release: %v: %v", record, err)
 			continue
 		}
+
+		if release.Labels, err = s.getReleaseCustomLabels(record.Key, record.Namespace); err != nil {
+			s.Log("failed to get release %s/%s custom labels: %v", record.Namespace, record.Key, err)
+			return nil, err
+		}
+
 		if filter(release) {
 			releases = append(releases, release)
 		}
@@ -275,7 +340,7 @@ func (s *SQL) List(filter func(*rspb.Release) bool) ([]*rspb.Release, error) {
 // Query returns the set of releases that match the provided set of labels.
 func (s *SQL) Query(labels map[string]string) ([]*rspb.Release, error) {
 	sb := s.statementBuilder.
-		Select(sqlReleaseTableBodyColumn).
+		Select(sqlReleaseTableKeyColumn, sqlReleaseTableNamespaceColumn, sqlReleaseTableBodyColumn).
 		From(sqlReleaseTableName)
 
 	keys := make([]string, 0, len(labels))
@@ -321,6 +386,12 @@ func (s *SQL) Query(labels map[string]string) ([]*rspb.Release, error) {
 			s.Log("list: failed to decode release: %v: %v", record, err)
 			continue
 		}
+
+		if release.Labels, err = s.getReleaseCustomLabels(record.Key, record.Namespace); err != nil {
+			s.Log("failed to get release %s/%s custom labels: %v", record.Namespace, record.Key, err)
+			return nil, err
+		}
+
 		releases = append(releases, release)
 	}
 
@@ -403,6 +474,36 @@ 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
 	}
+
+	// Filtering labels before insert cause in SQL storage driver system releases are stored in separate columns of release table
+	for k, v := range filterSystemLabels(rls.Labels) {
+		insertLabelsQuery, args, err := s.statementBuilder.
+			Insert(sqlCustomLabelsTableName).
+			Columns(
+				sqlCustomLabelsTableReleaseKeyColumn,
+				sqlCustomLabelsTableReleaseNamespaceColumn,
+				sqlCustomLabelsTableKeyColumn,
+				sqlCustomLabelsTableValueColumn,
+			).
+			Values(
+				key,
+				namespace,
+				k,
+				v,
+			).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
@@ -487,10 +588,51 @@ func (s *SQL) Delete(key string) (*rspb.Release, error) {
 		Where(sq.Eq{sqlReleaseTableNamespaceColumn: s.namespace}).
 		ToSql()
 	if err != nil {
-		s.Log("failed to build select query: %v", err)
+		s.Log("failed to build delete query: %v", err)
 		return nil, err
 	}
 
 	_, err = transaction.Exec(deleteQuery, args...)
+	if err != nil {
+		s.Log("failed perform delete query: %v", err)
 		return release, err
 	}
+
+	deleteCustomLabelsQuery, args, err := s.statementBuilder.
+		Delete(sqlCustomLabelsTableName).
+		Where(sq.Eq{sqlCustomLabelsTableReleaseKeyColumn: key}).
+		Where(sq.Eq{sqlCustomLabelsTableReleaseNamespaceColumn: s.namespace}).
+		ToSql()
+
+	if err != nil {
+		s.Log("failed to build delete Labels query: %v", err)
+		return nil, err
+	}
+	_, err = transaction.Exec(deleteCustomLabelsQuery, args...)
+	return release, err
+}
+
+// Get release custom labels from database
+func (s *SQL) getReleaseCustomLabels(key string, namespace string) (map[string]string, error) {
+	query, args, err := s.statementBuilder.
+		Select(sqlCustomLabelsTableKeyColumn, sqlCustomLabelsTableValueColumn).
+		From(sqlCustomLabelsTableName).
+		Where(sq.Eq{sqlCustomLabelsTableReleaseKeyColumn: key,
+			sqlCustomLabelsTableReleaseNamespaceColumn: s.namespace}).
+		ToSql()
+	if err != nil {
+		return nil, err
+	}
+
+	var labelsList = []SQLReleaseCustomLabelWrapper{}
+	if err := s.db.Select(&labelsList, query, args...); err != nil {
+		return nil, err
+	}
+
+	labelsMap := make(map[string]string)
+	for _, i := range labelsList {
+		labelsMap[i.Key] = i.Value
+	}
+
+	return filterSystemLabels(labelsMap), nil
+}

pkg/storage/driver/sql_test.go

@@ -62,6 +62,8 @@ func TestSQLGet(t *testing.T) {
 			),
 		).RowsWillBeClosed()
 
+	mockGetReleaseCustomLabels(mock, key, namespace, rel.Labels)
+
 	got, err := sqlDriver.Get(key)
 	if err != nil {
 		t.Fatalf("Failed to get release: %v", err)
@@ -77,38 +79,42 @@ func TestSQLGet(t *testing.T) {
 }
 
 func TestSQLList(t *testing.T) {
-	body1, _ := encodeRelease(releaseStub("key-1", 1, "default", rspb.StatusUninstalled))
-	body2, _ := encodeRelease(releaseStub("key-2", 1, "default", rspb.StatusUninstalled))
-	body3, _ := encodeRelease(releaseStub("key-3", 1, "default", rspb.StatusDeployed))
-	body4, _ := encodeRelease(releaseStub("key-4", 1, "default", rspb.StatusDeployed))
-	body5, _ := encodeRelease(releaseStub("key-5", 1, "default", rspb.StatusSuperseded))
-	body6, _ := encodeRelease(releaseStub("key-6", 1, "default", rspb.StatusSuperseded))
+	releases := []*rspb.Release{}
+	releases = append(releases, releaseStub("key-1", 1, "default", rspb.StatusUninstalled))
+	releases = append(releases, releaseStub("key-2", 1, "default", rspb.StatusUninstalled))
+	releases = append(releases, releaseStub("key-3", 1, "default", rspb.StatusDeployed))
+	releases = append(releases, releaseStub("key-4", 1, "default", rspb.StatusDeployed))
+	releases = append(releases, releaseStub("key-5", 1, "default", rspb.StatusSuperseded))
+	releases = append(releases, releaseStub("key-6", 1, "default", rspb.StatusSuperseded))
 
 	sqlDriver, mock := newTestFixtureSQL(t)
 
 	for i := 0; i < 3; i++ {
 		query := fmt.Sprintf(
-			"SELECT %s FROM %s WHERE %s = $1 AND %s = $2",
+			"SELECT %s, %s, %s FROM %s WHERE %s = $1 AND %s = $2",
+			sqlReleaseTableKeyColumn,
+			sqlReleaseTableNamespaceColumn,
 			sqlReleaseTableBodyColumn,
 			sqlReleaseTableName,
 			sqlReleaseTableOwnerColumn,
 			sqlReleaseTableNamespaceColumn,
 		)
 
+		rows := mock.NewRows([]string{
+			sqlReleaseTableBodyColumn,
+		})
+		for _, r := range releases {
+			body, _ := encodeRelease(r)
+			rows.AddRow(body)
+		}
 		mock.
 			ExpectQuery(regexp.QuoteMeta(query)).
 			WithArgs(sqlReleaseDefaultOwner, sqlDriver.namespace).
-			WillReturnRows(
-				mock.NewRows([]string{
-					sqlReleaseTableBodyColumn,
-				}).
-					AddRow(body1).
-					AddRow(body2).
-					AddRow(body3).
-					AddRow(body4).
-					AddRow(body5).
-					AddRow(body6),
-			).RowsWillBeClosed()
+			WillReturnRows(rows).RowsWillBeClosed()
+
+		for _, r := range releases {
+			mockGetReleaseCustomLabels(mock, "", r.Namespace, r.Labels)
+		}
 	}
 
 	// list all deleted releases
@@ -181,6 +187,21 @@ func TestSqlCreate(t *testing.T) {
 		ExpectExec(regexp.QuoteMeta(query)).
 		WithArgs(key, sqlReleaseDefaultType, body, rel.Name, rel.Namespace, int(rel.Version), rel.Info.Status.String(), sqlReleaseDefaultOwner, int(time.Now().Unix())).
 		WillReturnResult(sqlmock.NewResult(1, 1))
+
+	labelsQuery := fmt.Sprintf(
+		"INSERT INTO %s (%s,%s,%s,%s) VALUES ($1,$2,$3,$4)",
+		sqlCustomLabelsTableName,
+		sqlCustomLabelsTableReleaseKeyColumn,
+		sqlCustomLabelsTableReleaseNamespaceColumn,
+		sqlCustomLabelsTableKeyColumn,
+		sqlCustomLabelsTableValueColumn,
+	)
+	for k, v := range rel.Labels {
+		mock.
+			ExpectExec(regexp.QuoteMeta(labelsQuery)).
+			WithArgs(key, rel.Namespace, k, v).
+			WillReturnResult(sqlmock.NewResult(1, 1))
+	}
 	mock.ExpectCommit()
 
 	if err := sqlDriver.Create(key, rel); err != nil {
@@ -316,7 +337,9 @@ func TestSqlQuery(t *testing.T) {
 	sqlDriver, mock := newTestFixtureSQL(t)
 
 	query := fmt.Sprintf(
-		"SELECT %s FROM %s WHERE %s = $1 AND %s = $2 AND %s = $3 AND %s = $4",
+		"SELECT %s, %s, %s FROM %s WHERE %s = $1 AND %s = $2 AND %s = $3 AND %s = $4",
+		sqlReleaseTableKeyColumn,
+		sqlReleaseTableNamespaceColumn,
 		sqlReleaseTableBodyColumn,
 		sqlReleaseTableName,
 		sqlReleaseTableNameColumn,
@@ -345,8 +368,12 @@ func TestSqlQuery(t *testing.T) {
 			),
 		).RowsWillBeClosed()
 
+	mockGetReleaseCustomLabels(mock, "", deployedRelease.Namespace, deployedRelease.Labels)
+
 	query = fmt.Sprintf(
-		"SELECT %s FROM %s WHERE %s = $1 AND %s = $2 AND %s = $3",
+		"SELECT %s, %s, %s FROM %s WHERE %s = $1 AND %s = $2 AND %s = $3",
+		sqlReleaseTableKeyColumn,
+		sqlReleaseTableNamespaceColumn,
 		sqlReleaseTableBodyColumn,
 		sqlReleaseTableName,
 		sqlReleaseTableNameColumn,
@@ -367,6 +394,9 @@ func TestSqlQuery(t *testing.T) {
 			),
 		).RowsWillBeClosed()
 
+	mockGetReleaseCustomLabels(mock, "", supersededRelease.Namespace, supersededRelease.Labels)
+	mockGetReleaseCustomLabels(mock, "", deployedRelease.Namespace, deployedRelease.Labels)
+
 	_, err := sqlDriver.Query(labelSetUnknown)
 	if err == nil {
 		t.Errorf("Expected error {%v}, got nil", ErrReleaseNotFound)
@@ -447,6 +477,18 @@ func TestSqlDelete(t *testing.T) {
 		ExpectExec(regexp.QuoteMeta(deleteQuery)).
 		WithArgs(key, namespace).
 		WillReturnResult(sqlmock.NewResult(0, 1))
+
+	deleteLabelsQuery := fmt.Sprintf(
+		"DELETE FROM %s WHERE %s = $1 AND %s = $2",
+		sqlCustomLabelsTableName,
+		sqlCustomLabelsTableReleaseKeyColumn,
+		sqlCustomLabelsTableReleaseNamespaceColumn,
+	)
+	mock.
+		ExpectExec(regexp.QuoteMeta(deleteLabelsQuery)).
+		WithArgs(key, namespace).
+		WillReturnResult(sqlmock.NewResult(0, 1))
+
 	mock.ExpectCommit()
 
 	deletedRelease, err := sqlDriver.Delete(key)
@@ -461,3 +503,26 @@ func TestSqlDelete(t *testing.T) {
 		t.Errorf("Expected release {%v}, got {%v}", rel, deletedRelease)
 	}
 }
+
+func mockGetReleaseCustomLabels(mock sqlmock.Sqlmock, key string, namespace string, labels map[string]string) {
+	query := fmt.Sprintf(
+		regexp.QuoteMeta("SELECT %s, %s FROM %s WHERE %s = $1 AND %s = $2"),
+		sqlCustomLabelsTableKeyColumn,
+		sqlCustomLabelsTableValueColumn,
+		sqlCustomLabelsTableName,
+		sqlCustomLabelsTableReleaseKeyColumn,
+		sqlCustomLabelsTableReleaseNamespaceColumn,
+	)
+
+	eq := mock.ExpectQuery(query).
+		WithArgs(key, namespace)
+
+	returnRows := mock.NewRows([]string{
+		sqlCustomLabelsTableKeyColumn,
+		sqlCustomLabelsTableValueColumn,
+	})
+	for k, v := range labels {
+		returnRows.AddRow(k, v)
+	}
+	eq.WillReturnRows(returnRows).RowsWillBeClosed()
+}