some fixes

Author: James Cor

enginetest/engine_only_test.go

@@ -293,6 +293,73 @@ b (2/6 partitions)
 	require.ElementsMatch(expected, rows)
 }
 
+// TODO: this was an analyzer test, but we don't have a mock process list for it to use, so it has to be here
+func TestTrackProcess(t *testing.T) {
+	require := require.New(t)
+	db := memory.NewDatabase("db")
+	provider := memory.NewDBProvider(db)
+	a := analyzer.NewDefault(provider)
+	sess := memory.NewSession(sql.NewBaseSession(), provider)
+
+	node := plan.NewInnerJoin(
+		plan.NewResolvedTable(&nonIndexableTable{memory.NewPartitionedTable(db.BaseDatabase, "foo", sql.PrimaryKeySchema{}, nil, 2)}, nil, nil),
+		plan.NewResolvedTable(memory.NewPartitionedTable(db.BaseDatabase, "bar", sql.PrimaryKeySchema{}, nil, 4), nil, nil),
+		expression.NewLiteral(int64(1), types.Int64),
+	)
+
+	pl := sqle.NewProcessList()
+
+	ctx := sql.NewContext(context.Background(), sql.WithPid(1), sql.WithProcessList(pl), sql.WithSession(sess))
+	pl.AddConnection(ctx.Session.ID(), "localhost")
+	pl.ConnectionReady(ctx.Session)
+	ctx, err := ctx.ProcessList.BeginQuery(ctx, "SELECT foo")
+	require.NoError(err)
+
+	rule := getRuleFrom(analyzer.OnceAfterAll, analyzer.TrackProcessId)
+	result, _, err := rule.Apply(ctx, a, node, nil, analyzer.DefaultRuleSelector, nil)
+	require.NoError(err)
+
+	processes := ctx.ProcessList.Processes()
+	require.Len(processes, 1)
+	require.Equal("SELECT foo", processes[0].Query)
+	require.Equal(
+		map[string]sql.TableProgress{
+			"foo": {
+				Progress:           sql.Progress{Name: "foo", Done: 0, Total: 2},
+				PartitionsProgress: map[string]sql.PartitionProgress{},
+			},
+			"bar": {
+				Progress:           sql.Progress{Name: "bar", Done: 0, Total: 4},
+				PartitionsProgress: map[string]sql.PartitionProgress{},
+			},
+		},
+		processes[0].Progress)
+
+	join, ok := result.(*plan.JoinNode)
+	require.True(ok)
+	require.Equal(join.JoinType(), plan.JoinTypeInner)
+
+	lhs, ok := join.Left().(*plan.ResolvedTable)
+	require.True(ok)
+	_, ok = lhs.Table.(*plan.ProcessTable)
+	require.True(ok)
+
+	rhs, ok := join.Right().(*plan.ResolvedTable)
+	require.True(ok)
+	_, ok = rhs.Table.(*plan.ProcessTable)
+	require.True(ok)
+
+	iter, err := rowexec.DefaultBuilder.Build(ctx, result, nil)
+	require.NoError(err)
+	_, err = sql.RowIterToRows(ctx, iter)
+	require.NoError(err)
+
+	procs := ctx.ProcessList.Processes()
+	require.Len(procs, 1)
+	require.Equal(sql.ProcessCommandSleep, procs[0].Command)
+	require.Error(ctx.Err())
+}
+
 func TestConcurrentProcessList(t *testing.T) {
 	enginetest.TestConcurrentProcessList(t, enginetest.NewDefaultMemoryHarness())
 }

enginetest/evaluation.go

@@ -1080,8 +1080,6 @@ func assertSchemasEqualWithDefaults(t *testing.T, expected, actual sql.Schema) b
 
 func ExtractQueryNode(node sql.Node) sql.Node {
 	switch node := node.(type) {
-	case *plan.QueryProcess:
-		return ExtractQueryNode(node.Child())
 	case *plan.Releaser:
 		return ExtractQueryNode(node.Child)
 	default:

sql/rowexec/process_test.go

@@ -26,49 +26,6 @@ import (
 	"github.com/dolthub/go-mysql-server/sql/types"
 )
 
-func TestQueryProcess(t *testing.T) {
-	require := require.New(t)
-
-	db := memory.NewDatabase("test")
-	pro := memory.NewDBProvider(db)
-	ctx := newContext(pro)
-
-	table := memory.NewTable(db.BaseDatabase, "foo", sql.NewPrimaryKeySchema(sql.Schema{
-		{Name: "a", Type: types.Int64},
-	}), nil)
-
-	table.Insert(ctx, sql.NewRow(int64(1)))
-	table.Insert(ctx, sql.NewRow(int64(2)))
-
-	var notifications int
-
-	node := plan.NewQueryProcess(
-		plan.NewProject(
-			[]sql.Expression{
-				expression.NewGetField(0, types.Int64, "a", false),
-			},
-			plan.NewResolvedTable(table, nil, nil),
-		),
-		func() {
-			notifications++
-		},
-	)
-
-	iter, err := DefaultBuilder.Build(ctx, node, nil)
-	require.NoError(err)
-
-	rows, err := sql.RowIterToRows(ctx, iter)
-	require.NoError(err)
-
-	expected := []sql.Row{
-		{int64(1)},
-		{int64(2)},
-	}
-
-	require.ElementsMatch(expected, rows)
-	require.Equal(1, notifications)
-}
-
 func TestProcessTable(t *testing.T) {
 	require := require.New(t)