From d57b0e95c9d4e8f9ded07576b89f59d9a205d91d Mon Sep 17 00:00:00 2001 From: Sumit Solanki Date: Thu, 26 Mar 2026 16:44:24 +0530 Subject: [PATCH] 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 --- pkg/storage/driver/sql.go | 103 +++++++++++++++++++++++++++------ pkg/storage/driver/sql_test.go | 53 ++++++++++++++--- 2 files changed, 131 insertions(+), 25 deletions(-) diff --git a/pkg/storage/driver/sql.go b/pkg/storage/driver/sql.go index 21d9f6679..80768e78f 100644 --- a/pkg/storage/driver/sql.go +++ b/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 { - 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 - } + release.Labels = customLabelsByRelease[record.Namespace][record.Key] 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 { - 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 - } + release.Labels = customLabelsByRelease[record.Namespace][record.Key] 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{ diff --git a/pkg/storage/driver/sql_test.go b/pkg/storage/driver/sql_test.go index 4b4686b66..1c011c2d7 100644 --- a/pkg/storage/driver/sql_test.go +++ b/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 { - mockGetReleaseCustomLabels(mock, "", r.Namespace, r.Labels) - } + mockGetReleaseCustomLabelsBatch(mock, releases) } // 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) - mockGetReleaseCustomLabels(mock, "", deployedRelease.Namespace, deployedRelease.Labels) + mockGetReleaseCustomLabelsBatch(mock, []*rspb.Release{supersededRelease, deployedRelease}) _, 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