go sql: add option to panic on non-indexed queries

In this commit we add the ability to instrument a connection
with a boolean that, when enabled, will have a query panic
when no index is present in a sql explain for that query.

We also add the ability to bypass this check on any given query
in SelectOptions. Additionally we add some thin wrappers to our
query functions to set this option.

This will permit us to enable this check in
tests and have people explicitly make full table scan queries
when they fully intend to.

Implementation detail:
We check both key and possible keys since there are times where key
will be non-null (eg PRIMARY) while possible keys is null, and also
times when key is null, but possible keys is not null (eg when
mysql chooses a full scan despite the presence of an index due to small
table size)

livesql/live.go

@@ -184,6 +184,22 @@ func (ldb *LiveDB) Query(ctx context.Context, result interface{}, filter sqlgen.
 	return sqlgen.CopySlice(result, rows)
 }
 
+// FullScanQuery bypasses any index checking on a query.
+// Normal LiveDB.Query will check during tests if the query uses an index and will fail tests if not. This function
+// will skip those checks.
+// There are cases where we explicitly want to support full table scans such as
+// 1. During tests to verify results (eg get all)
+// 2. Some rare operations are infrequent and its better to have no index and instead perform full table scans
+//    when that query is run.
+func (ldb *LiveDB) FullScanQuery(ctx context.Context, result interface{}, filter sqlgen.Filter, options *sqlgen.SelectOptions) error {
+	if options == nil {
+		options = &sqlgen.SelectOptions{}
+	}
+	options.AllowNoIndex = true
+
+	return ldb.Query(ctx, result, filter, options)
+}
+
 // QueryRow fetches a single row from the database and will invalidate ctx when
 // the query result changes
 //

sqlgen/db.go

@@ -3,9 +3,11 @@ package sqlgen
 import (
 	"context"
 	"database/sql"
+	"encoding/json"
 	"errors"
 	"fmt"
 
+	"github.com/samsarahq/go/oops"
 	"github.com/samsarahq/thunder/batch"
 )
 
@@ -20,6 +22,8 @@ type DB struct {
 	shardLimit Filter
 
 	dynamicLimit DynamicLimit
+
+	panicOnNoIndex bool
 }
 
 type DynamicLimitFilterCallback func(context.Context, string) Filter
@@ -48,8 +52,9 @@ type DynamicLimit struct {
 
 func NewDB(conn *sql.DB, schema *Schema) *DB {
 	db := &DB{
-		Conn:   conn,
-		Schema: schema,
+		Conn:           conn,
+		Schema:         schema,
+		panicOnNoIndex: false,
 	}
 
 	db.batchFetch = &batch.Func{
@@ -134,6 +139,19 @@ func (db *DB) WithShardLimit(shardLimit Filter) (*DB, error) {
 	return &dbCopy, nil
 }
 
+// WithPanicOnNoIndex will configure this db connection to run an
+// explain on every query and panic when no index is found. This setting is
+// recommended only for use in testing so that you can find non-indexed
+// queries during tests.
+func (db *DB) WithPanicOnNoIndex() (*DB, error) {
+	if db.panicOnNoIndex {
+		return nil, errors.New("already is set panic on no index")
+	}
+
+	db.panicOnNoIndex = true
+	return db, nil
+}
+
 // WithDynamicLimit is similar to WithShardLimit except that it supports
 // dynamically computing what filter to enforce as a limit, based on a user-
 // provide callback. It may be used together with a shard limit.
@@ -232,6 +250,64 @@ func (db *DB) checkColumnValuesAgainstLimits(ctx context.Context, query SQLQuery
 	return nil
 }
 
+type ExplainResultRow struct {
+	Id           int64
+	SelectType   string `sql:"select_type"`
+	Table        string
+	TypeColumn   string  `sql:"type"`
+	PossibleKeys *string `sql:"possible_keys"`
+	Key          *string
+	KeyLen       *string `sql:"key_len"`
+	Ref          *string
+	Rows         int64
+	Extra        *string `sql:"Extra"`
+}
+
+func parseExplainResults(resultRows *sql.Rows) (results []ExplainResultRow, err error) {
+	for resultRows.Next() {
+		var explainType ExplainResultRow
+		err := resultRows.Scan(
+			&explainType.Id, &explainType.SelectType, &explainType.Table, &explainType.TypeColumn, &explainType.PossibleKeys, &explainType.Key,
+			&explainType.KeyLen, &explainType.Ref, &explainType.Rows, &explainType.Extra)
+		if err != nil {
+			return nil, oops.Wrapf(err, "Failed to parse explain results")
+		}
+		results = append(results, explainType)
+	}
+	return results, nil
+}
+
+func (db *DB) runExplainQuery(ctx context.Context, clause string, args []interface{}) error {
+	// We run an explain first and panic if there's no index
+	explained := "EXPLAIN " + clause
+	res, err := db.QueryExecer(ctx).QueryContext(ctx, explained, args...)
+	if err != nil {
+		return oops.Wrapf(err, "Failed to run explain on the query")
+	}
+	defer res.Close()
+
+	explainRes, err := parseExplainResults(res)
+	if err != nil {
+		return oops.Wrapf(err, "failed to parse explain results")
+	}
+
+	for _, explain := range explainRes {
+		if explain.Key == nil && explain.PossibleKeys == nil {
+			explainJSON, _ := json.Marshal(explain)
+			helpMsg := "If you get this message, either check your indices or you can explicitly use a FullScanQuery knowing you're performing a full table scan."
+			panic(fmt.Sprintf(
+				"A sql query was used that misses indexes. %s\n\n%s\n\nwith args\n%s\n\n%s",
+				helpMsg,
+				clause,
+				args,
+				string(explainJSON),
+			))
+		}
+	}
+
+	return nil
+}
+
 func (db *DB) BaseQuery(ctx context.Context, query *BaseSelectQuery) ([]interface{}, error) {
 	selectQuery, err := query.MakeSelectQuery()
 	if err != nil {
@@ -252,6 +328,13 @@ func (db *DB) BaseQuery(ctx context.Context, query *BaseSelectQuery) ([]interfac
 
 	clause, args := selectQuery.ToSQL()
 
+	if db.panicOnNoIndex && (query.Options == nil || !query.Options.AllowNoIndex) {
+		err = db.runExplainQuery(ctx, clause, args)
+		if err != nil {
+			return nil, oops.Wrapf(err, "Failed to run explain query")
+		}
+	}
+
 	res, err := db.QueryExecer(ctx).QueryContext(ctx, clause, args...)
 	if err != nil {
 		return nil, err
@@ -320,6 +403,22 @@ func (db *DB) Query(ctx context.Context, result interface{}, filter Filter, opti
 	return CopySlice(result, rows)
 }
 
+// FullScanQuery bypasses any index checking on a query.
+// Normal LiveDB.Query will check during tests if the query uses an index and will fail tests if not. This function
+// will skip those checks.
+// There are cases where we explicitly want to support full table scans such as
+// 1. During tests to verify results (eg get all)
+// 2. Some rare operations are infrequent and its better to have no index and instead perform full table scans
+//    when that query is run.
+func (db *DB) FullScanQuery(ctx context.Context, result interface{}, filter Filter, options *SelectOptions) error {
+	if options == nil {
+		options = &SelectOptions{}
+	}
+	options.AllowNoIndex = true
+
+	return db.Query(ctx, result, filter, options)
+}
+
 // QueryRow fetches a single row from the database
 //
 // result should be a pointer to a pointer to a struct, for example:

sqlgen/db_test.go

@@ -9,6 +9,45 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+func TestPanicOnNoIndex(t *testing.T) {
+	defer func() {
+		if r := recover(); r == nil {
+			t.Errorf("We expect non-indexed queries to panic.")
+		}
+	}()
+
+	tdb, db, err := setup()
+	assert.NoError(t, err)
+	defer tdb.Close()
+
+	db, err = db.WithPanicOnNoIndex()
+
+	// A second set will error
+	_, err = db.WithPanicOnNoIndex()
+	assert.EqualError(t, err, "already is set panic on no index")
+
+	// Querying users without any filters should panic (checked at the top of the test)
+	var users []*User
+	db.Query(context.Background(), &users, nil, nil)
+}
+
+func TestPanicOnNoIndexOverride(t *testing.T) {
+	tdb, db, err := setup()
+	assert.NoError(t, err)
+	defer tdb.Close()
+
+	db, err = db.WithPanicOnNoIndex()
+
+	// A second set will error
+	_, err = db.WithPanicOnNoIndex()
+	assert.EqualError(t, err, "already is set panic on no index")
+
+	// Querying users without any filters but with a full scan should not error.
+	var users []*User
+	err = db.FullScanQuery(context.Background(), &users, nil, nil)
+	assert.NoError(t, err)
+}
+
 type limitTestcase struct {
 	title           string
 	withFilterLimit func(*DB, Filter) *DB

sqlgen/reflect.go

@@ -23,6 +23,8 @@ type SelectOptions struct {
 
 	OrderBy string
 	Limit   int
+
+	AllowNoIndex bool
 }
 
 func (s *SelectOptions) IncludeFilter(table *Table, filter Filter) error {
@@ -381,8 +383,8 @@ func (b *baseCountQuery) makeCountQuery() (*countQuery, error) {
 	}
 
 	return &countQuery{
-		Table:   b.Table.Name,
-		Where:   where,
+		Table: b.Table.Name,
+		Where: where,
 	}, nil
 }