Provide feedback

.gitignore

@@ -25,6 +25,7 @@ _testmain.go
 *.exe
 *.test
 *.prof
+*.gz
 
 # Tpc-h tbl files
 benchmark/*tbl

enginetest/mysqlshim/table.go

@@ -15,7 +15,9 @@
 package mysqlshim
 
 import (
+	"errors"
 	"fmt"
+	"math"
 	"math/rand"
 	"strings"
 
@@ -143,7 +145,16 @@ func (t Table) Truncate(ctx *sql.Context) (int, error) {
 		return 0, err
 	}
 	err = t.db.shim.Exec("", fmt.Sprintf("TRUNCATE TABLE `%s`;", t.name))
-	return int(rowCount.(int64)), err
+	if int64RowCount, ok := rowCount.(int64); ok {
+		if int64RowCount >= math.MinInt && int64RowCount <= math.MaxInt {
+			return int(int64RowCount), err
+		}
+	} else if uint64RowCount, ok := rowCount.(uint64); ok {
+		if uint64RowCount <= math.MaxInt {
+			return int(uint64RowCount), err
+		}
+	}
+	return 0, errors.New("unable to convers row count")
 }
 
 // AddColumn implements the interface sql.AlterableTable.

go.mod

@@ -8,7 +8,7 @@ require (
 	github.com/go-sql-driver/mysql v1.6.0
 	github.com/gocraft/dbr/v2 v2.7.2
 	github.com/google/flatbuffers v2.0.6+incompatible
-	github.com/google/uuid v1.2.0
+	github.com/google/uuid v1.3.1
 	github.com/hashicorp/golang-lru v0.5.4
 	github.com/lestrrat-go/strftime v1.0.4
 	github.com/mitchellh/hashstructure v1.1.0
@@ -20,27 +20,31 @@ require (
 	github.com/stretchr/testify v1.7.1
 	go.opentelemetry.io/otel v1.7.0
 	go.opentelemetry.io/otel/trace v1.7.0
-	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
-	golang.org/x/text v0.3.7
-	golang.org/x/tools v0.1.10
+	golang.org/x/sync v0.14.0
+	golang.org/x/text v0.25.0
+	golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d
 	gopkg.in/src-d/go-errors.v1 v1.0.0
 )
 
+require (
+	golang.org/x/crypto v0.38.0 // indirect
+	golang.org/x/net v0.40.0 // indirect
+	google.golang.org/genproto/googleapis/rpc v0.0.0-20230822172742-b8732ec3820d // indirect
+)
+
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
-	github.com/golang/protobuf v1.5.2 // indirect
+	github.com/golang/protobuf v1.5.3 // indirect
 	github.com/kr/text v0.2.0 // indirect
 	github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e // indirect
-	golang.org/x/mod v0.6.0-dev.0.20220106191415-9b9b3d81d5e3 // indirect
-	golang.org/x/sys v0.0.0-20211019181941-9d821ace8654 // indirect
-	golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 // indirect
-	google.golang.org/genproto v0.0.0-20210506142907-4a47615972c2 // indirect
-	google.golang.org/grpc v1.37.0 // indirect
-	google.golang.org/protobuf v1.27.1 // indirect
+	golang.org/x/mod v0.17.0 // indirect
+	golang.org/x/sys v0.33.0 // indirect
+	google.golang.org/grpc v1.59.0 // indirect
+	google.golang.org/protobuf v1.33.0 // indirect
 	gopkg.in/check.v1 v1.0.0-20200902074654-038fdea0a05b // indirect
-	gopkg.in/yaml.v3 v3.0.0-20200615113413-eeeca48fe776 // indirect
+	gopkg.in/yaml.v3 v3.0.0 // indirect
 )
 
 replace github.com/oliveagle/jsonpath => github.com/dolthub/jsonpath v0.0.0-20210609232853-d49537a30474
 
-go 1.19
+go 1.23.0

memory/table.go

@@ -23,6 +23,7 @@ import (
 	"sort"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/dolthub/vitess/go/sqltypes"
@@ -87,6 +88,7 @@ type Table struct {
 	columns         []int
 
 	// Data storage
+	partitionMux  sync.RWMutex
 	partitions    map[string][]sql.Row
 	partitionKeys [][]byte
 
@@ -170,6 +172,7 @@ func NewPartitionedTableWithCollation(name string, schema sql.PrimaryKeySchema,
 		schema:        schema,
 		fkColl:        fkColl,
 		collation:     collation,
+		partitionMux:  sync.RWMutex{},
 		partitions:    partitions,
 		partitionKeys: keys,
 		autoIncVal:    autoIncVal,
@@ -178,7 +181,7 @@ func NewPartitionedTableWithCollation(name string, schema sql.PrimaryKeySchema,
 }
 
 // Name implements the sql.Table interface.
-func (t Table) Name() string {
+func (t *Table) Name() string {
 	return t.name
 }
 
@@ -196,7 +199,9 @@ func (t *Table) Collation() sql.CollationID {
 }
 
 func (t *Table) GetPartition(key string) []sql.Row {
+	t.partitionMux.RLock()
 	rows, ok := t.partitions[string(key)]
+	t.partitionMux.RUnlock()
 	if ok {
 		return rows
 	}
@@ -208,9 +213,11 @@ func (t *Table) GetPartition(key string) []sql.Row {
 func (t *Table) Partitions(ctx *sql.Context) (sql.PartitionIter, error) {
 	var keys [][]byte
 	for _, k := range t.partitionKeys {
+		t.partitionMux.RLock()
 		if rows, ok := t.partitions[string(k)]; ok && len(rows) > 0 {
 			keys = append(keys, k)
 		}
+		t.partitionMux.RUnlock()
 	}
 	return &partitionIter{keys: keys}, nil
 }
@@ -245,7 +252,11 @@ type rangePartition struct {
 
 // PartitionCount implements the sql.PartitionCounter interface.
 func (t *Table) PartitionCount(ctx *sql.Context) (int64, error) {
-	return int64(len(t.partitions)), nil
+	t.partitionMux.RLock()
+	partLen := len(t.partitions)
+	t.partitionMux.RUnlock()
+
+	return int64(partLen), nil
 }
 
 // PartitionRows implements the sql.PartitionRows interface.
@@ -256,14 +267,17 @@ func (t *Table) PartitionRows(ctx *sql.Context, partition sql.Partition) (sql.Ro
 		filters = append(t.filters, r.rang)
 	}
 
+	t.partitionMux.RLock()
 	rows, ok := t.partitions[string(partition.Key())]
 	if !ok {
+		t.partitionMux.RUnlock()
 		return nil, sql.ErrPartitionNotFound.New(partition.Key())
 	}
 	// The slice could be altered by other operations taking place during iteration (such as deletion or insertion), so
 	// make a copy of the values as they exist when execution begins.
 	rowsCopy := make([]sql.Row, len(rows))
 	copy(rowsCopy, rows)
+	t.partitionMux.RUnlock()
 
 	return &tableIter{
 		rows:    rowsCopy,
@@ -274,9 +288,11 @@ func (t *Table) PartitionRows(ctx *sql.Context, partition sql.Partition) (sql.Ro
 
 func (t *Table) numRows(ctx *sql.Context) (uint64, error) {
 	var count uint64 = 0
+	t.partitionMux.RLock()
 	for _, rows := range t.partitions {
 		count += uint64(len(rows))
 	}
+	t.partitionMux.RUnlock()
 
 	return count, nil
 }
@@ -570,10 +586,12 @@ func (t *Table) newTableEditor() *tableEditor {
 
 func (t *Table) Truncate(ctx *sql.Context) (int, error) {
 	count := 0
+	t.partitionMux.Lock()
 	for key := range t.partitions {
 		count += len(t.partitions[key])
 		t.partitions[key] = nil
 	}
+	t.partitionMux.Unlock()
 	return count, nil
 }
 
@@ -715,6 +733,7 @@ func (t *Table) addColumnToSchema(ctx *sql.Context, newCol *sql.Column, order *s
 		t.autoIncVal = 0
 
 		if newColIdx < len(t.schema.Schema) {
+			t.partitionMux.RLock()
 			for _, p := range t.partitions {
 				for _, row := range p {
 					if row[newColIdx] == nil {
@@ -736,6 +755,7 @@ func (t *Table) addColumnToSchema(ctx *sql.Context, newCol *sql.Column, order *s
 					}
 				}
 			}
+			t.partitionMux.RUnlock()
 		} else {
 			t.autoIncVal = 0
 		}
@@ -758,8 +778,10 @@ func (t *Table) addColumnToSchema(ctx *sql.Context, newCol *sql.Column, order *s
 }
 
 func (t *Table) insertValueInRows(ctx *sql.Context, idx int, colDefault *sql.ColumnDefaultValue) error {
+	t.partitionMux.RLock()
 	for k, p := range t.partitions {
-		newP := make([]sql.Row, len(p))
+		partLen := len(p)
+		newP := make([]sql.Row, partLen)
 		for i, row := range p {
 			var newRow sql.Row
 			newRow = append(newRow, row[:idx]...)
@@ -776,13 +798,19 @@ func (t *Table) insertValueInRows(ctx *sql.Context, idx int, colDefault *sql.Col
 			}
 			newP[i] = newRow
 		}
+		t.partitionMux.RUnlock()
+		t.partitionMux.Lock()
 		t.partitions[k] = newP
+		t.partitionMux.Unlock()
+		t.partitionMux.RLock()
 	}
+	t.partitionMux.RUnlock()
 	return nil
 }
 
 func (t *Table) DropColumn(ctx *sql.Context, columnName string) error {
 	droppedCol := t.dropColumnFromSchema(ctx, columnName)
+	t.partitionMux.RLock()
 	for k, p := range t.partitions {
 		newP := make([]sql.Row, len(p))
 		for i, row := range p {
@@ -791,8 +819,13 @@ func (t *Table) DropColumn(ctx *sql.Context, columnName string) error {
 			newRow = append(newRow, row[droppedCol+1:]...)
 			newP[i] = newRow
 		}
+		t.partitionMux.RUnlock()
+		t.partitionMux.Lock()
 		t.partitions[k] = newP
+		t.partitionMux.Unlock()
+		t.partitionMux.RLock()
 	}
+	t.partitionMux.RUnlock()
 	return nil
 }
 
@@ -857,6 +890,7 @@ func (t *Table) ModifyColumn(ctx *sql.Context, columnName string, column *sql.Co
 		}
 	}
 
+	t.partitionMux.RLock()
 	for k, p := range t.partitions {
 		newP := make([]sql.Row, len(p))
 		for i, row := range p {
@@ -876,8 +910,14 @@ func (t *Table) ModifyColumn(ctx *sql.Context, columnName string, column *sql.Co
 			newRow = append(newRow, oldRowWithoutVal[newIdx:]...)
 			newP[i] = newRow
 		}
+
+		t.partitionMux.RUnlock()
+		t.partitionMux.Lock()
 		t.partitions[k] = newP
+		t.partitionMux.Unlock()
+		t.partitionMux.RLock()
 	}
+	t.partitionMux.RUnlock()
 
 	pkNameToOrdIdx := make(map[string]int)
 	for i, ord := range t.schema.PkOrdinals {
@@ -1306,6 +1346,8 @@ func (t *Table) errIfDuplicateEntryExist(cols []string, idxName string) error {
 		return err
 	}
 	unique := make(map[uint64]struct{})
+	t.partitionMux.RLock()
+	defer t.partitionMux.RUnlock()
 	for _, partition := range t.partitions {
 		for _, row := range partition {
 			idxPrefixKey := projectOnRow(columnMapping, row)
@@ -1516,14 +1558,18 @@ func (t *Table) sortRows() {
 	}
 
 	var idx []partidx
+	t.partitionMux.RLock()
 	for _, k := range t.partitionKeys {
 		p := t.partitions[string(k)]
 		for i := 0; i < len(p); i++ {
 			idx = append(idx, partidx{string(k), i})
 		}
 	}
+	t.partitionMux.RUnlock()
 
-	sort.Sort(partitionssort{t.partitions, idx, less})
+	t.partitionMux.Lock()
+	defer t.partitionMux.Unlock()
+	sort.Sort(partitionssort{&t.partitionMux, t.partitions, idx, less})
 }
 
 type partidx struct {
@@ -1532,9 +1578,10 @@ type partidx struct {
 }
 
 type partitionssort struct {
-	ps   map[string][]sql.Row
-	idx  []partidx
-	less func(l, r sql.Row) bool
+	partitionMux *sync.RWMutex
+	ps           map[string][]sql.Row
+	idx          []partidx
+	less         func(l, r sql.Row) bool
 }
 
 func (ps partitionssort) Len() int {
@@ -1576,10 +1623,14 @@ func copyschema(sch sql.Schema) sql.Schema {
 }
 
 func newTable(t *Table, newSch sql.PrimaryKeySchema) (*Table, error) {
-	newTable := NewPartitionedTableWithCollation(t.name, newSch, t.fkColl, len(t.partitions), t.collation)
+	partitionCount, _ := t.PartitionCount(nil)
+
+	newTable := NewPartitionedTableWithCollation(t.name, newSch, t.fkColl, int(partitionCount), t.collation)
 	for _, partition := range t.partitions {
 		for _, partitionRow := range partition {
+			t.partitionMux.Lock()
 			err := newTable.Insert(sql.NewEmptyContext(), partitionRow)
+			t.partitionMux.Unlock()
 			if err != nil {
 				return nil, err
 			}