diff --git a/go/libraries/doltcore/doltdb/table.go b/go/libraries/doltcore/doltdb/table.go index f2b147cc448..f76565bec93 100644 --- a/go/libraries/doltcore/doltdb/table.go +++ b/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) diff --git a/go/libraries/doltcore/merge/merge_prolly_rows.go b/go/libraries/doltcore/merge/merge_prolly_rows.go index c26bdc84dc6..109212b2619 100644 --- a/go/libraries/doltcore/merge/merge_prolly_rows.go +++ b/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) } diff --git a/go/libraries/doltcore/sqle/cluster/controller.go b/go/libraries/doltcore/sqle/cluster/controller.go index 218c354da0d..05502049cf8 100644 --- a/go/libraries/doltcore/sqle/cluster/controller.go +++ b/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) } } diff --git a/go/libraries/doltcore/sqle/database.go b/go/libraries/doltcore/sqle/database.go index a287c78926d..194dc85ff90 100644 --- a/go/libraries/doltcore/sqle/database.go +++ b/go/libraries/doltcore/sqle/database.go @@ -1877,7 +1877,7 @@ func (db Database) dropTable(ctx *sql.Context, tableName string) error { if schema.HasAutoIncrement(sch) { ddb, _ := ds.GetDoltDB(ctx, db.RevisionQualifiedName()) - err = db.removeTableFromAutoIncrementTracker(ctx, tableName, ddb, ws.Ref()) + err = db.removeTableFromAutoIncrementTracker(ctx, tblName, ddb, ws.Ref()) if err != nil { return err } @@ -1892,7 +1892,7 @@ func (db Database) dropTable(ctx *sql.Context, tableName string) error { // otherwise. This operation is expensive if the func (db Database) removeTableFromAutoIncrementTracker( ctx *sql.Context, - tableName string, + tableName doltdb.TableName, ddb *doltdb.DoltDB, ws ref.WorkingSetRef, ) error { @@ -1924,12 +1924,12 @@ 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 } - err = ait.DropTable(ctx, tableName, wses...) + err = ait.DropRelation(ctx, tableName, wses...) 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.AddNewRelation(tableName, 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,14 @@ 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 + } + err = ait.AddNewRelation(tableName, doltdb.AutoIncrementState(1)) if err != nil { return err } - ait.AddNewTable(tableName.Name) } return db.createDoltTable(ctx, tableName.Name, tableName.Schema, root, doltSch) diff --git a/go/libraries/doltcore/sqle/dsess/auto_increment_tracker.go b/go/libraries/doltcore/sqle/dsess/auto_increment_tracker.go new file mode 100644 index 00000000000..9d0dd53415e --- /dev/null +++ b/go/libraries/doltcore/sqle/dsess/auto_increment_tracker.go @@ -0,0 +1,68 @@ +// Copyright 2026 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/dolt/go/libraries/doltcore/schema" + "iter" + + "github.com/dolthub/go-mysql-server/sql" + + "github.com/dolthub/dolt/go/libraries/doltcore/doltdb" + "github.com/dolthub/dolt/go/libraries/doltcore/sqle/globalstate" +) + +// DoltDBRelationSource implements RelationSource +// Specializations of SequenceTracker (such as AutoIncrementTracker) require an interface to read relations of a specified +// type out of a doltdb.RootValue. DoltDBRelationSource provides the ability to read values of type doltdb.Table. +type DoltDBRelationSource struct{} + +// GetRelation implements RelationSource +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) +} + +// IterRelations implements RelationSource +func (s DoltDBRelationSource) IterRelations(ctx context.Context, root doltdb.RootValue) iter.Seq2[doltdb.TableName, *doltdb.Table] { + return func(yield func(doltdb.TableName, *doltdb.Table) bool) { + _ = root.IterTables(ctx, func(name doltdb.TableName, table *doltdb.Table, sch schema.Schema) (stop bool, err error) { + if !yield(name, table) { + return true, nil + } + return false, nil + }) + } +} + +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...) +} + +// GetAutoIncrementTracker returns the AutoIncrementTracker stored within the global state. +func GetAutoIncrementTracker(ctx *sql.Context, gs globalstate.GlobalState) (*AutoIncrementTracker, error) { + return GetSequenceTracker(ctx, gs, autoIncrementTrackerKey) +} + +// autoIncrementTrackerKey is the key used to store the AutoIncrementTracker in the GlobalState's map of SequenceTrackers +var autoIncrementTrackerKey = TrackerKey[*AutoIncrementTracker]{} diff --git a/go/libraries/doltcore/sqle/dsess/auto_increment_tracker_test.go b/go/libraries/doltcore/sqle/dsess/auto_increment_tracker_test.go index 03f03b87b55..a4f460bdc3c 100644 --- a/go/libraries/doltcore/sqle/dsess/auto_increment_tracker_test.go +++ b/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{}), diff --git a/go/libraries/doltcore/sqle/dsess/autoincrement_tracker.go b/go/libraries/doltcore/sqle/dsess/autoincrement_tracker.go deleted file mode 100644 index 60b55ef63f9..00000000000 --- a/go/libraries/doltcore/sqle/dsess/autoincrement_tracker.go +++ /dev/null @@ -1,703 +0,0 @@ -// 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" - "errors" - "fmt" - "io" - "math" - "strings" - "sync" - "time" - - "github.com/dolthub/go-mysql-server/sql" - gmstypes "github.com/dolthub/go-mysql-server/sql/types" - "golang.org/x/sync/errgroup" - - "github.com/dolthub/dolt/go/libraries/doltcore/doltdb" - "github.com/dolthub/dolt/go/libraries/doltcore/doltdb/durable" - "github.com/dolthub/dolt/go/libraries/doltcore/doltdb/gcctx" - "github.com/dolthub/dolt/go/libraries/doltcore/ref" - "github.com/dolthub/dolt/go/libraries/doltcore/schema" - "github.com/dolthub/dolt/go/libraries/doltcore/sqle/dsess/mutexmap" - "github.com/dolthub/dolt/go/libraries/doltcore/sqle/globalstate" - "github.com/dolthub/dolt/go/store/prolly/tree" -) - -type LockMode int64 - -var ( - LockMode_Traditional LockMode = 0 - LockMode_Concurrent LockMode = 1 - LockMode_Interleaved LockMode = 2 -) - -type AutoIncrementTracker struct { - initErr error - sequences *sync.Map // map[string]uint64 - mm *mutexmap.MutexMap - // AutoIncrementTracker is lazily initialized by loading - // tracker state for every given |root|. On first access, we - // block on initialization being completed and we terminally - // return |initErr| if there was any error initializing. - init chan struct{} - // To clean up effectively we need to stop all access to - // storage. As part of that, we have the possibility to cancel - // async initialization and block on the process completing. - cancelInit chan struct{} - dbName string - // lockMode is the effective @@innodb_autoinc_lock_mode at the time of AutoIncrementTracker initialization. - // This value can only be set by config and cannot be changed in a running server. - lockMode LockMode -} - -// currentLockMode returns the effective @@innodb_autoinc_lock_mode stored in global server vars -func currentLockMode() LockMode { - _, i, _ := sql.SystemVariables.GetGlobal("innodb_autoinc_lock_mode") - if mode, ok := i.(int64); ok { - return LockMode(mode) - } - return LockMode_Interleaved -} - -var _ globalstate.AutoIncrementTracker = &AutoIncrementTracker{} - -// 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) { - ait := AutoIncrementTracker{ - dbName: dbName, - sequences: &sync.Map{}, - mm: mutexmap.NewMutexMap(), - init: make(chan struct{}), - cancelInit: make(chan struct{}), - } - gcSafepointController := getGCSafepointController(ctx) - ctx = context.Background() - if gcSafepointController != nil { - ctx = gcctx.WithGCSafepointController(ctx, gcSafepointController) - } - go func() { - if gcSafepointController != nil { - defer gcctx.SessionEnd(ctx) - gcctx.SessionCommandBegin(ctx) - defer gcctx.SessionCommandEnd(ctx) - } - ait.initWithRoots(ctx, roots...) - }() - return &ait, nil -} - -func getGCSafepointController(ctx context.Context) *gcctx.GCSafepointController { - if sqlCtx, ok := ctx.(*sql.Context); ok { - return DSessFromSess(sqlCtx.Session).GCSafepointController() - } - return gcctx.GetGCSafepointController(ctx) -} - -func loadAutoIncValue(sequences *sync.Map, tableName string) (current uint64, hasCurrent bool) { - tableName = strings.ToLower(tableName) - stored, hasCurrent := sequences.Load(tableName) - if !hasCurrent { - return 0, false - } - return stored.(uint64), true -} - -func (a *AutoIncrementTracker) initializeTableAutoIncrement(ctx *sql.Context, tableName string) (uint64, bool, error) { - sess := DSessFromSess(ctx.Session) - ws, err := sess.WorkingSet(ctx, a.dbName) - if err != nil { - return 0, false, err - } - - table, _, ok, err := doltdb.GetTableInsensitive(ctx, ws.WorkingRoot(), doltdb.TableName{Name: tableName}) - if err != nil || !ok { - return 0, false, err - } - - sch, err := table.GetSchema(ctx) - if err != nil { - return 0, false, err - } - if !schema.HasAutoIncrement(sch) { - return 0, false, nil - } - - seq, err := table.GetAutoIncrementValue(ctx) - if err != nil { - return 0, false, err - } - - table, err = a.deepSet(ctx, tableName, table, ws.Ref(), seq) - if err != nil { - return 0, false, err - } - - seq, ok = loadAutoIncValue(a.sequences, tableName) - if ok { - return seq, true, nil - } - - seq, err = table.GetAutoIncrementValue(ctx) - if err != nil { - return 0, false, err - } - a.sequences.Store(strings.ToLower(tableName), seq) - return seq, true, nil -} - -func (a *AutoIncrementTracker) Close() { - close(a.cancelInit) - <-a.init -} - -// Current returns the next value to be generated in the auto increment sequence for |tableName|. -func (a *AutoIncrementTracker) Current(tableName string) (uint64, error) { - err := a.waitForInit() - if err != nil { - return 0, err - } - seq, ok := loadAutoIncValue(a.sequences, tableName) - if !ok { - return 0, nil - } - return seq, nil -} - -// Next returns the next auto increment value for |tbl| using |insertVal| from an insert. If |insertVal| is -// null or 0, it is generated from the sequence. -func (a *AutoIncrementTracker) Next(ctx *sql.Context, tbl string, insertVal interface{}) (uint64, error) { - err := a.waitForInit() - if err != nil { - return 0, err - } - - tbl = strings.ToLower(tbl) - - given, err := CoerceAutoIncrementValue(ctx, insertVal) - if err != nil { - return 0, err - } - - // The read-modify-write of the sequence below must be atomic across concurrent inserters. In - // interleaved lock mode (the default) the engine holds no statement-level lock, so we take a - // short per-table lock here. - locked := false - if a.lockMode == LockMode_Interleaved { - release := a.mm.Lock(tbl) - defer release() - locked = true - } - - curr, ok := loadAutoIncValue(a.sequences, tbl) - if !ok { - // Missing tracker state after initialization can happen when a running sql-server discovers a database - // restored after startup, so initialize it here. - if !locked { - if a.lockMode == LockMode_Interleaved { - release := a.mm.Lock(tbl) - defer release() - locked = true - } - - curr, ok = loadAutoIncValue(a.sequences, tbl) - } - - if !ok { - curr, ok, err = a.initializeTableAutoIncrement(ctx, tbl) - if err != nil { - return 0, err - } - if !ok { - return 0, fmt.Errorf("autoIncrementTracker: unable to find sequence for table %s", tbl) - } - } - } - - if given == 0 { - // |given| is 0 or NULL - a.sequences.Store(tbl, curr+1) - return curr, nil - } - - if given >= curr { - // Check if the given value is valid for this column type - if !a.validateAutoIncrementBounds(ctx, tbl, given, false) { - return given, nil // Out of bounds, don't update sequence - } - - // Value is valid, determine next sequence value - nextVal := given - if a.validateAutoIncrementBounds(ctx, tbl, given, true) { - nextVal++ - } - a.sequences.Store(tbl, nextVal) - return given, nil - } - - // |given| < curr - return given, nil -} - -func (a *AutoIncrementTracker) CoerceAutoIncrementValue(ctx *sql.Context, val interface{}) (uint64, error) { - err := a.waitForInit() - if err != nil { - return 0, err - } - return CoerceAutoIncrementValue(ctx, val) -} - -// 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 -} - -// Set sets the auto increment value for the table named, if it's greater than the one already registered for this -// table. Otherwise, the update is silently disregarded. So far this matches the MySQL behavior, but Dolt uses the -// maximum value for this table across all branches. -func (a *AutoIncrementTracker) Set(ctx *sql.Context, tableName string, table *doltdb.Table, ws ref.WorkingSetRef, newAutoIncVal uint64) (*doltdb.Table, error) { - err := a.waitForInit() - if err != nil { - return nil, err - } - - tableName = strings.ToLower(tableName) - - release := a.mm.Lock(tableName) - defer release() - - existing, ok := loadAutoIncValue(a.sequences, tableName) - if !ok { - existing = 0 - } - if newAutoIncVal > existing && a.validateAutoIncrementBounds(ctx, tableName, newAutoIncVal, true) { - a.sequences.Store(tableName, newAutoIncVal) - return table.SetAutoIncrementValue(ctx, newAutoIncVal) - } else if newAutoIncVal > existing { - // Value is greater but out of bounds, don't update - return table, nil - } - // Value is not greater than current, do deep check across branches - return a.deepSet(ctx, tableName, table, ws, newAutoIncVal) -} - -// deepSet sets the auto increment value for the table named, if it's greater than the one on any branch head for this -// database, ignoring the current in-memory tracker value -func (a *AutoIncrementTracker) deepSet(ctx *sql.Context, tableName string, table *doltdb.Table, ws ref.WorkingSetRef, newAutoIncVal uint64) (*doltdb.Table, error) { - sess := DSessFromSess(ctx.Session) - db, ok := sess.Provider().BaseDatabase(ctx, a.dbName) - - // just give up if we can't find this db for any reason, or it's a non-versioned DB - if !ok || !db.Versioned() { - return table, nil - } - - // 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 nil, err - } - - aiCol, ok := schema.GetAutoIncrementColumn(sch) - if !ok { - return nil, nil - } - - var indexData durable.Index - aiIndex, ok := sch.Indexes().GetIndexByColumnNames(aiCol.Name) - if ok { - indexes, err := table.GetIndexSet(ctx) - if err != nil { - return nil, err - } - - indexData, err = indexes.GetIndex(ctx, sch, nil, aiIndex.Name()) - if err != nil { - return nil, err - } - } else { - indexData, err = table.GetRowData(ctx) - if err != nil { - return nil, err - } - } - - currentMax, err := getMaxIndexValue(ctx, indexData) - if err != nil { - return nil, err - } - - // If the given value is less than the current one, the operation is a no-op, bail out early - if newAutoIncVal <= currentMax { - return table, nil - } - - table, err = table.SetAutoIncrementValue(ctx, newAutoIncVal) - if err != nil { - return nil, err - } - - // Now that we have established the current max for this table, reset the global max accordingly - maxAutoInc := newAutoIncVal - doltdbs := db.DoltDatabases() - for _, db := range doltdbs { - branches, err := db.GetBranches(ctx) - if err != nil { - return nil, err - } - - remotes, err := db.GetRemoteRefs(ctx) - if err != nil { - return nil, err - } - - rootRefs := make([]ref.DoltRef, 0, len(branches)+len(remotes)) - rootRefs = append(rootRefs, branches...) - rootRefs = append(rootRefs, remotes...) - - for _, b := range rootRefs { - var rootish doltdb.Rootish - switch b.GetType() { - case ref.BranchRefType: - wsRef, err := ref.WorkingSetRefForHead(b) - if err != nil { - return nil, err - } - - if wsRef == ws { - // we don't need to check the working set we're updating - continue - } - - ws, err := db.ResolveWorkingSet(ctx, wsRef) - if err == doltdb.ErrWorkingSetNotFound { - // use the branch head if there isn't a working set for it - cm, err := db.ResolveCommitRef(ctx, b) - if err != nil { - return nil, err - } - rootish = cm - } else if err != nil { - return nil, err - } else { - rootish = ws - } - case ref.RemoteRefType: - cm, err := db.ResolveCommitRef(ctx, b) - if err != nil { - return nil, err - } - rootish = cm - } - - root, err := rootish.ResolveRootValue(ctx) - if err != nil { - return nil, err - } - - table, _, ok, err := doltdb.GetTableInsensitive(ctx, root, doltdb.TableName{Name: tableName}) - if err != nil { - return nil, err - } - if !ok { - continue - } - - sch, err := table.GetSchema(ctx) - if err != nil { - return nil, err - } - - if !schema.HasAutoIncrement(sch) { - continue - } - - tableName = strings.ToLower(tableName) - seq, err := table.GetAutoIncrementValue(ctx) - if err != nil { - return nil, err - } - - if seq > maxAutoInc { - maxAutoInc = seq - } - } - } - - if a.validateAutoIncrementBounds(ctx, tableName, maxAutoInc, true) { - a.sequences.Store(tableName, maxAutoInc) - } - return table, nil -} - -func getMaxIndexValue(ctx *sql.Context, indexData durable.Index) (uint64, 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? - field, err := tree.GetField(ctx, kd, 0, k, idx.NodeStore()) - if err != nil { - return 0, err - } - - maxVal, err := CoerceAutoIncrementValue(ctx, field) - if err != nil { - return 0, err - } - - return maxVal, nil -} - -// AddNewTable initializes a new table with an auto increment column to the tracker, as necessary -func (a *AutoIncrementTracker) AddNewTable(tableName string) error { - err := a.waitForInit() - if err != nil { - return err - } - - tableName = strings.ToLower(tableName) - // only initialize the sequence for this table if no other branch has such a table - a.sequences.LoadOrStore(tableName, uint64(1)) - return nil -} - -// DropTable drops the table with the name given. -// To establish the new auto increment value, callers must also pass all other working sets in scope that may include -// a table with the same name, omitting the working set that just deleted the table named. -func (a *AutoIncrementTracker) DropTable(ctx *sql.Context, tableName string, wses ...*doltdb.WorkingSet) error { - err := a.waitForInit() - if err != nil { - return err - } - - tableName = strings.ToLower(tableName) - - release := a.mm.Lock(tableName) - defer release() - - newHighestValue := uint64(1) - - // Get the new highest value from all tables in the working sets given - for _, ws := range wses { - table, _, exists, err := doltdb.GetTableInsensitive(ctx, ws.WorkingRoot(), doltdb.TableName{Name: tableName}) - if err != nil { - return err - } - - if !exists { - continue - } - - sch, err := table.GetSchema(ctx) - if err != nil { - return err - } - - if schema.HasAutoIncrement(sch) { - seq, err := table.GetAutoIncrementValue(ctx) - if err != nil { - return err - } - - if seq > newHighestValue { - newHighestValue = seq - } - } - } - - a.sequences.Store(tableName, newHighestValue) - - return nil -} - -func (a *AutoIncrementTracker) AcquireTableLock(ctx *sql.Context, tableName string) (func(), error) { - err := a.waitForInit() - if err != nil { - return nil, err - } - - if a.lockMode == LockMode_Interleaved { - // This shouldn't be possible, it's a serious programming error if it happens - panic("Attempted to acquire AutoInc lock for entire insert operation, but lock mode was set to Interleaved") - } - return a.mm.Lock(tableName), nil -} - -func (a *AutoIncrementTracker) waitForInit() error { - select { - case <-a.init: - return a.initErr - case <-time.After(5 * time.Minute): - return errors.New("failed to initialize autoincrement tracker") - } -} - -// This method will initialize the AutoIncrementTracker state with all -// data from the tables found in |roots|. This method closes the -// |a.init| channel when it completes. It is meant to be run in a -// goroutine, as in `go a.initWithRoots(...)`. When running this method, -// a newly allocated |a.init| channel should exist. -// -// It is the caller's responsibility to ensure that whatever |ctx| -// |initWithRoots| is called with appropriately outlives the end of -// the method and that it participates in GC lifecycle callbacks -// appropriately, if that is necessary. -func (a *AutoIncrementTracker) initWithRoots(ctx context.Context, roots ...doltdb.Rootish) { - defer close(a.init) - - // Cancel the parent context so that the errgroup work will - // complete with an error if we see cancelInit closed. - finishedCh := make(chan struct{}) - defer close(finishedCh) - ctx, cancel := context.WithCancelCause(ctx) - go func() { - select { - case <-a.cancelInit: - cancel(errors.New("initialization canceled. did not complete successfully.")) - case <-finishedCh: - } - }() - - eg, ctx := errgroup.WithContext(ctx) - eg.SetLimit(128) - - for _, root := range roots { - eg.Go(func() error { - if ctx.Err() != nil { - return context.Cause(ctx) - } - - r, err := root.ResolveRootValue(ctx) - if err != nil { - return err - } - - return r.IterTables(ctx, func(tableName doltdb.TableName, table *doltdb.Table, sch schema.Schema) (bool, error) { - if !schema.HasAutoIncrement(sch) { - return false, nil - } - - seq, err := table.GetAutoIncrementValue(ctx) - if err != nil { - return true, err - } - - tableNameStr := tableName.ToLower().Name - if oldValue, loaded := a.sequences.LoadOrStore(tableNameStr, seq); loaded { - old := oldValue.(uint64) - for seq > old && !a.sequences.CompareAndSwap(tableNameStr, old, seq) { - oldValue, _ = a.sequences.Load(tableNameStr) - old = oldValue.(uint64) - - } - } - - return false, nil - }) - }) - } - - a.lockMode = currentLockMode() - a.initErr = eg.Wait() -} - -// validateAutoIncrementBounds checks if a value (or value+1 if checkIncrement) is valid for the auto-increment column type -func (a *AutoIncrementTracker) validateAutoIncrementBounds(ctx *sql.Context, tbl string, val uint64, checkIncrement bool) bool { - sess := DSessFromSess(ctx.Session) - db, ok := sess.Provider().BaseDatabase(ctx, a.dbName) - if !ok || !db.Versioned() { - return true // fail-open for infrastructure errors - } - - ws, err := sess.WorkingSet(ctx, a.dbName) - if err != nil { - return true - } - - table, _, ok, err := doltdb.GetTableInsensitive(ctx, ws.WorkingRoot(), doltdb.TableName{Name: tbl}) - if err != nil || !ok { - return true - } - - sch, err := table.GetSchema(ctx) - if err != nil { - return true - } - - aiCol, ok := schema.GetAutoIncrementColumn(sch) - if !ok { - return true - } - - sqlType := aiCol.TypeInfo.ToSqlType() - - testVal := val - if checkIncrement { - // Check if incrementing would overflow - nextVal := val + 1 - if nextVal < val { - return false // uint64 overflow - } - testVal = nextVal - } - - _, inRange, err := sqlType.Convert(ctx, testVal) - return err == nil && inRange == sql.InRange -} - -func (a *AutoIncrementTracker) InitWithRoots(ctx context.Context, roots ...doltdb.Rootish) error { - err := a.waitForInit() - if err != nil { - return err - } - a.init = make(chan struct{}) - go a.initWithRoots(ctx, roots...) - return a.waitForInit() -} diff --git a/go/libraries/doltcore/sqle/dsess/globalstate.go b/go/libraries/doltcore/sqle/dsess/globalstate.go index b28c8e6e666..a2f153bd7c5 100644 --- a/go/libraries/doltcore/sqle/dsess/globalstate.go +++ b/go/libraries/doltcore/sqle/dsess/globalstate.go @@ -16,7 +16,6 @@ package dsess import ( "context" - "sync" "github.com/dolthub/go-mysql-server/sql" "golang.org/x/sync/errgroup" @@ -24,17 +23,25 @@ import ( "github.com/dolthub/dolt/go/libraries/doltcore/doltdb" "github.com/dolthub/dolt/go/libraries/doltcore/ref" "github.com/dolthub/dolt/go/libraries/doltcore/sqle/globalstate" + "github.com/dolthub/dolt/go/libraries/doltcore/sqle/globalstate/sequences" ) -func NewGlobalStateStoreForDb(ctx context.Context, dbName string, db *doltdb.DoltDB) (GlobalStateImpl, error) { +// TrackerKey is a type used as a key into the GlobalStateImpl's map of SequenceTrackers +type TrackerKey[TrackerType globalstate.SequenceTrackerBase] struct{} + +func NewSequenceTracker[ + RelationType sequences.SequencedRelation[RelationType, ValueType, StateType], + StateType sequences.SequenceState[StateType, ValueType], + ValueType comparable, +](ctx context.Context, dbName string, db *doltdb.DoltDB, relationSource RelationSource[RelationType, StateType, ValueType]) (*SequenceTracker[RelationType, StateType, ValueType], error) { branches, err := db.GetBranches(ctx) if err != nil { - return GlobalStateImpl{}, err + return nil, err } remotes, err := db.GetRemoteRefs(ctx) if err != nil { - return GlobalStateImpl{}, err + return nil, err } rootRefs := make([]ref.DoltRef, 0, len(branches)+len(remotes)) @@ -85,31 +92,59 @@ func NewGlobalStateStoreForDb(ctx context.Context, dbName string, db *doltdb.Dol err = eg.Wait() if err != nil { - return GlobalStateImpl{}, err + return nil, err } - tracker, err := NewAutoIncrementTracker(ctx, dbName, roots...) + return NewSequenceTrackerFromRoots(ctx, dbName, relationSource, roots...) +} + +func NewGlobalStateStoreForDb(ctx context.Context, dbName string, db *doltdb.DoltDB) (GlobalStateImpl, error) { + autoIncrementTracker, err := NewSequenceTracker(ctx, dbName, db, DoltDBRelationSource{}) if err != nil { return GlobalStateImpl{}, err } - return GlobalStateImpl{ - aiTracker: tracker, - mu: &sync.Mutex{}, + sequenceTrackers: map[interface{}]globalstate.SequenceTrackerBase{autoIncrementTrackerKey: autoIncrementTracker}, }, nil } type GlobalStateImpl struct { - aiTracker *AutoIncrementTracker - mu *sync.Mutex + sequenceTrackers map[interface{}]globalstate.SequenceTrackerBase } var _ globalstate.GlobalState = GlobalStateImpl{} -func (g GlobalStateImpl) AutoIncrementTracker(ctx *sql.Context) (globalstate.AutoIncrementTracker, error) { - return g.aiTracker, nil +func (g GlobalStateImpl) GetSequenceTracker(ctx *sql.Context, key interface{}) (globalstate.SequenceTrackerBase, error) { + return g.sequenceTrackers[key], nil +} + +func (g GlobalStateImpl) AddSequenceTracker(ctx *sql.Context, key interface{}, tracker globalstate.SequenceTrackerBase) error { + g.sequenceTrackers[key] = tracker + return nil } func (g GlobalStateImpl) Close() { - g.aiTracker.Close() + for _, tracker := range g.sequenceTrackers { + tracker.Close() + } +} + +func (g GlobalStateImpl) InitWithRoots(ctx *sql.Context, roots ...doltdb.Rootish) error { + for _, tracker := range g.sequenceTrackers { + err := tracker.InitWithRoots(ctx, roots...) + if err != nil { + return err + } + } + return nil +} + +// GetSequenceTracker returns a SequenceTracker held by the globalstate.GlobalState, keyed by the provided key. +// This function performs the necessary cast so that the caller doesn't have to cast the result. +func GetSequenceTracker[T globalstate.SequenceTrackerBase](ctx *sql.Context, gs globalstate.GlobalState, key TrackerKey[T]) (result T, err error) { + aiti, err := gs.GetSequenceTracker(ctx, key) + if err != nil { + return result, err + } + return aiti.(T), nil } diff --git a/go/libraries/doltcore/sqle/dsess/sequence_tracker.go b/go/libraries/doltcore/sqle/dsess/sequence_tracker.go new file mode 100644 index 00000000000..4f21313fc81 --- /dev/null +++ b/go/libraries/doltcore/sqle/dsess/sequence_tracker.go @@ -0,0 +1,647 @@ +// Copyright 2026 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" + "errors" + "fmt" + "iter" + "time" + + "github.com/dolthub/go-mysql-server/sql" + "golang.org/x/sync/errgroup" + + "github.com/dolthub/dolt/go/libraries/doltcore/doltdb" + "github.com/dolthub/dolt/go/libraries/doltcore/doltdb/gcctx" + "github.com/dolthub/dolt/go/libraries/doltcore/ref" + "github.com/dolthub/dolt/go/libraries/doltcore/sqle/dsess/mutexmap" + "github.com/dolthub/dolt/go/libraries/doltcore/sqle/globalstate" + "github.com/dolthub/dolt/go/libraries/doltcore/sqle/globalstate/sequences" +) + +type LockMode int64 + +var ( + LockMode_Traditional LockMode = 0 + LockMode_Concurrent LockMode = 1 + LockMode_Interleaved LockMode = 2 +) + +// A RelationSource maps table names to relations (which may be tables or root objects) at a supplied RootValue +type RelationSource[ + RelationType sequences.SequencedRelation[RelationType, ValueType, StateType], + StateType sequences.SequenceState[StateType, ValueType], + ValueType comparable, +] interface { + // GetRelation gets a relation at a specific doltdb.RootValue + GetRelation(ctx context.Context, root doltdb.RootValue, tName doltdb.TableName) (relation RelationType, resolvedName string, found bool, err error) + IterRelations(ctx context.Context, root doltdb.RootValue) iter.Seq2[doltdb.TableName, RelationType] +} + +type SequenceTracker[ + RelationType sequences.SequencedRelation[RelationType, ValueType, StateType], + StateType sequences.SequenceState[StateType, ValueType], + ValueType comparable, +] struct { + initErr error + sequences *SyncMap[doltdb.TableName, StateType] + mm *mutexmap.MutexMap + // SequenceTracker is lazily initialized by loading + // tracker state for every given |root|. On first access, we + // block on initialization being completed and we terminally + // return |initErr| if there was any error initializing. + init chan struct{} + // To clean up effectively we need to stop all access to + // storage. As part of that, we have the possibility to cancel + // async initialization and block on the process completing. + cancelInit chan struct{} + dbName string + // lockMode is the effective @@innodb_autoinc_lock_mode at the time of SequenceTracker initialization. + // This value can only be set by config and cannot be changed in a running server. + lockMode LockMode + // relationSource is how the tracker reads objects from a RootValue. + // It may read tables or RootObjects. + relationSource RelationSource[RelationType, StateType, ValueType] +} + +// currentLockMode returns the effective @@innodb_autoinc_lock_mode stored in global server vars +func currentLockMode() LockMode { + _, i, _ := sql.SystemVariables.GetGlobal("innodb_autoinc_lock_mode") + if mode, ok := i.(int64); ok { + return LockMode(mode) + } + return LockMode_Interleaved +} + +// staticAssertTypes contains compile-time assertions that SequenceTracker implements interfaces. +// It does not need to be called. +func (a *SequenceTracker[RelationType, StateType, ValueType]) staticAssertTypes() { + var _ globalstate.SequenceTracker[RelationType, StateType, ValueType] = a +} + +// NewSequenceTrackerFromRoots creates and initializes a new SequenceTracker by querying |relationSource| for +// objects that need to be globally tracked, and computing a single tracked global state for each object by merging +// the state at every root in |roots| +func NewSequenceTrackerFromRoots[ + RelationType sequences.SequencedRelation[RelationType, ValueType, StateType], + StateType sequences.SequenceState[StateType, ValueType], + ValueType comparable, +](ctx context.Context, dbName string, relationSource RelationSource[RelationType, StateType, ValueType], roots ...doltdb.Rootish) (*SequenceTracker[RelationType, StateType, ValueType], error) { + ait := SequenceTracker[RelationType, StateType, ValueType]{ + dbName: dbName, + sequences: &SyncMap[doltdb.TableName, StateType]{}, + mm: mutexmap.NewMutexMap(), + init: make(chan struct{}), + cancelInit: make(chan struct{}), + relationSource: relationSource, + } + gcSafepointController := getGCSafepointController(ctx) + ctx = context.Background() + if gcSafepointController != nil { + ctx = gcctx.WithGCSafepointController(ctx, gcSafepointController) + } + go func() { + if gcSafepointController != nil { + defer gcctx.SessionEnd(ctx) + gcctx.SessionCommandBegin(ctx) + defer gcctx.SessionCommandEnd(ctx) + } + ait.initWithRoots(ctx, roots...) + }() + return &ait, nil +} + +func getGCSafepointController(ctx context.Context) *gcctx.GCSafepointController { + if sqlCtx, ok := ctx.(*sql.Context); ok { + return DSessFromSess(sqlCtx.Session).GCSafepointController() + } + return gcctx.GetGCSafepointController(ctx) +} + +func loadSequenceState[StateType sequences.SequenceState[StateType, ValueType], ValueType comparable](sequences *SyncMap[doltdb.TableName, StateType], relationName doltdb.TableName) (current StateType, hasCurrent bool) { + return sequences.Load(relationName.ToLower()) +} + +func (a *SequenceTracker[RelationType, StateType, ValueType]) initializeSequenceState(ctx *sql.Context, relationName doltdb.TableName, initialValue interface{}) (state StateType, hasState bool, err error) { + sess := DSessFromSess(ctx.Session) + ws, err := sess.WorkingSet(ctx, a.dbName) + if err != nil { + return state, false, err + } + + table, _, ok, err := a.relationSource.GetRelation(ctx, ws.WorkingRoot(), relationName) + if err != nil || !ok { + return state, false, err + } + + hasAutoIncrement, err := table.HasSequenceState(ctx) + if err != nil { + return state, false, err + } + + var seq StateType + if !hasAutoIncrement { + // Create a new state based on the provided value + // TODO: This could cause problems when we need to create the more + // complicated sequence types for Doltgres + seq, err = state.WithSQLValue(ctx, initialValue) + if err != nil { + return state, false, err + } + } else { + seq, err = table.GetSequenceState(ctx) + if err != nil { + return state, false, err + } + } + + relation, err := a.deepSet(ctx, relationName, table, ws.Ref(), seq) + if err != nil { + return state, false, err + } + + state, ok = loadSequenceState(a.sequences, relationName) + if ok { + return state, true, nil + } + + seq, err = relation.GetSequenceState(ctx) + if err != nil { + return state, false, err + } + a.sequences.Store(relationName.ToLower(), seq) + return state, true, nil +} + +func (a *SequenceTracker[RelationType, StateType, ValueType]) Close() { + close(a.cancelInit) + <-a.init +} + +// Current returns the next value to be generated in the auto increment sequence for |relationName|. +func (a *SequenceTracker[RelationType, StateType, ValueType]) Current(relationName doltdb.TableName) (current StateType, err error) { + err = a.waitForInit() + if err != nil { + return current, err + } + seq, ok := loadSequenceState(a.sequences, relationName) + if !ok { + return current, nil + } + return seq, nil +} + +// Next returns the next auto increment value for |relationName| using |insertVal| from an insert. If |insertVal| is +// null or 0, it is generated from the sequence. +func (a *SequenceTracker[RelationType, StateType, ValueType]) Next(ctx *sql.Context, relationName doltdb.TableName, insertVal interface{}) (nextValue ValueType, err error) { + err = a.waitForInit() + if err != nil { + return nextValue, err + } + + relationName = relationName.ToLower() + + // The read-modify-write of the sequence below must be atomic across concurrent inserters. In + // interleaved lock mode (the default) the engine holds no statement-level lock, so we take a + // short per-table lock here. + locked := false + if a.lockMode == LockMode_Interleaved { + release := a.mm.Lock(relationName) + defer release() + locked = true + } + + currState, ok := loadSequenceState(a.sequences, relationName) + if !ok { + // Missing tracker state after initialization can happen when a running sql-server discovers a database + // restored after startup, so initialize it here. + if !locked { + if a.lockMode == LockMode_Interleaved { + release := a.mm.Lock(relationName) + defer release() + locked = true + } + + currState, ok = loadSequenceState(a.sequences, relationName) + } + + if !ok { + currState, ok, err = a.initializeSequenceState(ctx, relationName, insertVal) + if err != nil { + return nextValue, err + } + if !ok { + return nextValue, fmt.Errorf("autoIncrementTracker: unable to find sequence for table %s", relationName.Name) + } + } + } + + if insertVal == nil { + // |given| is 0 or NULL + currentVal, _, nextState, err := currState.Next() + if err != nil { + return nextValue, err + } + a.sequences.Store(relationName, nextState) + return currentVal, nil + } + + givenState, err := currState.WithSQLValue(ctx, insertVal) + if err != nil { + return nextValue, err + } + given := givenState.CurrentValue() + + if !currState.GreaterThan(givenState) { + // Check if the given value is valid for this column type + if !a.validateBounds(ctx, relationName, givenState, false) { + return givenState.CurrentValue(), nil // Out of bounds, don't update sequence + } + + // Value is valid, determine next sequence value + if a.validateBounds(ctx, relationName, givenState, true) { + _, _, givenState, err = givenState.Next() + if err != nil { + return nextValue, err + } + } + a.sequences.Store(relationName, givenState) + return given, nil + } + + return given, nil +} + +// Set sets the auto increment value for the table named, if it's greater than the one already registered for this +// table. Otherwise, the update is silently disregarded. So far this matches the MySQL behavior, but Dolt uses the +// maximum value for this table across all branches. +func (a *SequenceTracker[RelationType, StateType, ValueType]) Set(ctx *sql.Context, relationName doltdb.TableName, table RelationType, ws ref.WorkingSetRef, newSequenceState StateType) (newRelation RelationType, err error) { + err = a.waitForInit() + if err != nil { + return newRelation, err + } + + relationName = relationName.ToLower() + + release := a.mm.Lock(relationName) + defer release() + + existing, ok := loadSequenceState(a.sequences, relationName) + if !ok { + a.sequences.Store(relationName, newSequenceState) + return table.SetSequenceState(ctx, newSequenceState) + } + gt := newSequenceState.GreaterThan(existing) + if gt && a.validateBounds(ctx, relationName, newSequenceState, false) { + a.sequences.Store(relationName, newSequenceState) + return table.SetSequenceState(ctx, newSequenceState) + } else if gt { + // Value is greater but out of bounds, don't update + return table, nil + } + // Value is not greater than current, do deep check across branches + return a.deepSet(ctx, relationName, table, ws, newSequenceState) +} + +// deepSet sets the sequence state for the table named, if it's greater than the one on any branch head for this +// database, ignoring the current in-memory tracker value +func (a *SequenceTracker[RelationType, StateType, ValueType]) deepSet(ctx *sql.Context, relationName doltdb.TableName, table RelationType, ws ref.WorkingSetRef, newAutoIncVal StateType) (newRelation RelationType, err error) { + sess := DSessFromSess(ctx.Session) + db, ok := sess.Provider().BaseDatabase(ctx, a.dbName) + + // just give up if we can't find this db for any reason, or it's a non-versioned DB + if !ok || !db.Versioned() { + return table, nil + } + + table, success, err := table.TrySetSequenceState(ctx, newAutoIncVal) + if err != nil { + return newRelation, err + } + if !success { + return table, nil + } + + // Now that we have established the current max for this table, reset the global max accordingly + maxAutoInc := newAutoIncVal + doltdbs := db.DoltDatabases() + for _, db := range doltdbs { + branches, err := db.GetBranches(ctx) + if err != nil { + return newRelation, err + } + + remotes, err := db.GetRemoteRefs(ctx) + if err != nil { + return newRelation, err + } + + rootRefs := make([]ref.DoltRef, 0, len(branches)+len(remotes)) + rootRefs = append(rootRefs, branches...) + rootRefs = append(rootRefs, remotes...) + + for _, b := range rootRefs { + var rootish doltdb.Rootish + switch b.GetType() { + case ref.BranchRefType: + wsRef, err := ref.WorkingSetRefForHead(b) + if err != nil { + return newRelation, err + } + + if wsRef == ws { + // we don't need to check the working set we're updating + continue + } + + ws, err := db.ResolveWorkingSet(ctx, wsRef) + if err == doltdb.ErrWorkingSetNotFound { + // use the branch head if there isn't a working set for it + cm, err := db.ResolveCommitRef(ctx, b) + if err != nil { + return newRelation, err + } + rootish = cm + } else if err != nil { + return newRelation, err + } else { + rootish = ws + } + case ref.RemoteRefType: + cm, err := db.ResolveCommitRef(ctx, b) + if err != nil { + return newRelation, err + } + rootish = cm + } + + root, err := rootish.ResolveRootValue(ctx) + if err != nil { + return newRelation, err + } + + table, _, ok, err := a.relationSource.GetRelation(ctx, root, relationName) + if err != nil { + return newRelation, err + } + if !ok { + continue + } + + hasAutoIncrement, err := table.HasSequenceState(ctx) + if err != nil { + return newRelation, err + } + + if !hasAutoIncrement { + continue + } + + seq, err := table.GetSequenceState(ctx) + if err != nil { + return newRelation, err + } + + var mergeOk bool + maxAutoInc, mergeOk = maxAutoInc.Merge(seq) + if !mergeOk { + // TODO: This can't happen with AUTO INCREMENT but needs to be + // handled for sequences. + } + } + } + + if a.validateBounds(ctx, relationName, maxAutoInc, false) { + a.sequences.Store(relationName, maxAutoInc) + } + return table, nil +} + +// AddNewRelation initializes a new table with an auto increment column to the tracker, as necessary +func (a *SequenceTracker[RelationType, StateType, ValueType]) AddNewRelation(relationName doltdb.TableName, initialState StateType) error { + err := a.waitForInit() + if err != nil { + return err + } + + // only initialize the sequence for this table if no other branch has such a table + a.sequences.LoadOrStore(relationName.ToLower(), initialState) + return nil +} + +// DropRelation drops the table with the name given. +// To establish the new auto increment value, callers must also pass all other working sets in scope that may include +// a table with the same name, omitting the working set that just deleted the table named. +func (a *SequenceTracker[RelationType, StateType, ValueType]) DropRelation(ctx *sql.Context, relationName doltdb.TableName, wses ...*doltdb.WorkingSet) error { + err := a.waitForInit() + if err != nil { + return err + } + + relationName = relationName.ToLower() + + release := a.mm.Lock(relationName) + defer release() + + var newHighestValue *StateType + + // Get the new highest value from all tables in the working sets given + for _, ws := range wses { + table, _, exists, err := a.relationSource.GetRelation(ctx, ws.WorkingRoot(), relationName) + if err != nil { + return err + } + + if !exists { + continue + } + + hasAutoIncrement, err := table.HasSequenceState(ctx) + if err != nil { + return err + } + if hasAutoIncrement { + seq, err := table.GetSequenceState(ctx) + if err != nil { + return err + } + if newHighestValue == nil { + newHighestValue = &seq + } else { + var ok bool + *newHighestValue, ok = (*newHighestValue).Merge(seq) + if !ok { + // TODO: This can't happen with AUTO INCREMENT but needs to be + // handled for sequences. + } + } + + } + } + + if newHighestValue != nil { + a.sequences.Store(relationName, *newHighestValue) + } else { + a.sequences.Delete(relationName) + } + + return nil +} + +func (a *SequenceTracker[RelationType, StateType, ValueType]) AcquireLock(ctx *sql.Context, relationName doltdb.TableName) (func(), error) { + err := a.waitForInit() + if err != nil { + return nil, err + } + + if a.lockMode == LockMode_Interleaved { + // This shouldn't be possible, it's a serious programming error if it happens + panic("Attempted to acquire AutoInc lock for entire insert operation, but lock mode was set to Interleaved") + } + return a.mm.Lock(relationName), nil +} + +func (a *SequenceTracker[RelationType, StateType, ValueType]) waitForInit() error { + select { + case <-a.init: + return a.initErr + case <-time.After(5 * time.Minute): + return errors.New("failed to initialize autoincrement tracker") + } +} + +// This method will initialize the SequenceTracker state with all +// data from the tables found in |roots|. This method closes the +// |a.init| channel when it completes. It is meant to be run in a +// goroutine, as in `go a.initWithRoots(...)`. When running this method, +// a newly allocated |a.init| channel should exist. +// +// It is the caller's responsibility to ensure that whatever |ctx| +// |initWithRoots| is called with appropriately outlives the end of +// the method and that it participates in GC lifecycle callbacks +// appropriately, if that is necessary. +func (a *SequenceTracker[RelationType, StateType, ValueType]) initWithRoots(ctx context.Context, roots ...doltdb.Rootish) { + defer close(a.init) + + // Cancel the parent context so that the errgroup work will + // complete with an error if we see cancelInit closed. + finishedCh := make(chan struct{}) + defer close(finishedCh) + ctx, cancel := context.WithCancelCause(ctx) + go func() { + select { + case <-a.cancelInit: + cancel(errors.New("initialization canceled. did not complete successfully.")) + case <-finishedCh: + } + }() + + eg, ctx := errgroup.WithContext(ctx) + eg.SetLimit(128) + + for _, root := range roots { + eg.Go(func() error { + if ctx.Err() != nil { + return context.Cause(ctx) + } + + r, err := root.ResolveRootValue(ctx) + if err != nil { + return err + } + + for relationName, relation := range a.relationSource.IterRelations(ctx, r) { + hasSequenceState, err := relation.HasSequenceState(ctx) + if err != nil { + return err + } + if !hasSequenceState { + continue + } + seq, err := relation.GetSequenceState(ctx) + if err != nil { + return err + } + + key := relationName.ToLower() + if oldValue, loaded := a.sequences.LoadOrStore(key, seq); loaded { + for seq.GreaterThan(oldValue) && !a.sequences.CompareAndSwap(key, oldValue, seq) { + oldValue, _ = a.sequences.Load(key) + } + } + } + return nil + }) + } + + a.lockMode = currentLockMode() + a.initErr = eg.Wait() +} + +// validateAutoIncrementBounds checks if a value (or value+1 if checkIncrement) is valid for the auto-increment column type +func (a *SequenceTracker[RelationType, StateType, ValueType]) validateBounds(ctx *sql.Context, relationName doltdb.TableName, val StateType, checkIncrement bool) bool { + sess := DSessFromSess(ctx.Session) + db, ok := sess.Provider().BaseDatabase(ctx, a.dbName) + if !ok || !db.Versioned() { + return true // fail-open for infrastructure errors + } + + ws, err := sess.WorkingSet(ctx, a.dbName) + if err != nil { + return true + } + + table, _, ok, err := a.relationSource.GetRelation(ctx, ws.WorkingRoot(), relationName) + if err != nil || !ok { + return true + } + + hasSequenceState, err := table.HasSequenceState(ctx) + if !hasSequenceState { + // fail-open because the table writer could be in the process of adding auto-increment to a column + return true + } + + sqlType, ok, err := table.GetSequenceSqlType(ctx) + if err != nil || !ok { + return true + } + + testVal := val + if checkIncrement { + // TODO: Remove error parameter? + _, hasNext, nextVal, _ := val.Next() + // SequenceState can only error if there is no next value. + // Consider changing this to a separate |ok| return value. + if !hasNext { + return false + } + testVal = nextVal + } + + _, inRange, err := sqlType.Convert(ctx, testVal.CurrentValue()) + return err == nil && inRange == sql.InRange +} + +func (a *SequenceTracker[RelationType, StateType, ValueType]) InitWithRoots(ctx context.Context, roots ...doltdb.Rootish) error { + err := a.waitForInit() + if err != nil { + return err + } + a.init = make(chan struct{}) + go a.initWithRoots(ctx, roots...) + return a.waitForInit() +} diff --git a/go/libraries/doltcore/sqle/dsess/session.go b/go/libraries/doltcore/sqle/dsess/session.go index 5d4894ac4b8..9adea55325d 100644 --- a/go/libraries/doltcore/sqle/dsess/session.go +++ b/go/libraries/doltcore/sqle/dsess/session.go @@ -1144,12 +1144,7 @@ func (d *DoltSession) ResetGlobals(ctx *sql.Context, dbName string, root doltdb. return err } - tracker, err := sessionState.dbState.globalState.AutoIncrementTracker(ctx) - if err != nil { - return err - } - - err = tracker.InitWithRoots(ctx, root) + err = sessionState.dbState.globalState.InitWithRoots(ctx, root) if err != nil { return err } @@ -1412,7 +1407,7 @@ func (d *DoltSession) addDB(ctx *sql.Context, db SqlDatabase) error { } sessionState.globalState = stateProvider.GetGlobalState() - tracker, err := sessionState.globalState.AutoIncrementTracker(ctx) + tracker, err := GetAutoIncrementTracker(ctx, sessionState.globalState) if err != nil { return err } @@ -2080,4 +2075,4 @@ func DefaultHead(ctx *sql.Context, baseName string, db SqlDatabase) (string, err // WriteSessFunc is a constructor that session builders use to // create fresh table editors. // The indirection avoids a writer/dsess package import cycle. -type WriteSessFunc func(dbName string, ws *doltdb.WorkingSet, aiTracker globalstate.AutoIncrementTracker, setter SessionRootSetter, opts editor.Options) WriteSession +type WriteSessFunc func(dbName string, ws *doltdb.WorkingSet, aiTracker *AutoIncrementTracker, setter SessionRootSetter, opts editor.Options) WriteSession diff --git a/go/libraries/doltcore/sqle/dsess/sync_map.go b/go/libraries/doltcore/sqle/dsess/sync_map.go new file mode 100644 index 00000000000..791c0fb151b --- /dev/null +++ b/go/libraries/doltcore/sqle/dsess/sync_map.go @@ -0,0 +1,51 @@ +// Copyright 2026 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 "sync" + +// SyncMap is a simple generic wrapper around sync.Map, designed for type safety around callsites. +// Using this type instead of sync.Map removes the need for casts. +type SyncMap[Key any, Value any] struct { + m sync.Map +} + +func (a *SyncMap[Key, Value]) Store(key Key, val Value) { + a.m.Store(key, val) +} + +func (a *SyncMap[Key, Value]) Load(key Key) (val Value, ok bool) { + v, ok := a.m.Load(key) + if !ok { + return val, ok + } + return v.(Value), ok +} + +func (a *SyncMap[Key, Value]) Delete(key Key) { + a.m.Delete(key) +} + +func (a *SyncMap[Key, Value]) LoadOrStore(key Key, val Value) (actual Value, loaded bool) { + act, loaded := a.m.LoadOrStore(key, val) + if !loaded { + return actual, loaded + } + return act.(Value), loaded +} + +func (a *SyncMap[Key, Value]) CompareAndSwap(key Key, old Value, new Value) (swapped bool) { + return a.m.CompareAndSwap(key, old, new) +} diff --git a/go/libraries/doltcore/sqle/globalstate/auto_increment_tracker.go b/go/libraries/doltcore/sqle/globalstate/auto_increment_tracker.go deleted file mode 100644 index f8aed7a5498..00000000000 --- a/go/libraries/doltcore/sqle/globalstate/auto_increment_tracker.go +++ /dev/null @@ -1,49 +0,0 @@ -// Copyright 2021 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 globalstate - -import ( - "context" - - "github.com/dolthub/go-mysql-server/sql" - - "github.com/dolthub/dolt/go/libraries/doltcore/doltdb" - "github.com/dolthub/dolt/go/libraries/doltcore/ref" -) - -// AutoIncrementTracker knows how to get and set the current auto increment value for a table. It's defined as an -// interface here because implementations need to reach into session state, requiring a dependency on this package. -type AutoIncrementTracker interface { - // Current returns the current auto increment value for the given table. - Current(tableName string) (uint64, error) - // Next returns the next auto increment value for the given table, and increments the current value. - Next(ctx *sql.Context, tbl string, insertVal interface{}) (uint64, error) - // AddNewTable adds a new table to the tracker, initializing the auto increment value to 1. - AddNewTable(tableName string) error - // DropTable removes a table from the tracker. - DropTable(ctx *sql.Context, tableName string, wses ...*doltdb.WorkingSet) error - // CoerceAutoIncrementValue coerces the given value to a uint64, returning an error if it can't be done. - CoerceAutoIncrementValue(ctx *sql.Context, val interface{}) (uint64, error) - // Set sets the auto increment value for the given table. This operation may silently do nothing if this value is - // below the current value for this table. The table in the provided working set is assumed to already have the value - // given, so the new global maximum is computed without regard for its value in that working set. - Set(ctx *sql.Context, tableName string, table *doltdb.Table, ws ref.WorkingSetRef, newAutoIncVal uint64) (*doltdb.Table, error) - // AcquireTableLock acquires the auto increment lock on a table, and returns a callback function to release the lock. - // Depending on the value of the `innodb_autoinc_lock_mode` system variable, the engine may need to acquire and hold - // the lock for the duration of an insert statement. - AcquireTableLock(ctx *sql.Context, tableName string) (func(), error) - // InitWithRoots fills the AutoIncrementTracker with values pulled from each root in order. - InitWithRoots(ctx context.Context, roots ...doltdb.Rootish) error -} diff --git a/go/libraries/doltcore/sqle/globalstate/global_state.go b/go/libraries/doltcore/sqle/globalstate/global_state.go index 720e6e901b0..cf08580ef2d 100644 --- a/go/libraries/doltcore/sqle/globalstate/global_state.go +++ b/go/libraries/doltcore/sqle/globalstate/global_state.go @@ -14,13 +14,21 @@ package globalstate -import "github.com/dolthub/go-mysql-server/sql" +import ( + "github.com/dolthub/go-mysql-server/sql" + + "github.com/dolthub/dolt/go/libraries/doltcore/doltdb" +) // GlobalState is just a holding interface for pieces of global state, of which the auto increment tracking info is // the only example at the moment. type GlobalState interface { - // AutoIncrementTracker returns the auto increment tracker for this global state. - AutoIncrementTracker(ctx *sql.Context) (AutoIncrementTracker, error) + // GetSequenceTracker returns the auto increment tracker for this global state. + GetSequenceTracker(ctx *sql.Context, key interface{}) (SequenceTrackerBase, error) + // AddSequenceTracker adds a new SequenceTracker to the GlobalState, accessible by the provided key. + AddSequenceTracker(ctx *sql.Context, key interface{}, value SequenceTrackerBase) error + // InitWithRoots initializes all of the state's SequenceTrackers + InitWithRoots(ctx *sql.Context, roots ...doltdb.Rootish) error } // GlobalStateProvider is an optional interface for databases that provide global state tracking diff --git a/go/libraries/doltcore/sqle/globalstate/sequence_tracker.go b/go/libraries/doltcore/sqle/globalstate/sequence_tracker.go new file mode 100644 index 00000000000..454c61b9e04 --- /dev/null +++ b/go/libraries/doltcore/sqle/globalstate/sequence_tracker.go @@ -0,0 +1,59 @@ +// Copyright 2026 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 globalstate + +import ( + "context" + + "github.com/dolthub/go-mysql-server/sql" + + "github.com/dolthub/dolt/go/libraries/doltcore/doltdb" + "github.com/dolthub/dolt/go/libraries/doltcore/ref" + "github.com/dolthub/dolt/go/libraries/doltcore/sqle/globalstate/sequences" +) + +// SequenceTrackerBase is the non-generic base interface for SequenceTracker +// It is useful for establishing an upper bound on type parameters in some circumstances. +type SequenceTrackerBase interface { + // AcquireLock acquires the lock on the global state for a specific relation, and returns a callback function to release the lock. + // Depending on the value of the `innodb_autoinc_lock_mode` system variable, the engine may need to acquire and hold + // the lock for the duration of an insert statement. + AcquireLock(ctx *sql.Context, tableName doltdb.TableName) (func(), error) + // DropRelation removes a relation from the tracker. + DropRelation(ctx *sql.Context, tableName doltdb.TableName, wses ...*doltdb.WorkingSet) error + // InitWithRoots fills the SequenceTracker with values pulled from each root in order. + InitWithRoots(ctx context.Context, roots ...doltdb.Rootish) error + Close() +} + +// SequenceTracker knows how to get and set the current sequence state for a relation (a table or a root object). +// It's defined as an interface here because implementations need to reach into session state, requiring a dependency on this package. +type SequenceTracker[ + RelationType sequences.SequencedRelation[RelationType, ValueType, StateType], + StateType sequences.SequenceState[StateType, ValueType], + ValueType comparable, +] interface { + SequenceTrackerBase + // Current returns the current sequence state for the given relation. + Current(tableName doltdb.TableName) (StateType, error) + // Next returns the next SQL value produced by the given relation, and advances that relation's state. + Next(ctx *sql.Context, tableName doltdb.TableName, insertVal interface{}) (ValueType, error) + // AddNewRelation adds a new table to the tracker, initializing the sequence state to the provided |initialState|. + AddNewRelation(tableName doltdb.TableName, initialState StateType) error + // Set sets the sequence state for the given relation. This operation may silently do nothing if this value is + // below the current value for this relation. The relation in the provided working set is assumed to already have the value + // given, so the new global maximum is computed without regard for its value in that working set. + Set(ctx *sql.Context, tableName doltdb.TableName, table RelationType, ws ref.WorkingSetRef, newSequenceState StateType) (RelationType, error) +} diff --git a/go/libraries/doltcore/sqle/globalstate/sequences/doc.go b/go/libraries/doltcore/sqle/globalstate/sequences/doc.go new file mode 100644 index 00000000000..9cf7290992b --- /dev/null +++ b/go/libraries/doltcore/sqle/globalstate/sequences/doc.go @@ -0,0 +1,18 @@ +// Copyright 2026 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. + +// This package is a collection of interfaces that describe state machines that are controlled by the global state. +// It is a separate package from the parent package globalstate so that doltdb can depend on it. + +package sequences diff --git a/go/libraries/doltcore/sqle/globalstate/sequences/state.go b/go/libraries/doltcore/sqle/globalstate/sequences/state.go new file mode 100644 index 00000000000..026475bcda7 --- /dev/null +++ b/go/libraries/doltcore/sqle/globalstate/sequences/state.go @@ -0,0 +1,64 @@ +// Copyright 2026 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 sequences + +import ( + "context" + + "github.com/dolthub/go-mysql-server/sql" +) + +// A SequenceState is an incrementing state that must be shared across all branches and transactions. +// It corresponds to a table or root object in the database. +// |Self| should always be the same type as the implementation. +// It produces a sequence of |ValueType| each time Next() is called. +type SequenceState[Self any, ValueType comparable] interface { + // Next advances the state of the sequence, producing a SQL value and the next SequenceState. + Next() (sqlVal ValueType, hasNext bool, nextState Self, err error) + // CurrentValue returns the next SQL value in the sequence without advancing it. + CurrentValue() (sqlVal ValueType) + // WithValue sets the state such that the next SQL value will be the one provided, then returns the new SequenceState. + WithValue(sqlVal ValueType) Self + // WithSQLValue coerces the input into the required type, then sets the state such that the next SQL value will be the one provided, then returns the new SequenceState. + WithSQLValue(ctx *sql.Context, v interface{}) (Self, error) + // GreaterThan compares two states and determines which one has advanced further. + GreaterThan(other Self) bool + // AtEnd returns whether the sequence is at its end. If true, then subsequent calls to SequenceState will + // no longer advance the sequence and may return an error. + AtEnd() bool + // Merge combines two states, usually returning the one that's further along. If the states can't be combined, + // |ok| is false. + Merge(other Self) (merged Self, ok bool) +} + +// A SequencedRelation is a table or root object with a state that must be shared across all branches +// and transactions in order to ensure global uniqueness of created rows. It wraps a SequenceState. +// Because some implementations are immutable (doltdb.Table), the methods that set the state must return a new +// value instead of self-mutating. +type SequencedRelation[Self any, ValueType comparable, StateType SequenceState[StateType, ValueType]] interface { + // GetSequenceState returns the current SequenceState of the object. + GetSequenceState(ctx context.Context) (StateType, error) + // HasSequenceState returns whether the relation wraps a sequence. + // (This may be false, for instance, for tables that do not have an AUTO INCREMENT column) + HasSequenceState(ctx context.Context) (bool, error) + // GetSequenceSqlType returns the SQL type generated by the sequence. + GetSequenceSqlType(ctx context.Context) (sql.Type, bool, error) + // SetSequenceState unconditionally sets the SequenceState for the object. + SetSequenceState(ctx context.Context, val StateType) (Self, error) + // TrySetSequenceState attempts to set the SequenceState for the object, but may fail if the provided state is + // incompatible with the other state of the object. Returns a bool indicating success. + // (For example, setting a table's AUTO INCREMENT value to be less than a value currently written to the table will fail.) + TrySetSequenceState(ctx *sql.Context, val StateType) (Self, bool, error) +} diff --git a/go/libraries/doltcore/sqle/tables.go b/go/libraries/doltcore/sqle/tables.go index d4d0902d62a..365fae4f32f 100644 --- a/go/libraries/doltcore/sqle/tables.go +++ b/go/libraries/doltcore/sqle/tables.go @@ -455,7 +455,11 @@ func (t *DoltTable) PeekNextAutoIncrementValue(ctx *sql.Context) (uint64, error) if err != nil { return 0, err } - return table.GetAutoIncrementValue(ctx) + seq, err := table.GetAutoIncrementValue(ctx) + if err != nil { + return 0, err + } + return uint64(seq), err } // Name returns the name of the table. @@ -1069,7 +1073,7 @@ func (t *WritableDoltTable) truncate( if schema.HasAutoIncrement(sch) { ddb, _ := sess.GetDoltDB(ctx, t.db.RevisionQualifiedName()) - err = t.db.removeTableFromAutoIncrementTracker(ctx, t.Name(), ddb, ws.Ref()) + err = t.db.removeTableFromAutoIncrementTracker(ctx, t.TableName(), ddb, ws.Ref()) if err != nil { return nil, err } @@ -1570,11 +1574,14 @@ func (t *AlterableDoltTable) AddColumn(ctx *sql.Context, column *sql.Column, ord } if column.AutoIncrement { - ait, err := t.db.gs.AutoIncrementTracker(ctx) + ait, err := dsess.GetAutoIncrementTracker(ctx, t.db.gs) + if err != nil { + return err + } + err = ait.AddNewRelation(t.TableName(), doltdb.AutoIncrementState(1)) if err != nil { return err } - ait.AddNewTable(t.tableName) } newRoot, err := root.PutTable(ctx, t.TableName(), updatedTable) @@ -1797,7 +1804,7 @@ func (t *AlterableDoltTable) RewriteInserter( // TODO: figure out locking. Other DBs automatically lock a table during this kind of operation, we should probably // do the same. We're messing with global auto-increment values here and it's not safe. - ait, err := t.db.gs.AutoIncrementTracker(ctx) + ait, err := dsess.GetAutoIncrementTracker(ctx, t.db.gs) if err != nil { return nil, err } @@ -1854,7 +1861,7 @@ func fullTextRewriteEditor( // TODO: figure out locking. Other DBs automatically lock a table during this kind of operation, we should probably // do the same. We're messing with global auto-increment values here and it's not safe. - ait, err := t.db.gs.AutoIncrementTracker(ctx) + ait, err := dsess.GetAutoIncrementTracker(ctx, t.db.gs) if err != nil { return nil, err } @@ -2271,21 +2278,24 @@ func (t *AlterableDoltTable) ModifyColumn(ctx *sql.Context, columnName string, c return err } - updatedTable, err = updatedTable.SetAutoIncrementValue(ctx, seq) + updatedTable, err = updatedTable.SetAutoIncrementValue(ctx, doltdb.AutoIncrementState(seq)) if err != nil { return err } - ait, err := t.db.gs.AutoIncrementTracker(ctx) + ait, err := dsess.GetAutoIncrementTracker(ctx, t.db.gs) if err != nil { return err } // TODO: this isn't transactional, and it should be (but none of the auto increment tracking is) - ait.AddNewTable(t.tableName) + err = ait.AddNewRelation(t.TableName(), doltdb.AutoIncrementState(1)) + if err != nil { + return err + } // Since this is a new auto increment table, we don't need to exclude the current working set from consideration // when computing its new sequence value, hence the empty ref - _, err = ait.Set(ctx, t.tableName, updatedTable, ref.WorkingSetRef{}, seq) + _, err = ait.Set(ctx, t.TableName(), updatedTable, ref.WorkingSetRef{}, doltdb.AutoIncrementState(seq)) if err != nil { return err } @@ -2296,7 +2306,7 @@ func (t *AlterableDoltTable) ModifyColumn(ctx *sql.Context, columnName string, c // TODO: this isn't transactional, and it should be sess := dsess.DSessFromSess(ctx.Session) ddb, _ := sess.GetDoltDB(ctx, t.db.RevisionQualifiedName()) - err = t.db.removeTableFromAutoIncrementTracker(ctx, t.Name(), ddb, ws.Ref()) + err = t.db.removeTableFromAutoIncrementTracker(ctx, t.TableName(), ddb, ws.Ref()) if err != nil { return err } @@ -2361,7 +2371,7 @@ func (t *AlterableDoltTable) getFirstAutoIncrementValue( } } - seq, err := dsess.CoerceAutoIncrementValue(ctx, initialValue) + seq, err := doltdb.CoerceAutoIncrementValue(ctx, initialValue) if err != nil { return 0, err } diff --git a/go/libraries/doltcore/sqle/temp_table.go b/go/libraries/doltcore/sqle/temp_table.go index 3e96e281ba2..6dc48d676e2 100644 --- a/go/libraries/doltcore/sqle/temp_table.go +++ b/go/libraries/doltcore/sqle/temp_table.go @@ -484,7 +484,11 @@ func temporaryDoltSchema(ctx context.Context, pkSch sql.PrimaryKeySchema, tags [ } func (t *TempTable) PeekNextAutoIncrementValue(ctx *sql.Context) (uint64, error) { - return t.table.GetAutoIncrementValue(ctx) + autoIncState, err := t.table.GetAutoIncrementValue(ctx) + if err != nil { + return 0, err + } + return autoIncState.CurrentValue(), nil } func (t *TempTable) GetNextAutoIncrementValue(ctx *sql.Context, insertVal interface{}) (uint64, error) { diff --git a/go/libraries/doltcore/sqle/writer/prolly_table_writer.go b/go/libraries/doltcore/sqle/writer/prolly_table_writer.go index cf3853ded0d..e633ce6e98b 100644 --- a/go/libraries/doltcore/sqle/writer/prolly_table_writer.go +++ b/go/libraries/doltcore/sqle/writer/prolly_table_writer.go @@ -23,7 +23,6 @@ import ( "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/dsess" - "github.com/dolthub/dolt/go/libraries/doltcore/sqle/globalstate" "github.com/dolthub/dolt/go/libraries/doltcore/sqle/index" "github.com/dolthub/dolt/go/store/hash" "github.com/dolthub/dolt/go/store/pool" @@ -52,7 +51,7 @@ type prollyTableWriter struct { writeSess dsess.WriteSession aiCol schema.Column - aiTracker globalstate.AutoIncrementTracker + aiTracker *dsess.AutoIncrementTracker aiAlterVal uint64 aiAltered bool // True when an ALTER TABLE affects the auto increment value aiSet bool // True when an INSERT/UPDATE affects the auto increment value @@ -176,7 +175,6 @@ func (w *prollyTableWriter) Insert(ctx *sql.Context, sqlRow sql.Row) (err error) // TODO: need schema name in ai tracker w.aiSet = true - w.aiTracker.Next(ctx, w.tblName.Name, sqlRow) return nil } @@ -273,12 +271,16 @@ func (w *prollyTableWriter) PreciseMatch() bool { // GetNextAutoIncrementValue implements TableWriter. func (w *prollyTableWriter) GetNextAutoIncrementValue(ctx *sql.Context, insertVal interface{}) (uint64, error) { - return w.aiTracker.Next(ctx, w.tblName.Name, insertVal) + v, err := w.aiTracker.Next(ctx, w.tblName, insertVal) + if err != nil { + return 0, err + } + return v, nil } // SetAutoIncrementValue implements AutoIncrementSetter. func (w *prollyTableWriter) SetAutoIncrementValue(ctx *sql.Context, val uint64) error { - seq, err := w.aiTracker.CoerceAutoIncrementValue(ctx, val) + seq, err := doltdb.CoerceAutoIncrementValue(ctx, val) if err != nil { return err } @@ -292,7 +294,7 @@ func (w *prollyTableWriter) SetAutoIncrementValue(ctx *sql.Context, val uint64) // AcquireAutoIncrementLock implements AutoIncrementSetter. func (w *prollyTableWriter) AcquireAutoIncrementLock(ctx *sql.Context) (func(), error) { - return w.aiTracker.AcquireTableLock(ctx, w.tblName.Name) + return w.aiTracker.AcquireLock(ctx, w.tblName) } // Close implements Closer @@ -396,13 +398,12 @@ func (w *prollyTableWriter) table(ctx *sql.Context) (tbl *doltdb.Table, err erro if w.aiCol.AutoIncrement { if w.aiAltered { - tbl, err = w.aiTracker.Set(ctx, w.tblName.Name, tbl, w.writeSess.GetWorkingSet().Ref(), w.aiAlterVal) + tbl, err = w.aiTracker.Set(ctx, w.tblName, tbl, w.writeSess.GetWorkingSet().Ref(), doltdb.AutoIncrementState(w.aiAlterVal)) if err != nil { return nil, err } } else if w.aiSet { - var aiVal uint64 - aiVal, err = w.aiTracker.Current(w.tblName.Name) + aiVal, err := w.aiTracker.Current(w.tblName) if err != nil { return nil, err } diff --git a/go/libraries/doltcore/sqle/writer/prolly_write_session.go b/go/libraries/doltcore/sqle/writer/prolly_write_session.go index 6afd83c2802..579521447c1 100644 --- a/go/libraries/doltcore/sqle/writer/prolly_write_session.go +++ b/go/libraries/doltcore/sqle/writer/prolly_write_session.go @@ -23,7 +23,6 @@ import ( "github.com/dolthub/dolt/go/libraries/doltcore/doltdb" "github.com/dolthub/dolt/go/libraries/doltcore/sqle/dsess" - "github.com/dolthub/dolt/go/libraries/doltcore/sqle/globalstate" "github.com/dolthub/dolt/go/libraries/doltcore/table/editor" "github.com/dolthub/dolt/go/store/hash" ) @@ -31,7 +30,7 @@ import ( // NewWriteSession creates and returns a WriteSession. Inserting a nil root is not an error, as there are // locations that do not have a root at the time of this call. However, a root must be set through SetWorkingRoot before any // table editors are returned. -func NewWriteSession(dbName string, ws *doltdb.WorkingSet, aiTracker globalstate.AutoIncrementTracker, setter dsess.SessionRootSetter, opts editor.Options) dsess.WriteSession { +func NewWriteSession(dbName string, ws *doltdb.WorkingSet, aiTracker *dsess.AutoIncrementTracker, setter dsess.SessionRootSetter, opts editor.Options) dsess.WriteSession { return &prollyWriteSession{ dbName: dbName, tables: make(map[doltdb.TableName]*prollyTableWriter), @@ -47,7 +46,7 @@ func NewWriteSession(dbName string, ws *doltdb.WorkingSet, aiTracker globalstate type prollyWriteSession struct { dbName string tables map[doltdb.TableName]*prollyTableWriter - aiTracker globalstate.AutoIncrementTracker + aiTracker *dsess.AutoIncrementTracker workingSet *doltdb.WorkingSet setter dsess.SessionRootSetter targetStaging bool