github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/virtual_table_test.go (about)

     1  // Copyright 2020 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package sql
    12  
    13  import (
    14  	"context"
    15  	"testing"
    16  
    17  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    18  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    19  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    20  	"github.com/cockroachdb/errors"
    21  	"github.com/stretchr/testify/require"
    22  )
    23  
    24  func TestVirtualTableGenerators(t *testing.T) {
    25  	defer leaktest.AfterTest(t)()
    26  	t.Run("test cleanup", func(t *testing.T) {
    27  		ctx := context.Background()
    28  		worker := func(pusher rowPusher) error {
    29  			if err := pusher.pushRow(tree.NewDInt(1)); err != nil {
    30  				return err
    31  			}
    32  			if err := pusher.pushRow(tree.NewDInt(2)); err != nil {
    33  				return err
    34  			}
    35  			return nil
    36  		}
    37  
    38  		next, cleanup := setupGenerator(ctx, worker)
    39  		d, err := next()
    40  		if err != nil {
    41  			t.Fatal(err)
    42  		}
    43  		require.Equal(t, tree.Datums{tree.NewDInt(1)}, d)
    44  
    45  		// Check that we can safely cleanup in the middle of execution.
    46  		cleanup()
    47  	})
    48  
    49  	t.Run("test worker error", func(t *testing.T) {
    50  		// Test that if the worker returns an error we catch it.
    51  		ctx := context.Background()
    52  		worker := func(pusher rowPusher) error {
    53  			if err := pusher.pushRow(tree.NewDInt(1)); err != nil {
    54  				return err
    55  			}
    56  			if err := pusher.pushRow(tree.NewDInt(2)); err != nil {
    57  				return err
    58  			}
    59  			return errors.New("dummy error")
    60  		}
    61  		next, cleanup := setupGenerator(ctx, worker)
    62  		_, err := next()
    63  		require.NoError(t, err)
    64  		_, err = next()
    65  		require.NoError(t, err)
    66  		_, err = next()
    67  		require.Error(t, err)
    68  		cleanup()
    69  	})
    70  
    71  	t.Run("test no next", func(t *testing.T) {
    72  		ctx := context.Background()
    73  		// Test we don't leak anything if we call cleanup before next.
    74  		worker := func(pusher rowPusher) error {
    75  			return nil
    76  		}
    77  		_, cleanup := setupGenerator(ctx, worker)
    78  		cleanup()
    79  	})
    80  
    81  	t.Run("test context cancellation", func(t *testing.T) {
    82  		ctx, cancel := context.WithCancel(context.Background())
    83  		// Test cancellation before asking for any rows.
    84  		worker := func(pusher rowPusher) error {
    85  			if err := pusher.pushRow(tree.NewDInt(1)); err != nil {
    86  				return err
    87  			}
    88  			if err := pusher.pushRow(tree.NewDInt(2)); err != nil {
    89  				return err
    90  			}
    91  			return nil
    92  		}
    93  		next, cleanup := setupGenerator(ctx, worker)
    94  		cancel()
    95  		_, err := next()
    96  		// There is a small chance that we race and don't return
    97  		// a query canceled here. So, only check the error if
    98  		// it is non-nil.
    99  		if err != nil {
   100  			require.Equal(t, sqlbase.QueryCanceledError, err)
   101  		}
   102  		cleanup()
   103  
   104  		// Test cancellation after asking for a row.
   105  		ctx, cancel = context.WithCancel(context.Background())
   106  		next, cleanup = setupGenerator(ctx, worker)
   107  		row, err := next()
   108  		require.NoError(t, err)
   109  		require.Equal(t, tree.Datums{tree.NewDInt(1)}, row)
   110  		cancel()
   111  		_, err = next()
   112  		require.Equal(t, sqlbase.QueryCanceledError, err)
   113  		cleanup()
   114  
   115  		// Test cancellation after asking for all the rows.
   116  		ctx, cancel = context.WithCancel(context.Background())
   117  		next, cleanup = setupGenerator(ctx, worker)
   118  		_, err = next()
   119  		require.NoError(t, err)
   120  		_, err = next()
   121  		require.NoError(t, err)
   122  		cancel()
   123  		cleanup()
   124  	})
   125  }
   126  
   127  func BenchmarkVirtualTableGenerators(b *testing.B) {
   128  	defer leaktest.AfterTest(b)()
   129  	ctx := context.Background()
   130  	worker := func(pusher rowPusher) error {
   131  		for {
   132  			if err := pusher.pushRow(tree.NewDInt(tree.DInt(1))); err != nil {
   133  				return err
   134  			}
   135  		}
   136  	}
   137  	b.Run("bench read", func(b *testing.B) {
   138  		next, cleanup := setupGenerator(ctx, worker)
   139  		b.ResetTimer()
   140  		for i := 0; i < b.N; i++ {
   141  			_, err := next()
   142  			require.NoError(b, err)
   143  		}
   144  		cleanup()
   145  	})
   146  }