Draft: Create a SequenceTracker for Postgres sequences.

core/sequences/collection.go

@@ -18,18 +18,24 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"iter"
 	"math"
 	"sort"
 	"strings"
 
 	"github.com/cockroachdb/errors"
 	"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/sequences"
 	"github.com/dolthub/dolt/go/store/hash"
 	"github.com/dolthub/dolt/go/store/prolly"
 	"github.com/dolthub/dolt/go/store/prolly/tree"
+	"github.com/dolthub/go-mysql-server/sql"
+	"github.com/dolthub/go-mysql-server/sql/types"
 
 	"github.com/dolthub/doltgresql/core/id"
 	"github.com/dolthub/doltgresql/core/rootobject/objinterface"
+	"github.com/dolthub/doltgresql/utils"
 )
 
 // Collection contains a collection of sequences.
@@ -48,11 +54,8 @@ const (
 	Persistence_Unlogged  Persistence = 2
 )
 
-// Sequence represents a single sequence within the pg_sequence table.
-type Sequence struct {
+type SequenceState struct {
 	Id            id.Sequence
-	DataTypeID    id.Type
-	Persistence   Persistence
 	Start         int64
 	Current       int64
 	Increment     int64
@@ -62,13 +65,163 @@ type Sequence struct {
 	Cycle         bool
 	IsAtEnd       bool
 	HasBeenCalled bool
-	OwnerTable    id.Table
-	OwnerColumn   string
+}
+
+type SequenceTracker = dsess.SequenceTracker[*Sequence, SequenceState, int64]
+
+// SequenceTrackerKey is the key to identify the SequenceTracker in the globalstate.GlobalState
+var SequenceTrackerKey dsess.TrackerKey[*SequenceTracker] = struct{}{}
+
+func (sequence SequenceState) Merge(otherSequenceState SequenceState) (merged SequenceState) {
+	newSequenceState := sequence
+	thisIsIncrementing := sequence.Increment > 0
+	otherIsIncrementing := otherSequenceState.Increment > 0
+	if thisIsIncrementing != otherIsIncrementing {
+		// These states can't be merged.
+		// A zero-valued state is the "invalid" state.
+		return SequenceState{}
+	}
+	if thisIsIncrementing {
+		newSequenceState.Increment = utils.Min(sequence.Increment, otherSequenceState.Increment)
+		newSequenceState.Start = utils.Min(sequence.Start, otherSequenceState.Start)
+	} else {
+		newSequenceState.Increment = utils.Max(sequence.Increment, otherSequenceState.Increment)
+		newSequenceState.Start = utils.Max(sequence.Start, otherSequenceState.Start)
+	}
+	if sequence.GreaterThan(otherSequenceState) {
+		newSequenceState.Current = sequence.Current
+	} else {
+		newSequenceState.Current = otherSequenceState.Current
+	}
+	newSequenceState.Minimum = utils.Min(sequence.Minimum, otherSequenceState.Minimum)
+	newSequenceState.Maximum = utils.Max(sequence.Maximum, otherSequenceState.Maximum)
+	newSequenceState.Cycle = sequence.Cycle || otherSequenceState.Cycle
+	newSequenceState.IsAtEnd = sequence.IsAtEnd && otherSequenceState.IsAtEnd
+	newSequenceState.HasBeenCalled = sequence.HasBeenCalled || otherSequenceState.HasBeenCalled
+	return newSequenceState
+}
+
+var _ sequences.SequenceState[SequenceState, int64] = SequenceState{}
+
+func (sequence SequenceState) CurrentValue() int64 {
+	return sequence.Current
+}
+
+func (sequence SequenceState) WithValue(v int64) SequenceState {
+	sequence.Current = v
+	sequence.IsAtEnd = false
+	return sequence
+}
+
+func (sequence SequenceState) WithSQLValue(ctx *sql.Context, v interface{}) (SequenceState, error) {
+	// TODO: Coercing happens here, based on the type of the sequence
+	return sequence.WithValue(v.(int64)), nil
+}
+
+func (sequence SequenceState) GreaterThan(other SequenceState) bool {
+	// A sequence that has wrapped around is further along than a sequence that hasn't.
+	// Otherwise, we see which sequence is further alone, in the direction that it's incrementing.
+	if sequence.Increment > 0 {
+		hasWrapped := sequence.Current < sequence.Start
+		otherHasWrapped := other.Current < sequence.Start
+		if hasWrapped == otherHasWrapped {
+			return sequence.Current > other.Current
+		} else {
+			// Exactly one of the sequences has wrapped around. That sequence is greater.
+			return hasWrapped
+		}
+	} else {
+		hasWrapped := sequence.Current > sequence.Start
+		otherHasWrapped := other.Current > sequence.Start
+		if hasWrapped == otherHasWrapped {
+			return sequence.Current < other.Current
+		} else {
+			// Exactly one of the sequences has wrapped around. That sequence is greater.
+			return hasWrapped
+		}
+	}
+}
+
+func (sequence SequenceState) AtEnd() bool {
+	return sequence.IsAtEnd
+}
+
+func (sequence SequenceState) Next() (sqlVal int64, hasNext bool, nextState SequenceState, err error) {
+	// First we'll check if we've reached the end, and cycle or error as necessary
+	sequence.HasBeenCalled = true
+	if sequence.IsAtEnd {
+		if !sequence.Cycle {
+			if sequence.Increment > 0 {
+				return 0, false, SequenceState{}, errors.Errorf(`nextval: reached maximum value of sequence "%s" (%d)`, sequence.Id, sequence.Maximum)
+			} else {
+				return 0, false, SequenceState{}, errors.Errorf(`nextval: reached minimum value of sequence "%s" (%d)`, sequence.Id, sequence.Minimum)
+			}
+		}
+		sequence.IsAtEnd = false
+		if sequence.Increment > 0 {
+			sequence.Current = sequence.Minimum
+		} else {
+			sequence.Current = sequence.Maximum
+		}
+	}
+	// We'll return the current value, so everything after this sets the value for the next call
+	valueToReturn := sequence.Current
+	// Increment the current value
+	if sequence.Increment > 0 {
+		// Check for overflow or crossing the maximum, meaning we're at the end
+		if sequence.Current > math.MaxInt64-sequence.Increment || sequence.Current+sequence.Increment > sequence.Maximum {
+			sequence.IsAtEnd = true
+		} else {
+			sequence.Current += sequence.Increment
+		}
+	} else {
+		// Check for underflow or crossing the minimum, meaning we're at the end
+		if sequence.Current < math.MinInt64-sequence.Increment || sequence.Current+sequence.Increment < sequence.Minimum {
+			sequence.IsAtEnd = true
+		} else {
+			sequence.Current += sequence.Increment
+		}
+	}
+	return valueToReturn, true, sequence, nil
+}
+
+// Sequence represents a single sequence within the pg_sequence table.
+type Sequence struct {
+	DataTypeID  id.Type
+	Persistence Persistence
+	SequenceState
+	OwnerTable  id.Table
+	OwnerColumn string
+}
+
+func (sequence *Sequence) GetSequenceState(ctx context.Context) (SequenceState, error) {
+	return sequence.SequenceState, nil
+}
+
+func (sequence *Sequence) HasSequenceState(ctx context.Context) (bool, error) {
+	return true, nil
+}
+
+func (sequence *Sequence) SetSequenceState(ctx context.Context, newSequenceState SequenceState) (*Sequence, error) {
+	newSequence := sequence
+	newSequence.SequenceState = newSequenceState
+	return newSequence, nil
+}
+
+func (sequence *Sequence) GetSequenceSqlType(ctx context.Context) (sql.Type, bool, error) {
+	// TODO: Return the actual correct type here
+	return types.Int64, true, nil
+}
+
+func (sequence *Sequence) TrySetSequenceState(ctx *sql.Context, val SequenceState) (*Sequence, bool, error) {
+	newSequence, err := sequence.SetSequenceState(ctx, val)
+	return newSequence, true, err
 }
 
 var _ objinterface.Collection = (*Collection)(nil)
 var _ objinterface.RootObject = (*Sequence)(nil)
 var _ doltdb.RootObject = (*Sequence)(nil)
+var _ sequences.SequencedRelation[*Sequence, int64, SequenceState] = (*Sequence)(nil)
 
 // GetSequence returns the sequence with the given schema and name. Returns nil if the sequence cannot be found.
 func (pgs *Collection) GetSequence(ctx context.Context, name id.Sequence) (*Sequence, error) {
@@ -284,7 +437,7 @@ func (pgs *Collection) NextVal(ctx context.Context, name id.Sequence) (int64, er
 }
 
 // SetVal sets the sequence to the
-func (pgs *Collection) SetVal(ctx context.Context, name id.Sequence, newValue int64, autoAdvance bool) error {
+func (pgs *Collection) SetVal(ctx context.Context, name id.Sequence, newValue int64, hasBeenCalled bool, autoAdvance bool) error {
 	seq, err := pgs.getSequence(ctx, name)
 	if err != nil {
 		return err
@@ -298,7 +451,7 @@ func (pgs *Collection) SetVal(ctx context.Context, name id.Sequence, newValue in
 	}
 	seq.Current = newValue
 	seq.IsAtEnd = false
-	seq.HasBeenCalled = false
+	seq.HasBeenCalled = hasBeenCalled
 	if autoAdvance {
 		_, err := seq.nextValForSequence()
 		return err
@@ -435,40 +588,39 @@ func (pgs *Collection) writeCache(ctx context.Context) (err error) {
 
 // nextValForSequence increments the calling sequence.
 func (sequence *Sequence) nextValForSequence() (int64, error) {
-	// First we'll check if we've reached the end, and cycle or error as necessary
-	if sequence.IsAtEnd {
-		if !sequence.Cycle {
-			if sequence.Increment > 0 {
-				return 0, errors.Errorf(`nextval: reached maximum value of sequence "%s" (%d)`, sequence.Id, sequence.Maximum)
-			} else {
-				return 0, errors.Errorf(`nextval: reached minimum value of sequence "%s" (%d)`, sequence.Id, sequence.Minimum)
-			}
-		}
-		sequence.IsAtEnd = false
-		if sequence.Increment > 0 {
-			sequence.Current = sequence.Minimum
-		} else {
-			sequence.Current = sequence.Maximum
-		}
+	result, _, newSequence, err := sequence.Next()
+	if err != nil {
+		return 0, err
 	}
-	// We'll return the current value, so everything after this sets the value for the next call
-	sequence.HasBeenCalled = true
-	valueToReturn := sequence.Current
-	// Increment the current value
-	if sequence.Increment > 0 {
-		// Check for overflow or crossing the maximum, meaning we're at the end
-		if sequence.Current > math.MaxInt64-sequence.Increment || sequence.Current+sequence.Increment > sequence.Maximum {
-			sequence.IsAtEnd = true
-		} else {
-			sequence.Current += sequence.Increment
-		}
-	} else {
-		// Check for underflow or crossing the minimum, meaning we're at the end
-		if sequence.Current < math.MinInt64-sequence.Increment || sequence.Current+sequence.Increment < sequence.Minimum {
-			sequence.IsAtEnd = true
-		} else {
-			sequence.Current += sequence.Increment
-		}
+	sequence.SequenceState = newSequence
+	return result, nil
+}
+
+// SequenceSource reads relations from a RootValue by reading its RootObjects
+type SequenceSource struct{}
+
+var _ doltdb.RelationSource[*Sequence] = SequenceSource{}
+
+func (s SequenceSource) GetRelation(ctx context.Context, root doltdb.RootValue, tName doltdb.TableName) (relation *Sequence, resolvedName string, found bool, err error) {
+	obj, found, err := root.GetRootObject(ctx, tName)
+	if !found || err != nil {
+		return nil, "", found, err
+	}
+	if seq, ok := obj.(*Sequence); ok {
+		return seq, tName.Name, true, nil
+	}
+	return nil, "", found, nil
+}
+
+func (s SequenceSource) IterRelations(ctx context.Context, root doltdb.RootValue) iter.Seq2[doltdb.TableName, *Sequence] {
+	return func(yield func(doltdb.TableName, *Sequence) bool) {
+		_ = root.IterRootObjects(ctx, func(name doltdb.TableName, obj doltdb.RootObject) (stop bool, err error) {
+			if seq, ok := obj.(*Sequence); ok {
+				if !yield(name, seq) {
+					return true, nil
+				}
+			}
+			return false, nil
+		})
 	}
-	return valueToReturn, nil
 }

core/sequences/collection_test.go

@@ -89,13 +89,15 @@ func newTestCollection(t *testing.T, ns tree.NodeStore) *Collection {
 // through Serialize and Deserialize. The exact values are not significant.
 func newTestSequence(schema, name string) *Sequence {
 	return &Sequence{
-		Id:        id.NewSequence(schema, name),
-		Start:     1,
-		Current:   1,
-		Increment: 1,
-		Minimum:   1,
-		Maximum:   math.MaxInt64,
-		Cache:     1,
+		SequenceState: SequenceState{
+			Id:        id.NewSequence(schema, name),
+			Start:     1,
+			Current:   1,
+			Increment: 1,
+			Minimum:   1,
+			Maximum:   math.MaxInt64,
+			Cache:     1,
+		},
 	}
 }
 

core/sequences/root_object.go

@@ -325,6 +325,7 @@ func (pgs *Collection) RenameRootObject(ctx context.Context, oldName id.Id, newN
 	if !oldName.IsValid() || !newName.IsValid() || oldName.Section() != newName.Section() || oldName.Section() != id.Section_Sequence {
 		return errors.New("cannot rename sequence due to invalid name")
 	}
+	// TODO: Update ait
 	oldSeqName := id.Sequence(oldName)
 	newSeqName := id.Sequence(newName)
 	seq, err := pgs.GetSequence(ctx, oldSeqName)

go.mod

@@ -6,7 +6,7 @@ require (
 	github.com/PuerkitoBio/goquery v1.8.1
 	github.com/cockroachdb/apd/v3 v3.2.3
 	github.com/cockroachdb/errors v1.7.5
-	github.com/dolthub/dolt/go v0.40.5-0.20260804000445-86daffc60fe6
+	github.com/dolthub/dolt/go v0.40.5-0.20260804160803-748773e6b7c6
 	github.com/dolthub/eventsapi_schema v0.0.0-20260715220557-d9b4a1c6b4d4
 	github.com/dolthub/flatbuffers/v23 v23.3.3-dh.2
 	github.com/dolthub/go-mysql-server v0.20.1-0.20260803224759-f6896710fd7d

go.sum

@@ -248,6 +248,10 @@ github.com/dolthub/dolt-mcp v0.3.4 h1:AyG5cw+fNWXDHXujtQnqUPZrpWtPg6FN6yYtjv1pP4
 github.com/dolthub/dolt-mcp v0.3.4/go.mod h1:bCZ7KHvDYs+M0e+ySgmGiNvLhcwsN7bbf5YCyillLrk=
 github.com/dolthub/dolt/go v0.40.5-0.20260804000445-86daffc60fe6 h1:avF4Wlos4AlJ1UbrCKfQrVAdLlP2ppLTcGSTvJkhYm8=
 github.com/dolthub/dolt/go v0.40.5-0.20260804000445-86daffc60fe6/go.mod h1:tu5+NUsslUw+3i+7b57DumjuAegSkbZSpFccQz0u0nk=
+github.com/dolthub/dolt/go v0.40.5-0.20260804155236-475bfe5f5ba4 h1:30ZdVM5S1gCWnNhcwpvZAJ8fcBM4Q3cKYl4CmdeYUOI=
+github.com/dolthub/dolt/go v0.40.5-0.20260804155236-475bfe5f5ba4/go.mod h1:azS/FhEQSpp0L9ARSwB6A98nqpRdKlvSYaEoQh2Grvw=
+github.com/dolthub/dolt/go v0.40.5-0.20260804160803-748773e6b7c6 h1:Um4xgEcROlfq0g25sMVdzBkXf1s/x6o8TUl85MRIUS8=
+github.com/dolthub/dolt/go v0.40.5-0.20260804160803-748773e6b7c6/go.mod h1:tu5+NUsslUw+3i+7b57DumjuAegSkbZSpFccQz0u0nk=
 github.com/dolthub/eventsapi_schema v0.0.0-20260715220557-d9b4a1c6b4d4 h1:0mg9QEFdkkBwJMxvz1tCjHYmfG2iIC6aShj1InDq9/M=
 github.com/dolthub/eventsapi_schema v0.0.0-20260715220557-d9b4a1c6b4d4/go.mod h1:SSLraQS/jGLYFgff3vuZ+JbVUct6vyEeMzjLBqWqoyM=
 github.com/dolthub/flatbuffers/v23 v23.3.3-dh.2 h1:u3PMzfF8RkKd3lB9pZ2bfn0qEG+1Gms9599cr0REMww=

server/analyzer/serial.go

@@ -138,17 +138,19 @@ func ReplaceSerial(ctx *sql.Context, a *analyzer.Analyzer, node sql.Node, scope
 		}
 
 		ctSequences = append(ctSequences, pgnodes.NewCreateSequence(false, "", false, &sequences.Sequence{
-			Id:          id.NewSequence("", sequenceName),
 			DataTypeID:  col.Type.(*pgtypes.DoltgresType).ID,
 			Persistence: sequences.Persistence_Permanent,
-			Start:       1,
-			Current:     1,
-			Increment:   1,
-			Minimum:     1,
-			Maximum:     maxValue,
-			Cache:       1,
-			Cycle:       false,
-			IsAtEnd:     false,
+			SequenceState: sequences.SequenceState{
+				Id:        id.NewSequence("", sequenceName),
+				Start:     1,
+				Current:   1,
+				Increment: 1,
+				Minimum:   1,
+				Maximum:   maxValue,
+				Cache:     1,
+				Cycle:     false,
+				IsAtEnd:   false,
+			},
 			OwnerTable:  id.NewTable("", createTable.Name()),
 			OwnerColumn: col.Name,
 		}))