package inmem import ( "sort" "github.com/kolide/fleet/server/kolide" "github.com/pkg/errors" ) func (d *Datastore) NewQuery(query *kolide.Query, opts ...kolide.OptionalArg) (*kolide.Query, error) { d.mtx.Lock() defer d.mtx.Unlock() newQuery := *query for _, q := range d.queries { if query.Name == q.Name { return nil, alreadyExists("Query", q.ID) } } newQuery.ID = d.nextID(newQuery) newQuery.Packs = []kolide.Pack{} d.queries[newQuery.ID] = &newQuery return &newQuery, nil } func (d *Datastore) QueryByName(name string, opts ...kolide.OptionalArg) (*kolide.Query, bool, error) { d.mtx.Lock() defer d.mtx.Unlock() for _, q := range d.queries { if name == q.Name { return q, true, nil } } return nil, false, nil } func (d *Datastore) SaveQuery(query *kolide.Query) error { d.mtx.Lock() defer d.mtx.Unlock() if _, ok := d.queries[query.ID]; !ok { return notFound("Query").WithID(query.ID) } d.queries[query.ID] = query return nil } func (d *Datastore) DeleteQuery(qid uint) error { d.mtx.Lock() defer d.mtx.Unlock() if _, ok := d.queries[qid]; !ok { return notFound("Query").WithID(qid) } delete(d.queries, qid) return nil } // DeleteQueries (soft) deletes the existing query objects with the provided // IDs. The number of deleted queries is returned along with any error. func (d *Datastore) DeleteQueries(ids []uint) (uint, error) { d.mtx.Lock() defer d.mtx.Unlock() deleted := uint(0) for _, id := range ids { if _, ok := d.queries[id]; ok { delete(d.queries, id) deleted++ } } return deleted, nil } func (d *Datastore) getUserNameByID(id uint) string { if u, ok := d.users[id]; ok { return u.Name } return "" } func (d *Datastore) Query(id uint) (*kolide.Query, error) { d.mtx.Lock() defer d.mtx.Unlock() query, ok := d.queries[id] if !ok { return nil, notFound("Query").WithID(id) } query.AuthorName = d.getUserNameByID(query.AuthorID) if err := d.loadPacksForQueries([]*kolide.Query{query}); err != nil { return nil, errors.Wrap(err, "error fetching query by id") } return query, nil } func (d *Datastore) ListQueries(opt kolide.ListOptions) ([]*kolide.Query, error) { d.mtx.Lock() defer d.mtx.Unlock() // We need to sort by keys to provide reliable ordering keys := []int{} for k, _ := range d.queries { keys = append(keys, int(k)) } sort.Ints(keys) queries := []*kolide.Query{} for _, k := range keys { q := d.queries[uint(k)] if q.Saved { q.AuthorName = d.getUserNameByID(q.AuthorID) queries = append(queries, q) } } // Apply ordering if opt.OrderKey != "" { var fields = map[string]string{ "id": "ID", "created_at": "CreatedAt", "updated_at": "UpdatedAt", "name": "Name", "query": "Query", "interval": "Interval", "snapshot": "Snapshot", "differential": "Differential", "platform": "Platform", "version": "Version", } if err := sortResults(queries, opt, fields); err != nil { return nil, err } } // Apply limit/offset low, high := d.getLimitOffsetSliceBounds(opt, len(queries)) queries = queries[low:high] if err := d.loadPacksForQueries(queries); err != nil { return nil, errors.Wrap(err, "error listing queries") } return queries, nil } // loadPacksForQueries loads the packs associated with the provided queries func (d *Datastore) loadPacksForQueries(queries []*kolide.Query) error { for _, q := range queries { q.Packs = make([]kolide.Pack, 0) for _, sq := range d.scheduledQueries { if sq.QueryID == q.ID { q.Packs = append(q.Packs, *d.packs[sq.PackID]) } } } return nil }