mirror of
https://github.com/fleetdm/fleet
synced 2026-05-23 17:08:53 +00:00
change recordlabelexecutions to use map[uint]bool instead of map[string]bool (#965)
Closes #345
This commit is contained in:
parent
a2fded1b2c
commit
6135f90f57
9 changed files with 29 additions and 51 deletions
|
|
@ -396,11 +396,10 @@ func testDistributedQueriesForHost(t *testing.T, ds kolide.Datastore) {
|
|||
Query: "query1",
|
||||
})
|
||||
require.Nil(t, err)
|
||||
l1ID := fmt.Sprintf("%d", l1.ID)
|
||||
|
||||
// Add hosts to label
|
||||
for _, h := range []*kolide.Host{h1, h2} {
|
||||
err = ds.RecordLabelQueryExecutions(h, map[string]bool{l1ID: true}, time.Now())
|
||||
err = ds.RecordLabelQueryExecutions(h, map[uint]bool{l1.ID: true}, time.Now())
|
||||
require.Nil(t, err)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -89,7 +89,7 @@ func testLabels(t *testing.T, db kolide.Datastore) {
|
|||
assert.Empty(t, labels)
|
||||
|
||||
// Record a query execution
|
||||
err = db.RecordLabelQueryExecutions(host, map[string]bool{"1": true}, baseTime)
|
||||
err = db.RecordLabelQueryExecutions(host, map[uint]bool{1: true}, baseTime)
|
||||
assert.Nil(t, err)
|
||||
|
||||
// Use a 10 minute interval, so the query we just added should show up
|
||||
|
|
@ -99,14 +99,14 @@ func testLabels(t *testing.T, db kolide.Datastore) {
|
|||
assert.Equal(t, expectQueries, queries)
|
||||
|
||||
// Record an old query execution -- Shouldn't change the return
|
||||
err = db.RecordLabelQueryExecutions(host, map[string]bool{"2": true}, baseTime.Add(-1*time.Hour))
|
||||
err = db.RecordLabelQueryExecutions(host, map[uint]bool{2: true}, baseTime.Add(-1*time.Hour))
|
||||
assert.Nil(t, err)
|
||||
queries, err = db.LabelQueriesForHost(host, time.Now().Add(-(10 * time.Minute)))
|
||||
assert.Nil(t, err)
|
||||
assert.Equal(t, expectQueries, queries)
|
||||
|
||||
// Record a newer execution for that query and another
|
||||
err = db.RecordLabelQueryExecutions(host, map[string]bool{"2": false, "3": true}, baseTime)
|
||||
err = db.RecordLabelQueryExecutions(host, map[uint]bool{2: false, 3: true}, baseTime)
|
||||
assert.Nil(t, err)
|
||||
|
||||
// Now these should no longer show up in the necessary to run queries
|
||||
|
|
@ -275,7 +275,6 @@ func testListHostsInLabel(t *testing.T, db kolide.Datastore) {
|
|||
})
|
||||
require.Nil(t, err)
|
||||
require.NotZero(t, l1.ID)
|
||||
l1ID := fmt.Sprintf("%d", l1.ID)
|
||||
|
||||
{
|
||||
|
||||
|
|
@ -285,7 +284,7 @@ func testListHostsInLabel(t *testing.T, db kolide.Datastore) {
|
|||
}
|
||||
|
||||
for _, h := range []*kolide.Host{h1, h2, h3} {
|
||||
err = db.RecordLabelQueryExecutions(h, map[string]bool{l1ID: true}, time.Now())
|
||||
err = db.RecordLabelQueryExecutions(h, map[uint]bool{l1.ID: true}, time.Now())
|
||||
assert.Nil(t, err)
|
||||
}
|
||||
|
||||
|
|
@ -329,7 +328,6 @@ func testListUniqueHostsInLabels(t *testing.T, db kolide.Datastore) {
|
|||
})
|
||||
require.Nil(t, err)
|
||||
require.NotZero(t, l1.ID)
|
||||
l1ID := fmt.Sprintf("%d", l1.ID)
|
||||
|
||||
l2, err := db.NewLabel(&kolide.Label{
|
||||
Name: "label bar",
|
||||
|
|
@ -337,15 +335,14 @@ func testListUniqueHostsInLabels(t *testing.T, db kolide.Datastore) {
|
|||
})
|
||||
require.Nil(t, err)
|
||||
require.NotZero(t, l2.ID)
|
||||
l2ID := fmt.Sprintf("%d", l2.ID)
|
||||
|
||||
for i := 0; i < 3; i++ {
|
||||
err = db.RecordLabelQueryExecutions(hosts[i], map[string]bool{l1ID: true}, time.Now())
|
||||
err = db.RecordLabelQueryExecutions(hosts[i], map[uint]bool{l1.ID: true}, time.Now())
|
||||
assert.Nil(t, err)
|
||||
}
|
||||
// host 2 executes twice
|
||||
for i := 2; i < len(hosts); i++ {
|
||||
err = db.RecordLabelQueryExecutions(hosts[i], map[string]bool{l2ID: true}, time.Now())
|
||||
err = db.RecordLabelQueryExecutions(hosts[i], map[uint]bool{l2.ID: true}, time.Now())
|
||||
assert.Nil(t, err)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,6 @@
|
|||
package datastore
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
|
||||
"github.com/WatchBeam/clock"
|
||||
|
|
@ -90,7 +89,7 @@ func testGetHostsInPack(t *testing.T, ds kolide.Datastore) {
|
|||
|
||||
err = ds.RecordLabelQueryExecutions(
|
||||
h1,
|
||||
map[string]bool{fmt.Sprintf("%d", l1.ID): true},
|
||||
map[uint]bool{l1.ID: true},
|
||||
mockClock.Now(),
|
||||
)
|
||||
require.Nil(t, err)
|
||||
|
|
@ -107,7 +106,7 @@ func testGetHostsInPack(t *testing.T, ds kolide.Datastore) {
|
|||
|
||||
err = ds.RecordLabelQueryExecutions(
|
||||
h2,
|
||||
map[string]bool{fmt.Sprintf("%d", l1.ID): true},
|
||||
map[uint]bool{l1.ID: true},
|
||||
mockClock.Now(),
|
||||
)
|
||||
require.Nil(t, err)
|
||||
|
|
|
|||
|
|
@ -67,29 +67,12 @@ func (d *Datastore) LabelQueriesForHost(host *kolide.Host, cutoff time.Time) (ma
|
|||
return queries, nil
|
||||
}
|
||||
|
||||
func (d *Datastore) getLabelByIDString(id string) (*kolide.Label, error) {
|
||||
labelID, err := strconv.Atoi(id)
|
||||
if err != nil {
|
||||
return nil, errors.New("non-int label ID")
|
||||
}
|
||||
|
||||
d.mtx.Lock()
|
||||
label, ok := d.labels[uint(labelID)]
|
||||
d.mtx.Unlock()
|
||||
|
||||
if !ok {
|
||||
return nil, errors.New("label ID not found: " + string(labelID))
|
||||
}
|
||||
|
||||
return label, nil
|
||||
}
|
||||
|
||||
func (d *Datastore) RecordLabelQueryExecutions(host *kolide.Host, results map[string]bool, t time.Time) error {
|
||||
func (d *Datastore) RecordLabelQueryExecutions(host *kolide.Host, results map[uint]bool, t time.Time) error {
|
||||
// Record executions
|
||||
for strLabelID, matches := range results {
|
||||
label, err := d.getLabelByIDString(strLabelID)
|
||||
if err != nil {
|
||||
return err
|
||||
for labelID, matches := range results {
|
||||
label, ok := d.labels[labelID]
|
||||
if !ok {
|
||||
return notFound("Label").WithID(labelID)
|
||||
}
|
||||
|
||||
updated := false
|
||||
|
|
|
|||
|
|
@ -108,7 +108,7 @@ func (d *Datastore) LabelQueriesForHost(host *kolide.Host, cutoff time.Time) (ma
|
|||
|
||||
}
|
||||
|
||||
func (d *Datastore) RecordLabelQueryExecutions(host *kolide.Host, results map[string]bool, updated time.Time) error {
|
||||
func (d *Datastore) RecordLabelQueryExecutions(host *kolide.Host, results map[uint]bool, updated time.Time) error {
|
||||
sqlStatement := `
|
||||
INSERT INTO label_query_executions (updated_at, matches, label_id, host_id) VALUES
|
||||
|
||||
|
|
|
|||
|
|
@ -24,7 +24,7 @@ type LabelStore interface {
|
|||
// results map is a map of label id -> whether or not the label
|
||||
// matches. The time parameter is the timestamp to save with the query
|
||||
// execution.
|
||||
RecordLabelQueryExecutions(host *Host, results map[string]bool, t time.Time) error
|
||||
RecordLabelQueryExecutions(host *Host, results map[uint]bool, t time.Time) error
|
||||
|
||||
// LabelsForHost returns the labels that the given host is in.
|
||||
ListLabelsForHost(hid uint) ([]Label, error)
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ import (
|
|||
hostctx "github.com/kolide/kolide-ose/server/contexts/host"
|
||||
"github.com/kolide/kolide-ose/server/kolide"
|
||||
"github.com/kolide/kolide-ose/server/pubsub"
|
||||
"github.com/pkg/errors"
|
||||
"golang.org/x/net/context"
|
||||
)
|
||||
|
||||
|
|
@ -403,11 +404,15 @@ func (svc service) ingestDetailQuery(host *kolide.Host, name string, rows []map[
|
|||
}
|
||||
|
||||
// ingestLabelQuery records the results of label queries run by a host
|
||||
func (svc service) ingestLabelQuery(host kolide.Host, query string, rows []map[string]string, results map[string]bool) error {
|
||||
func (svc service) ingestLabelQuery(host kolide.Host, query string, rows []map[string]string, results map[uint]bool) error {
|
||||
trimmedQuery := strings.TrimPrefix(query, hostLabelQueryPrefix)
|
||||
trimmedQueryNum, err := strconv.Atoi(trimmedQuery)
|
||||
if err != nil {
|
||||
return errors.Wrap(err, "converting query from string to int")
|
||||
}
|
||||
// A label query matches if there is at least one result for that
|
||||
// query. We must also store negative results.
|
||||
results[trimmedQuery] = len(rows) > 0
|
||||
results[uint(trimmedQueryNum)] = len(rows) > 0
|
||||
return nil
|
||||
}
|
||||
|
||||
|
|
@ -487,7 +492,7 @@ func (svc service) SubmitDistributedQueryResults(ctx context.Context, results ko
|
|||
}
|
||||
|
||||
detailUpdated := false
|
||||
labelResults := map[string]bool{}
|
||||
labelResults := map[uint]bool{}
|
||||
for query, rows := range results {
|
||||
switch {
|
||||
case strings.HasPrefix(query, hostDetailQueryPrefix):
|
||||
|
|
|
|||
|
|
@ -4,7 +4,6 @@ import (
|
|||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
|
@ -454,7 +453,7 @@ func TestGetClientConfig(t *testing.T) {
|
|||
|
||||
err = ds.RecordLabelQueryExecutions(
|
||||
host,
|
||||
map[string]bool{fmt.Sprintf("%d", mysqlLabel.ID): true},
|
||||
map[uint]bool{mysqlLabel.ID: true},
|
||||
mockClock.Now(),
|
||||
)
|
||||
assert.Nil(t, err)
|
||||
|
|
@ -645,7 +644,6 @@ func TestDistributedQueries(t *testing.T) {
|
|||
Query: &q,
|
||||
})
|
||||
require.Nil(t, err)
|
||||
labelId := strconv.Itoa(int(label.ID))
|
||||
|
||||
// Record match with label
|
||||
ctx = viewer.NewContext(ctx, viewer.Viewer{
|
||||
|
|
@ -653,7 +651,7 @@ func TestDistributedQueries(t *testing.T) {
|
|||
ID: 0,
|
||||
},
|
||||
})
|
||||
err = ds.RecordLabelQueryExecutions(host, map[string]bool{labelId: true}, mockClock.Now())
|
||||
err = ds.RecordLabelQueryExecutions(host, map[uint]bool{label.ID: true}, mockClock.Now())
|
||||
require.Nil(t, err)
|
||||
err = ds.MarkHostSeen(host, mockClock.Now())
|
||||
require.Nil(t, err)
|
||||
|
|
|
|||
|
|
@ -111,7 +111,6 @@ func TestCountHostsInTargets(t *testing.T) {
|
|||
})
|
||||
require.Nil(t, err)
|
||||
require.NotZero(t, l1.ID)
|
||||
l1ID := fmt.Sprintf("%d", l1.ID)
|
||||
|
||||
l2, err := ds.NewLabel(&kolide.Label{
|
||||
Name: "label bar",
|
||||
|
|
@ -119,15 +118,14 @@ func TestCountHostsInTargets(t *testing.T) {
|
|||
})
|
||||
require.Nil(t, err)
|
||||
require.NotZero(t, l2.ID)
|
||||
l2ID := fmt.Sprintf("%d", l2.ID)
|
||||
|
||||
for _, h := range []*kolide.Host{h1, h2, h3, h6} {
|
||||
err = ds.RecordLabelQueryExecutions(h, map[string]bool{l1ID: true}, mockClock.Now())
|
||||
err = ds.RecordLabelQueryExecutions(h, map[uint]bool{l1.ID: true}, mockClock.Now())
|
||||
assert.Nil(t, err)
|
||||
}
|
||||
|
||||
for _, h := range []*kolide.Host{h3, h4, h5} {
|
||||
err = ds.RecordLabelQueryExecutions(h, map[string]bool{l2ID: true}, mockClock.Now())
|
||||
err = ds.RecordLabelQueryExecutions(h, map[uint]bool{l2.ID: true}, mockClock.Now())
|
||||
assert.Nil(t, err)
|
||||
}
|
||||
|
||||
|
|
@ -269,10 +267,9 @@ func TestSearchHostsInLabels(t *testing.T) {
|
|||
})
|
||||
require.Nil(t, err)
|
||||
require.NotZero(t, l1.ID)
|
||||
l1ID := fmt.Sprintf("%d", l1.ID)
|
||||
|
||||
for _, h := range []*kolide.Host{h1, h2, h3} {
|
||||
err = ds.RecordLabelQueryExecutions(h, map[string]bool{l1ID: true}, time.Now())
|
||||
err = ds.RecordLabelQueryExecutions(h, map[uint]bool{l1.ID: true}, time.Now())
|
||||
assert.Nil(t, err)
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue