change recordlabelexecutions to use map[uint]bool instead of map[string]bool (#965)

Closes #345
This commit is contained in:
Victor Vrantchan 2017-01-17 01:03:51 -05:00 committed by GitHub
parent a2fded1b2c
commit 6135f90f57
9 changed files with 29 additions and 51 deletions

View file

@ -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)
}

View file

@ -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)
}

View file

@ -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)

View file

@ -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

View file

@ -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

View file

@ -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)

View file

@ -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):

View file

@ -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)

View file

@ -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)
}