Refactor AutoIncrementTracker into generitc type with interfaces

go/libraries/doltcore/doltdb/table.go

@@ -18,10 +18,16 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
+	"math"
 	"unicode"
 
+	"github.com/dolthub/go-mysql-server/sql"
+	gmstypes "github.com/dolthub/go-mysql-server/sql/types"
+
 	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb/durable"
 	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
+	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/globalstate/sequences"
 	"github.com/dolthub/dolt/go/store/hash"
 	"github.com/dolthub/dolt/go/store/prolly/tree"
 	"github.com/dolthub/dolt/go/store/types"
@@ -63,6 +69,8 @@ type Table struct {
 	overriddenSchema schema.Schema
 }
 
+var _ sequences.SequencedRelation[*Table, uint64, AutoIncrementState] = &Table{}
+
 // NewTable creates a durable object which stores row data, index data, and schema.
 func NewTable(ctx context.Context, vrw types.ValueReadWriter, ns tree.NodeStore, sch schema.Schema, rows durable.Index, indexes durable.IndexSet, autoIncVal types.Value) (*Table, error) {
 	dt, err := durable.NewTable(ctx, vrw, ns, sch, rows, indexes, autoIncVal)
@@ -406,22 +414,198 @@ func (t *Table) RenameIndexRowData(ctx context.Context, oldIndexName, newIndexNa
 	return t.SetIndexSet(ctx, indexes)
 }
 
+// HasSequenceState returns whether the table has an AUTO_INCREMENT value.
+func (t *Table) HasSequenceState(ctx context.Context) (bool, error) {
+	sch, err := t.GetSchema(ctx)
+	if err != nil {
+		return false, err
+	}
+	return schema.HasAutoIncrement(sch), nil
+}
+
+// GetSequenceState implements SequencedRelation
+func (t *Table) GetSequenceState(ctx context.Context) (AutoIncrementState, error) {
+	return t.GetAutoIncrementValue(ctx)
+}
+
 // GetAutoIncrementValue returns the current AUTO_INCREMENT value for this table.
-func (t *Table) GetAutoIncrementValue(ctx context.Context) (uint64, error) {
-	return t.table.GetAutoIncrement(ctx)
+func (t *Table) GetAutoIncrementValue(ctx context.Context) (AutoIncrementState, error) {
+	currentValue, err := t.table.GetAutoIncrement(ctx)
+	if err != nil {
+		return 0, err
+	}
+	return AutoIncrementState(currentValue), nil
+}
+
+func (t *Table) GetSequenceSqlType(ctx context.Context) (sql.Type, bool, error) {
+	sch, err := t.table.GetSchema(ctx)
+	if err != nil {
+		return nil, false, err
+	}
+
+	aiCol, ok := schema.GetAutoIncrementColumn(sch)
+	if !ok {
+		return nil, false, nil
+	}
+
+	return aiCol.TypeInfo.ToSqlType(), true, nil
 }
 
 // SetAutoIncrementValue sets the current AUTO_INCREMENT value for this table. This method does not verify that the
 // value given is greater than current table values. Setting it lower than current table values will result in
 // incorrect key generation on future inserts, causing duplicate key errors.
-func (t *Table) SetAutoIncrementValue(ctx context.Context, val uint64) (*Table, error) {
-	table, err := t.table.SetAutoIncrement(ctx, val)
+func (t *Table) SetAutoIncrementValue(ctx context.Context, val AutoIncrementState) (*Table, error) {
+	table, err := t.table.SetAutoIncrement(ctx, val.CurrentValue())
 	if err != nil {
 		return nil, err
 	}
 	return &Table{table: table}, nil
 }
 
+// SetSequenceState implements sequences.SequencedRelation
+func (t *Table) SetSequenceState(ctx context.Context, val AutoIncrementState) (*Table, error) {
+	return t.SetAutoIncrementValue(ctx, val)
+}
+
+// SetSequenceState implements sequences.SequencedRelation
+func (t *Table) TrySetSequenceState(ctx *sql.Context, newAutoIncVal AutoIncrementState) (*Table, bool, error) {
+	currentMax, err := getMaxAutoIncrementValue(ctx, t)
+	if err != nil {
+		return nil, false, err
+	}
+	currentMaxVal := AutoIncrementState(currentMax)
+
+	if !newAutoIncVal.GreaterThan(currentMaxVal) {
+		return t, false, nil
+	}
+
+	newTable, err := t.SetSequenceState(ctx, newAutoIncVal)
+	return newTable, true, err
+}
+
+// CoerceAutoIncrementValue converts |val| into an AUTO_INCREMENT sequence value
+func CoerceAutoIncrementValue(ctx *sql.Context, val interface{}) (uint64, error) {
+	switch typ := val.(type) {
+	case float32:
+		val = math.Round(float64(typ))
+	case float64:
+		val = math.Round(typ)
+	}
+
+	var err error
+	val, _, err = gmstypes.Uint64.Convert(ctx, val)
+	if err != nil {
+		return 0, err
+	}
+	if val == nil || val == uint64(0) {
+		return 0, nil
+	}
+	return val.(uint64), nil
+}
+
+type AutoIncrementState uint64
+
+func (s AutoIncrementState) Next() (aiVal uint64, ok bool, nextState AutoIncrementState, err error) {
+	if s == math.MaxUint64 {
+		return uint64(math.MaxUint64), false, s, nil
+	}
+	return uint64(s), true, s + 1, nil
+}
+
+func (s AutoIncrementState) CurrentValue() uint64 {
+	return uint64(s)
+}
+
+func (s AutoIncrementState) WithValue(v uint64) AutoIncrementState {
+	return AutoIncrementState(v)
+}
+
+func (s AutoIncrementState) WithSQLValue(ctx *sql.Context, v interface{}) (AutoIncrementState, error) {
+	given, err := CoerceAutoIncrementValue(ctx, v)
+	if err != nil {
+		return s, err
+	}
+	return AutoIncrementState(given), nil
+}
+
+func (s AutoIncrementState) GreaterThan(other AutoIncrementState) bool {
+	return s > other
+}
+
+func (s AutoIncrementState) Merge(other AutoIncrementState) (AutoIncrementState, bool) {
+	if s > other {
+		return s, true
+	}
+	return other, true
+}
+
+func (s AutoIncrementState) AtEnd() bool {
+	return s == math.MaxUint64
+}
+
+// getMaxAutoIncrementValue gets the highest value in a table's AUTO INCREMENT column
+func getMaxAutoIncrementValue(ctx *sql.Context, table *Table) (uint64, error) {
+	// First, establish whether to update this table based on the given value and its current max value.
+	sch, err := table.GetSchema(ctx)
+	if err != nil {
+		return 0, err
+	}
+
+	aiCol, ok := schema.GetAutoIncrementColumn(sch)
+	if !ok {
+		return 0, nil
+	}
+
+	var indexData durable.Index
+	aiIndex, ok := sch.Indexes().GetIndexByColumnNames(aiCol.Name)
+	if ok {
+		indexes, err := table.GetIndexSet(ctx)
+		if err != nil {
+			return 0, err
+		}
+
+		indexData, err = indexes.GetIndex(ctx, sch, nil, aiIndex.Name())
+		if err != nil {
+			return 0, err
+		}
+	} else {
+		indexData, err = table.GetRowData(ctx)
+		if err != nil {
+			return 0, err
+		}
+	}
+
+	maxValue, err := getMaxIndexValue(ctx, indexData)
+	if err != nil {
+		return 0, err
+	}
+	return CoerceAutoIncrementValue(ctx, maxValue)
+}
+
+// getMaxIndexValue reads the highest value for the first column in an index.
+func getMaxIndexValue(ctx *sql.Context, indexData durable.Index) (interface{}, error) {
+	idx, err := durable.ProllyMapFromIndex(indexData)
+	if err != nil {
+		return 0, err
+	}
+
+	iter, err := idx.IterAllReverse(ctx)
+	if err != nil {
+		return 0, err
+	}
+
+	kd, _ := idx.Descriptors()
+	k, _, err := iter.Next(ctx)
+	if err == io.EOF {
+		return 0, nil
+	} else if err != nil {
+		return 0, err
+	}
+
+	// TODO: is the auto-inc column always the first column in the index?
+	return tree.GetField(ctx, kd, 0, k, idx.NodeStore())
+}
+
 // AddColumnToRows adds the column named to row data as necessary and returns the resulting table.
 func (t *Table) AddColumnToRows(ctx context.Context, newCol string, newSchema schema.Schema) (*Table, error) {
 	idx, err := t.table.GetTableRows(ctx)

go/libraries/doltcore/merge/merge_prolly_rows.go

@@ -140,8 +140,8 @@ func mergeAutoIncrementValues(ctx context.Context, tbl, otherTbl, resultTbl *dol
 	if err != nil {
 		return nil, err
 	}
-	if autoVal < mergeAutoVal {
-		autoVal = mergeAutoVal
+	if mergeAutoVal.GreaterThan(autoVal) {
+		return resultTbl.SetAutoIncrementValue(ctx, mergeAutoVal)
 	}
 	return resultTbl.SetAutoIncrementValue(ctx, autoVal)
 }

go/libraries/doltcore/sqle/cluster/controller.go

@@ -976,10 +976,6 @@ func (c *Controller) refreshAutoIncrementTrackersForSessionDatabases() error {
 			// Non-versioned DBs don't participate in AUTO_INCREMENT global state
 			continue
 		}
-		ai, err := gsp.GetGlobalState().AutoIncrementTracker(sqlCtx)
-		if err != nil {
-			return fmt.Errorf("cluster/controller: auto-inc refresh: %s: tracker: %w", name, err)
-		}
 
 		// Get working set roots only
 		state, ok, err := sess.LookupDbState(sqlCtx, name)
@@ -990,7 +986,7 @@ func (c *Controller) refreshAutoIncrementTrackersForSessionDatabases() error {
 			// Not loaded in session; defer to lazy initialization on first use
 			continue
 		}
-		if err := ai.InitWithRoots(sqlCtx, state.WorkingSet()); err != nil {
+		if err := gsp.GetGlobalState().InitWithRoots(sqlCtx, state.WorkingSet()); err != nil {
 			return fmt.Errorf("cluster/controller: auto-inc refresh: %s: init: %w", name, err)
 		}
 	}

go/libraries/doltcore/sqle/database.go

@@ -1924,7 +1924,7 @@ func (db Database) removeTableFromAutoIncrementTracker(
 		wses = append(wses, ws)
 	}
 
-	ait, err := db.gs.AutoIncrementTracker(ctx)
+	ait, err := dsess.GetAutoIncrementTracker(ctx, db.gs)
 	if err != nil {
 		return err
 	}
@@ -2084,11 +2084,14 @@ func (db Database) createSqlTable(ctx *sql.Context, table string, schemaName str
 	// Prevent any tables that use BINARY, CHAR, VARBINARY, VARCHAR prefixes
 
 	if schema.HasAutoIncrement(doltSch) {
-		ait, err := db.gs.AutoIncrementTracker(ctx)
+		ait, err := dsess.GetAutoIncrementTracker(ctx, db.gs)
+		if err != nil {
+			return err
+		}
+		err = ait.AddNewTable(tableName.Name, doltdb.AutoIncrementState(1))
 		if err != nil {
 			return err
 		}
-		ait.AddNewTable(tableName.Name)
 	}
 
 	return db.createDoltTable(ctx, tableName.Name, tableName.Schema, root, doltSch)
@@ -2144,11 +2147,11 @@ func (db Database) createIndexedSqlTable(ctx *sql.Context, table string, schemaN
 	}
 
 	if schema.HasAutoIncrement(doltSch) {
-		ait, err := db.gs.AutoIncrementTracker(ctx)
+		ait, err := dsess.GetAutoIncrementTracker(ctx, db.gs)
 		if err != nil {
 			return err
 		}
-		ait.AddNewTable(tableName.Name)
+		ait.AddNewTable(tableName.Name, doltdb.AutoIncrementState(1))
 	}
 
 	return db.createDoltTable(ctx, tableName.Name, tableName.Schema, root, doltSch)

go/libraries/doltcore/sqle/dsess/auto_increment_tracker.go

@@ -0,0 +1,56 @@
+// Copyright 2023 Dolthub, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package dsess
+
+import (
+	"context"
+
+	"github.com/dolthub/go-mysql-server/sql"
+
+	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
+	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
+	"github.com/dolthub/dolt/go/libraries/doltcore/sqle/globalstate"
+)
+
+type DoltDBRelationSource struct{}
+
+func (s DoltDBRelationSource) GetRelation(ctx context.Context, root doltdb.RootValue, tName doltdb.TableName) (relation *doltdb.Table, resolvedName string, found bool, err error) {
+	return doltdb.GetTableInsensitive(ctx, root, tName)
+}
+
+func (s DoltDBRelationSource) GetRelations(ctx context.Context, root doltdb.RootValue, cb func(doltdb.TableName, *doltdb.Table) (bool, error)) error {
+	return root.IterTables(ctx, func(name doltdb.TableName, table *doltdb.Table, sch schema.Schema) (stop bool, err error) {
+		return cb(name, table)
+	})
+}
+
+var _ RelationSource[*doltdb.Table, doltdb.AutoIncrementState, uint64] = (*DoltDBRelationSource)(nil)
+
+type AutoIncrementTracker = SequenceTracker[*doltdb.Table, doltdb.AutoIncrementState, uint64]
+
+// NewAutoIncrementTracker returns a new autoincrement tracker for the roots given. All roots sets must be
+// considered because the auto increment value for a table is tracked globally, across all branches.
+// Roots provided should be the working sets when available, or the branches when they are not (e.g. for remote
+// branches that don't have a local working set)
+func NewAutoIncrementTracker(ctx context.Context, dbName string, roots ...doltdb.Rootish) (*AutoIncrementTracker, error) {
+	return NewSequenceTrackerFromRoots(ctx, dbName, DoltDBRelationSource{}, roots...)
+}
+
+func GetAutoIncrementTracker(ctx *sql.Context, gs globalstate.GlobalState) (*AutoIncrementTracker, error) {
+	return GetSequenceTracker(ctx, gs, autoIncrementTrackerKey)
+}
+
+// autoIncrementTrackerKey is the key used to store the AutoIncrmenetTracker in the GlobalState's map of SequenceTrackers
+var autoIncrementTrackerKey = TrackerKey[*AutoIncrementTracker]{}

go/libraries/doltcore/sqle/dsess/auto_increment_tracker_test.go

@@ -17,7 +17,6 @@ package dsess
 import (
 	"context"
 	"fmt"
-	"sync"
 	"testing"
 
 	"github.com/dolthub/go-mysql-server/sql"
@@ -68,7 +67,7 @@ func TestCoerceAutoIncrementValue(t *testing.T) {
 	for _, test := range tests {
 		name := fmt.Sprintf("Coerce %v to %v", test.val, test.exp)
 		t.Run(name, func(t *testing.T) {
-			act, err := CoerceAutoIncrementValue(ctx, test.val)
+			act, err := doltdb.CoerceAutoIncrementValue(ctx, test.val)
 			if test.err {
 				assert.Error(t, err)
 			} else {
@@ -83,7 +82,7 @@ func TestInitWithRoots(t *testing.T) {
 	t.Run("EmptyRoots", func(t *testing.T) {
 		ait := AutoIncrementTracker{
 			dbName:     "test_database",
-			sequences:  &sync.Map{},
+			sequences:  &SyncMap[string, doltdb.AutoIncrementState]{},
 			mm:         mutexmap.NewMutexMap(),
 			init:       make(chan struct{}),
 			cancelInit: make(chan struct{}),
@@ -94,7 +93,7 @@ func TestInitWithRoots(t *testing.T) {
 	t.Run("CloseCancelsInit", func(t *testing.T) {
 		ait := AutoIncrementTracker{
 			dbName:     "test_database",
-			sequences:  &sync.Map{},
+			sequences:  &SyncMap[string, doltdb.AutoIncrementState]{},
 			mm:         mutexmap.NewMutexMap(),
 			init:       make(chan struct{}),
 			cancelInit: make(chan struct{}),