github.com/dolthub/go-mysql-server@v0.18.0/enginetest/parallelism.go (about)

     1  package enginetest
     2  
     3  import (
     4  	"fmt"
     5  	"testing"
     6  
     7  	"github.com/stretchr/testify/require"
     8  
     9  	"github.com/dolthub/go-mysql-server/enginetest/queries"
    10  	"github.com/dolthub/go-mysql-server/enginetest/scriptgen/setup"
    11  	"github.com/dolthub/go-mysql-server/sql"
    12  	"github.com/dolthub/go-mysql-server/sql/plan"
    13  	"github.com/dolthub/go-mysql-server/sql/transform"
    14  )
    15  
    16  func TestParallelismQueries(t *testing.T, harness Harness) {
    17  	harness.Setup(setup.XySetup...)
    18  	e := mustNewEngine(t, harness)
    19  	defer e.Close()
    20  	for _, tt := range queries.ParallelismTests {
    21  		t.Run(tt.Query, func(t *testing.T) {
    22  			evalParallelismTest(t, harness, e, tt.Query, tt.Parallel)
    23  		})
    24  	}
    25  }
    26  
    27  func evalParallelismTest(t *testing.T, harness Harness, e QueryEngine, query string, parallel bool) {
    28  	ctx := NewContext(harness)
    29  	ctx = ctx.WithQuery(query)
    30  	a, err := analyzeQuery(ctx, e, query)
    31  	require.NoError(t, err)
    32  	require.Equal(t, parallel, findExchange(a), fmt.Sprintf("expected exchange: %t\nplan:\n%s", parallel, sql.DebugString(a)))
    33  }
    34  
    35  func findExchange(n sql.Node) bool {
    36  	return transform.InspectUp(n, func(n sql.Node) bool {
    37  		if n == nil {
    38  			return false
    39  		}
    40  		_, ok := n.(*plan.Exchange)
    41  		if ok {
    42  			return true
    43  		}
    44  
    45  		if ex, ok := n.(sql.Expressioner); ok {
    46  			for _, e := range ex.Expressions() {
    47  				found := transform.InspectExpr(e, func(e sql.Expression) bool {
    48  					sq, ok := e.(*plan.Subquery)
    49  					if !ok {
    50  						return false
    51  					}
    52  					return findExchange(sq.Query)
    53  				})
    54  				if found {
    55  					return true
    56  				}
    57  			}
    58  		}
    59  		return false
    60  	})
    61  }