Expressions: Sql expressions with Duckdb (#81666)

duckdb temp storage of dataframes using parquet and querying from sql expressions
---------

Co-authored-by: Ryan McKinley <ryantxu@gmail.com>

pkg/expr/commands.go

@@ -324,6 +324,8 @@ const (
 	TypeClassicConditions
 	// TypeThreshold is the CMDType for checking if a threshold has been crossed
 	TypeThreshold
+	// TypeSQL is the CMDType for running SQL expressions
+	TypeSQL
 )
 
 func (gt CommandType) String() string {
@@ -336,6 +338,8 @@ func (gt CommandType) String() string {
 		return "resample"
 	case TypeClassicConditions:
 		return "classic_conditions"
+	case TypeSQL:
+		return "sql"
 	default:
 		return "unknown"
 	}
@@ -354,6 +358,8 @@ func ParseCommandType(s string) (CommandType, error) {
 		return TypeClassicConditions, nil
 	case "threshold":
 		return TypeThreshold, nil
+	case "sql":
+		return TypeSQL, nil
 	default:
 		return TypeUnknown, fmt.Errorf("'%v' is not a recognized expression type", s)
 	}

pkg/expr/graph.go

@@ -75,6 +75,8 @@ func (dp *DataPipeline) execute(c context.Context, now time.Time, s *Service) (m
 		executeDSNodesGrouped(c, now, vars, s, dsNodes)
 	}
 
+	s.allowLongFrames = hasSqlExpression(*dp)
+
 	for _, node := range *dp {
 		if groupByDSFlag && node.NodeType() == TypeDatasourceNode {
 			continue // already executed via executeDSNodesGrouped
@@ -266,6 +268,10 @@ func buildGraphEdges(dp *simple.DirectedGraph, registry map[string]Node) error {
 		for _, neededVar := range cmdNode.Command.NeedsVars() {
 			neededNode, ok := registry[neededVar]
 			if !ok {
+				_, ok := cmdNode.Command.(*SQLCommand)
+				if ok {
+					continue
+				}
 				return fmt.Errorf("unable to find dependent node '%v'", neededVar)
 			}
 
@@ -312,3 +318,37 @@ func GetCommandsFromPipeline[T Command](pipeline DataPipeline) []T {
 	}
 	return results
 }
+
+func hasSqlExpression(dp DataPipeline) bool {
+	for _, node := range dp {
+		if node.NodeType() == TypeCMDNode {
+			cmdNode := node.(*CMDNode)
+			_, ok := cmdNode.Command.(*SQLCommand)
+			if ok {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+// func graphHasSqlExpresssion(dp *simple.DirectedGraph) bool {
+// 	node := dp.Nodes()
+// 	for node.Next() {
+// 		if cmdNode, ok := node.Node().(*CMDNode); ok {
+// 			// res[dpNode.RefID()] = dpNode
+// 			_, ok := cmdNode.Command.(*SQLCommand)
+// 			if ok {
+// 				return true
+// 			}
+// 		}
+// 		// if node.NodeType() == TypeCMDNode {
+// 		// 	cmdNode := node.(*CMDNode)
+// 		// 	_, ok := cmdNode.Command.(*SQLCommand)
+// 		// 	if ok {
+// 		// 		return true
+// 		// 	}
+// 		// }
+// 	}
+// 	return false
+// }

pkg/expr/mathexp/parse/node.go

@@ -405,6 +405,8 @@ const (
 	TypeVariantSet
 	// TypeNoData is a no data response without a known data type.
 	TypeNoData
+	// TypeTableData is a tabular data response.
+	TypeTableData
 )
 
 // String returns a string representation of the ReturnType.
@@ -422,6 +424,8 @@ func (f ReturnType) String() string {
 		return "variant"
 	case TypeNoData:
 		return "noData"
+	case TypeTableData:
+		return "tableData"
 	default:
 		return "unknown"
 	}

pkg/expr/mathexp/types.go

@@ -246,3 +246,48 @@ func (s NoData) New() NoData {
 func NewNoData() NoData {
 	return NoData{data.NewFrame("no data")}
 }
+
+// TableData is an untyped no data response.
+type TableData struct{ Frame *data.Frame }
+
+// Type returns the Value type and allows it to fulfill the Value interface.
+func (s TableData) Type() parse.ReturnType { return parse.TypeTableData }
+
+// Value returns the actual value allows it to fulfill the Value interface.
+func (s TableData) Value() any { return s }
+
+func (s TableData) GetLabels() data.Labels { return nil }
+
+func (s TableData) SetLabels(ls data.Labels) {}
+
+func (s TableData) GetMeta() any {
+	return s.Frame.Meta.Custom
+}
+
+func (s TableData) SetMeta(v any) {
+	m := s.Frame.Meta
+	if m == nil {
+		m = &data.FrameMeta{}
+		s.Frame.SetMeta(m)
+	}
+	m.Custom = v
+}
+
+func (s TableData) AddNotice(notice data.Notice) {
+	m := s.Frame.Meta
+	if m == nil {
+		m = &data.FrameMeta{}
+		s.Frame.SetMeta(m)
+	}
+	m.Notices = append(m.Notices, notice)
+}
+
+func (s TableData) AsDataFrame() *data.Frame { return s.Frame }
+
+func (s TableData) New() TableData {
+	return NewTableData()
+}
+
+func NewTableData() TableData {
+	return TableData{data.NewFrame("")}
+}

pkg/expr/models.go

@@ -24,6 +24,9 @@ const (
 
 	// Threshold
 	QueryTypeThreshold QueryType = "threshold"
+
+	// SQL query via DuckDB
+	QueryTypeSQL QueryType = "sql"
 )
 
 type MathQuery struct {
@@ -69,6 +72,11 @@ type ClassicQuery struct {
 	Conditions []classic.ConditionJSON `json:"conditions"`
 }
 
+// SQLQuery requires the sqlExpression feature flag
+type SQLExpression struct {
+	Expression string `json:"expression" jsonschema:"minLength=1,example=SELECT * FROM A LIMIT 1"`
+}
+
 //-------------------------------
 // Non-query commands
 //-------------------------------

pkg/expr/nodes.go

@@ -155,6 +155,8 @@ func buildCMDNode(rn *rawNode, toggles featuremgmt.FeatureToggles) (*CMDNode, er
 		node.Command, err = classic.UnmarshalConditionsCmd(rn.Query, rn.RefID)
 	case TypeThreshold:
 		node.Command, err = UnmarshalThresholdCommand(rn, toggles)
+	case TypeSQL:
+		node.Command, err = UnmarshalSQLCommand(rn)
 	default:
 		return nil, fmt.Errorf("expression command type '%v' in expression '%v' not implemented", commandType, rn.RefID)
 	}
@@ -471,7 +473,8 @@ func convertDataFramesToResults(ctx context.Context, frames data.Frames, datasou
 			logger.Warn("Ignoring InfluxDB data frame due to missing numeric fields")
 			continue
 		}
-		if schema.Type != data.TimeSeriesTypeWide {
+
+		if schema.Type != data.TimeSeriesTypeWide && !s.allowLongFrames {
 			return "", mathexp.Results{}, fmt.Errorf("input data must be a wide series but got type %s (input refid)", schema.Type)
 		}
 		filtered = append(filtered, frame)
@@ -484,20 +487,29 @@ func convertDataFramesToResults(ctx context.Context, frames data.Frames, datasou
 
 	maybeFixerFn := checkIfSeriesNeedToBeFixed(filtered, datasourceType)
 
-	vals := make([]mathexp.Value, 0, totalLen)
-	for _, frame := range filtered {
-		series, err := WideToMany(frame, maybeFixerFn)
-		if err != nil {
-			return "", mathexp.Results{}, err
-		}
-		for _, ser := range series {
-			vals = append(vals, ser)
-		}
-	}
 	dataType := "single frame series"
 	if len(filtered) > 1 {
 		dataType = "multi frame series"
 	}
+
+	vals := make([]mathexp.Value, 0, totalLen)
+	for _, frame := range filtered {
+		schema := frame.TimeSeriesSchema()
+		if schema.Type == data.TimeSeriesTypeWide {
+			series, err := WideToMany(frame, maybeFixerFn)
+			if err != nil {
+				return "", mathexp.Results{}, err
+			}
+			for _, ser := range series {
+				vals = append(vals, ser)
+			}
+		} else {
+			v := mathexp.TableData{Frame: frame}
+			vals = append(vals, v)
+			dataType = "single frame"
+		}
+	}
+
 	return dataType, mathexp.Results{
 		Values: vals,
 	}, nil

pkg/expr/reader.go

@@ -30,6 +30,7 @@ func NewExpressionQueryReader(features featuremgmt.FeatureToggles) (*ExpressionQ
 }
 
 // ReadQuery implements query.TypedQueryHandler.
+// nolint:gocyclo
 func (h *ExpressionQueryReader) ReadQuery(
 	// Properties that have been parsed off the same node
 	common *rawNode, // common query.CommonQueryProperties
@@ -102,6 +103,13 @@ func (h *ExpressionQueryReader) ReadQuery(
 			eq.Command, err = classic.NewConditionCmd(common.RefID, q.Conditions)
 		}
 
+	case QueryTypeSQL:
+		q := &SQLExpression{}
+		err = iter.ReadVal(q)
+		if err == nil {
+			eq.Command, err = NewSQLCommand(common.RefID, q.Expression, common.TimeRange)
+		}
+
 	case QueryTypeThreshold:
 		q := &ThresholdQuery{}
 		err = iter.ReadVal(q)

pkg/expr/service.go

@@ -63,8 +63,9 @@ type Service struct {
 
 	pluginsClient backend.CallResourceHandler
 
-	tracer  tracing.Tracer
-	metrics *metrics
+	tracer          tracing.Tracer
+	metrics         *metrics
+	allowLongFrames bool
 }
 
 type pluginContextProvider interface {

pkg/expr/sql/parser.go

@@ -0,0 +1,99 @@
+package sql
+
+import (
+	"errors"
+	"strings"
+
+	parser "github.com/krasun/gosqlparser"
+	"github.com/xwb1989/sqlparser"
+)
+
+// TablesList returns a list of tables for the sql statement
+func TablesList(rawSQL string) ([]string, error) {
+	stmt, err := sqlparser.Parse(rawSQL)
+	if err != nil {
+		tables, err := parse(rawSQL)
+		if err != nil {
+			return parseTables(rawSQL)
+		}
+		return tables, nil
+	}
+
+	tables := []string{}
+	switch kind := stmt.(type) {
+	case *sqlparser.Select:
+		for _, t := range kind.From {
+			buf := sqlparser.NewTrackedBuffer(nil)
+			t.Format(buf)
+			table := buf.String()
+			if table != "dual" {
+				tables = append(tables, buf.String())
+			}
+		}
+	default:
+		return nil, errors.New("not a select statement")
+	}
+	return tables, nil
+}
+
+// uses a simple tokenizer
+func parse(rawSQL string) ([]string, error) {
+	query, err := parser.Parse(rawSQL)
+	if err != nil {
+		return nil, err
+	}
+	if query.GetType() == parser.StatementSelect {
+		sel, ok := query.(*parser.Select)
+		if ok {
+			return []string{sel.Table}, nil
+		}
+	}
+	return nil, err
+}
+
+func parseTables(rawSQL string) ([]string, error) {
+	checkSql := strings.ToUpper(rawSQL)
+	if strings.HasPrefix(checkSql, "SELECT") || strings.HasPrefix(rawSQL, "WITH") {
+		tables := []string{}
+		tokens := strings.Split(rawSQL, " ")
+		checkNext := false
+		takeNext := false
+		for _, t := range tokens {
+			t = strings.ToUpper(t)
+			t = strings.TrimSpace(t)
+
+			if takeNext {
+				tables = append(tables, t)
+				checkNext = false
+				takeNext = false
+				continue
+			}
+			if checkNext {
+				if strings.Contains(t, "(") {
+					checkNext = false
+					continue
+				}
+				if strings.Contains(t, ",") {
+					values := strings.Split(t, ",")
+					for _, v := range values {
+						v := strings.TrimSpace(v)
+						if v != "" {
+							tables = append(tables, v)
+						} else {
+							takeNext = true
+							break
+						}
+					}
+					continue
+				}
+				tables = append(tables, t)
+				checkNext = false
+			}
+			if t == "FROM" {
+				checkNext = true
+			}
+		}
+		return tables, nil
+	}
+	return nil, errors.New("not a select statement")
+}

pkg/expr/sql/parser_test.go

@@ -0,0 +1,58 @@
+package sql
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestParse(t *testing.T) {
+	sql := "select * from foo"
+	tables, err := parseTables((sql))
+	assert.Nil(t, err)
+
+	assert.Equal(t, "FOO", tables[0])
+}
+
+func TestParseWithComma(t *testing.T) {
+	sql := "select * from foo,bar"
+	tables, err := parseTables((sql))
+	assert.Nil(t, err)
+
+	assert.Equal(t, "FOO", tables[0])
+	assert.Equal(t, "BAR", tables[1])
+}
+
+func TestParseWithCommas(t *testing.T) {
+	sql := "select * from foo,bar,baz"
+	tables, err := parseTables((sql))
+	assert.Nil(t, err)
+
+	assert.Equal(t, "FOO", tables[0])
+	assert.Equal(t, "BAR", tables[1])
+	assert.Equal(t, "BAZ", tables[2])
+}
+
+func TestArray(t *testing.T) {
+	sql := "SELECT array_value(1, 2, 3)"
+	tables, err := TablesList((sql))
+	assert.Nil(t, err)
+
+	assert.Equal(t, 0, len(tables))
+}
+
+func TestArray2(t *testing.T) {
+	sql := "SELECT array_value(1, 2, 3)[2]"
+	tables, err := TablesList((sql))
+	assert.Nil(t, err)
+
+	assert.Equal(t, 0, len(tables))
+}
+
+func TestXxx(t *testing.T) {
+	sql := "SELECT [3, 2, 1]::INT[3];"
+	tables, err := TablesList((sql))
+	assert.Nil(t, err)
+
+	assert.Equal(t, 0, len(tables))
+}

pkg/expr/sql_command.go

@@ -0,0 +1,107 @@
+package expr
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/grafana/grafana-plugin-sdk-go/data"
+	"github.com/scottlepp/go-duck/duck"
+
+	"github.com/grafana/grafana/pkg/expr/mathexp"
+	"github.com/grafana/grafana/pkg/expr/sql"
+	"github.com/grafana/grafana/pkg/infra/tracing"
+	"github.com/grafana/grafana/pkg/util/errutil"
+)
+
+// SQLCommand is an expression to run SQL over results
+type SQLCommand struct {
+	query       string
+	varsToQuery []string
+	timeRange   TimeRange
+	refID       string
+}
+
+// NewSQLCommand creates a new SQLCommand.
+func NewSQLCommand(refID, rawSQL string, tr TimeRange) (*SQLCommand, error) {
+	if rawSQL == "" {
+		return nil, errutil.BadRequest("sql-missing-query",
+			errutil.WithPublicMessage("missing SQL query"))
+	}
+	tables, err := sql.TablesList(rawSQL)
+	if err != nil {
+		logger.Warn("invalid sql query", "sql", rawSQL, "error", err)
+		return nil, errutil.BadRequest("sql-invalid-sql",
+			errutil.WithPublicMessage("error reading SQL command"),
+		)
+	}
+	return &SQLCommand{
+		query:       rawSQL,
+		varsToQuery: tables,
+		timeRange:   tr,
+		refID:       refID,
+	}, nil
+}
+
+// UnmarshalSQLCommand creates a SQLCommand from Grafana's frontend query.
+func UnmarshalSQLCommand(rn *rawNode) (*SQLCommand, error) {
+	if rn.TimeRange == nil {
+		return nil, fmt.Errorf("time range must be specified for refID %s", rn.RefID)
+	}
+
+	expressionRaw, ok := rn.Query["expression"]
+	if !ok {
+		return nil, errors.New("no expression in the query")
+	}
+	expression, ok := expressionRaw.(string)
+	if !ok {
+		return nil, fmt.Errorf("expected sql expression to be type string, but got type %T", expressionRaw)
+	}
+
+	return NewSQLCommand(rn.RefID, expression, rn.TimeRange)
+}
+
+// NeedsVars returns the variable names (refIds) that are dependencies
+// to execute the command and allows the command to fulfill the Command interface.
+func (gr *SQLCommand) NeedsVars() []string {
+	return gr.varsToQuery
+}
+
+// Execute runs the command and returns the results or an error if the command
+// failed to execute.
+func (gr *SQLCommand) Execute(ctx context.Context, now time.Time, vars mathexp.Vars, tracer tracing.Tracer) (mathexp.Results, error) {
+	_, span := tracer.Start(ctx, "SSE.ExecuteSQL")
+	defer span.End()
+
+	allFrames := []*data.Frame{}
+	for _, ref := range gr.varsToQuery {
+		results := vars[ref]
+		frames := results.Values.AsDataFrames(ref)
+		allFrames = append(allFrames, frames...)
+	}
+
+	rsp := mathexp.Results{}
+
+	duckDB := duck.NewInMemoryDB()
+	var frame = &data.Frame{}
+	err := duckDB.QueryFramesInto(gr.refID, gr.query, allFrames, frame)
+	if err != nil {
+		rsp.Error = err
+		return rsp, nil
+	}
+
+	frame.RefID = gr.refID
+
+	if frame.Rows() == 0 {
+		rsp.Values = mathexp.Values{
+			mathexp.NoData{Frame: frame},
+		}
+	}
+
+	rsp.Values = mathexp.Values{
+		mathexp.TableData{Frame: frame},
+	}
+
+	return rsp, nil
+}

pkg/expr/sql_command_test.go

@@ -0,0 +1,26 @@
+package expr
+
+import (
+	"strings"
+	"testing"
+)
+
+func TestNewCommand(t *testing.T) {
+	cmd, err := NewSQLCommand("a", "select a from foo, bar", nil)
+	if err != nil && strings.Contains(err.Error(), "feature is not enabled") {
+		return
+	}
+
+	if err != nil {
+		t.Fail()
+		return
+	}
+
+	for _, v := range cmd.varsToQuery {
+		if strings.Contains("foo bar", v) {
+			continue
+		}
+		t.Fail()
+		return
+	}
+}