From 691eb55cf25ca29aea0093926f608472bcd8e8d9 Mon Sep 17 00:00:00 2001 From: Zachary Wasserman Date: Tue, 6 Dec 2016 10:22:28 -0800 Subject: [PATCH] Return packs with queries (#575) - New datastore method for loading packs associated with a query - ListQueries and Query datastore methods now load packs Addresses #388 --- server/datastore/datastore_campaigns_test.go | 90 +-------------- server/datastore/datastore_queries_test.go | 46 ++++++++ server/datastore/datastore_test.go | 1 + server/datastore/inmem/queries.go | 22 ++++ server/datastore/mysql/packs.go | 4 +- server/datastore/mysql/queries.go | 52 +++++++++ server/datastore/test_util.go | 114 +++++++++++++++++++ server/kolide/queries.go | 12 +- 8 files changed, 253 insertions(+), 88 deletions(-) create mode 100644 server/datastore/test_util.go diff --git a/server/datastore/datastore_campaigns_test.go b/server/datastore/datastore_campaigns_test.go index 223a498d1b..f7fc0fa309 100644 --- a/server/datastore/datastore_campaigns_test.go +++ b/server/datastore/datastore_campaigns_test.go @@ -11,86 +11,6 @@ import ( "github.com/stretchr/testify/require" ) -func newQuery(t *testing.T, ds kolide.Datastore, name, q string) *kolide.Query { - query, err := ds.NewQuery(&kolide.Query{ - Name: name, - Query: q, - }) - require.Nil(t, err) - - return query -} - -func newCampaign(t *testing.T, ds kolide.Datastore, queryID uint, status kolide.DistributedQueryStatus, now time.Time) *kolide.DistributedQueryCampaign { - campaign, err := ds.NewDistributedQueryCampaign(&kolide.DistributedQueryCampaign{ - UpdateCreateTimestamps: kolide.UpdateCreateTimestamps{ - CreateTimestamp: kolide.CreateTimestamp{ - CreatedAt: now, - }, - }, - QueryID: queryID, - Status: status, - }) - require.Nil(t, err) - - return campaign -} - -func newExecution(t *testing.T, ds kolide.Datastore, campaignID uint, hostID uint) *kolide.DistributedQueryExecution { - execution, err := ds.NewDistributedQueryExecution(&kolide.DistributedQueryExecution{ - HostID: hostID, - DistributedQueryCampaignID: campaignID, - }) - require.Nil(t, err) - - return execution -} - -func newHost(t *testing.T, ds kolide.Datastore, name, ip, key, uuid string, now time.Time) *kolide.Host { - h, err := ds.NewHost(&kolide.Host{ - HostName: name, - NodeKey: key, - UUID: uuid, - DetailUpdateTime: now, - }) - - require.Nil(t, err) - require.NotZero(t, h.ID) - require.Nil(t, ds.MarkHostSeen(h, now)) - - return h -} - -func newLabel(t *testing.T, ds kolide.Datastore, name, query string) *kolide.Label { - l, err := ds.NewLabel(&kolide.Label{Name: name, Query: query}) - - require.Nil(t, err) - require.NotZero(t, l.ID) - - return l -} - -func addHost(t *testing.T, ds kolide.Datastore, campaignID, hostID uint) { - _, err := ds.NewDistributedQueryCampaignTarget( - &kolide.DistributedQueryCampaignTarget{ - Type: kolide.TargetHost, - TargetID: hostID, - DistributedQueryCampaignID: campaignID, - }) - require.Nil(t, err) - -} - -func addLabel(t *testing.T, ds kolide.Datastore, campaignID, labelID uint) { - _, err := ds.NewDistributedQueryCampaignTarget( - &kolide.DistributedQueryCampaignTarget{ - Type: kolide.TargetLabel, - TargetID: labelID, - DistributedQueryCampaignID: campaignID, - }) - require.Nil(t, err) -} - func checkTargets(t *testing.T, ds kolide.Datastore, campaignID uint, expectedHostIDs []uint, expectedLabelIDs []uint) { hostIDs, labelIDs, err := ds.DistributedQueryCampaignTargetIDs(campaignID) require.Nil(t, err) @@ -127,17 +47,17 @@ func testDistributedQueryCampaign(t *testing.T, ds kolide.Datastore) { checkTargets(t, ds, campaign.ID, []uint{}, []uint{}) - addHost(t, ds, campaign.ID, h1.ID) + addHostToCampaign(t, ds, campaign.ID, h1.ID) checkTargets(t, ds, campaign.ID, []uint{h1.ID}, []uint{}) - addLabel(t, ds, campaign.ID, l1.ID) + addLabelToCampaign(t, ds, campaign.ID, l1.ID) checkTargets(t, ds, campaign.ID, []uint{h1.ID}, []uint{l1.ID}) - addLabel(t, ds, campaign.ID, l2.ID) + addLabelToCampaign(t, ds, campaign.ID, l2.ID) checkTargets(t, ds, campaign.ID, []uint{h1.ID}, []uint{l1.ID, l2.ID}) - addHost(t, ds, campaign.ID, h2.ID) - addHost(t, ds, campaign.ID, h3.ID) + addHostToCampaign(t, ds, campaign.ID, h2.ID) + addHostToCampaign(t, ds, campaign.ID, h3.ID) checkTargets(t, ds, campaign.ID, []uint{h1.ID, h2.ID, h3.ID}, []uint{l1.ID, l2.ID}) diff --git a/server/datastore/datastore_queries_test.go b/server/datastore/datastore_queries_test.go index b33d4ff8ea..177711d33e 100644 --- a/server/datastore/datastore_queries_test.go +++ b/server/datastore/datastore_queries_test.go @@ -5,6 +5,7 @@ import ( "testing" "github.com/kolide/kolide-ose/server/kolide" + "github.com/patrickmn/sortutil" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -69,3 +70,48 @@ func testListQuery(t *testing.T, ds kolide.Datastore) { assert.Nil(t, err) assert.Equal(t, 10, len(results)) } + +func checkPacks(t *testing.T, expected []kolide.Pack, actual []kolide.Pack) { + sortutil.AscByField(expected, "ID") + sortutil.AscByField(actual, "ID") + assert.Equal(t, expected, actual) +} + +func testLoadPacksForQueries(t *testing.T, ds kolide.Datastore) { + q1 := newQuery(t, ds, "q1", "select * from time") + q2 := newQuery(t, ds, "q2", "select * from osquery_info") + + p1 := newPack(t, ds, "p1") + p2 := newPack(t, ds, "p2") + p3 := newPack(t, ds, "p3") + + var err error + + addQueryToPack(t, ds, q1.ID, p2.ID) + + q1, err = ds.Query(q1.ID) + require.Nil(t, err) + q2, err = ds.Query(q2.ID) + require.Nil(t, err) + checkPacks(t, []kolide.Pack{*p2}, q1.Packs) + checkPacks(t, []kolide.Pack{}, q2.Packs) + + addQueryToPack(t, ds, q2.ID, p1.ID) + addQueryToPack(t, ds, q2.ID, p3.ID) + + q1, err = ds.Query(q1.ID) + require.Nil(t, err) + q2, err = ds.Query(q2.ID) + require.Nil(t, err) + checkPacks(t, []kolide.Pack{*p2}, q1.Packs) + checkPacks(t, []kolide.Pack{*p1, *p3}, q2.Packs) + + addQueryToPack(t, ds, q1.ID, p3.ID) + + q1, err = ds.Query(q1.ID) + require.Nil(t, err) + q2, err = ds.Query(q2.ID) + require.Nil(t, err) + checkPacks(t, []kolide.Pack{*p2, *p3}, q1.Packs) + checkPacks(t, []kolide.Pack{*p1, *p3}, q2.Packs) +} diff --git a/server/datastore/datastore_test.go b/server/datastore/datastore_test.go index ec4a0965e0..817b496024 100644 --- a/server/datastore/datastore_test.go +++ b/server/datastore/datastore_test.go @@ -49,4 +49,5 @@ var testFunctions = [...]func(*testing.T, kolide.Datastore){ testDistributedQueryCampaign, testCleanupDistributedQueryCampaigns, testBuiltInLabels, + testLoadPacksForQueries, } diff --git a/server/datastore/inmem/queries.go b/server/datastore/inmem/queries.go index 3fb6f58225..b9caf0b8bd 100644 --- a/server/datastore/inmem/queries.go +++ b/server/datastore/inmem/queries.go @@ -58,6 +58,10 @@ func (orm *Datastore) Query(id uint) (*kolide.Query, error) { return nil, errors.ErrNotFound } + if err := orm.loadPacksForQueries([]*kolide.Query{query}); err != nil { + return nil, errors.DatabaseError(err) + } + return query, nil } @@ -102,5 +106,23 @@ func (orm *Datastore) ListQueries(opt kolide.ListOptions) ([]*kolide.Query, erro low, high := orm.getLimitOffsetSliceBounds(opt, len(queries)) queries = queries[low:high] + if err := orm.loadPacksForQueries(queries); err != nil { + return nil, errors.DatabaseError(err) + } + return queries, nil } + +// loadPacksForQueries loads the packs associated with the provided queries +func (orm *Datastore) loadPacksForQueries(queries []*kolide.Query) error { + for _, q := range queries { + q.Packs = make([]kolide.Pack, 0) + for _, pq := range orm.packQueries { + if pq.QueryID == q.ID { + q.Packs = append(q.Packs, *orm.packs[pq.PackID]) + } + } + } + + return nil +} diff --git a/server/datastore/mysql/packs.go b/server/datastore/mysql/packs.go index 7b4aa8e9a5..4a783110de 100644 --- a/server/datastore/mysql/packs.go +++ b/server/datastore/mysql/packs.go @@ -86,10 +86,10 @@ func (d *Datastore) ListPacks(opt kolide.ListOptions) ([]*kolide.Pack, error) { // AddQueryToPack associates a kolide.Query with a kolide.Pack func (d *Datastore) AddQueryToPack(qid uint, pid uint) error { sql := ` - INSERT INTO pack_queries ( pack_id, query_id) + INSERT INTO pack_queries (query_id, pack_id) VALUES (?, ?) ` - if _, err := d.db.Exec(sql, pid, qid); err != nil { + if _, err := d.db.Exec(sql, qid, pid); err != nil { return errors.DatabaseError(err) } diff --git a/server/datastore/mysql/queries.go b/server/datastore/mysql/queries.go index 554b26bc4a..8b8124e9f5 100644 --- a/server/datastore/mysql/queries.go +++ b/server/datastore/mysql/queries.go @@ -1,6 +1,7 @@ package mysql import ( + "github.com/jmoiron/sqlx" "github.com/kolide/kolide-ose/server/errors" "github.com/kolide/kolide-ose/server/kolide" ) @@ -70,6 +71,10 @@ func (d *Datastore) Query(id uint) (*kolide.Query, error) { return nil, errors.DatabaseError(err) } + if err := d.loadPacksForQueries([]*kolide.Query{query}); err != nil { + return nil, errors.DatabaseError(err) + } + return query, nil } @@ -88,6 +93,53 @@ func (d *Datastore) ListQueries(opt kolide.ListOptions) ([]*kolide.Query, error) return nil, errors.DatabaseError(err) } + if err := d.loadPacksForQueries(results); err != nil { + return nil, errors.DatabaseError(err) + } + return results, nil } + +// loadPacksForQueries loads the packs associated with the provided queries +func (d *Datastore) loadPacksForQueries(queries []*kolide.Query) error { + sql := ` + SELECT p.*, pq.query_id AS query_id + FROM packs p + JOIN pack_queries pq + ON p.id = pq.pack_id + WHERE query_id IN (?) + ` + + // Used to map the results + id_queries := map[uint]*kolide.Query{} + // Used for the IN clause + ids := []uint{} + for _, q := range queries { + q.Packs = make([]kolide.Pack, 0) + ids = append(ids, q.ID) + id_queries[q.ID] = q + } + + query, args, err := sqlx.In(sql, ids) + if err != nil { + return errors.DatabaseError(err) + } + + rows := []struct { + QueryID uint `db:"query_id"` + kolide.Pack + }{} + + err = d.db.Select(&rows, query, args...) + if err != nil { + return errors.DatabaseError(err) + } + + for _, row := range rows { + q := id_queries[row.QueryID] + q.Packs = append(q.Packs, row.Pack) + } + + return nil +} diff --git a/server/datastore/test_util.go b/server/datastore/test_util.go new file mode 100644 index 0000000000..7040d6aa40 --- /dev/null +++ b/server/datastore/test_util.go @@ -0,0 +1,114 @@ +package datastore + +import ( + "testing" + "time" + + "github.com/kolide/kolide-ose/server/kolide" + "github.com/stretchr/testify/require" +) + +func newQuery(t *testing.T, ds kolide.Datastore, name, q string) *kolide.Query { + query, err := ds.NewQuery(&kolide.Query{ + Name: name, + Query: q, + }) + require.Nil(t, err) + + // Loading gives us the timestamps + query, err = ds.Query(query.ID) + require.Nil(t, err) + + return query +} + +func newPack(t *testing.T, ds kolide.Datastore, name string) *kolide.Pack { + pack, err := ds.NewPack(&kolide.Pack{ + Name: name, + }) + require.Nil(t, err) + + // Loading gives us the timestamps + pack, err = ds.Pack(pack.ID) + require.Nil(t, err) + + return pack +} + +func addQueryToPack(t *testing.T, ds kolide.Datastore, queryID, packID uint) { + err := ds.AddQueryToPack(queryID, packID) + require.Nil(t, err) +} + +func newCampaign(t *testing.T, ds kolide.Datastore, queryID uint, status kolide.DistributedQueryStatus, now time.Time) *kolide.DistributedQueryCampaign { + campaign, err := ds.NewDistributedQueryCampaign(&kolide.DistributedQueryCampaign{ + UpdateCreateTimestamps: kolide.UpdateCreateTimestamps{ + CreateTimestamp: kolide.CreateTimestamp{ + CreatedAt: now, + }, + }, + QueryID: queryID, + Status: status, + }) + require.Nil(t, err) + + // Loading gives us the timestamps + campaign, err = ds.DistributedQueryCampaign(campaign.ID) + require.Nil(t, err) + + return campaign +} + +func addHostToCampaign(t *testing.T, ds kolide.Datastore, campaignID, hostID uint) { + _, err := ds.NewDistributedQueryCampaignTarget( + &kolide.DistributedQueryCampaignTarget{ + Type: kolide.TargetHost, + TargetID: hostID, + DistributedQueryCampaignID: campaignID, + }) + require.Nil(t, err) +} + +func addLabelToCampaign(t *testing.T, ds kolide.Datastore, campaignID, labelID uint) { + _, err := ds.NewDistributedQueryCampaignTarget( + &kolide.DistributedQueryCampaignTarget{ + Type: kolide.TargetLabel, + TargetID: labelID, + DistributedQueryCampaignID: campaignID, + }) + require.Nil(t, err) +} + +func newExecution(t *testing.T, ds kolide.Datastore, campaignID uint, hostID uint) *kolide.DistributedQueryExecution { + execution, err := ds.NewDistributedQueryExecution(&kolide.DistributedQueryExecution{ + HostID: hostID, + DistributedQueryCampaignID: campaignID, + }) + require.Nil(t, err) + + return execution +} + +func newHost(t *testing.T, ds kolide.Datastore, name, ip, key, uuid string, now time.Time) *kolide.Host { + h, err := ds.NewHost(&kolide.Host{ + HostName: name, + NodeKey: key, + UUID: uuid, + DetailUpdateTime: now, + }) + + require.Nil(t, err) + require.NotZero(t, h.ID) + require.Nil(t, ds.MarkHostSeen(h, now)) + + return h +} + +func newLabel(t *testing.T, ds kolide.Datastore, name, query string) *kolide.Label { + l, err := ds.NewLabel(&kolide.Label{Name: name, Query: query}) + + require.Nil(t, err) + require.NotZero(t, l.ID) + + return l +} diff --git a/server/kolide/queries.go b/server/kolide/queries.go index f74386cdb0..bf563f0d32 100644 --- a/server/kolide/queries.go +++ b/server/kolide/queries.go @@ -7,11 +7,18 @@ import ( ) type QueryStore interface { - // Query methods + // NewQuery creates a new query object in thie datastore. The returned + // query should have the ID updated. NewQuery(query *Query) (*Query, error) + // SaveQuery saves changes to an existing query object. SaveQuery(query *Query) error + // DeleteQuery (soft) deletes an existing query object. DeleteQuery(query *Query) error + // Query returns the query associated with the provided ID. Associated + // packs should also be loaded. Query(id uint) (*Query, error) + // ListQueries returns a list of queries with the provided sorting and + // paging options. Associated packs should also be loaded. ListQueries(opt ListOptions) ([]*Query, error) } @@ -50,6 +57,9 @@ type Query struct { Differential bool `json:"differential"` Platform string `json:"platform"` Version string `json:"version"` + // Packs is loaded when retrieving queries, but is stored in a join + // table in the MySQL backend. + Packs []Pack `json:"packs" db:"-"` } type Option struct {