Fixed query and list tests for SQL storage driver

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

pkg/storage/driver/labels.go

@@ -49,7 +49,7 @@ func (lbs *labels) fromMap(kvs map[string]string) {
 
 func (lbs *labels) mergeUserLabels(list map[string]string) {
 	for k, v := range list {
-		if ! isSystemLabel(k) {
+		if !isSystemLabel(k) {
 			lbs.set(k, v)
 		}
 	}

pkg/storage/driver/mock_test.go

@@ -40,10 +40,14 @@ func releaseStub(name string, vers int, namespace string, status rspb.Status) *r
 		Version:   vers,
 		Namespace: namespace,
 		Info:      &rspb.Info{Status: status},
-		Labels: map[string]string{"key": "value"},
+		Labels:    defaultTestLabels(),
 	}
 }
 
+func defaultTestLabels() map[string]string {
+	return map[string]string{"key1": "value1", "key2": "value2"}
+}
+
 func testKey(name string, vers int) string {
 	return fmt.Sprintf("%s.v%d", name, vers)
 }

pkg/storage/driver/sql.go

@@ -34,7 +34,7 @@ import (
 
 var _ Driver = (*SQL)(nil)
 
-var LabelMap = map[string]struct{}{
+var labelMap = map[string]struct{}{
 	"modifiedAt": {},
 	"createdAt":  {},
 	"version":    {},
@@ -208,7 +208,7 @@ type SQLReleaseWrapper struct {
 	// The rspb.Release body, as a base64-encoded string
 	Body string `db:"body"`
 
-	// Release "Labels" that can be used as filters in the storage.Query(Labels map[string]string)
+	// Release "labels" that can be used as filters in the storage.Query(labels map[string]string)
 	// we implemented. Note that allowing Helm users to filter against new dimensions will require a
 	// new migration to be added, and the Create and/or update functions to be updated accordingly.
 	Name       string `db:"name"`
@@ -277,36 +277,18 @@ func (s *SQL) Get(key string) (*rspb.Release, error) {
 		return nil, err
 	}
 
-	LabelsQuery, args, err := s.statementBuilder.
+	if release.Labels, err = s.getReleaseLabels(key, s.namespace); err != nil {
-		Select(sqlLabelTableKeyColumn, sqlLabelTableValueColumn).
+		s.Log("failed to get release %s labels: %v", key, err)
-		From(sqlLabelTableName).
-		Where(sq.Eq{sqlLabelTableReleaseKeyColumn: key,
-			sqlLabelTableReleaseNamespaceColumn: s.namespace}).
-		ToSql()
-	if err != nil {
-		s.Log("failed to build query: %v", err)
 		return nil, err
 	}
 
-	var LabelsList = []SQLReleaseLabelWrapper{}
-	if err := s.db.Select(&LabelsList, LabelsQuery, args...); err != nil {
-		s.Log("get: failed to get release Labels: %v", err)
-		return nil, err
-	}
-
-	LabelsMap := make(map[string]string)
-	for _, i := range LabelsList {
-		LabelsMap[i.Key] = i.Value
-	}
-	release.Labels = filterSystemLabels(LabelsMap)
-
 	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, sqlReleaseTableBodyColumn).
 		From(sqlReleaseTableName).
 		Where(sq.Eq{sqlReleaseTableOwnerColumn: sqlReleaseDefaultOwner})
 
@@ -334,6 +316,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.getReleaseLabels(record.Key, release.Namespace); err != nil {
+			s.Log("failed to get release %s labels: %v", record.Key, err)
+			return nil, err
+		}
+
 		if filter(release) {
 			releases = append(releases, release)
 		}
@@ -342,23 +330,23 @@ func (s *SQL) List(filter func(*rspb.Release) bool) ([]*rspb.Release, error) {
 	return releases, nil
 }
 
-// Query returns the set of releases that match the provided set of Labels.
+// Query returns the set of releases that match the provided set of labels.
-func (s *SQL) Query(Labels map[string]string) ([]*rspb.Release, error) {
+func (s *SQL) Query(labels map[string]string) ([]*rspb.Release, error) {
 	sb := s.statementBuilder.
-		Select(sqlReleaseTableBodyColumn).
+		Select(sqlReleaseTableKeyColumn, sqlReleaseTableBodyColumn).
 		From(sqlReleaseTableName)
 
-	keys := make([]string, 0, len(Labels))
+	keys := make([]string, 0, len(labels))
-	for key := range Labels {
+	for key := range labels {
 		keys = append(keys, key)
 	}
 	sort.Strings(keys)
 	for _, key := range keys {
-		if _, ok := LabelMap[key]; ok {
+		if _, ok := labelMap[key]; ok {
-			sb = sb.Where(sq.Eq{key: Labels[key]})
+			sb = sb.Where(sq.Eq{key: labels[key]})
 		} else {
 			s.Log("unknown Label %s", key)
-			return nil, fmt.Errorf("unknow Label %s", key)
+			return nil, fmt.Errorf("unknown label %s", key)
 		}
 	}
 
@@ -376,7 +364,7 @@ func (s *SQL) Query(Labels map[string]string) ([]*rspb.Release, error) {
 
 	var records = []SQLReleaseWrapper{}
 	if err := s.db.Select(&records, query, args...); err != nil {
-		s.Log("list: failed to query with Labels: %v", err)
+		s.Log("list: failed to query with labels: %v", err)
 		return nil, err
 	}
 
@@ -387,6 +375,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.getReleaseLabels(record.Key, release.Namespace); err != nil {
+			s.Log("failed to get release %s labels: %v", record.Key, err)
+			return nil, err
+		}
+
 		releases = append(releases, release)
 	}
 
@@ -576,28 +570,12 @@ func (s *SQL) Delete(key string) (*rspb.Release, error) {
 		return nil, err
 	}
 
-	LabelsQuery, args, err := s.statementBuilder.
+	if release.Labels, err = s.getReleaseLabels(key, s.namespace); err != nil {
-		Select(sqlLabelTableKeyColumn, sqlLabelTableValueColumn).
+		s.Log("failed to get release %s labels: %v", key, err)
-		From(sqlLabelTableName).
+		transaction.Rollback()
-		Where(sq.Eq{sqlLabelTableReleaseKeyColumn: key,
-			sqlLabelTableReleaseNamespaceColumn: s.namespace}).
-		ToSql()
-	if err != nil {
-		s.Log("failed to build query: %v", err)
-		return nil, err
-	}
-
-	var LabelsList = []SQLReleaseLabelWrapper{}
-	if err := s.db.Select(&LabelsList, LabelsQuery, args...); err != nil {
-		s.Log("get: failed to get release Labels: %v", err)
 		return nil, err
 	}
 
-	LabelsMap := make(map[string]string)
-	for _, i := range LabelsList {
-		LabelsMap[i.Key] = i.Value
-	}
-	release.Labels = filterSystemLabels(LabelsMap)
 	defer transaction.Commit()
 
 	deleteQuery, args, err := s.statementBuilder.
@@ -628,3 +606,28 @@ func (s *SQL) Delete(key string) (*rspb.Release, error) {
 	_, err = transaction.Exec(deleteLabelsQuery, args...)
 	return release, err
 }
+
+// Fetches release labels from database
+func (s *SQL) getReleaseLabels(key string, namespace string) (map[string]string, error) {
+	query, args, err := s.statementBuilder.
+		Select(sqlLabelTableKeyColumn, sqlLabelTableValueColumn).
+		From(sqlLabelTableName).
+		Where(sq.Eq{sqlLabelTableReleaseKeyColumn: key,
+			sqlLabelTableReleaseNamespaceColumn: s.namespace}).
+		ToSql()
+	if err != nil {
+		return nil, err
+	}
+
+	var labelsList = []SQLReleaseLabelWrapper{}
+	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,26 +62,7 @@ func TestSQLGet(t *testing.T) {
 			),
 		).RowsWillBeClosed()
 
-	queryLabels := fmt.Sprintf(
+	mockLabelsFetch(mock, key, namespace, defaultTestLabels())
-		regexp.QuoteMeta("SELECT %s, %s FROM %s WHERE %s = $1 AND %s = $2"),
-		sqlLabelTableKeyColumn,
-		sqlLabelTableValueColumn,
-		sqlLabelTableName,
-		sqlLabelTableReleaseKeyColumn,
-		sqlLabelTableReleaseNamespaceColumn,
-	)
-
-	eq := mock.ExpectQuery(queryLabels).
-		WithArgs(key, namespace)
-
-	returnRows := mock.NewRows([]string{
-		sqlLabelTableKeyColumn,
-		sqlLabelTableValueColumn,
-	})
-	for k, v := range rel.Labels {
-		returnRows.AddRow(k, v)
-	}
-	eq.WillReturnRows(returnRows).RowsWillBeClosed()
 
 	got, err := sqlDriver.Get(key)
 	if err != nil {
@@ -109,7 +90,8 @@ func TestSQLList(t *testing.T) {
 
 	for i := 0; i < 3; i++ {
 		query := fmt.Sprintf(
-			"SELECT %s FROM %s WHERE %s = $1 AND %s = $2",
+			"SELECT %s, %s FROM %s WHERE %s = $1 AND %s = $2",
+			sqlReleaseTableKeyColumn,
 			sqlReleaseTableBodyColumn,
 			sqlReleaseTableName,
 			sqlReleaseTableOwnerColumn,
@@ -130,6 +112,10 @@ func TestSQLList(t *testing.T) {
 					AddRow(body5).
 					AddRow(body6),
 			).RowsWillBeClosed()
+
+		for j := 1; j <= 6; j++ {
+			mockLabelsFetch(mock, "", sqlDriver.namespace, defaultTestLabels())
+		}
 	}
 
 	// list all deleted releases
@@ -347,7 +333,8 @@ 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 FROM %s WHERE %s = $1 AND %s = $2 AND %s = $3 AND %s = $4",
+		sqlReleaseTableKeyColumn,
 		sqlReleaseTableBodyColumn,
 		sqlReleaseTableName,
 		sqlReleaseTableNameColumn,
@@ -367,8 +354,11 @@ func TestSqlQuery(t *testing.T) {
 			),
 		).RowsWillBeClosed()
 
+	mockLabelsFetch(mock, "", "default", defaultTestLabels())
+
 	query = fmt.Sprintf(
-		"SELECT %s FROM %s WHERE %s = $1 AND %s = $2 AND %s = $3",
+		"SELECT %s, %s FROM %s WHERE %s = $1 AND %s = $2 AND %s = $3",
+		sqlReleaseTableKeyColumn,
 		sqlReleaseTableBodyColumn,
 		sqlReleaseTableName,
 		sqlReleaseTableNameColumn,
@@ -389,6 +379,9 @@ func TestSqlQuery(t *testing.T) {
 			),
 		).RowsWillBeClosed()
 
+	mockLabelsFetch(mock, "", "default", defaultTestLabels())
+	mockLabelsFetch(mock, "", "default", defaultTestLabels())
+
 	results, err := sqlDriver.Query(labelSetDeployed)
 	if err != nil {
 		t.Fatalf("failed to query for deployed smug-pigeon release: %v", err)
@@ -451,26 +444,7 @@ func TestSqlDelete(t *testing.T) {
 			),
 		).RowsWillBeClosed()
 
-	queryLabels := fmt.Sprintf(
+	mockLabelsFetch(mock, key, namespace, defaultTestLabels())
-		regexp.QuoteMeta("SELECT %s, %s FROM %s WHERE %s = $1 AND %s = $2"),
-		sqlLabelTableKeyColumn,
-		sqlLabelTableValueColumn,
-		sqlLabelTableName,
-		sqlLabelTableReleaseKeyColumn,
-		sqlLabelTableReleaseNamespaceColumn,
-	)
-
-	eq := mock.ExpectQuery(queryLabels).
-		WithArgs(key, namespace)
-
-	returnRows := mock.NewRows([]string{
-		sqlLabelTableKeyColumn,
-		sqlLabelTableValueColumn,
-	})
-	for k, v := range rel.Labels {
-		returnRows.AddRow(k, v)
-	}
-	eq.WillReturnRows(returnRows).RowsWillBeClosed()
 
 	deleteQuery := fmt.Sprintf(
 		"DELETE FROM %s WHERE %s = $1 AND %s = $2",
@@ -478,6 +452,7 @@ func TestSqlDelete(t *testing.T) {
 		sqlReleaseTableKeyColumn,
 		sqlReleaseTableNamespaceColumn,
 	)
+
 	mock.
 		ExpectExec(regexp.QuoteMeta(deleteQuery)).
 		WithArgs(key, namespace).
@@ -508,3 +483,27 @@ func TestSqlDelete(t *testing.T) {
 		t.Errorf("Expected release {%v}, got {%v}", rel, deletedRelease)
 	}
 }
+
+// Mocks labels fetch queries
+func mockLabelsFetch(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"),
+		sqlLabelTableKeyColumn,
+		sqlLabelTableValueColumn,
+		sqlLabelTableName,
+		sqlLabelTableReleaseKeyColumn,
+		sqlLabelTableReleaseNamespaceColumn,
+	)
+
+	eq := mock.ExpectQuery(query).
+		WithArgs(key, namespace)
+
+	returnRows := mock.NewRows([]string{
+		sqlLabelTableKeyColumn,
+		sqlLabelTableValueColumn,
+	})
+	for k, v := range labels {
+		returnRows.AddRow(k, v)
+	}
+	eq.WillReturnRows(returnRows).RowsWillBeClosed()
+}

pkg/storage/driver/util.go

@@ -85,7 +85,7 @@ func decodeRelease(data string) (*rspb.Release, error) {
 }
 
 // Returns array of system labels' keys
-func systemLablesKeys() []string {
+func systemLabelsKeys() []string {
 	return []string{"name", "owner", "status", "version", "createdAt", "modifiedAt"}
 }
 
@@ -93,7 +93,7 @@ func systemLablesKeys() []string {
 func filterSystemLabels(lbs map[string]string) map[string]string {
 	result := make(map[string]string)
 	for k, v := range lbs {
-		if !isSystemLabel(k){
+		if !isSystemLabel(k) {
 			result[k] = v
 		}
 	}
@@ -101,7 +101,7 @@ func filterSystemLabels(lbs map[string]string) map[string]string {
 }
 
 func isSystemLabel(key string) bool {
-	for _, v := range systemLablesKeys() {
+	for _, v := range systemLabelsKeys() {
 		if key == v {
 			return true
 		}