github.com/bingoohuang/gg@v0.0.0-20240325092523-45da7dee9335/pkg/sqlparse/tidbparser/dependency/sessionctx/stmtctx/stmtctx.go (about)

     1  // Copyright 2017 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package stmtctx
    15  
    16  import (
    17  	"math"
    18  	"sync"
    19  	"time"
    20  
    21  	"github.com/bingoohuang/gg/pkg/sqlparse/tidbparser/dependency/mysql"
    22  	"github.com/bingoohuang/gg/pkg/sqlparse/tidbparser/dependency/util/memory"
    23  )
    24  
    25  // StatementContext contains variables for a statement.
    26  // It should be reset before executing a statement.
    27  type StatementContext struct {
    28  	// Set the following variables before execution
    29  
    30  	InInsertStmt           bool
    31  	InUpdateOrDeleteStmt   bool
    32  	InSelectStmt           bool
    33  	IgnoreTruncate         bool
    34  	IgnoreZeroInDate       bool
    35  	DividedByZeroAsWarning bool
    36  	TruncateAsWarning      bool
    37  	OverflowAsWarning      bool
    38  	InShowWarning          bool
    39  	UseCache               bool
    40  	PadCharToFullLength    bool
    41  	BatchCheck             bool
    42  
    43  	// mu struct holds variables that change during execution.
    44  	mu struct {
    45  		sync.Mutex
    46  		affectedRows      uint64
    47  		foundRows         uint64
    48  		warnings          []error
    49  		histogramsNotLoad bool
    50  	}
    51  
    52  	// Copied from SessionVars.TimeZone.
    53  	TimeZone     *time.Location
    54  	Priority     mysql.PriorityEnum
    55  	NotFillCache bool
    56  	MemTracker   *memory.Tracker
    57  }
    58  
    59  // AddAffectedRows adds affected rows.
    60  func (sc *StatementContext) AddAffectedRows(rows uint64) {
    61  	sc.mu.Lock()
    62  	sc.mu.affectedRows += rows
    63  	sc.mu.Unlock()
    64  }
    65  
    66  // AffectedRows gets affected rows.
    67  func (sc *StatementContext) AffectedRows() uint64 {
    68  	sc.mu.Lock()
    69  	rows := sc.mu.affectedRows
    70  	sc.mu.Unlock()
    71  	return rows
    72  }
    73  
    74  // FoundRows gets found rows.
    75  func (sc *StatementContext) FoundRows() uint64 {
    76  	sc.mu.Lock()
    77  	rows := sc.mu.foundRows
    78  	sc.mu.Unlock()
    79  	return rows
    80  }
    81  
    82  // AddFoundRows adds found rows.
    83  func (sc *StatementContext) AddFoundRows(rows uint64) {
    84  	sc.mu.Lock()
    85  	sc.mu.foundRows += rows
    86  	sc.mu.Unlock()
    87  }
    88  
    89  // GetWarnings gets warnings.
    90  func (sc *StatementContext) GetWarnings() []error {
    91  	sc.mu.Lock()
    92  	warns := make([]error, len(sc.mu.warnings))
    93  	copy(warns, sc.mu.warnings)
    94  	sc.mu.Unlock()
    95  	return warns
    96  }
    97  
    98  // WarningCount gets warning count.
    99  func (sc *StatementContext) WarningCount() uint16 {
   100  	if sc.InShowWarning {
   101  		return 0
   102  	}
   103  	sc.mu.Lock()
   104  	wc := uint16(len(sc.mu.warnings))
   105  	sc.mu.Unlock()
   106  	return wc
   107  }
   108  
   109  // SetWarnings sets warnings.
   110  func (sc *StatementContext) SetWarnings(warns []error) {
   111  	sc.mu.Lock()
   112  	sc.mu.warnings = warns
   113  	sc.mu.Unlock()
   114  }
   115  
   116  // AppendWarning appends a warning.
   117  func (sc *StatementContext) AppendWarning(warn error) {
   118  	sc.mu.Lock()
   119  	if len(sc.mu.warnings) < math.MaxUint16 {
   120  		sc.mu.warnings = append(sc.mu.warnings, warn)
   121  	}
   122  	sc.mu.Unlock()
   123  }
   124  
   125  // SetHistogramsNotLoad sets histogramsNotLoad.
   126  func (sc *StatementContext) SetHistogramsNotLoad() {
   127  	sc.mu.Lock()
   128  	sc.mu.histogramsNotLoad = true
   129  	sc.mu.Unlock()
   130  }
   131  
   132  // HistogramsNotLoad gets histogramsNotLoad.
   133  func (sc *StatementContext) HistogramsNotLoad() bool {
   134  	sc.mu.Lock()
   135  	notLoad := sc.mu.histogramsNotLoad
   136  	sc.mu.Unlock()
   137  	return notLoad
   138  }
   139  
   140  // HandleTruncate ignores or returns the error based on the StatementContext state.
   141  func (sc *StatementContext) HandleTruncate(err error) error {
   142  	// TODO: At present we have not checked whether the error can be ignored or treated as warning.
   143  	// We will do that later, and then append WarnDataTruncated instead of the error itself.
   144  	if err == nil {
   145  		return nil
   146  	}
   147  	if sc.IgnoreTruncate {
   148  		return nil
   149  	}
   150  	if sc.TruncateAsWarning {
   151  		sc.AppendWarning(err)
   152  		return nil
   153  	}
   154  	return err
   155  }
   156  
   157  // HandleOverflow treats ErrOverflow as warnings or returns the error based on the StmtCtx.OverflowAsWarning state.
   158  func (sc *StatementContext) HandleOverflow(err error, warnErr error) error {
   159  	if err == nil {
   160  		return nil
   161  	}
   162  
   163  	if sc.OverflowAsWarning {
   164  		sc.AppendWarning(warnErr)
   165  		return nil
   166  	}
   167  	return err
   168  }
   169  
   170  // ResetForRetry resets the changed states during execution.
   171  func (sc *StatementContext) ResetForRetry() {
   172  	sc.mu.Lock()
   173  	sc.mu.affectedRows = 0
   174  	sc.mu.foundRows = 0
   175  	sc.mu.warnings = nil
   176  	sc.mu.Unlock()
   177  }