github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/pkg/sqlmodel/where_handle.go (about)

     1  // Copyright 2022 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 sqlmodel
    15  
    16  import (
    17  	"github.com/pingcap/log"
    18  	"github.com/pingcap/tidb/pkg/parser/model"
    19  	"github.com/pingcap/tidb/pkg/parser/mysql"
    20  	"github.com/pingcap/tidb/pkg/types"
    21  )
    22  
    23  // WhereHandle is used to generate a WHERE clause in SQL.
    24  type WhereHandle struct {
    25  	UniqueNotNullIdx *model.IndexInfo
    26  	// If the index and columns have no NOT NULL constraint, but all data is NOT
    27  	// NULL, we can still use it.
    28  	// every index that is UNIQUE should be added to UniqueIdxs, even for
    29  	// PK and NOT NULL.
    30  	UniqueIdxs []*model.IndexInfo
    31  }
    32  
    33  // GetWhereHandle calculates a WhereHandle by source/target TableInfo's indices,
    34  // columns and state. Other component can cache the result.
    35  func GetWhereHandle(source, target *model.TableInfo) *WhereHandle {
    36  	ret := WhereHandle{}
    37  	indices := make([]*model.IndexInfo, 0, len(target.Indices)+1)
    38  	indices = append(indices, target.Indices...)
    39  	if idx := getPKIsHandleIdx(target); target.PKIsHandle && idx != nil {
    40  		indices = append(indices, idx)
    41  	}
    42  
    43  	for _, idx := range indices {
    44  		if !idx.Unique {
    45  			continue
    46  		}
    47  		// when the tableInfo is from CDC, it may contain some index that is
    48  		// creating.
    49  		if idx.State != model.StatePublic {
    50  			continue
    51  		}
    52  
    53  		rewritten := rewriteColsOffset(idx, source)
    54  		if rewritten == nil {
    55  			continue
    56  		}
    57  		ret.UniqueIdxs = append(ret.UniqueIdxs, rewritten)
    58  
    59  		if rewritten.Primary {
    60  			// PK is prior to UNIQUE NOT NULL for better performance
    61  			ret.UniqueNotNullIdx = rewritten
    62  			continue
    63  		}
    64  		// use downstream columns to check NOT NULL constraint
    65  		if ret.UniqueNotNullIdx == nil && allColsNotNull(idx, target.Columns) {
    66  			ret.UniqueNotNullIdx = rewritten
    67  			continue
    68  		}
    69  	}
    70  	return &ret
    71  }
    72  
    73  // rewriteColsOffset rewrites index columns offset to those from source table.
    74  // Returns nil when any column does not represent in source.
    75  func rewriteColsOffset(index *model.IndexInfo, source *model.TableInfo) *model.IndexInfo {
    76  	if index == nil || source == nil {
    77  		return nil
    78  	}
    79  
    80  	columns := make([]*model.IndexColumn, 0, len(index.Columns))
    81  	for _, key := range index.Columns {
    82  		sourceColumn := model.FindColumnInfo(source.Columns, key.Name.L)
    83  		if sourceColumn == nil {
    84  			return nil
    85  		}
    86  		column := &model.IndexColumn{
    87  			Name:   key.Name,
    88  			Offset: sourceColumn.Offset,
    89  			Length: key.Length,
    90  		}
    91  		columns = append(columns, column)
    92  	}
    93  	clone := *index
    94  	clone.Columns = columns
    95  	return &clone
    96  }
    97  
    98  func getPKIsHandleIdx(ti *model.TableInfo) *model.IndexInfo {
    99  	if pk := ti.GetPkColInfo(); pk != nil {
   100  		return &model.IndexInfo{
   101  			Table:   ti.Name,
   102  			Unique:  true,
   103  			Primary: true,
   104  			State:   model.StatePublic,
   105  			Tp:      model.IndexTypeBtree,
   106  			Columns: []*model.IndexColumn{{
   107  				Name:   pk.Name,
   108  				Offset: pk.Offset,
   109  				Length: types.UnspecifiedLength,
   110  			}},
   111  		}
   112  	}
   113  	return nil
   114  }
   115  
   116  func allColsNotNull(idx *model.IndexInfo, cols []*model.ColumnInfo) bool {
   117  	for _, idxCol := range idx.Columns {
   118  		col := cols[idxCol.Offset]
   119  		if !mysql.HasNotNullFlag(col.GetFlag()) {
   120  			return false
   121  		}
   122  	}
   123  	return true
   124  }
   125  
   126  // getWhereIdxByData returns the index that is identical to a row change, it
   127  // may be
   128  // - a PK, or
   129  // - an UNIQUE index whose columns are all NOT NULL, or
   130  // - an UNIQUE index and the data are all NOT NULL.
   131  // For the last case, last used index is swapped to front.
   132  func (h *WhereHandle) getWhereIdxByData(data []interface{}) *model.IndexInfo {
   133  	if h == nil {
   134  		log.L().DPanic("WhereHandle is nil")
   135  		return nil
   136  	}
   137  	if h.UniqueNotNullIdx != nil {
   138  		return h.UniqueNotNullIdx
   139  	}
   140  	for i, idx := range h.UniqueIdxs {
   141  		ok := true
   142  		for _, idxCol := range idx.Columns {
   143  			if data[idxCol.Offset] == nil {
   144  				ok = false
   145  				break
   146  			}
   147  		}
   148  		if ok {
   149  			h.UniqueIdxs[0], h.UniqueIdxs[i] = h.UniqueIdxs[i], h.UniqueIdxs[0]
   150  			return idx
   151  		}
   152  	}
   153  	return nil
   154  }