featurebase/dax/queryer/queryer.go

// Package queryer provides the core query-related structs.
package queryer

import (
	"context"
	"fmt"
	"net/http"
	"strings"
	"sync"
	"time"

	featurebase "github.com/molecula/featurebase/v3"
	fbcontext "github.com/molecula/featurebase/v3/context"
	"github.com/molecula/featurebase/v3/dax"
	"github.com/molecula/featurebase/v3/encoding/proto"
	"github.com/molecula/featurebase/v3/errors"
	idkmds "github.com/molecula/featurebase/v3/idk/mds"
	"github.com/molecula/featurebase/v3/logger"
	featurebase_pql "github.com/molecula/featurebase/v3/pql"
	fbproto "github.com/molecula/featurebase/v3/proto"
	"github.com/molecula/featurebase/v3/server"
	"github.com/molecula/featurebase/v3/sql3/parser"
	"github.com/molecula/featurebase/v3/sql3/planner"
	plannertypes "github.com/molecula/featurebase/v3/sql3/planner/types"
	"github.com/molecula/featurebase/v3/stats"
	"github.com/molecula/featurebase/v3/systemlayer"
	uuid "github.com/satori/go.uuid"
)

// Queryer represents the query layer in a Molecula implementation. The idea is
// that the externally-facing Molecula API would proxy query requests to a pool
// of "Queryer" nodes, which handle incoming query requests.
type Queryer struct {
	mu            sync.RWMutex
	orchestrators map[dax.TableQualifier]*qualifiedOrchestrator

	fbClient *featurebase.InternalClient

	noder   dax.Noder
	schemar dax.Schemar

	logger logger.Logger
}

// New returns a new instance of Queryer.
func New(cfg Config) *Queryer {
	q := &Queryer{
		noder:         dax.NewNopNoder(),
		schemar:       dax.NewNopSchemar(),
		orchestrators: make(map[dax.TableQualifier]*qualifiedOrchestrator),
		logger:        logger.NopLogger,
	}

	if cfg.Logger != nil {
		q.logger = cfg.Logger
	}

	return q
}

// Orchestrator gets (or creates) an instance of qualifiedOrchestrator based on
// the provided dax.TableQualifier.
func (q *Queryer) Orchestrator(qual dax.TableQualifier) *qualifiedOrchestrator {
	// Try to get orchestrator under a read lock first.
	if orch := func() *qualifiedOrchestrator {
		q.mu.RLock()
		defer q.mu.RUnlock()
		if orch, ok := q.orchestrators[qual]; ok {
			return orch
		}
		return nil
	}(); orch != nil {
		return orch
	}

	// Since we didn't find an orchestrator under a read lock, obtain a write
	// lock and try a read/write.
	q.mu.Lock()
	defer q.mu.Unlock()
	if orch, ok := q.orchestrators[qual]; ok {
		return orch
	}

	sapi := newQualifiedSchemaAPI(qual, q.schemar)

	orch := &orchestrator{
		schema:   sapi,
		trans:    NewMDSTranslator(q.noder, q.schemar),
		topology: &MDSTopology{noder: q.noder},
		// TODO(jaffee) using default http.Client probably bad... need to set some timeouts.
		client: q.fbClient,
		stats:  stats.NopStatsClient,
		logger: q.logger,
	}

	qorch := newQualifiedOrchestrator(orch, qual)
	q.orchestrators[qual] = qorch

	return qorch
}

func (q *Queryer) SetNoder(noder dax.Noder) error {
	q.noder = noder
	return nil
}

func (q *Queryer) SetSchemar(schemar dax.Schemar) error {
	q.schemar = schemar
	return nil
}

func (q *Queryer) Start() error {
	if q.noder == nil {
		return errors.New(errors.ErrUncoded, "queryer requires noder to be configured")
	} else if q.schemar == nil {
		return errors.New(errors.ErrUncoded, "queryer requires schemar to be configured")
	}

	// fbClient is an instance of internal client. It's used in one place in the
	// orchestrator (o.client.QueryNode()), but in that case, the host is
	// replaces with the actual host (another computer node) to connect to.
	// That's why we set it up with a dummy host here.
	fbClient, err := featurebase.NewInternalClient("fakehostname:8080",
		&http.Client{},
		featurebase.WithSerializer(proto.Serializer{}),
		featurebase.WithPathPrefix("should-not-be-used"),
	)
	if err != nil {
		return errors.Wrap(err, "setting up internal client")
	}
	q.fbClient = fbClient

	return nil
}

func (q *Queryer) QuerySQL(ctx context.Context, qual dax.TableQualifier, sql string) (*featurebase.WireQueryResponse, error) {
	start := time.Now()

	if len(sql) > 0 && sql[0] == '[' {
		return q.parseAndQueryPQL(ctx, qual, sql)
	}
	ret := &featurebase.WireQueryResponse{}

	applyExecutionTime := func() {
		ret.ExecutionTime = time.Since(start).Microseconds()
	}

	applyError := func(e error) {
		ret.Error = e.Error()
		applyExecutionTime()
	}

	// Create a requestID and add it to the context.
	requestID, err := uuid.NewV4()
	if err != nil {
		applyError(errors.Wrap(err, "creating requestID"))
		return ret, nil
	}
	// put the requestId in the context
	ctx = fbcontext.WithRequestID(ctx, requestID.String())

	st, err := parser.NewParser(strings.NewReader(sql)).ParseStatement()
	if err != nil {
		applyError(errors.Wrap(err, "parsing sql"))
		return ret, nil
	}

	// SchemaAPI
	sapi := newQualifiedSchemaAPI(qual, q.schemar)

	// Importer
	imp := idkmds.NewImporter(q.noder, q.schemar, qual, nil)

	// TODO(tlt): this obviously doesn't work; we don't have an API here. We
	// need a dax-compatible implementation of the SystemAPI (or at least a
	// no-op implementation).
	sysapi := &featurebase.FeatureBaseSystemAPI{API: nil}

	systemLayer := systemlayer.NewSystemLayer()

	pl := planner.NewExecutionPlanner(q.Orchestrator(qual), sapi, sysapi, systemLayer, imp, q.logger, sql)

	planOp, err := pl.CompilePlan(ctx, st)
	if err != nil {
		applyError(errors.Wrap(err, "compiling plan"))
		return ret, nil
	}

	// Get a query iterator.
	iter, err := planOp.Iterator(ctx, nil)
	if err != nil {
		applyError(errors.Wrap(err, "getting iterator"))
		return ret, nil
	}

	// Read schema.
	columns := planOp.Schema()
	schema := featurebase.WireQuerySchema{
		Fields: make([]*featurebase.WireQueryField, len(columns)),
	}
	for i, col := range columns {
		btype, err := dax.BaseTypeFromString(col.Type.BaseTypeName())
		if err != nil {
			applyError(errors.Wrap(err, "getting fieldtype from string"))
			return ret, nil
		}
		schema.Fields[i] = &featurebase.WireQueryField{
			Name:     dax.FieldName(col.ColumnName),
			Type:     strings.ToLower(col.Type.TypeDescription()), // TODO(tlt): remove this once sql3 uses BaseTypes.
			BaseType: btype,
			TypeInfo: col.Type.TypeInfo(),
		}
	}

	// Read rows.
	data := make([][]interface{}, 0)
	var currentRow plannertypes.Row
	for currentRow, err = iter.Next(ctx); err == nil; currentRow, err = iter.Next(ctx) {
		data = append(data, currentRow)
	}
	if err != nil && err != plannertypes.ErrNoMoreRows {
		applyError(errors.Wrap(err, "getting row"))
		return ret, nil
	}

	ret.Schema = schema
	ret.Data = data
	applyExecutionTime()

	return ret, nil
}

func (q *Queryer) parseAndQueryPQL(ctx context.Context, qual dax.TableQualifier, sql string) (*featurebase.WireQueryResponse, error) {
	var i int
	for i = 1; sql[i] != ']'; i++ {
		if i == len(sql)-1 {
			return nil, errors.Errorf("couldn't parse table name out of '%s'", sql)
		}
	}
	table := sql[1:i]
	query := sql[i+1:]
	fmt.Println("got table/query", table, query)

	return q.QueryPQL(ctx, qual, dax.TableName(table), query)
}

// convertIndex tries to covert any "index" specified in the call.Args map to a
// TableKeyer. Note, since the Call.CallIndex() method currently only looks for
// strings, we can't just set the value to a TableKeyer; we have to set it to
// the equivalent string and then parse it back out later. A TODO would be to
// modify Call.CallIndex() to be TableKeyer aware. I didn't do that along with
// these changes because I'm not sure if we want to introduce dax types into the
// pql package.
func (q *Queryer) convertIndex(ctx context.Context, qual dax.TableQualifier, call *featurebase_pql.Call) {
	if index := call.CallIndex(); index != "" {
		qtbl, err := q.schemar.TableByName(ctx, qual, dax.TableName(index))
		if err != nil {
			return
		}
		call.Args["index"] = string(qtbl.Key())
	}

	// Apply to children.
	for _, child := range call.Children {
		q.convertIndex(ctx, qual, child)
	}
}

func (q *Queryer) QueryPQL(ctx context.Context, qual dax.TableQualifier, table dax.TableName, pql string) (*featurebase.WireQueryResponse, error) {
	// Parse the pql into a pql.Query containing []pql.Call.
	qry, err := featurebase_pql.NewParser(strings.NewReader(pql)).Parse()
	if err != nil {
		return nil, errors.Wrap(err, "parsing pql")
	}
	if len(qry.Calls) != 1 {
		return nil, errors.Errorf("must have exactly 1 query, but got: %+v", qry.Calls)
	}

	// Replace any "index" arguments within the PQL with a TableKey.
	q.convertIndex(ctx, qual, qry.Calls[0])

	qtbl, err := q.schemar.TableByName(ctx, qual, dax.TableName(table))
	if err != nil {
		return nil, errors.Wrap(err, "converting index to qualified table")
	}

	results, err := q.Orchestrator(qual).Execute(ctx, qtbl, qry, nil, &featurebase.ExecOptions{})
	if err != nil {
		return nil, errors.Wrap(err, "orchestrator.Execute")
	}
	if len(results.Results) != 1 {
		return nil, errors.Errorf("expected single result but got %+v", results.Results)
	}

	return pqlResultToQueryResult(results.Results[0])
}

func pqlResultToQueryResult(pqlResult interface{}) (*featurebase.WireQueryResponse, error) {
	toTabler, err := server.ToTablerWrapper(pqlResult)
	if err != nil {
		return nil, errors.Wrap(err, "wrapping as type ToTabler")
	}
	table, err := toTabler.ToTable()
	if err != nil {
		return nil, errors.Wrap(err, "ToTable")
	}

	return tableResponseToQueryResult(table)
}

func tableResponseToQueryResult(t *fbproto.TableResponse) (*featurebase.WireQueryResponse, error) {
	qr := &featurebase.WireQueryResponse{
		Schema: featurebase.WireQuerySchema{Fields: make([]*featurebase.WireQueryField, len(t.Headers))},
		Data:   make([][]interface{}, len(t.Rows)),
	}
	for i, ci := range t.Headers {
		qr.Schema.Fields[i] = &featurebase.WireQueryField{
			Name:     dax.FieldName(ci.Name),
			Type:     string(datatypeToBaseType(ci.Datatype)), // TODO(tlt): this doesn't contain typeInfo
			BaseType: datatypeToBaseType(ci.Datatype),
		}
	}

	for i, row := range t.Rows {
		qr.Data[i] = rowToSliceInterface(t.Headers, row)
	}

	return qr, nil
}

func datatypeToBaseType(ciDatatype string) dax.BaseType {
	switch ciDatatype {
	case "string":
		return dax.BaseTypeString
	case "uint64":
		return dax.BaseTypeID
	case "float64":
		// ??
		panic("float64 doesn't have sql3 field type?")
	case "int64":
		return dax.BaseTypeInt
	case "bool":
		return dax.BaseTypeBool
	case "decimal":
		return dax.BaseTypeDecimal
	case "timestamp":
		return dax.BaseTypeTimestamp
	case "[]string":
		return dax.BaseTypeStringSet
	case "[]uint64":
		return dax.BaseTypeIDSet
	// TODO []byte??
	default:
		panic(fmt.Sprintf("unknown ColumnInfo Datatype: %s", ciDatatype))
	}
}

func rowToSliceInterface(header []*fbproto.ColumnInfo, row *fbproto.Row) []interface{} {
	ret := make([]interface{}, len(row.Columns))
	for i, col := range row.Columns {
		switch header[i].Datatype {
		case "string":
			ret[i] = col.GetStringVal()
		case "uint64":
			ret[i] = col.GetUint64Val()
		case "int64":
			ret[i] = col.GetInt64Val()
		case "bool":
			ret[i] = col.GetBoolVal()
		case "[]byte":
			ret[i] = col.GetBlobVal()
		case "[]uint64":
			ret[i] = col.GetUint64ArrayVal()
		case "[]string":
			ret[i] = col.GetStringArrayVal()
		case "float64":
			ret[i] = col.GetFloat64Val()
		case "decimal":
			dec := col.GetDecimalVal()
			ret[i] = featurebase_pql.NewDecimal(dec.Value, dec.Scale)
		case "timestamp":
			ret[i] = col.GetTimestampVal()
		default:
			panic(fmt.Sprintf("don't know how to get value for columninfo datatype %s, val: %+v, type: %[2]T", header[i].Datatype, col.ColumnVal))
		}
	}
	return ret
}