Merge pull request #1547 from niaow/external-lookup

[CORE-388] Add ExternalLookup query

.circleci/config.yml

@@ -108,6 +108,19 @@ jobs:
       - run:
           command: make << parameters.test_make_target >> SHARD_WIDTH=<< parameters.shard_width >> GOARCH=<< parameters.goarch >>
           no_output_timeout: 30m
+  test-external-lookup:
+    docker:
+      - image: circleci/golang:1.15.8
+      - image: circleci/postgres:13.2-ram
+        environment:
+          POSTGRES_PASSWORD=password
+    steps:
+      - checkout-plus
+      - run: sudo apt-get install postgresql-client
+      - run: (for i in `seq 1 20`; do pg_isready -h localhost && exit 0 || sleep 1; done; exit 1)
+      - run:
+          command: make test-external-lookup EXTERNAL_LOOKUP_DSN=postgresql://postgres:password@localhost/circle_test?sslmode=disable
+          no_output_timeout: 30m
   cluster-tests:
     executor:
       name: golang
@@ -226,6 +239,10 @@ workflows:
           resource_class: large
           requires:
             - setup
+      - test-external-lookup:
+          context: molecula
+          requires:
+            - setup
       - cluster-tests:
           context: molecula
           requires:

Makefile

@@ -341,3 +341,5 @@ test-txstore-rbf:
 test-txstore-rbf_bolt:
 	PILOSA_STORAGE_BACKEND=rbf_bolt $(MAKE) testv-race
 
+test-external-lookup:
+	$(GO) test . -tags='$(BUILD_TAGS) $(TEST_TAGS)' $(TESTFLAGS) -run ^TestExternalLookup$$ -externalLookupDSN $(EXTERNAL_LOOKUP_DSN)

ctl/server.go

@@ -66,6 +66,9 @@ func BuildServerFlags(cmd *cobra.Command, srv *server.Command) {
 	flags.StringVar(&srv.Config.Etcd.ClusterURL, "etcd.cluster-url", srv.Config.Etcd.ClusterURL, "Cluster URL to join.")
 	flags.StringVar(&srv.Config.Etcd.InitCluster, "etcd.initial-cluster", srv.Config.Etcd.InitCluster, "Initial cluster name1=apurl1,name2=apurl2")
 
+	// External postgres database for ExternalLookup
+	flags.StringVar(&srv.Config.LookupDBDSN, "lookup-db-dsn", "", "external (postgres) database DSN to use for ExternalLookup calls")
+
 	// AntiEntropy
 	flags.DurationVar((*time.Duration)(&srv.Config.AntiEntropy.Interval), "anti-entropy.interval", (time.Duration)(srv.Config.AntiEntropy.Interval), "Interval at which to run anti-entropy routine.")
 

executor.go

@@ -16,10 +16,12 @@ package pilosa
 
 import (
 	"context"
+	"database/sql"
 	"encoding/json"
 	"fmt"
 	"math"
 	"math/bits"
+	"reflect"
 	"sort"
 	"strings"
 	"sync"
@@ -28,6 +30,7 @@ import (
 
 	"golang.org/x/sync/errgroup"
 
+	"github.com/lib/pq"
 	"github.com/pilosa/pilosa/v2/disco"
 	"github.com/pilosa/pilosa/v2/pql"
 	"github.com/pilosa/pilosa/v2/proto"
@@ -799,6 +802,10 @@ func (e *executor) executeCall(ctx context.Context, qcx *Qcx, index string, c *p
 		statFn()
 		res, err := e.executeRows(ctx, qcx, index, c, shards, opt)
 		return res, errors.Wrapf(err, "executeRows %v", shardSlice(shards))
+	case "ExternalLookup":
+		statFn()
+		res, err := e.executeExternalLookup(ctx, qcx, index, c, shards, opt)
+		return res, errors.Wrapf(err, "executeExternalLookup %v", shardSlice(shards))
 	case "Extract":
 		statFn()
 		res, err := e.executeExtract(ctx, qcx, index, c, shards, opt)
@@ -3999,6 +4006,236 @@ func (e *ExtractedIDMatrix) Append(m ExtractedIDMatrix) {
 	}
 }
 
+var (
+	typeSQLNullString = reflect.TypeOf(sql.NullString{})
+	typeSQLNullBool   = reflect.TypeOf(sql.NullBool{})
+	typeSQLNullInt32  = reflect.TypeOf(sql.NullInt32{})
+	typeSQLNullInt64  = reflect.TypeOf(sql.NullInt64{})
+)
+
+func (e *executor) executeExternalLookup(ctx context.Context, qcx *Qcx, index string, c *pql.Call, shards []uint64, opt *execOptions) (res ExtractedTable, err error) {
+	if e.Holder.lookupDB == nil {
+		return ExtractedTable{}, errors.New("external DB connection is not configured")
+	}
+
+	idx := e.Holder.Index(index)
+	if idx == nil {
+		return ExtractedTable{}, errors.Errorf("index not found: %q", index)
+	}
+
+	query, ok, err := c.StringArg("query")
+	if err != nil {
+		return ExtractedTable{}, errors.Wrap(err, "looking up query")
+	}
+	if !ok {
+		return ExtractedTable{}, errors.New("missing query")
+	}
+
+	switch len(c.Children) {
+	case 0:
+		return ExtractedTable{}, errors.New("missing lookup input")
+	case 1:
+	default:
+		return ExtractedTable{}, errors.New("too many inputs to lookup query")
+	}
+
+	rawArg, err := e.executeCall(ctx, qcx, index, c.Children[0], shards, opt)
+	if err != nil {
+		return ExtractedTable{}, errors.Wrapf(err, "evaluating SQL argument call %q", c.String())
+	}
+
+	qr := []interface{}{rawArg}
+	err = e.translateResults(ctx, index, idx, c.Children, qr)
+	if err != nil {
+		return ExtractedTable{}, errors.Wrap(err, "translating query result")
+	}
+	argRow, ok := qr[0].(*Row)
+	if !ok {
+		return ExtractedTable{}, errors.Errorf("argument call result is a %T but expected a row", qr[0])
+	}
+	if !argRow.Any() {
+		return ExtractedTable{}, nil
+	}
+
+	var arg interface{}
+	if argRow.Keys != nil {
+		arg = argRow.Keys
+	} else {
+		arg = argRow.Columns()
+	}
+
+	result, err := e.Holder.lookupDB.QueryContext(ctx, query, pq.Array(arg))
+	if err != nil {
+		return ExtractedTable{}, errors.Wrapf(err, "SQL query failed")
+	}
+	defer func() {
+		cerr := result.Close()
+		if cerr != nil && err == nil {
+			err = cerr
+		}
+	}()
+	if !result.Next() {
+		return ExtractedTable{}, errors.Wrap(result.Err(), "reading SQL query result")
+	}
+
+	colTypes, err := result.ColumnTypes()
+	if err != nil {
+		return ExtractedTable{}, errors.Wrapf(err, "fetching SQL query result types")
+	}
+
+	scanSlots := make([]interface{}, len(colTypes))
+	scanMapper := make([]func() interface{}, len(colTypes))
+	header := make([]ExtractedTableField, len(colTypes))
+	for i, colType := range colTypes {
+		scanType := colType.ScanType()
+	setupScan:
+		if scanType.PkgPath() != "" {
+			// This is a named type.
+			switch scanType {
+			case typeSQLNullString:
+				var dst sql.NullString
+				scanSlots[i] = &dst
+				scanMapper[i] = func() interface{} {
+					if !dst.Valid {
+						return nil
+					}
+
+					return dst.String
+				}
+				header[i] = ExtractedTableField{
+					Name: colType.Name(),
+					Type: "string",
+				}
+
+			case typeSQLNullBool:
+				var dst sql.NullBool
+				scanSlots[i] = &dst
+				scanMapper[i] = func() interface{} {
+					if !dst.Valid {
+						return nil
+					}
+
+					return dst.Bool
+				}
+				header[i] = ExtractedTableField{
+					Name: colType.Name(),
+					Type: "bool",
+				}
+
+			case typeSQLNullInt32:
+				var dst sql.NullInt32
+				scanSlots[i] = &dst
+				scanMapper[i] = func() interface{} {
+					if !dst.Valid {
+						return nil
+					}
+
+					return int64(dst.Int32)
+				}
+				header[i] = ExtractedTableField{
+					Name: colType.Name(),
+					Type: "int64",
+				}
+
+			case typeSQLNullInt64:
+				var dst sql.NullInt64
+				scanSlots[i] = &dst
+				scanMapper[i] = func() interface{} {
+					if !dst.Valid {
+						return nil
+					}
+
+					return dst.Int64
+				}
+				header[i] = ExtractedTableField{
+					Name: colType.Name(),
+					Type: "int64",
+				}
+
+			default:
+				return ExtractedTable{}, errors.Errorf("unable to process result type %v from SQL query %q", scanType, query)
+			}
+			continue
+		}
+
+		switch scanType.Kind() {
+		case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
+			var dst uint64
+			scanSlots[i] = &dst
+			scanMapper[i] = func() interface{} { return dst }
+			header[i] = ExtractedTableField{
+				Name: colType.Name(),
+				Type: "uint64",
+			}
+
+		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
+			// The database may have lied and this may be nullable.
+			scanType = typeSQLNullInt64
+			goto setupScan
+
+		case reflect.String:
+			// The database may have lied and this may be nullable.
+			scanType = typeSQLNullString
+			goto setupScan
+
+		case reflect.Bool:
+			// The database may have lied and this may be nullable.
+			scanType = typeSQLNullBool
+			goto setupScan
+
+		default:
+			return ExtractedTable{}, errors.Errorf("unable to process result type %v from SQL query %q", scanType, query)
+		}
+	}
+
+	var columns []ExtractedTableColumn
+	for {
+		err := result.Scan(scanSlots...)
+		if err != nil {
+			return ExtractedTable{}, errors.Wrap(err, "scanning SQL result")
+		}
+
+		var col KeyOrID
+		switch v := scanMapper[0]().(type) {
+		case nil:
+			return ExtractedTable{}, errors.Errorf("missing primary key in result")
+		case uint64:
+			col.ID = v
+		case int64:
+			col.ID = uint64(v)
+		case string:
+			col.Keyed = true
+			col.Key = v
+		default:
+			return ExtractedTable{}, errors.Wrap(err, "cannot use %v of type %T as primary key for result table")
+		}
+
+		data := make([]interface{}, len(scanMapper)-1)
+		for i, m := range scanMapper[1:] {
+			data[i] = m()
+		}
+
+		columns = append(columns, ExtractedTableColumn{
+			Column: col,
+			Rows:   data,
+		})
+
+		if !result.Next() {
+			break
+		}
+	}
+
+	err = result.Err()
+	if err != nil {
+		return ExtractedTable{}, errors.Wrap(err, "reading SQL result")
+	}
+
+	return ExtractedTable{
+		Fields:  header[1:],
+		Columns: columns,
+	}, nil
+}
+
 func (e *executor) executeExtract(ctx context.Context, qcx *Qcx, index string, c *pql.Call, shards []uint64, opt *execOptions) (ExtractedIDMatrix, error) {
 	// Extract the column filter call.
 	if len(c.Children) < 1 {

executor_test.go

@@ -17,6 +17,7 @@ package pilosa_test
 import (
 	"bytes"
 	"context"
+	"database/sql"
 	"encoding/csv"
 	"encoding/json"
 	"flag"
@@ -7899,3 +7900,279 @@ func tableResponseToCSVString(m *proto.TableResponse) (string, error) {
 	}
 	return buf.String(), nil
 }
+
+var dbDSN string
+
+func init() {
+	flag.StringVar(&dbDSN, "externalLookupDSN", "", "SQL DSN to use for external database access")
+}
+
+func TestExternalLookup(t *testing.T) {
+	// Set up access to a SQL database.
+	if dbDSN == "" {
+		t.Skip("no database provided")
+	}
+	db, err := sql.Open("postgres", dbDSN)
+	if err != nil {
+		t.Fatalf("failed to set up test database: %v", err)
+	}
+	defer func() {
+		if cerr := db.Close(); cerr != nil {
+			t.Errorf("failed to close test database: %v", cerr)
+		}
+	}()
+
+	// Set up some data to use in the SQL DB.
+	// This creates 3 tables:
+	// - "lookup" - which stores an id->string mapping
+	// - "misc" - which has misc non-nullable fields
+	// - "nullable" - to test handling of nullable fields
+	func() {
+		// Set up a write transaction.
+		// This uses a function scope so that the defer is scoped appropriately.
+		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+		defer cancel()
+		tx, err := db.BeginTx(ctx, nil)
+		if err != nil {
+			t.Fatalf("failed to create transaction: %v", err)
+		}
+		var ok bool
+		defer func() {
+			if !ok {
+				rerr := tx.Rollback()
+				if rerr != nil {
+					t.Errorf("failed to roll-back transaction: %v", rerr)
+				}
+			}
+		}()
+
+		// Delete all tables so that we start fresh.
+		_, err = tx.Exec(`DROP INDEX IF EXISTS lookupIndex`)
+		if err != nil {
+			t.Fatalf("deleting lookup index: %v", err)
+		}
+		_, err = tx.Exec(`DROP TABLE IF EXISTS lookup`)
+		if err != nil {
+			t.Fatalf("deleting lookup table: %v", err)
+		}
+		_, err = tx.Exec(`DROP TABLE IF EXISTS misc`)
+		if err != nil {
+			t.Fatalf("deleting misc table: %v", err)
+		}
+		_, err = tx.Exec(`DROP TABLE IF EXISTS nullable`)
+		if err != nil {
+			t.Fatalf("deleting nullable table: %v", err)
+		}
+
+		_, err = tx.Exec(`CREATE TABLE lookup (
+			id int NOT NULL,
+			data text NOT NULL
+		)`)
+		if err != nil {
+			t.Fatalf("failed to create lookup table: %v", err)
+		}
+		_, err = tx.Exec(`INSERT INTO lookup (id, data) VALUES
+			(0, 'h'),
+			(1, 'xyzzy'),
+			(2, 'plugh'),
+			(3, 'wow')
+		`)
+		if err != nil {
+			t.Fatalf("failed to populate lookup table: %v", err)
+		}
+		_, err = tx.Exec(`CREATE UNIQUE INDEX lookupIndex ON lookup (id);`)
+		if err != nil {
+			t.Fatalf("failed to index lookup table: %v", err)
+		}
+
+		_, err = tx.Exec(`CREATE TABLE misc (
+			id int NOT NULL,
+			stringval text NOT NULL,
+			boolval boolean NOT NULL,
+			intval int NOT NULL
+		)`)
+		if err != nil {
+			t.Fatalf("failed to create misc table: %v", err)
+		}
+		_, err = tx.Exec(`INSERT INTO misc (id, stringval, boolval, intval) VALUES
+			(0, 'h', true, 4),
+			(1, 'y', false, 11)
+		`)
+		if err != nil {
+			t.Fatalf("failed to populate misc table: %v", err)
+		}
+
+		_, err = tx.Exec(`CREATE TABLE nullable (
+			id int NOT NULL,
+			stringval text,
+			boolval boolean,
+			intval int
+		)`)
+		if err != nil {
+			t.Fatalf("failed to create nullable table: %v", err)
+		}
+		_, err = tx.Exec(`INSERT INTO nullable (id, stringval, boolval, intval) VALUES
+			(0, 'h', true, 4),
+			(1, 'y', false, 11),
+			(2, null, null, null),
+			(3, null, true, null),
+			(4, 'plugh', null, 0)
+		;`)
+		if err != nil {
+			t.Fatalf("failed to populate nullable table: %v", err)
+		}
+
+		err = tx.Commit()
+		if err != nil {
+			t.Fatalf("failed to commit DB setup transaction: %v", err)
+		}
+
+		ok = true
+	}()
+
+	// Start up a Pilosa cluster with access to the DB.
+	c := test.MustRunCluster(t, 3, []server.CommandOption{server.OptCommandServerOptions(pilosa.OptServerLookupDB(dbDSN))})
+	defer c.Close()
+
+	// Populate a field with some data that can be used in queries.
+	c.CreateField(t, "i", pilosa.IndexOptions{}, "f")
+	c.ImportBits(t, "i", "f", [][2]uint64{
+		{1, 1},
+		{1, 3},
+		{2, 2},
+		{2, 3},
+	})
+
+	cases := []struct {
+		name   string
+		query  string
+		expect pilosa.QueryResponse
+	}{
+		{
+			name:  "Empty",
+			query: `ExternalLookup(Union(), query="select * from lookup where id = ANY($1)")`,
+			expect: pilosa.QueryResponse{
+				Results: []interface{}{
+					pilosa.ExtractedTable{},
+				},
+			},
+		},
+		{
+			name:  "ConstRowLookup",
+			query: `ExternalLookup(ConstRow(columns=[1, 3]), query="select id, data from lookup where id = ANY($1)")`,
+			expect: pilosa.QueryResponse{
+				Results: []interface{}{
+					pilosa.ExtractedTable{
+						Fields: []pilosa.ExtractedTableField{
+							{Name: "data", Type: "string"},
+						},
+						Columns: []pilosa.ExtractedTableColumn{
+							{
+								Column: pilosa.KeyOrID{ID: 1},
+								Rows:   []interface{}{"xyzzy"},
+							},
+							{
+								Column: pilosa.KeyOrID{ID: 3},
+								Rows:   []interface{}{"wow"},
+							},
+						},
+					},
+				},
+			},
+		},
+		{
+			name:  "ComputedFilter",
+			query: `ExternalLookup(Intersect(Row(f=1), Row(f=2)), query="select id, data from lookup where id = ANY($1)")`,
+			expect: pilosa.QueryResponse{
+				Results: []interface{}{
+					pilosa.ExtractedTable{
+						Fields: []pilosa.ExtractedTableField{
+							{Name: "data", Type: "string"},
+						},
+						Columns: []pilosa.ExtractedTableColumn{
+							{
+								Column: pilosa.KeyOrID{ID: 3},
+								Rows:   []interface{}{"wow"},
+							},
+						},
+					},
+				},
+			},
+		},
+		{
+			name:  "Misc",
+			query: `ExternalLookup(ConstRow(columns=[0, 1]), query="select id, stringval, boolval, intval from misc where id = ANY($1)")`,
+			expect: pilosa.QueryResponse{
+				Results: []interface{}{
+					pilosa.ExtractedTable{
+						Fields: []pilosa.ExtractedTableField{
+							{Name: "stringval", Type: "string"},
+							{Name: "boolval", Type: "bool"},
+							{Name: "intval", Type: "int64"},
+						},
+						Columns: []pilosa.ExtractedTableColumn{
+							{
+								Column: pilosa.KeyOrID{ID: 0},
+								Rows:   []interface{}{"h", true, int64(4)},
+							},
+							{
+								Column: pilosa.KeyOrID{ID: 1},
+								Rows:   []interface{}{"y", false, int64(11)},
+							},
+						},
+					},
+				},
+			},
+		},
+		{
+			name:  "Nullable",
+			query: `ExternalLookup(ConstRow(columns=[0, 1, 2, 3, 4]), query="select id, stringval, boolval, intval from nullable where id = ANY($1)")`,
+			expect: pilosa.QueryResponse{
+				Results: []interface{}{
+					pilosa.ExtractedTable{
+						Fields: []pilosa.ExtractedTableField{
+							{Name: "stringval", Type: "string"},
+							{Name: "boolval", Type: "bool"},
+							{Name: "intval", Type: "int64"},
+						},
+						Columns: []pilosa.ExtractedTableColumn{
+							{
+								Column: pilosa.KeyOrID{ID: 0},
+								Rows:   []interface{}{"h", true, int64(4)},
+							},
+							{
+								Column: pilosa.KeyOrID{ID: 1},
+								Rows:   []interface{}{"y", false, int64(11)},
+							},
+							{
+								Column: pilosa.KeyOrID{ID: 2},
+								Rows:   []interface{}{nil, nil, nil},
+							},
+							{
+								Column: pilosa.KeyOrID{ID: 3},
+								Rows:   []interface{}{nil, true, nil},
+							},
+							{
+								Column: pilosa.KeyOrID{ID: 4},
+								Rows:   []interface{}{"plugh", nil, int64(0)},
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+	t.Run("Query", func(t *testing.T) {
+		for _, tc := range cases {
+			tc := tc
+			t.Run(tc.name, func(t *testing.T) {
+				t.Parallel()
+
+				result := c.Query(t, "i", tc.query)
+				if !reflect.DeepEqual(result, tc.expect) {
+					t.Errorf("expected %v but got %v", tc.expect, result)
+				}
+			})
+		}
+	})
+}

holder.go

@@ -16,6 +16,7 @@ package pilosa
 
 import (
 	"context"
+	"database/sql"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -151,6 +152,8 @@ type Holder struct {
 
 	txf *TxFactory
 
+	lookupDB *sql.DB
+
 	// a separate lock out for indexes, to avoid the deadlock/race dilema
 	// on holding mu.
 	imu     sync.RWMutex
@@ -240,6 +243,8 @@ type HolderConfig struct {
 	StorageConfig       *storage.Config
 	RBFConfig           *rbfcfg.Config
 	AntiEntropyInterval time.Duration
+
+	LookupDBDSN string
 }
 
 func DefaultHolderConfig() *HolderConfig {
@@ -729,6 +734,17 @@ func (h *Holder) Open() error {
 
 	h.txf.blueGreenOnIfRunningBlueGreen()
 
+	if h.cfg.LookupDBDSN != "" {
+		h.Logger.Printf("connecting to lookup DB")
+
+		db, err := sql.Open("postgres", h.cfg.LookupDBDSN)
+		if err != nil {
+			return errors.Wrap(err, "connecting to lookup database")
+		}
+
+		h.lookupDB = db
+	}
+
 	h.Logger.Printf("open holder: complete")
 
 	return nil
@@ -839,6 +855,14 @@ func (h *Holder) Close() error {
 		h.SnapshotQueue = nil
 	}
 
+	if h.lookupDB != nil {
+		err := h.lookupDB.Close()
+		if err != nil {
+			return errors.Wrap(err, "closing DB")
+		}
+		h.lookupDB = nil
+	}
+
 	_ = testhook.Closed(h.Auditor, h, nil)
 
 	return nil

pql/ast.go

@@ -400,6 +400,12 @@ var callInfoByFunc = map[string]callInfo{
 	"Union":     {allowUnknown: false},
 	"UnionRows": {allowUnknown: false, callType: PrecallGlobal},
 	"Extract":   {allowUnknown: false},
+	"ExternalLookup": {
+		allowUnknown: false,
+		prototypes: map[string]interface{}{
+			"query": "",
+		},
+	},
 	"Limit": {
 		allowUnknown: false,
 		prototypes: map[string]interface{}{

server.go

@@ -40,6 +40,8 @@ import (
 	"github.com/pilosa/pilosa/v2/topology"
 	"github.com/pkg/errors"
 	"golang.org/x/sync/errgroup"
+
+	_ "github.com/lib/pq"
 )
 
 // Default server settings.
@@ -398,6 +400,14 @@ func OptServerDisCo(disCo disco.DisCo,
 	}
 }
 
+// OptServerLookupDB configures a connection to an external postgres database for ExternalLookup queries.
+func OptServerLookupDB(dsn string) ServerOption {
+	return func(s *Server) error {
+		s.holderConfig.LookupDBDSN = dsn
+		return nil
+	}
+}
+
 // NewServer returns a new instance of Server.
 func NewServer(opts ...ServerOption) (*Server, error) {
 	cluster := newCluster()

server/config.go

@@ -222,6 +222,9 @@ type Config struct {
 	// for the /query-history endpoint. This parameter is per-node, and the
 	// result combines the history from all nodes.
 	QueryHistoryLength int `toml:"query-history-length"`
+
+	// LookupDBDSN is an external database to connect to for `ExternalLookup` queries.
+	LookupDBDSN string `toml:"lookup-db-dsn"`
 }
 
 // MustValidate checks that all ports in a Config are unique and not zero.

server/server.go

@@ -429,6 +429,10 @@ func (m *Command) SetupServer() error {
 		discoOpt,
 	}
 
+	if m.Config.LookupDBDSN != "" {
+		serverOptions = append(serverOptions, pilosa.OptServerLookupDB(m.Config.LookupDBDSN))
+	}
+
 	serverOptions = append(serverOptions, m.serverOptions...)
 
 	m.Server, err = pilosa.NewServer(serverOptions...)