github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/sessiondata/sequence_state.go (about)

     1  // Copyright 2018 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package sessiondata
    12  
    13  import (
    14  	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
    15  	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
    16  	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
    17  )
    18  
    19  // SequenceState stores session-scoped state used by sequence builtins.
    20  //
    21  // All public methods of SequenceState are thread-safe, as the structure is
    22  // meant to be shared by statements executing in parallel on a session.
    23  type SequenceState struct {
    24  	mu struct {
    25  		syncutil.Mutex
    26  		// latestValues stores the last value obtained by nextval() in this session
    27  		// by descriptor id.
    28  		latestValues map[uint32]int64
    29  
    30  		// lastSequenceIncremented records the descriptor id of the last sequence
    31  		// nextval() was called on in this session.
    32  		lastSequenceIncremented uint32
    33  	}
    34  }
    35  
    36  // NewSequenceState creates a SequenceState.
    37  func NewSequenceState() *SequenceState {
    38  	ss := SequenceState{}
    39  	ss.mu.latestValues = make(map[uint32]int64)
    40  	return &ss
    41  }
    42  
    43  // NextVal ever called returns true if a sequence has ever been incremented on
    44  // this session.
    45  func (ss *SequenceState) nextValEverCalledLocked() bool {
    46  	return len(ss.mu.latestValues) > 0
    47  }
    48  
    49  // RecordValue records the latest manipulation of a sequence done by a session.
    50  func (ss *SequenceState) RecordValue(seqID uint32, val int64) {
    51  	ss.mu.Lock()
    52  	ss.mu.lastSequenceIncremented = seqID
    53  	ss.mu.latestValues[seqID] = val
    54  	ss.mu.Unlock()
    55  }
    56  
    57  // SetLastSequenceIncremented sets the id of the last incremented sequence.
    58  // Usually this id is set through RecordValue().
    59  func (ss *SequenceState) SetLastSequenceIncremented(seqID uint32) {
    60  	ss.mu.Lock()
    61  	ss.mu.lastSequenceIncremented = seqID
    62  	ss.mu.Unlock()
    63  }
    64  
    65  // GetLastValue returns the value most recently obtained by
    66  // nextval() for the last sequence for which RecordLatestVal() was called.
    67  func (ss *SequenceState) GetLastValue() (int64, error) {
    68  	ss.mu.Lock()
    69  	defer ss.mu.Unlock()
    70  
    71  	if !ss.nextValEverCalledLocked() {
    72  		return 0, pgerror.New(
    73  			pgcode.ObjectNotInPrerequisiteState, "lastval is not yet defined in this session")
    74  	}
    75  
    76  	return ss.mu.latestValues[ss.mu.lastSequenceIncremented], nil
    77  }
    78  
    79  // GetLastValueByID returns the value most recently obtained by nextval() for
    80  // the given sequence in this session.
    81  // The bool retval is false if RecordLatestVal() was never called on the
    82  // requested sequence.
    83  func (ss *SequenceState) GetLastValueByID(seqID uint32) (int64, bool) {
    84  	ss.mu.Lock()
    85  	defer ss.mu.Unlock()
    86  
    87  	val, ok := ss.mu.latestValues[seqID]
    88  	return val, ok
    89  }
    90  
    91  // Export returns a copy of the SequenceState's state - the latestValues and
    92  // lastSequenceIncremented.
    93  // lastSequenceIncremented is only defined if latestValues is non-empty.
    94  func (ss *SequenceState) Export() (map[uint32]int64, uint32) {
    95  	ss.mu.Lock()
    96  	defer ss.mu.Unlock()
    97  	res := make(map[uint32]int64, len(ss.mu.latestValues))
    98  	for k, v := range ss.mu.latestValues {
    99  		res[k] = v
   100  	}
   101  	return res, ss.mu.lastSequenceIncremented
   102  }