query cache (#12)

* query cache

* add info to the frame that the data is from cache

Author: scottlepp

README.md

@@ -5,9 +5,6 @@
 * Doesn't require CGO.
 * Requires duckdb cli to be in the path.
 
-## Dataframes
-* Allows querying dataframes using SQL, with all the aggregate/window/analytics functions from DuckDB.
-
 ## In Memory Database
 ```
 	db := NewInMemoryDB()
@@ -30,6 +27,10 @@
 	_, err := db.RunCommands(commands)
 	res, err := db.Query("SELECT * from t1;")
 ```
+
+## Dataframes
+* Allows querying dataframes using SQL, with all the aggregate/window/analytics functions from DuckDB.
+
 ## Query Dataframes
 ```
 	db := NewInMemoryDB()

duck/duckdb.go

@@ -17,30 +17,35 @@ import (
 
 var logger = log.DefaultLogger
 
+type Dirs map[string]string
+
 type DuckDB struct {
-	Name   string
-	mode   string
-	format string
-	exe    string
-	chunk  int
+	Name          string
+	mode          string
+	format        string
+	exe           string
+	chunk         int
+	cacheDuration int
+	cache         cache
 }
 
 type Opts struct {
-	Mode   string
-	Format string
-	Chunk  int
-	Exe    string
+	Mode          string
+	Format        string
+	Chunk         int
+	Exe           string
+	CacheDuration int
 }
 
 const newline = "\n"
 
 // NewInMemoryDB creates a new in-memory DuckDB
-func NewInMemoryDB(opts ...Opts) DuckDB {
+func NewInMemoryDB(opts ...Opts) *DuckDB {
 	return NewDuckDB("", opts...)
 }
 
 // NewDuckDB creates a new DuckDB
-func NewDuckDB(name string, opts ...Opts) DuckDB {
+func NewDuckDB(name string, opts ...Opts) *DuckDB {
 	db := DuckDB{
 		Name:   name,
 		mode:   "json",
@@ -59,6 +64,9 @@ func NewDuckDB(name string, opts ...Opts) DuckDB {
 		if opt.Chunk > 0 {
 			db.chunk = opt.Chunk
 		}
+		if opt.CacheDuration > 0 {
+			db.cacheDuration = opt.CacheDuration
+		}
 	}
 
 	// Find the executable if it is not configured
@@ -68,7 +76,8 @@ func NewDuckDB(name string, opts ...Opts) DuckDB {
 			db.exe = "/usr/local/bin/duckdb"
 		}
 	}
-	return db
+	db.cache = cache{}
+	return &db
 }
 
 // RunCommands runs a series of of sql commands against duckdb
@@ -108,51 +117,48 @@ func (d *DuckDB) Query(query string) (string, error) {
 }
 
 // QueryFrame will load a dataframe into a view named RefID, and run the query against that view
-func (d *DuckDB) QueryFrames(name string, query string, frames []*sdk.Frame) (string, error) {
-	dirs, err := data.ToParquet(frames, d.chunk)
-	if err != nil {
-		logger.Error("error converting to parquet", "error", err)
-		return "", err
+func (d *DuckDB) QueryFrames(name string, query string, frames []*sdk.Frame) (string, bool, error) {
+	data := FrameData{
+		cacheDuration: d.cacheDuration,
+		cache:         &d.cache,
+		db:            d,
 	}
 
-	defer func() {
-		for _, dir := range dirs {
-			err := os.RemoveAll(dir)
-			if err != nil {
-				logger.Error("failed to remove parquet files", "error", err)
-			}
-		}
-	}()
-
-	commands := []string{}
-	created := map[string]bool{}
-	logger.Debug("starting to create views from frames", "frames", len(frames))
-	for _, frame := range frames {
-		if created[frame.RefID] {
-			continue
-		}
-		cmd := fmt.Sprintf("CREATE VIEW %s AS (SELECT * from '%s/*.parquet');", frame.RefID, dirs[frame.RefID])
-		logger.Debug("creating view", "cmd", cmd)
-		commands = append(commands, cmd)
-		created[frame.RefID] = true
-	}
+	return data.Query(name, query, frames)
+}
 
-	commands = append(commands, query)
-	res, err := d.RunCommands(commands)
-	if err != nil {
-		logger.Error("error running commands", "error", err)
-		return "", err
+func wipe(dirs map[string]string) {
+	for _, dir := range dirs {
+		err := os.RemoveAll(dir)
+		if err != nil {
+			logger.Error("failed to remove parquet files", "error", err)
+		}
 	}
-	return res, nil
 }
 
 func (d *DuckDB) QueryFramesInto(name string, query string, frames []*sdk.Frame, f *sdk.Frame) error {
-	res, err := d.QueryFrames(name, query, frames)
+	res, cached, err := d.QueryFrames(name, query, frames)
 	if err != nil {
 		return err
 	}
 
-	return resultsToFrame(name, res, f, frames)
+	err = resultsToFrame(name, res, f, frames)
+	if err != nil {
+		return err
+	}
+	if cached {
+		for _, frame := range frames {
+			if frame.Meta == nil {
+				frame.Meta = &sdk.FrameMeta{}
+			}
+			notice := sdk.Notice{
+				Severity: sdk.NoticeSeverityInfo,
+				Text:     "Data retrieved from cache",
+			}
+			frame.Meta.Notices = append(frame.Meta.Notices, notice)
+		}
+	}
+	return nil
 }
 
 // Destroy will remove database files created by duckdb

duck/duckdb_test.go

@@ -52,12 +52,42 @@ func TestQueryFrame(t *testing.T) {
 	frame.RefID = "foo"
 	frames := []*data.Frame{frame}
 
-	res, err := db.QueryFrames("foo", "select * from foo", frames)
+	res, _, err := db.QueryFrames("foo", "select * from foo", frames)
 	assert.Nil(t, err)
 
 	assert.Contains(t, res, `[{"value":"test"}]`)
 }
 
+func TestQueryFrameCache(t *testing.T) {
+	opts := Opts{
+		CacheDuration: 5,
+	}
+	db := NewInMemoryDB(opts)
+
+	var values = []string{"test"}
+	frame := data.NewFrame("foo", data.NewField("value", nil, values))
+	frame.RefID = "foo"
+	frames := []*data.Frame{frame}
+
+	res, cached, err := db.QueryFrames("foo", "select * from foo", frames)
+	assert.Nil(t, err)
+	assert.False(t, cached)
+	assert.Contains(t, res, `[{"value":"test"}]`)
+
+	res, cached, err = db.QueryFrames("foo", "select * from foo", frames)
+	assert.Nil(t, err)
+	assert.True(t, cached)
+	assert.Contains(t, res, `[{"value":"test"}]`)
+
+	// wait for cache to expire
+	time.Sleep(6 * time.Second)
+
+	res, cached, err = db.QueryFrames("foo", "select * from foo", frames)
+	assert.Nil(t, err)
+	assert.False(t, cached)
+	assert.Contains(t, res, `[{"value":"test"}]`)
+}
+
 func TestQueryFrameWithDisplayName(t *testing.T) {
 
 	db := NewInMemoryDB()
@@ -71,7 +101,7 @@ func TestQueryFrameWithDisplayName(t *testing.T) {
 	frame.RefID = "foo"
 	frames := []*data.Frame{frame}
 
-	res, err := db.QueryFrames("foo", "select * from foo", frames)
+	res, _, err := db.QueryFrames("foo", "select * from foo", frames)
 	assert.Nil(t, err)
 
 	assert.Contains(t, res, `[{"some value":"test"}]`)
@@ -88,7 +118,7 @@ func TestQueryFrameChunks(t *testing.T) {
 	frame.RefID = "foo"
 	frames := []*data.Frame{frame}
 
-	res, err := db.QueryFrames("foo", "select * from foo", frames)
+	res, _, err := db.QueryFrames("foo", "select * from foo", frames)
 	assert.Nil(t, err)
 
 	assert.Contains(t, res, `test2`)

duck/frame-data.go

@@ -0,0 +1,151 @@
+package duck
+
+import (
+	"fmt"
+	"sync"
+	"time"
+
+	sdk "github.com/grafana/grafana-plugin-sdk-go/data"
+	"github.com/scottlepp/go-duck/duck/data"
+)
+
+type FrameData struct {
+	cacheDuration int
+	cache         *cache
+	db            *DuckDB
+}
+
+func (f *FrameData) Query(name string, query string, frames []*sdk.Frame) (string, bool, error) {
+	dirs, cached, err := f.data(name, query, frames)
+	if err != nil {
+		logger.Error("error converting to parquet", "error", err)
+		return "", cached, err
+	}
+
+	defer f.postProcess(name, query, dirs, cached)
+
+	// create a wait group to wait for the query to finish
+	// if the cache duration is exceeded, wait before deleting the cache ( parquet files )
+	var wg sync.WaitGroup
+	wg.Add(1)
+	f.cache.setWait(fmt.Sprintf("%s:%s", name, query), &wg)
+
+	var res string
+	var qerr error
+
+	go func() {
+		res, qerr = f.runQuery(query, dirs, frames)
+		wg.Done()
+	}()
+
+	wg.Wait()
+	f.cache.deleteWait(fmt.Sprintf("%s:%s", name, query))
+
+	if qerr != nil {
+		logger.Error("error running commands", "error", err)
+		return "", cached, err
+	}
+
+	key := fmt.Sprintf("%s:%s", name, query)
+	if f.cacheDuration > 0 && !cached {
+		f.cache.set(key, dirs)
+	}
+
+	return res, cached, nil
+}
+
+func (f *FrameData) runQuery(query string, dirs Dirs, frames []*sdk.Frame) (string, error) {
+	commands := createViews(frames, dirs)
+	commands = append(commands, query)
+	return f.db.RunCommands(commands)
+}
+
+func createViews(frames []*sdk.Frame, dirs Dirs) []string {
+	commands := []string{}
+	created := map[string]bool{}
+	logger.Debug("starting to create views from frames", "frames", len(frames))
+	for _, frame := range frames {
+		if created[frame.RefID] {
+			continue
+		}
+		cmd := fmt.Sprintf("CREATE VIEW %s AS (SELECT * from '%s/*.parquet');", frame.RefID, dirs[frame.RefID])
+		logger.Debug("creating view", "cmd", cmd)
+		commands = append(commands, cmd)
+		created[frame.RefID] = true
+	}
+	return commands
+}
+
+func (f *FrameData) data(name string, query string, frames []*sdk.Frame) (Dirs, bool, error) {
+	if f.cacheDuration > 0 {
+		// check the cache
+		key := fmt.Sprintf("%s:%s", name, query)
+		if d, ok := f.cache.get(key); ok {
+			return d, true, nil
+		}
+	}
+
+	dirs, err := data.ToParquet(frames, f.db.chunk)
+	return dirs, false, err
+}
+
+func (f *FrameData) postProcess(name string, query string, dirs Dirs, cached bool) {
+	go func() {
+		if f.cacheDuration == 0 {
+			wipe(dirs)
+			return
+		}
+		// if result is not cached, a new cache entry will be created
+		// delete the new cache entry after cacheDuration
+		if !cached {
+			time.Sleep(time.Duration(f.cacheDuration) * time.Second)
+			key := fmt.Sprintf("%s:%s", name, query)
+			f.cache.delete(key)
+			// if the query is running wait for the query to finish before deleting the parquet files
+			wg, wait := f.cache.getWait(key)
+			if wait {
+				wg.Wait()
+				wipe(dirs)
+				return
+			}
+			wipe(dirs)
+		}
+	}()
+}
+
+type cache struct {
+	store sync.Map
+	wait  sync.Map
+}
+
+func (c *cache) set(key string, value Dirs) {
+	c.store.Store(key, value)
+}
+
+func (c *cache) get(key string) (Dirs, bool) {
+	val, ok := c.store.Load(key)
+	if !ok {
+		return nil, false
+	}
+	return val.(Dirs), true
+}
+
+func (c *cache) delete(key string) {
+	c.store.Delete(key)
+}
+
+func (c *cache) setWait(key string, value *sync.WaitGroup) {
+	c.wait.Store(key, value)
+}
+
+func (c *cache) getWait(key string) (*sync.WaitGroup, bool) {
+	val, ok := c.wait.Load(key)
+	if !ok {
+		return nil, false
+	}
+	return val.(*sync.WaitGroup), true
+}
+
+func (c *cache) deleteWait(key string) {
+	c.wait.Delete(key)
+}