From bf98dd7742cfb65400e720e23a09c3cf45a69948 Mon Sep 17 00:00:00 2001 From: Sharon Date: Fri, 10 Jul 2026 14:21:32 -0400 Subject: [PATCH 1/3] Batch extension label-membership checks in GetOrbitConfig (#45320) Replace N+1 per-extension HostMemberOfAllLabels calls in filterExtensionsForHost with a single HostMembershipForLabels query that fetches all label memberships at once, then filters in memory. --- changes/batch-extension-label-check | 1 + server/datastore/mysql/labels.go | 29 ++++++ server/datastore/mysql/labels_test.go | 121 ++++++++++++++++++++++++++ server/fleet/datastore.go | 4 + server/mock/datastore_mock.go | 12 +++ server/service/orbit.go | 40 +++++++-- 6 files changed, 198 insertions(+), 9 deletions(-) create mode 100644 changes/batch-extension-label-check diff --git a/changes/batch-extension-label-check b/changes/batch-extension-label-check new file mode 100644 index 00000000000..4bc7af2bc90 --- /dev/null +++ b/changes/batch-extension-label-check @@ -0,0 +1 @@ +- Improved performance of Orbit config endpoint by batching extension label-membership checks into a single database query. diff --git a/server/datastore/mysql/labels.go b/server/datastore/mysql/labels.go index 4bfa8aadb57..9ab1a09ec54 100644 --- a/server/datastore/mysql/labels.go +++ b/server/datastore/mysql/labels.go @@ -1743,6 +1743,35 @@ func (ds *Datastore) HostMemberOfAllLabels(ctx context.Context, hostID uint, lab return ok, nil } +// HostMembershipForLabels returns the set of label names (from the provided list) that the host is a member of. +// Labels that do not exist are not included in the result. +func (ds *Datastore) HostMembershipForLabels(ctx context.Context, hostID uint, labelNames []string) (map[string]bool, error) { + if len(labelNames) == 0 { + return nil, nil + } + + sqlStatement := ` + SELECT l.name FROM labels l + JOIN label_membership lm ON l.id = lm.label_id + WHERE lm.host_id = ? AND l.name IN (?) + ` + sql, args, err := sqlx.In(sqlStatement, hostID, labelNames) + if err != nil { + return nil, ctxerr.Wrap(ctx, err, "building query for host label membership") + } + + var names []string + if err := sqlx.SelectContext(ctx, ds.reader(ctx), &names, sql, args...); err != nil { + return nil, ctxerr.Wrap(ctx, err, "get host label membership") + } + + result := make(map[string]bool, len(names)) + for _, n := range names { + result[n] = true + } + return result, nil +} + // AddLabelsToHost skips auth as it's only used in tests, and where label teams have already been validated. func (ds *Datastore) AddLabelsToHost(ctx context.Context, hostID uint, labelIDs []uint) error { if len(labelIDs) == 0 { diff --git a/server/datastore/mysql/labels_test.go b/server/datastore/mysql/labels_test.go index 623e8c0e282..352d0be048c 100644 --- a/server/datastore/mysql/labels_test.go +++ b/server/datastore/mysql/labels_test.go @@ -96,6 +96,7 @@ func TestLabels(t *testing.T) { {"ListHostsInLabelIssues", testListHostsInLabelIssues}, {"ListHostsInLabelDiskEncryptionStatus", testListHostsInLabelDiskEncryptionStatus}, {"HostMemberOfAllLabels", testHostMemberOfAllLabels}, + {"HostMembershipForLabels", testHostMembershipForLabels}, {"ListHostsInLabelOSSettings", testLabelsListHostsInLabelOSSettings}, {"AddDeleteLabelsToFromHost", testAddDeleteLabelsToFromHost}, {"ApplyLabelSpecSerialUUID", testApplyLabelSpecsForSerialUUID}, @@ -2106,6 +2107,126 @@ func testHostMemberOfAllLabels(t *testing.T, ds *Datastore) { } } +func testHostMembershipForLabels(t *testing.T, ds *Datastore) { + ctx := context.Background() + + // + // Setup test + // - h1 member of 'All hosts', 'Foobar' and 'Zoobar' + // - h2 member of 'All hosts' and 'Foobar' + // - h3 member of no labels + // + + allHostsLabel, err := ds.NewLabel(ctx, + &fleet.Label{ + Name: "All hosts", + Query: "SELECT 1", + LabelType: fleet.LabelTypeBuiltIn, + LabelMembershipType: fleet.LabelMembershipTypeDynamic, + }, + ) + require.NoError(t, err) + foobarLabel, err := ds.NewLabel(ctx, &fleet.Label{ + Name: "Foobar", + Query: "SELECT 1;", + LabelType: fleet.LabelTypeRegular, + LabelMembershipType: fleet.LabelMembershipTypeDynamic, + }) + require.NoError(t, err) + zoobarLabel, err := ds.NewLabel(ctx, &fleet.Label{ + Name: "Zoobar", + Query: "SELECT 2;", + LabelType: fleet.LabelTypeRegular, + LabelMembershipType: fleet.LabelMembershipTypeDynamic, + }) + require.NoError(t, err) + + newHostFunc := func(name string) *fleet.Host { + h, err := ds.NewHost(ctx, &fleet.Host{ + DetailUpdatedAt: time.Now(), + LabelUpdatedAt: time.Now(), + PolicyUpdatedAt: time.Now(), + SeenTime: time.Now(), + OsqueryHostID: ptr.String(name), + NodeKey: ptr.String(name), + UUID: name, + Hostname: "foo.local" + name, + }) + require.NoError(t, err) + return h + } + + h1 := newHostFunc("h1") + h2 := newHostFunc("h2") + h3 := newHostFunc("h3") + + err = ds.RecordLabelQueryExecutions(ctx, h1, map[uint]*bool{ + allHostsLabel.ID: ptr.Bool(true), + foobarLabel.ID: ptr.Bool(true), + zoobarLabel.ID: ptr.Bool(true), + }, time.Now(), false) + require.NoError(t, err) + err = ds.RecordLabelQueryExecutions(ctx, h2, map[uint]*bool{ + allHostsLabel.ID: ptr.Bool(true), + foobarLabel.ID: ptr.Bool(true), + }, time.Now(), false) + require.NoError(t, err) + + for _, tc := range []struct { + name string + hostID uint + labelNames []string + expectedResult map[string]bool + }{ + { + name: "empty label names returns nil", + hostID: h1.ID, + labelNames: nil, + expectedResult: nil, + }, + { + name: "h1 is member of all three labels", + hostID: h1.ID, + labelNames: []string{allHostsLabel.Name, foobarLabel.Name, zoobarLabel.Name}, + expectedResult: map[string]bool{allHostsLabel.Name: true, foobarLabel.Name: true, zoobarLabel.Name: true}, + }, + { + name: "h2 is member of some but not all labels", + hostID: h2.ID, + labelNames: []string{allHostsLabel.Name, foobarLabel.Name, zoobarLabel.Name}, + expectedResult: map[string]bool{allHostsLabel.Name: true, foobarLabel.Name: true}, + }, + { + name: "nonexistent labels are not included", + hostID: h1.ID, + labelNames: []string{allHostsLabel.Name, "nonexistent-label"}, + expectedResult: map[string]bool{allHostsLabel.Name: true}, + }, + { + name: "nonexistent host returns empty result", + hostID: 999, + labelNames: []string{allHostsLabel.Name}, + expectedResult: map[string]bool{}, + }, + { + name: "h3 is member of no labels", + hostID: h3.ID, + labelNames: []string{allHostsLabel.Name, foobarLabel.Name}, + expectedResult: map[string]bool{}, + }, + } { + t.Run(tc.name, func(t *testing.T) { + result, err := ds.HostMembershipForLabels(ctx, tc.hostID, tc.labelNames) + require.NoError(t, err) + if tc.expectedResult == nil { + require.Nil(t, result) + } else { + require.Equal(t, tc.expectedResult, result) + } + }) + } +} + func testLabelsListHostsInLabelOSSettings(t *testing.T, db *Datastore) { h1, err := db.NewHost(context.Background(), &fleet.Host{ DetailUpdatedAt: time.Now(), diff --git a/server/fleet/datastore.go b/server/fleet/datastore.go index 97e7bb4c122..c4cfd8d3dc2 100644 --- a/server/fleet/datastore.go +++ b/server/fleet/datastore.go @@ -337,6 +337,10 @@ type Datastore interface { // if labelNames is empty. HostMemberOfAllLabels(ctx context.Context, hostID uint, labelNames []string) (bool, error) + // HostMembershipForLabels returns the set of label names (from the provided list) that the host is a member of. + // Labels that do not exist are not included in the result. + HostMembershipForLabels(ctx context.Context, hostID uint, labelNames []string) (map[string]bool, error) + // TODO JUAN: Refactor this to use the Operating System type instead. // HostIDsByOSVersion retrieves the IDs of all host matching osVersion HostIDsByOSVersion(ctx context.Context, osVersion OSVersion, offset int, limit int) ([]uint, error) diff --git a/server/mock/datastore_mock.go b/server/mock/datastore_mock.go index 6ecb4d300d9..284150ab7ef 100644 --- a/server/mock/datastore_mock.go +++ b/server/mock/datastore_mock.go @@ -256,6 +256,8 @@ type HostIDsByOSIDFunc func(ctx context.Context, osID uint, offset int, limit in type HostMemberOfAllLabelsFunc func(ctx context.Context, hostID uint, labelNames []string) (bool, error) +type HostMembershipForLabelsFunc func(ctx context.Context, hostID uint, labelNames []string) (map[string]bool, error) + type HostIDsByOSVersionFunc func(ctx context.Context, osVersion fleet.OSVersion, offset int, limit int) ([]uint, error) type HostByIdentifierFunc func(ctx context.Context, identifier string) (*fleet.Host, error) @@ -2507,6 +2509,9 @@ type DataStore struct { HostMemberOfAllLabelsFunc HostMemberOfAllLabelsFunc HostMemberOfAllLabelsFuncInvoked bool + HostMembershipForLabelsFunc HostMembershipForLabelsFunc + HostMembershipForLabelsFuncInvoked bool + HostIDsByOSVersionFunc HostIDsByOSVersionFunc HostIDsByOSVersionFuncInvoked bool @@ -6175,6 +6180,13 @@ func (s *DataStore) HostMemberOfAllLabels(ctx context.Context, hostID uint, labe return s.HostMemberOfAllLabelsFunc(ctx, hostID, labelNames) } +func (s *DataStore) HostMembershipForLabels(ctx context.Context, hostID uint, labelNames []string) (map[string]bool, error) { + s.mu.Lock() + s.HostMembershipForLabelsFuncInvoked = true + s.mu.Unlock() + return s.HostMembershipForLabelsFunc(ctx, hostID, labelNames) +} + func (s *DataStore) HostIDsByOSVersion(ctx context.Context, osVersion fleet.OSVersion, offset int, limit int) ([]uint, error) { s.mu.Lock() s.HostIDsByOSVersionFuncInvoked = true diff --git a/server/service/orbit.go b/server/service/orbit.go index 791e33f52e7..9caf41599be 100644 --- a/server/service/orbit.go +++ b/server/service/orbit.go @@ -971,17 +971,39 @@ func (svc *Service) filterExtensionsForHost(ctx context.Context, extensions json // Filter the extensions by labels (premium only feature). if license, _ := license.FromContext(ctx); license != nil && license.IsPremium() { - for extensionName, extensionInfo := range extensionsInfo { - hostIsMemberOfAllLabels, err := svc.ds.HostMemberOfAllLabels(ctx, host.ID, extensionInfo.Labels) + // Collect all unique label names across all extensions. + allLabels := make(map[string]struct{}) + for _, extInfo := range extensionsInfo { + for _, l := range extInfo.Labels { + allLabels[l] = struct{}{} + } + } + + if len(allLabels) > 0 { + labelNames := make([]string, 0, len(allLabels)) + for l := range allLabels { + labelNames = append(labelNames, l) + } + + memberOf, err := svc.ds.HostMembershipForLabels(ctx, host.ID, labelNames) if err != nil { - return nil, ctxerr.Wrap(ctx, err, "check host labels") + return nil, ctxerr.Wrap(ctx, err, "check host label membership") } - if hostIsMemberOfAllLabels { - // Do not filter out, but there's no need to send the label names to the devices. - extensionInfo.Labels = nil - extensionsInfo[extensionName] = extensionInfo - } else { - delete(extensionsInfo, extensionName) + + for extensionName, extensionInfo := range extensionsInfo { + allMatch := true + for _, l := range extensionInfo.Labels { + if !memberOf[l] { + allMatch = false + break + } + } + if allMatch { + extensionInfo.Labels = nil + extensionsInfo[extensionName] = extensionInfo + } else { + delete(extensionsInfo, extensionName) + } } } } From 46f24bc88561ef2a2f71c670339efac75d654a67 Mon Sep 17 00:00:00 2001 From: Sharon Date: Fri, 10 Jul 2026 15:52:10 -0400 Subject: [PATCH 2/3] Fix lint: use map[string]struct{} for set, replace deprecated ptr.String - setboolcheck: HostMembershipForLabels return type changed from map[string]bool to map[string]struct{} (Fleet convention for sets) - staticcheck: replaced ptr.String/ptr.Bool with new() in test --- server/datastore/mysql/labels.go | 6 +++--- server/datastore/mysql/labels_test.go | 26 +++++++++++++------------- server/fleet/datastore.go | 2 +- server/mock/datastore_mock.go | 4 ++-- server/service/orbit.go | 2 +- 5 files changed, 20 insertions(+), 20 deletions(-) diff --git a/server/datastore/mysql/labels.go b/server/datastore/mysql/labels.go index 9ab1a09ec54..b423c5aa906 100644 --- a/server/datastore/mysql/labels.go +++ b/server/datastore/mysql/labels.go @@ -1745,7 +1745,7 @@ func (ds *Datastore) HostMemberOfAllLabels(ctx context.Context, hostID uint, lab // HostMembershipForLabels returns the set of label names (from the provided list) that the host is a member of. // Labels that do not exist are not included in the result. -func (ds *Datastore) HostMembershipForLabels(ctx context.Context, hostID uint, labelNames []string) (map[string]bool, error) { +func (ds *Datastore) HostMembershipForLabels(ctx context.Context, hostID uint, labelNames []string) (map[string]struct{}, error) { if len(labelNames) == 0 { return nil, nil } @@ -1765,9 +1765,9 @@ func (ds *Datastore) HostMembershipForLabels(ctx context.Context, hostID uint, l return nil, ctxerr.Wrap(ctx, err, "get host label membership") } - result := make(map[string]bool, len(names)) + result := make(map[string]struct{}, len(names)) for _, n := range names { - result[n] = true + result[n] = struct{}{} } return result, nil } diff --git a/server/datastore/mysql/labels_test.go b/server/datastore/mysql/labels_test.go index 352d0be048c..5500554f736 100644 --- a/server/datastore/mysql/labels_test.go +++ b/server/datastore/mysql/labels_test.go @@ -2147,8 +2147,8 @@ func testHostMembershipForLabels(t *testing.T, ds *Datastore) { LabelUpdatedAt: time.Now(), PolicyUpdatedAt: time.Now(), SeenTime: time.Now(), - OsqueryHostID: ptr.String(name), - NodeKey: ptr.String(name), + OsqueryHostID: new(name), + NodeKey: new(name), UUID: name, Hostname: "foo.local" + name, }) @@ -2161,14 +2161,14 @@ func testHostMembershipForLabels(t *testing.T, ds *Datastore) { h3 := newHostFunc("h3") err = ds.RecordLabelQueryExecutions(ctx, h1, map[uint]*bool{ - allHostsLabel.ID: ptr.Bool(true), - foobarLabel.ID: ptr.Bool(true), - zoobarLabel.ID: ptr.Bool(true), + allHostsLabel.ID: new(true), + foobarLabel.ID: new(true), + zoobarLabel.ID: new(true), }, time.Now(), false) require.NoError(t, err) err = ds.RecordLabelQueryExecutions(ctx, h2, map[uint]*bool{ - allHostsLabel.ID: ptr.Bool(true), - foobarLabel.ID: ptr.Bool(true), + allHostsLabel.ID: new(true), + foobarLabel.ID: new(true), }, time.Now(), false) require.NoError(t, err) @@ -2176,7 +2176,7 @@ func testHostMembershipForLabels(t *testing.T, ds *Datastore) { name string hostID uint labelNames []string - expectedResult map[string]bool + expectedResult map[string]struct{} }{ { name: "empty label names returns nil", @@ -2188,31 +2188,31 @@ func testHostMembershipForLabels(t *testing.T, ds *Datastore) { name: "h1 is member of all three labels", hostID: h1.ID, labelNames: []string{allHostsLabel.Name, foobarLabel.Name, zoobarLabel.Name}, - expectedResult: map[string]bool{allHostsLabel.Name: true, foobarLabel.Name: true, zoobarLabel.Name: true}, + expectedResult: map[string]struct{}{allHostsLabel.Name: {}, foobarLabel.Name: {}, zoobarLabel.Name: {}}, }, { name: "h2 is member of some but not all labels", hostID: h2.ID, labelNames: []string{allHostsLabel.Name, foobarLabel.Name, zoobarLabel.Name}, - expectedResult: map[string]bool{allHostsLabel.Name: true, foobarLabel.Name: true}, + expectedResult: map[string]struct{}{allHostsLabel.Name: {}, foobarLabel.Name: {}}, }, { name: "nonexistent labels are not included", hostID: h1.ID, labelNames: []string{allHostsLabel.Name, "nonexistent-label"}, - expectedResult: map[string]bool{allHostsLabel.Name: true}, + expectedResult: map[string]struct{}{allHostsLabel.Name: {}}, }, { name: "nonexistent host returns empty result", hostID: 999, labelNames: []string{allHostsLabel.Name}, - expectedResult: map[string]bool{}, + expectedResult: map[string]struct{}{}, }, { name: "h3 is member of no labels", hostID: h3.ID, labelNames: []string{allHostsLabel.Name, foobarLabel.Name}, - expectedResult: map[string]bool{}, + expectedResult: map[string]struct{}{}, }, } { t.Run(tc.name, func(t *testing.T) { diff --git a/server/fleet/datastore.go b/server/fleet/datastore.go index c4cfd8d3dc2..376b737f1c2 100644 --- a/server/fleet/datastore.go +++ b/server/fleet/datastore.go @@ -339,7 +339,7 @@ type Datastore interface { // HostMembershipForLabels returns the set of label names (from the provided list) that the host is a member of. // Labels that do not exist are not included in the result. - HostMembershipForLabels(ctx context.Context, hostID uint, labelNames []string) (map[string]bool, error) + HostMembershipForLabels(ctx context.Context, hostID uint, labelNames []string) (map[string]struct{}, error) // TODO JUAN: Refactor this to use the Operating System type instead. // HostIDsByOSVersion retrieves the IDs of all host matching osVersion diff --git a/server/mock/datastore_mock.go b/server/mock/datastore_mock.go index 284150ab7ef..7d841ac6238 100644 --- a/server/mock/datastore_mock.go +++ b/server/mock/datastore_mock.go @@ -256,7 +256,7 @@ type HostIDsByOSIDFunc func(ctx context.Context, osID uint, offset int, limit in type HostMemberOfAllLabelsFunc func(ctx context.Context, hostID uint, labelNames []string) (bool, error) -type HostMembershipForLabelsFunc func(ctx context.Context, hostID uint, labelNames []string) (map[string]bool, error) +type HostMembershipForLabelsFunc func(ctx context.Context, hostID uint, labelNames []string) (map[string]struct{}, error) type HostIDsByOSVersionFunc func(ctx context.Context, osVersion fleet.OSVersion, offset int, limit int) ([]uint, error) @@ -6180,7 +6180,7 @@ func (s *DataStore) HostMemberOfAllLabels(ctx context.Context, hostID uint, labe return s.HostMemberOfAllLabelsFunc(ctx, hostID, labelNames) } -func (s *DataStore) HostMembershipForLabels(ctx context.Context, hostID uint, labelNames []string) (map[string]bool, error) { +func (s *DataStore) HostMembershipForLabels(ctx context.Context, hostID uint, labelNames []string) (map[string]struct{}, error) { s.mu.Lock() s.HostMembershipForLabelsFuncInvoked = true s.mu.Unlock() diff --git a/server/service/orbit.go b/server/service/orbit.go index 9caf41599be..ce3faf16979 100644 --- a/server/service/orbit.go +++ b/server/service/orbit.go @@ -993,7 +993,7 @@ func (svc *Service) filterExtensionsForHost(ctx context.Context, extensions json for extensionName, extensionInfo := range extensionsInfo { allMatch := true for _, l := range extensionInfo.Labels { - if !memberOf[l] { + if _, ok := memberOf[l]; !ok { allMatch = false break } From 99b2fe79a4396f6deb98e3ef37e7a811e40116cc Mon Sep 17 00:00:00 2001 From: Sharon Date: Wed, 15 Jul 2026 09:39:51 -0400 Subject: [PATCH 3/3] Address review feedback: case-sensitivity fix, remove old method, add tests - Fix case-sensitivity bug: HostMembershipForLabels now keys the result map by the caller's requested label names, not the DB-stored names. MySQL's utf8mb4_unicode_ci collation can match differently-cased names, but Go map lookups are case-sensitive. - Remove HostMemberOfAllLabels (no remaining callers) - Remove redundant empty-input guard (caller already checks) - Use t.Context() instead of context.Background() - Add test case for subset of labels - Add test case for case-insensitive match preserving requested casing --- server/datastore/mysql/labels.go | 53 ++----- server/datastore/mysql/labels_test.go | 203 ++------------------------ server/fleet/datastore.go | 9 +- server/mock/datastore_mock.go | 12 -- 4 files changed, 30 insertions(+), 247 deletions(-) diff --git a/server/datastore/mysql/labels.go b/server/datastore/mysql/labels.go index b423c5aa906..8e47fb9687b 100644 --- a/server/datastore/mysql/labels.go +++ b/server/datastore/mysql/labels.go @@ -1715,41 +1715,10 @@ func (ds *Datastore) LabelsSummary(ctx context.Context, filter fleet.TeamFilter) return labelsSummary, nil } -// HostMemberOfAllLabels returns whether the given host is a member of all the provided labels. -// If the labels do not exist, then the host is considered not a member of the provided labels. -// A host will always be a member of an empty label set, so this method returns (true, nil) -// if labelNames is empty. -func (ds *Datastore) HostMemberOfAllLabels(ctx context.Context, hostID uint, labelNames []string) (bool, error) { - if len(labelNames) == 0 { - return true, nil - } - - sqlStatement := ` - SELECT COUNT(*) = ? FROM labels l - LEFT JOIN (SELECT label_id FROM label_membership WHERE host_id = ?) lm - ON l.id = lm.label_id - WHERE l.name IN (?) AND lm.label_id IS NOT NULL; - ` - sql, args, err := sqlx.In(sqlStatement, len(labelNames), hostID, labelNames) - if err != nil { - return false, ctxerr.Wrap(ctx, err, "building query to get label IDs") - } - - var ok bool - if err := sqlx.GetContext(ctx, ds.reader(ctx), &ok, sql, args...); err != nil { - return false, ctxerr.Wrap(ctx, err, "get label IDs") - } - - return ok, nil -} - // HostMembershipForLabels returns the set of label names (from the provided list) that the host is a member of. -// Labels that do not exist are not included in the result. +// Labels that do not exist are not included in the result. The returned map is keyed by the +// requested label names (preserving the caller's casing), not the DB-stored names. func (ds *Datastore) HostMembershipForLabels(ctx context.Context, hostID uint, labelNames []string) (map[string]struct{}, error) { - if len(labelNames) == 0 { - return nil, nil - } - sqlStatement := ` SELECT l.name FROM labels l JOIN label_membership lm ON l.id = lm.label_id @@ -1760,14 +1729,22 @@ func (ds *Datastore) HostMembershipForLabels(ctx context.Context, hostID uint, l return nil, ctxerr.Wrap(ctx, err, "building query for host label membership") } - var names []string - if err := sqlx.SelectContext(ctx, ds.reader(ctx), &names, sql, args...); err != nil { + var dbNames []string + if err := sqlx.SelectContext(ctx, ds.reader(ctx), &dbNames, sql, args...); err != nil { return nil, ctxerr.Wrap(ctx, err, "get host label membership") } - result := make(map[string]struct{}, len(names)) - for _, n := range names { - result[n] = struct{}{} + // The DB may return names with different casing than requested (utf8mb4_unicode_ci), + // so map DB names back to the caller's requested names for case-sensitive Go map lookups. + dbNameSet := make(map[string]struct{}, len(dbNames)) + for _, n := range dbNames { + dbNameSet[strings.ToLower(n)] = struct{}{} + } + result := make(map[string]struct{}, len(dbNames)) + for _, req := range labelNames { + if _, ok := dbNameSet[strings.ToLower(req)]; ok { + result[req] = struct{}{} + } } return result, nil } diff --git a/server/datastore/mysql/labels_test.go b/server/datastore/mysql/labels_test.go index 5500554f736..a5b30208494 100644 --- a/server/datastore/mysql/labels_test.go +++ b/server/datastore/mysql/labels_test.go @@ -95,7 +95,6 @@ func TestLabels(t *testing.T) { {"LabelsSummaryAndListTeamFiltering", testLabelsSummaryAndListTeamFiltering}, {"ListHostsInLabelIssues", testListHostsInLabelIssues}, {"ListHostsInLabelDiskEncryptionStatus", testListHostsInLabelDiskEncryptionStatus}, - {"HostMemberOfAllLabels", testHostMemberOfAllLabels}, {"HostMembershipForLabels", testHostMembershipForLabels}, {"ListHostsInLabelOSSettings", testLabelsListHostsInLabelOSSettings}, {"AddDeleteLabelsToFromHost", testAddDeleteLabelsToFromHost}, @@ -1925,190 +1924,8 @@ func testListHostsInLabelDiskEncryptionStatus(t *testing.T, ds *Datastore) { listHostsCheckCount(t, ds, fleet.TeamFilter{User: test.UserAdmin}, fleet.HostListOptions{MacOSSettingsDiskEncryptionFilter: fleet.DiskEncryptionRemovingEnforcement}, 1) } -func testHostMemberOfAllLabels(t *testing.T, ds *Datastore) { - ctx := context.Background() - - // - // Setup test - // - h1 member of 'All hosts', 'Foobar' and 'Zoobar' - // - h2 member of 'All hosts' and 'Foobar' - // - h3 member of 'All hosts' and 'Zoobar' - // - h4 member of 'All hosts' - // - h5 member of no labels - // - - allHostsLabel, err := ds.NewLabel(ctx, - &fleet.Label{ - Name: "All hosts", - Query: "SELECT 1", - LabelType: fleet.LabelTypeBuiltIn, - LabelMembershipType: fleet.LabelMembershipTypeDynamic, - }, - ) - require.NoError(t, err) - foobarLabel, err := ds.NewLabel(ctx, &fleet.Label{ - Name: "Foobar", - Query: "SELECT 1;", - LabelType: fleet.LabelTypeRegular, - LabelMembershipType: fleet.LabelMembershipTypeDynamic, - }) - require.NoError(t, err) - zoobarLabel, err := ds.NewLabel(ctx, &fleet.Label{ - Name: "Zoobar", - Query: "SELECT 2;", - LabelType: fleet.LabelTypeRegular, - LabelMembershipType: fleet.LabelMembershipTypeDynamic, - }) - require.NoError(t, err) - - newHostFunc := func(name string) *fleet.Host { - h, err := ds.NewHost(ctx, &fleet.Host{ - DetailUpdatedAt: time.Now(), - LabelUpdatedAt: time.Now(), - PolicyUpdatedAt: time.Now(), - SeenTime: time.Now(), - OsqueryHostID: ptr.String(name), - NodeKey: ptr.String(name), - UUID: name, - Hostname: "foo.local" + name, - }) - require.NoError(t, err) - return h - } - - h1 := newHostFunc("h1") - h2 := newHostFunc("h2") - h3 := newHostFunc("h3") - h4 := newHostFunc("h4") - h5 := newHostFunc("h5") - _ = h5 - - err = ds.RecordLabelQueryExecutions(ctx, h1, map[uint]*bool{ - allHostsLabel.ID: ptr.Bool(true), - foobarLabel.ID: ptr.Bool(true), - zoobarLabel.ID: ptr.Bool(true), - }, time.Now(), false) - require.NoError(t, err) - err = ds.RecordLabelQueryExecutions(ctx, h2, map[uint]*bool{ - allHostsLabel.ID: ptr.Bool(true), - foobarLabel.ID: ptr.Bool(true), - }, time.Now(), false) - require.NoError(t, err) - err = ds.RecordLabelQueryExecutions(ctx, h3, map[uint]*bool{ - allHostsLabel.ID: ptr.Bool(true), - zoobarLabel.ID: ptr.Bool(true), - }, time.Now(), false) - require.NoError(t, err) - err = ds.RecordLabelQueryExecutions(ctx, h4, map[uint]*bool{ - allHostsLabel.ID: ptr.Bool(true), - }, time.Now(), false) - require.NoError(t, err) - - // - // Run tests for HostMemberOfAllLabels - // - - for _, tc := range []struct { - name string - hostID uint - labelNames []string - expectedResult bool - }{ - { - name: "nonexistent host", - hostID: 999, - labelNames: []string{allHostsLabel.Name}, - expectedResult: false, - }, - { - name: "h1 does not belong to nonexistent label", - hostID: h1.ID, - labelNames: []string{"Non existent label"}, - expectedResult: false, - }, - { - name: "h1 does not belong to All hosts + nonexistent label", - hostID: h1.ID, - labelNames: []string{allHostsLabel.Name, "Non existent label"}, - expectedResult: false, - }, - { - name: "h1 belongs to the given subset of labels", - hostID: h1.ID, - labelNames: []string{allHostsLabel.Name, foobarLabel.Name}, - expectedResult: true, - }, - { - name: "h1 belongs to all the given labels", - hostID: h1.ID, - labelNames: []string{allHostsLabel.Name, foobarLabel.Name, zoobarLabel.Name}, - expectedResult: true, - }, - { - name: "h1 member of empty label set", - hostID: h1.ID, - labelNames: []string{}, - expectedResult: true, - }, - { - name: "h2 belongs to all the given labels", - hostID: h2.ID, - labelNames: []string{allHostsLabel.Name, foobarLabel.Name}, - expectedResult: true, - }, - { - name: "h2 does not belongs to all the given labels", - hostID: h2.ID, - labelNames: []string{allHostsLabel.Name, foobarLabel.Name, zoobarLabel.Name}, - expectedResult: false, - }, - { - name: "h2 belongs to the given label", - hostID: h2.ID, - labelNames: []string{foobarLabel.Name}, - expectedResult: true, - }, - { - name: "h2 does not belong to the given label", - hostID: h2.ID, - labelNames: []string{zoobarLabel.Name}, - expectedResult: false, - }, - { - name: "h3 belongs to all the given labels", - hostID: h3.ID, - labelNames: []string{allHostsLabel.Name, zoobarLabel.Name}, - expectedResult: true, - }, - { - name: "h4 belongs to all the given labels", - hostID: h4.ID, - labelNames: []string{allHostsLabel.Name}, - expectedResult: true, - }, - { - name: "h4 does not belong to the given labels", - hostID: h4.ID, - labelNames: []string{foobarLabel.Name}, - expectedResult: false, - }, - { - name: "h5 does not belong to the given labels", - hostID: h5.ID, - labelNames: []string{allHostsLabel.Name}, - expectedResult: false, - }, - } { - t.Run(tc.name, func(t *testing.T) { - v, err := ds.HostMemberOfAllLabels(ctx, tc.hostID, tc.labelNames) - require.NoError(t, err) - require.Equal(t, tc.expectedResult, v) - }) - } -} - func testHostMembershipForLabels(t *testing.T, ds *Datastore) { - ctx := context.Background() + ctx := t.Context() // // Setup test @@ -2178,18 +1995,18 @@ func testHostMembershipForLabels(t *testing.T, ds *Datastore) { labelNames []string expectedResult map[string]struct{} }{ - { - name: "empty label names returns nil", - hostID: h1.ID, - labelNames: nil, - expectedResult: nil, - }, { name: "h1 is member of all three labels", hostID: h1.ID, labelNames: []string{allHostsLabel.Name, foobarLabel.Name, zoobarLabel.Name}, expectedResult: map[string]struct{}{allHostsLabel.Name: {}, foobarLabel.Name: {}, zoobarLabel.Name: {}}, }, + { + name: "h1 is member of a subset of labels", + hostID: h1.ID, + labelNames: []string{allHostsLabel.Name, foobarLabel.Name}, + expectedResult: map[string]struct{}{allHostsLabel.Name: {}, foobarLabel.Name: {}}, + }, { name: "h2 is member of some but not all labels", hostID: h2.ID, @@ -2202,6 +2019,12 @@ func testHostMembershipForLabels(t *testing.T, ds *Datastore) { labelNames: []string{allHostsLabel.Name, "nonexistent-label"}, expectedResult: map[string]struct{}{allHostsLabel.Name: {}}, }, + { + name: "case-insensitive match preserves requested casing", + hostID: h1.ID, + labelNames: []string{"foobar", "ZOOBAR"}, + expectedResult: map[string]struct{}{"foobar": {}, "ZOOBAR": {}}, + }, { name: "nonexistent host returns empty result", hostID: 999, diff --git a/server/fleet/datastore.go b/server/fleet/datastore.go index 376b737f1c2..e30a73af817 100644 --- a/server/fleet/datastore.go +++ b/server/fleet/datastore.go @@ -331,14 +331,9 @@ type Datastore interface { // HostIDsByOSID retrieves the IDs of all host for the given OS ID HostIDsByOSID(ctx context.Context, osID uint, offset int, limit int) ([]uint, error) - // HostMemberOfAllLabels returns whether the given host is a member of all the provided labels. - // If a label name does not exist, then the host is considered not a member of the provided label. - // A host will always be a member of an empty label set, so this method returns (true, nil) - // if labelNames is empty. - HostMemberOfAllLabels(ctx context.Context, hostID uint, labelNames []string) (bool, error) - // HostMembershipForLabels returns the set of label names (from the provided list) that the host is a member of. - // Labels that do not exist are not included in the result. + // Labels that do not exist are not included in the result. The returned map is keyed by the + // requested label names (preserving the caller's casing), not the DB-stored names. HostMembershipForLabels(ctx context.Context, hostID uint, labelNames []string) (map[string]struct{}, error) // TODO JUAN: Refactor this to use the Operating System type instead. diff --git a/server/mock/datastore_mock.go b/server/mock/datastore_mock.go index 7d841ac6238..6dfb5294491 100644 --- a/server/mock/datastore_mock.go +++ b/server/mock/datastore_mock.go @@ -254,8 +254,6 @@ type HostIDsByIdentifierFunc func(ctx context.Context, filter fleet.TeamFilter, type HostIDsByOSIDFunc func(ctx context.Context, osID uint, offset int, limit int) ([]uint, error) -type HostMemberOfAllLabelsFunc func(ctx context.Context, hostID uint, labelNames []string) (bool, error) - type HostMembershipForLabelsFunc func(ctx context.Context, hostID uint, labelNames []string) (map[string]struct{}, error) type HostIDsByOSVersionFunc func(ctx context.Context, osVersion fleet.OSVersion, offset int, limit int) ([]uint, error) @@ -2506,9 +2504,6 @@ type DataStore struct { HostIDsByOSIDFunc HostIDsByOSIDFunc HostIDsByOSIDFuncInvoked bool - HostMemberOfAllLabelsFunc HostMemberOfAllLabelsFunc - HostMemberOfAllLabelsFuncInvoked bool - HostMembershipForLabelsFunc HostMembershipForLabelsFunc HostMembershipForLabelsFuncInvoked bool @@ -6173,13 +6168,6 @@ func (s *DataStore) HostIDsByOSID(ctx context.Context, osID uint, offset int, li return s.HostIDsByOSIDFunc(ctx, osID, offset, limit) } -func (s *DataStore) HostMemberOfAllLabels(ctx context.Context, hostID uint, labelNames []string) (bool, error) { - s.mu.Lock() - s.HostMemberOfAllLabelsFuncInvoked = true - s.mu.Unlock() - return s.HostMemberOfAllLabelsFunc(ctx, hostID, labelNames) -} - func (s *DataStore) HostMembershipForLabels(ctx context.Context, hostID uint, labelNames []string) (map[string]struct{}, error) { s.mu.Lock() s.HostMembershipForLabelsFuncInvoked = true