EntityStore: support filtering by labels (#59905)

* label filtering

* filtering in sql

* filtering in sql group by

* label is unique - no need for distinct

* capitalize

* fix capitalization

pkg/services/store/entity/sqlstash/querybuilder.go

@@ -17,6 +17,11 @@ func (q *selectQuery) addWhere(f string, val interface{}) {
 	q.where = append(q.where, f+"=?")
 }
 
+func (q *selectQuery) addWhereInSubquery(f string, subquery string, subqueryArgs []interface{}) {
+	q.args = append(q.args, subqueryArgs...)
+	q.where = append(q.where, f+" IN ("+subquery+")")
+}
+
 func (q *selectQuery) addWhereIn(f string, vals []string) {
 	count := len(vals)
 	if count > 1 {

pkg/services/store/entity/sqlstash/sql_storage_server.go

@@ -623,7 +623,7 @@ func (s *sqlEntityServer) Search(ctx context.Context, r *entity.EntitySearchRequ
 		return nil, fmt.Errorf("missing user in context")
 	}
 
-	if r.NextPageToken != "" || len(r.Sort) > 0 || len(r.Labels) > 0 {
+	if r.NextPageToken != "" || len(r.Sort) > 0 {
 		return nil, fmt.Errorf("not yet supported")
 	}
 
@@ -637,6 +637,7 @@ func (s *sqlEntityServer) Search(ctx context.Context, r *entity.EntitySearchRequ
 	if r.WithBody {
 		fields = append(fields, "body")
 	}
+
 	if r.WithLabels {
 		fields = append(fields, "labels")
 	}
@@ -644,25 +645,40 @@ func (s *sqlEntityServer) Search(ctx context.Context, r *entity.EntitySearchRequ
 		fields = append(fields, "fields")
 	}
 
-	selectQuery := selectQuery{
+	entityQuery := selectQuery{
 		fields:   fields,
 		from:     "entity", // the table
 		args:     []interface{}{},
 		limit:    int(r.Limit),
 		oneExtra: true, // request one more than the limit (and show next token if it exists)
 	}
-	selectQuery.addWhere("tenant_id", user.OrgID)
+	entityQuery.addWhere("tenant_id", user.OrgID)
 
 	if len(r.Kind) > 0 {
-		selectQuery.addWhereIn("kind", r.Kind)
+		entityQuery.addWhereIn("kind", r.Kind)
 	}
 
 	// Folder UID or OID?
 	if r.Folder != "" {
-		selectQuery.addWhere("folder", r.Folder)
+		entityQuery.addWhere("folder", r.Folder)
+	}
+
+	if len(r.Labels) > 0 {
+		var args []interface{}
+		var conditions []string
+		for labelKey, labelValue := range r.Labels {
+			args = append(args, labelKey)
+			args = append(args, labelValue)
+			conditions = append(conditions, "(label = ? AND value = ?)")
+		}
+		joinedConditions := strings.Join(conditions, " OR ")
+		query := "SELECT grn FROM entity_labels WHERE " + joinedConditions + " GROUP BY grn HAVING COUNT(label) = ?"
+		args = append(args, len(r.Labels))
+
+		entityQuery.addWhereInSubquery("grn", query, args)
 	}
 
-	query, args := selectQuery.toQuery()
+	query, args := entityQuery.toQuery()
 
 	fmt.Printf("\n\n-------------\n")
 	fmt.Printf("%s\n", query)
@@ -704,7 +720,7 @@ func (s *sqlEntityServer) Search(ctx context.Context, r *entity.EntitySearchRequ
 		}
 
 		// found one more than requested
-		if len(rsp.Results) >= selectQuery.limit {
+		if len(rsp.Results) >= entityQuery.limit {
 			// TODO? should this encode start+offset?
 			rsp.NextPageToken = oid
 			break
@@ -732,5 +748,6 @@ func (s *sqlEntityServer) Search(ctx context.Context, r *entity.EntitySearchRequ
 
 		rsp.Results = append(rsp.Results, result)
 	}
+
 	return rsp, err
 }

pkg/services/store/entity/tests/server_integration_test.go

@@ -1,6 +1,7 @@
 package entity_server_tests
 
 import (
+	_ "embed"
 	"encoding/json"
 	"fmt"
 	"reflect"
@@ -15,6 +16,13 @@ import (
 	"google.golang.org/grpc/metadata"
 )
 
+var (
+	//go:embed testdata/dashboard-with-tags-b-g.json
+	dashboardWithTagsBlueGreen string
+	//go:embed testdata/dashboard-with-tags-r-g.json
+	dashboardWithTagsRedGreen string
+)
+
 type rawEntityMatcher struct {
 	grn          *entity.GRN
 	createdRange []time.Time
@@ -379,4 +387,88 @@ func TestIntegrationEntityServer(t *testing.T) {
 			w2.Entity.Version,
 		}, version)
 	})
+
+	t.Run("should be able to filter objects based on their labels", func(t *testing.T) {
+		kind := models.StandardKindDashboard
+		_, err := testCtx.client.Write(ctx, &entity.WriteEntityRequest{
+			GRN: &entity.GRN{
+				Kind: kind,
+				UID:  "blue-green",
+			},
+			Body: []byte(dashboardWithTagsBlueGreen),
+		})
+		require.NoError(t, err)
+
+		_, err = testCtx.client.Write(ctx, &entity.WriteEntityRequest{
+			GRN: &entity.GRN{
+				Kind: kind,
+				UID:  "red-green",
+			},
+			Body: []byte(dashboardWithTagsRedGreen),
+		})
+		require.NoError(t, err)
+
+		search, err := testCtx.client.Search(ctx, &entity.EntitySearchRequest{
+			Kind:       []string{kind},
+			WithBody:   false,
+			WithLabels: true,
+			Labels: map[string]string{
+				"red": "",
+			},
+		})
+		require.NoError(t, err)
+		require.NotNil(t, search)
+		require.Len(t, search.Results, 1)
+		require.Equal(t, search.Results[0].GRN.UID, "red-green")
+
+		search, err = testCtx.client.Search(ctx, &entity.EntitySearchRequest{
+			Kind:       []string{kind},
+			WithBody:   false,
+			WithLabels: true,
+			Labels: map[string]string{
+				"red":   "",
+				"green": "",
+			},
+		})
+		require.NoError(t, err)
+		require.NotNil(t, search)
+		require.Len(t, search.Results, 1)
+		require.Equal(t, search.Results[0].GRN.UID, "red-green")
+
+		search, err = testCtx.client.Search(ctx, &entity.EntitySearchRequest{
+			Kind:       []string{kind},
+			WithBody:   false,
+			WithLabels: true,
+			Labels: map[string]string{
+				"red": "invalid",
+			},
+		})
+		require.NoError(t, err)
+		require.NotNil(t, search)
+		require.Len(t, search.Results, 0)
+
+		search, err = testCtx.client.Search(ctx, &entity.EntitySearchRequest{
+			Kind:       []string{kind},
+			WithBody:   false,
+			WithLabels: true,
+			Labels: map[string]string{
+				"green": "",
+			},
+		})
+		require.NoError(t, err)
+		require.NotNil(t, search)
+		require.Len(t, search.Results, 2)
+
+		search, err = testCtx.client.Search(ctx, &entity.EntitySearchRequest{
+			Kind:       []string{kind},
+			WithBody:   false,
+			WithLabels: true,
+			Labels: map[string]string{
+				"yellow": "",
+			},
+		})
+		require.NoError(t, err)
+		require.NotNil(t, search)
+		require.Len(t, search.Results, 0)
+	})
 }

pkg/services/store/entity/tests/testdata/dashboard-with-tags-b-g.json

@@ -0,0 +1,40 @@
+{
+  "tags": [
+    "blue",
+    "green"
+  ],
+  "editable": true,
+  "fiscalYearStartMonth": 0,
+  "graphTooltip": 0,
+  "id": 221,
+  "links": [],
+  "liveNow": false,
+  "panels": [
+    {
+      "gridPos": {
+        "h": 1,
+        "w": 24,
+        "x": 0,
+        "y": 8
+      },
+      "id": 8,
+      "title": "Row title",
+      "type": "row"
+    }
+  ],
+  "schemaVersion": 36,
+  "style": "dark",
+  "templating": {
+    "list": []
+  },
+  "time": {
+    "from": "now-6h",
+    "to": "now"
+  },
+  "timepicker": {},
+  "timezone": "",
+  "title": "special ds",
+  "uid": "mocpwtR4k",
+  "version": 1,
+  "weekStart": ""
+}

pkg/services/store/entity/tests/testdata/dashboard-with-tags-r-g.json

@@ -0,0 +1,40 @@
+{
+  "tags": [
+    "red",
+    "green"
+  ],
+  "editable": true,
+  "fiscalYearStartMonth": 0,
+  "graphTooltip": 0,
+  "id": 221,
+  "links": [],
+  "liveNow": false,
+  "panels": [
+    {
+      "gridPos": {
+        "h": 1,
+        "w": 24,
+        "x": 0,
+        "y": 8
+      },
+      "id": 8,
+      "title": "Row title",
+      "type": "row"
+    }
+  ],
+  "schemaVersion": 36,
+  "style": "dark",
+  "templating": {
+    "list": []
+  },
+  "time": {
+    "from": "now-6h",
+    "to": "now"
+  },
+  "timepicker": {},
+  "timezone": "",
+  "title": "special ds",
+  "uid": "mocpwtR4k",
+  "version": 1,
+  "weekStart": ""
+}