storage/sql: batch custom label lookups for List and Query

Avoid N+1 SQL queries when listing/querying releases by loading custom labels in a single batch query and mapping them back per release. This reduces database round-trips and improves release query performance at scale while preserving existing behavior.

Signed-off-by: Sumit Solanki <sumit.solanki@ibm.com>

pkg/storage/driver/sql.go

@@ -366,6 +366,12 @@ func (s *SQL) List(filter func(release.Releaser) bool) ([]release.Releaser, erro
 		return nil, err
 	}
 
+	customLabelsByRelease, err := s.getReleaseCustomLabelsMap(records)
+	if err != nil {
+		s.Logger().Debug("failed to list release custom labels", slog.Any("error", err))
+		return nil, err
+	}
+
 	var releases []release.Releaser
 	for _, record := range records {
 		release, err := decodeRelease(record.Body)
@@ -374,15 +380,7 @@ func (s *SQL) List(filter func(release.Releaser) bool) ([]release.Releaser, erro
 			continue
 		}
 
-		if release.Labels, err = s.getReleaseCustomLabels(record.Key, record.Namespace); err != nil {
+		release.Labels = customLabelsByRelease[record.Namespace][record.Key]
-			s.Logger().Debug(
-				"failed to get release custom labels",
-				slog.String("namespace", record.Namespace),
-				slog.String("key", record.Key),
-				slog.Any("error", err),
-			)
-			return nil, err
-		}
 		maps.Copy(release.Labels, getReleaseSystemLabels(release))
 
 		if filter(release) {
@@ -435,6 +433,12 @@ func (s *SQL) Query(labels map[string]string) ([]release.Releaser, error) {
 		return nil, ErrReleaseNotFound
 	}
 
+	customLabelsByRelease, err := s.getReleaseCustomLabelsMap(records)
+	if err != nil {
+		s.Logger().Debug("failed to list release custom labels", slog.Any("error", err))
+		return nil, err
+	}
+
 	var releases []release.Releaser
 	for _, record := range records {
 		release, err := decodeRelease(record.Body)
@@ -443,15 +447,7 @@ func (s *SQL) Query(labels map[string]string) ([]release.Releaser, error) {
 			continue
 		}
 
-		if release.Labels, err = s.getReleaseCustomLabels(record.Key, record.Namespace); err != nil {
+		release.Labels = customLabelsByRelease[record.Namespace][record.Key]
-			s.Logger().Debug(
-				"failed to get release custom labels",
-				slog.String("namespace", record.Namespace),
-				slog.String("key", record.Key),
-				slog.Any("error", err),
-			)
-			return nil, err
-		}
 
 		releases = append(releases, release)
 	}
@@ -716,6 +712,77 @@ func (s *SQL) getReleaseCustomLabels(key string, _ string) (map[string]string, e
 	return filterSystemLabels(labelsMap), nil
 }
 
+func (s *SQL) getReleaseCustomLabelsMap(records []SQLReleaseWrapper) (map[string]map[string]map[string]string, error) {
+	labelsByNamespaceAndKey := make(map[string]map[string]map[string]string)
+	if len(records) == 0 {
+		return labelsByNamespaceAndKey, nil
+	}
+
+	keys := make([]string, 0, len(records))
+	namespaces := make([]string, 0, len(records))
+	seen := make(map[string]struct{}, len(records))
+	for _, record := range records {
+		id := record.Namespace + "\x00" + record.Key
+		if _, ok := seen[id]; ok {
+			continue
+		}
+		seen[id] = struct{}{}
+		keys = append(keys, record.Key)
+		namespaces = append(namespaces, record.Namespace)
+		if labelsByNamespaceAndKey[record.Namespace] == nil {
+			labelsByNamespaceAndKey[record.Namespace] = make(map[string]map[string]string)
+		}
+		labelsByNamespaceAndKey[record.Namespace][record.Key] = make(map[string]string)
+	}
+
+	query, args, err := s.statementBuilder.
+		Select(
+			sqlCustomLabelsTableReleaseNamespaceColumn,
+			sqlCustomLabelsTableReleaseKeyColumn,
+			sqlCustomLabelsTableKeyColumn,
+			sqlCustomLabelsTableValueColumn,
+		).
+		From(sqlCustomLabelsTableName).
+		Where(sq.Eq{sqlCustomLabelsTableReleaseNamespaceColumn: namespaces}).
+		Where(sq.Eq{sqlCustomLabelsTableReleaseKeyColumn: keys}).
+		ToSql()
+	if err != nil {
+		return nil, err
+	}
+
+	type customLabelRecord struct {
+		ReleaseNamespace string `db:"releaseNamespace"`
+		ReleaseKey       string `db:"releaseKey"`
+		Key              string `db:"key"`
+		Value            string `db:"value"`
+	}
+
+	var labelsList []customLabelRecord
+	if err := s.db.Select(&labelsList, query, args...); err != nil {
+		return nil, err
+	}
+
+	for _, item := range labelsList {
+		nsMap, ok := labelsByNamespaceAndKey[item.ReleaseNamespace]
+		if !ok {
+			continue
+		}
+		keyMap, ok := nsMap[item.ReleaseKey]
+		if !ok {
+			continue
+		}
+		keyMap[item.Key] = item.Value
+	}
+
+	for namespace, perKey := range labelsByNamespaceAndKey {
+		for key, labels := range perKey {
+			labelsByNamespaceAndKey[namespace][key] = filterSystemLabels(labels)
+		}
+	}
+
+	return labelsByNamespaceAndKey, nil
+}
+
 // Rebuild system labels from release object
 func getReleaseSystemLabels(rls *rspb.Release) map[string]string {
 	return map[string]string{

pkg/storage/driver/sql_test.go

@@ -133,20 +133,20 @@ func TestSQLList(t *testing.T) {
 		)
 
 		rows := mock.NewRows([]string{
+			sqlReleaseTableKeyColumn,
+			sqlReleaseTableNamespaceColumn,
 			sqlReleaseTableBodyColumn,
 		})
 		for _, r := range releases {
 			body, _ := encodeRelease(r)
-			rows.AddRow(body)
+			rows.AddRow("", r.Namespace, body)
 		}
 		mock.
 			ExpectQuery(regexp.QuoteMeta(query)).
 			WithArgs(sqlReleaseDefaultOwner, sqlDriver.namespace).
 			WillReturnRows(rows).RowsWillBeClosed()
 
-		for _, r := range releases {
+		mockGetReleaseCustomLabelsBatch(mock, releases)
-			mockGetReleaseCustomLabels(mock, "", r.Namespace, r.Labels)
-		}
 	}
 
 	// list all deleted releases
@@ -401,6 +401,8 @@ func TestSqlQuery(t *testing.T) {
 		WithArgs("smug-pigeon", sqlReleaseDefaultOwner, "unknown", "default").
 		WillReturnRows(
 			mock.NewRows([]string{
+				sqlReleaseTableKeyColumn,
+				sqlReleaseTableNamespaceColumn,
 				sqlReleaseTableBodyColumn,
 			}),
 		).RowsWillBeClosed()
@@ -410,13 +412,17 @@ func TestSqlQuery(t *testing.T) {
 		WithArgs("smug-pigeon", sqlReleaseDefaultOwner, "deployed", "default").
 		WillReturnRows(
 			mock.NewRows([]string{
+				sqlReleaseTableKeyColumn,
+				sqlReleaseTableNamespaceColumn,
 				sqlReleaseTableBodyColumn,
 			}).AddRow(
+				"",
+				deployedRelease.Namespace,
 				deployedReleaseBody,
 			),
 		).RowsWillBeClosed()
 
-	mockGetReleaseCustomLabels(mock, "", deployedRelease.Namespace, deployedRelease.Labels)
+	mockGetReleaseCustomLabelsBatch(mock, []*rspb.Release{deployedRelease})
 
 	query = fmt.Sprintf(
 		"SELECT %s, %s, %s FROM %s WHERE %s = $1 AND %s = $2 AND %s = $3",
@@ -434,16 +440,21 @@ func TestSqlQuery(t *testing.T) {
 		WithArgs("smug-pigeon", sqlReleaseDefaultOwner, "default").
 		WillReturnRows(
 			mock.NewRows([]string{
+				sqlReleaseTableKeyColumn,
+				sqlReleaseTableNamespaceColumn,
 				sqlReleaseTableBodyColumn,
 			}).AddRow(
+				"",
+				supersededRelease.Namespace,
 				supersededReleaseBody,
 			).AddRow(
+				"",
+				deployedRelease.Namespace,
 				deployedReleaseBody,
 			),
 		).RowsWillBeClosed()
 
-	mockGetReleaseCustomLabels(mock, "", supersededRelease.Namespace, supersededRelease.Labels)
+	mockGetReleaseCustomLabelsBatch(mock, []*rspb.Release{supersededRelease, deployedRelease})
-	mockGetReleaseCustomLabels(mock, "", deployedRelease.Namespace, deployedRelease.Labels)
 
 	_, err := sqlDriver.Query(labelSetUnknown)
 	if err == nil {
@@ -577,6 +588,34 @@ func mockGetReleaseCustomLabels(mock sqlmock.Sqlmock, key string, namespace stri
 	eq.WillReturnRows(returnRows).RowsWillBeClosed()
 }
 
+func mockGetReleaseCustomLabelsBatch(mock sqlmock.Sqlmock, releases []*rspb.Release) {
+	query := fmt.Sprintf(
+		regexp.QuoteMeta("SELECT %s, %s, %s, %s FROM %s WHERE %s IN (")+".*"+regexp.QuoteMeta(") AND %s IN (")+".*"+regexp.QuoteMeta(")"),
+		sqlCustomLabelsTableReleaseNamespaceColumn,
+		sqlCustomLabelsTableReleaseKeyColumn,
+		sqlCustomLabelsTableKeyColumn,
+		sqlCustomLabelsTableValueColumn,
+		sqlCustomLabelsTableName,
+		sqlCustomLabelsTableReleaseNamespaceColumn,
+		sqlCustomLabelsTableReleaseKeyColumn,
+	)
+
+	eq := mock.ExpectQuery(query)
+
+	returnRows := mock.NewRows([]string{
+		sqlCustomLabelsTableReleaseNamespaceColumn,
+		sqlCustomLabelsTableReleaseKeyColumn,
+		sqlCustomLabelsTableKeyColumn,
+		sqlCustomLabelsTableValueColumn,
+	})
+	for _, rel := range releases {
+		for k, v := range rel.Labels {
+			returnRows.AddRow(rel.Namespace, "", k, v)
+		}
+	}
+	eq.WillReturnRows(returnRows).RowsWillBeClosed()
+}
+
 func TestSqlCheckAppliedMigrations(t *testing.T) {
 	cases := []struct {
 		migrationsToApply    []*migrate.Migration