github.com/dolthub/go-mysql-server@v0.18.0/sql/rowexec/resolved_table_test.go (about)

     1  // Copyright 2020-2021 Dolthub, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package rowexec
    16  
    17  import (
    18  	"context"
    19  	"io"
    20  	"testing"
    21  
    22  	"github.com/stretchr/testify/require"
    23  
    24  	"github.com/dolthub/go-mysql-server/sql"
    25  	. "github.com/dolthub/go-mysql-server/sql/plan"
    26  	"github.com/dolthub/go-mysql-server/sql/planbuilder"
    27  	"github.com/dolthub/go-mysql-server/sql/types"
    28  )
    29  
    30  func TestResolvedTable(t *testing.T) {
    31  	var require = require.New(t)
    32  
    33  	table := NewResolvedTable(newTableTest("test"), nil, nil)
    34  	require.NotNil(table)
    35  
    36  	ctx := sql.NewEmptyContext()
    37  	iter, err := DefaultBuilder.Build(ctx, table, nil)
    38  	require.NoError(err)
    39  
    40  	rows, err := sql.RowIterToRows(ctx, iter)
    41  	require.NoError(err)
    42  	require.Len(rows, 9)
    43  
    44  	tableTest, ok := table.Table.(*dummyTable)
    45  	require.True(ok)
    46  
    47  	for i, row := range rows {
    48  		expected := tableTest.rows[i]
    49  		require.ElementsMatch(expected, row)
    50  	}
    51  }
    52  
    53  func TestResolvedTableCancelled(t *testing.T) {
    54  	var require = require.New(t)
    55  
    56  	table := NewResolvedTable(newTableTest("test"), nil, nil)
    57  	require.NotNil(table)
    58  
    59  	octx, cancel := context.WithCancel(context.Background())
    60  	cancel()
    61  	ctx := sql.NewContext(octx)
    62  
    63  	iter, err := DefaultBuilder.Build(ctx, table, nil)
    64  	require.NoError(err)
    65  
    66  	_, err = iter.Next(ctx)
    67  	require.Equal(context.Canceled, err)
    68  }
    69  
    70  func newTableTest(source string) sql.Table {
    71  	schema := []*sql.Column{
    72  		{Name: "col1", Type: types.Int32, Source: source, Default: planbuilder.MustStringToColumnDefaultValue(sql.NewEmptyContext(), "0", types.Int32, false), Nullable: false},
    73  		{Name: "col2", Type: types.Int64, Source: source, Default: planbuilder.MustStringToColumnDefaultValue(sql.NewEmptyContext(), "0", types.Int64, false), Nullable: false},
    74  		{Name: "col3", Type: types.Text, Source: source, Default: planbuilder.MustStringToColumnDefaultValue(sql.NewEmptyContext(), `""`, types.Text, false), Nullable: false},
    75  	}
    76  
    77  	keys := [][]byte{
    78  		[]byte("partition1"),
    79  		[]byte("partition2"),
    80  		[]byte("partition3"),
    81  	}
    82  
    83  	rows := []sql.Row{
    84  		sql.NewRow(int32(1), int64(9), "one, nine"),
    85  		sql.NewRow(int32(2), int64(8), "two, eight"),
    86  		sql.NewRow(int32(3), int64(7), "three, seven"),
    87  		sql.NewRow(int32(4), int64(6), "four, six"),
    88  		sql.NewRow(int32(5), int64(5), "five, five"),
    89  		sql.NewRow(int32(6), int64(4), "six, four"),
    90  		sql.NewRow(int32(7), int64(3), "seven, three"),
    91  		sql.NewRow(int32(8), int64(2), "eight, two"),
    92  		sql.NewRow(int32(9), int64(1), "nine, one"),
    93  	}
    94  
    95  	partitions := map[string][]sql.Row{
    96  		"partition1": []sql.Row{rows[0], rows[1], rows[2]},
    97  		"partition2": []sql.Row{rows[3], rows[4], rows[5]},
    98  		"partition3": []sql.Row{rows[6], rows[7], rows[8]},
    99  	}
   100  
   101  	return &dummyTable{schema, keys, partitions, rows}
   102  }
   103  
   104  type dummyTable struct {
   105  	schema     sql.Schema
   106  	keys       [][]byte
   107  	partitions map[string][]sql.Row
   108  	rows       []sql.Row
   109  }
   110  
   111  var _ sql.Table = (*dummyTable)(nil)
   112  
   113  func (t *dummyTable) Name() string { return "dummy" }
   114  
   115  func (t *dummyTable) String() string {
   116  	panic("not implemented")
   117  }
   118  
   119  func (*dummyTable) Insert(*sql.Context, sql.Row) error {
   120  	panic("not implemented")
   121  }
   122  
   123  func (t *dummyTable) Schema() sql.Schema { return t.schema }
   124  
   125  func (t *dummyTable) Collation() sql.CollationID { return sql.Collation_Default }
   126  
   127  func (t *dummyTable) Partitions(ctx *sql.Context) (sql.PartitionIter, error) {
   128  	return &partitionIter{keys: t.keys}, nil
   129  }
   130  
   131  func (t *dummyTable) PartitionRows(ctx *sql.Context, partition sql.Partition) (sql.RowIter, error) {
   132  	rows, ok := t.partitions[string(partition.Key())]
   133  	if !ok {
   134  		return nil, sql.ErrPartitionNotFound.New(partition.Key())
   135  	}
   136  
   137  	return sql.RowsToRowIter(rows...), nil
   138  }
   139  
   140  type partition struct {
   141  	key []byte
   142  }
   143  
   144  func (p *partition) Key() []byte { return p.key }
   145  
   146  type partitionIter struct {
   147  	keys [][]byte
   148  	pos  int
   149  }
   150  
   151  func (p *partitionIter) Next(ctx *sql.Context) (sql.Partition, error) {
   152  	if p.pos >= len(p.keys) {
   153  		return nil, io.EOF
   154  	}
   155  
   156  	key := p.keys[p.pos]
   157  	p.pos++
   158  	return &partition{key}, nil
   159  }
   160  
   161  func (p *partitionIter) Close(_ *sql.Context) error { return nil }