fleet/server/datastore/cached_mysql/cached_mysql.go
Lucas Manuel Rodriguez 371c533bfc
Improved Datastore usage of osquery hosts requests (#3601)
* WIP

* Amend tests

* Do not load aggregated stats for packs

* Add option to host lite

* Fix remaining TODOs

* Fix osquery_utils tests

* Fix SQL

* Fix SQL (bis)

* Restore AuthenticateHost to load once

* Code improvements and re-add deferred host save

* More fixes to the PR

* Wrap users table update on tx

* Add caching to ListPacksForHost and ListScheduledQueriesInPack

* Remove SaveHostSoftware (replaced by UpdateHostSoftware)

* Add unit tests for new functionality

* Add changes file

* Fix scheduled queries test
2022-01-17 22:52:09 -03:00

134 lines
3.0 KiB
Go

package cached_mysql
import (
"context"
"fmt"
"time"
"github.com/fleetdm/fleet/v4/server/fleet"
"github.com/patrickmn/go-cache"
)
type cachedMysql struct {
fleet.Datastore
c *cache.Cache
packsExp time.Duration
scheduledQueriesExp time.Duration
}
const (
appConfigKey = "AppConfig"
packsKey = "Packs"
scheduledQueriesKey = "ScheduledQueries"
defaultAppConfigExpiration = 1 * time.Second
defaultPacksExpiration = 1 * time.Minute
defaultScheduledQueriesExpiration = 1 * time.Minute
)
type Option func(*cachedMysql)
func WithPacksExpiration(t time.Duration) Option {
return func(o *cachedMysql) {
o.packsExp = t
}
}
func WithScheduledQueriesExpiration(t time.Duration) Option {
return func(o *cachedMysql) {
o.scheduledQueriesExp = t
}
}
func New(ds fleet.Datastore, opts ...Option) fleet.Datastore {
c := &cachedMysql{
Datastore: ds,
c: cache.New(5*time.Minute, 10*time.Minute),
packsExp: defaultPacksExpiration,
scheduledQueriesExp: defaultScheduledQueriesExpiration,
}
for _, fn := range opts {
fn(c)
}
return c
}
func (ds *cachedMysql) NewAppConfig(ctx context.Context, info *fleet.AppConfig) (*fleet.AppConfig, error) {
ac, err := ds.Datastore.NewAppConfig(ctx, info)
if err != nil {
return nil, err
}
ds.c.Set(appConfigKey, ac, defaultAppConfigExpiration)
return ac.Clone()
}
func (ds *cachedMysql) AppConfig(ctx context.Context) (*fleet.AppConfig, error) {
cachedAc, found := ds.c.Get(appConfigKey)
if found {
return cachedAc.(*fleet.AppConfig).Clone()
}
ac, err := ds.Datastore.AppConfig(ctx)
if err != nil {
return nil, err
}
ds.c.Set(appConfigKey, ac, defaultAppConfigExpiration)
return ac.Clone()
}
func (ds *cachedMysql) SaveAppConfig(ctx context.Context, info *fleet.AppConfig) error {
err := ds.Datastore.SaveAppConfig(ctx, info)
if err != nil {
return err
}
ds.c.Set(appConfigKey, info, defaultAppConfigExpiration)
return nil
}
func (ds *cachedMysql) ListPacksForHost(ctx context.Context, hid uint) ([]*fleet.Pack, error) {
key := fmt.Sprintf("%s_%d", packsKey, hid)
cachedPacks, found := ds.c.Get(key)
if found && cachedPacks != nil {
casted, ok := cachedPacks.([]*fleet.Pack)
if ok {
return casted, nil
}
}
packs, err := ds.Datastore.ListPacksForHost(ctx, hid)
if err != nil {
return nil, err
}
ds.c.Set(key, packs, ds.packsExp)
return packs, nil
}
func (ds *cachedMysql) ListScheduledQueriesInPack(ctx context.Context, id uint) ([]*fleet.ScheduledQuery, error) {
key := fmt.Sprintf("%s_%d", scheduledQueriesKey, id)
cachedScheduledQueries, found := ds.c.Get(key)
if found && cachedScheduledQueries != nil {
casted, ok := cachedScheduledQueries.([]*fleet.ScheduledQuery)
if ok {
return casted, nil
}
}
scheduledQueries, err := ds.Datastore.ListScheduledQueriesInPack(ctx, id)
if err != nil {
return nil, err
}
ds.c.Set(key, scheduledQueries, ds.scheduledQueriesExp)
return scheduledQueries, nil
}