Elgato_dark/fleet
1
0
Fork 0
mirror of https://github.com/fleetdm/fleet synced 2026-05-23 08:58:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 111

Return packs with queries (#575)

Browse source 
- New datastore method for loading packs associated with a query
- ListQueries and Query datastore methods now load packs

Addresses #388
This commit is contained in:
Zachary Wasserman 2016-12-06 10:22:28 -08:00 committed by GitHub
parent bc6109fb56
commit 691eb55cf2
8 changed files with 253 additions and 88 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

90
server/datastore/datastore_campaigns_test.go
View file

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

46
server/datastore/datastore_queries_test.go
View file

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

1
server/datastore/datastore_test.go
View file

@ -49,4 +49,5 @@ var testFunctions = [...]func(*testing.T, kolide.Datastore){
testDistributedQueryCampaign,
testCleanupDistributedQueryCampaigns,
testBuiltInLabels,
testLoadPacksForQueries,
}

22
server/datastore/inmem/queries.go
View file

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

4
server/datastore/mysql/packs.go
View file

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

52
server/datastore/mysql/queries.go
View file

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

114
server/datastore/test_util.go Normal file
View file

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

12
server/kolide/queries.go
View file

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