vitess.io/vitess@v0.16.2/go/vt/vtgate/vindexes/lookup_internal.go (about)

     1  /*
     2  Copyright 2019 The Vitess Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package vindexes
    18  
    19  import (
    20  	"bytes"
    21  	"context"
    22  	"fmt"
    23  	"sort"
    24  	"strconv"
    25  	"strings"
    26  
    27  	"vitess.io/vitess/go/vt/vterrors"
    28  
    29  	"vitess.io/vitess/go/sqltypes"
    30  
    31  	querypb "vitess.io/vitess/go/vt/proto/query"
    32  	vtgatepb "vitess.io/vitess/go/vt/proto/vtgate"
    33  	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
    34  )
    35  
    36  var (
    37  	readLockExclusive = "exclusive"
    38  	readLockShared    = "shared"
    39  	readLockNone      = "none"
    40  	readLockDefault   = readLockExclusive
    41  
    42  	readLockExprs = map[string]string{
    43  		readLockExclusive: "for update",
    44  		readLockShared:    "lock in share mode",
    45  		readLockNone:      "",
    46  	}
    47  )
    48  
    49  // lookupInternal implements the functions for the Lookup vindexes.
    50  type lookupInternal struct {
    51  	Table                   string   `json:"table"`
    52  	FromColumns             []string `json:"from_columns"`
    53  	To                      string   `json:"to"`
    54  	Autocommit              bool     `json:"autocommit,omitempty"`
    55  	MultiShardAutocommit    bool     `json:"multi_shard_autocommit,omitempty"`
    56  	Upsert                  bool     `json:"upsert,omitempty"`
    57  	IgnoreNulls             bool     `json:"ignore_nulls,omitempty"`
    58  	BatchLookup             bool     `json:"batch_lookup,omitempty"`
    59  	ReadLock                string   `json:"read_lock,omitempty"`
    60  	sel, selTxDml, ver, del string   // sel: map query, ver: verify query, del: delete query
    61  }
    62  
    63  func (lkp *lookupInternal) Init(lookupQueryParams map[string]string, autocommit, upsert, multiShardAutocommit bool) error {
    64  	lkp.Table = lookupQueryParams["table"]
    65  	lkp.To = lookupQueryParams["to"]
    66  	var fromColumns []string
    67  	for _, from := range strings.Split(lookupQueryParams["from"], ",") {
    68  		fromColumns = append(fromColumns, strings.TrimSpace(from))
    69  	}
    70  	lkp.FromColumns = fromColumns
    71  
    72  	var err error
    73  	lkp.IgnoreNulls, err = boolFromMap(lookupQueryParams, "ignore_nulls")
    74  	if err != nil {
    75  		return err
    76  	}
    77  	lkp.BatchLookup, err = boolFromMap(lookupQueryParams, "batch_lookup")
    78  	if err != nil {
    79  		return err
    80  	}
    81  	if readLock, ok := lookupQueryParams["read_lock"]; ok {
    82  		if _, valid := readLockExprs[readLock]; !valid {
    83  			return fmt.Errorf("invalid read_lock value: %s", readLock)
    84  		}
    85  		lkp.ReadLock = readLock
    86  	}
    87  
    88  	lkp.Autocommit = autocommit
    89  	lkp.Upsert = upsert
    90  	if multiShardAutocommit {
    91  		lkp.Autocommit = true
    92  		lkp.MultiShardAutocommit = true
    93  	}
    94  
    95  	// TODO @rafael: update sel and ver to support multi column vindexes. This will be done
    96  	// as part of face 2 of https://github.com/vitessio/vitess/issues/3481
    97  	// For now multi column behaves as a single column for Map and Verify operations
    98  	lkp.sel = fmt.Sprintf("select %s, %s from %s where %s in ::%s", lkp.FromColumns[0], lkp.To, lkp.Table, lkp.FromColumns[0], lkp.FromColumns[0])
    99  	if lkp.ReadLock != readLockNone {
   100  		lockExpr, ok := readLockExprs[lkp.ReadLock]
   101  		if !ok {
   102  			lockExpr = readLockExprs[readLockDefault]
   103  		}
   104  		lkp.selTxDml = fmt.Sprintf("%s %s", lkp.sel, lockExpr)
   105  	} else {
   106  		lkp.selTxDml = lkp.sel
   107  	}
   108  	lkp.ver = fmt.Sprintf("select %s from %s where %s = :%s and %s = :%s", lkp.FromColumns[0], lkp.Table, lkp.FromColumns[0], lkp.FromColumns[0], lkp.To, lkp.To)
   109  	lkp.del = lkp.initDelStmt()
   110  	return nil
   111  }
   112  
   113  // Lookup performs a lookup for the ids.
   114  func (lkp *lookupInternal) Lookup(ctx context.Context, vcursor VCursor, ids []sqltypes.Value, co vtgatepb.CommitOrder) ([]*sqltypes.Result, error) {
   115  	if vcursor == nil {
   116  		return nil, fmt.Errorf("cannot perform lookup: no vcursor provided")
   117  	}
   118  	results := make([]*sqltypes.Result, 0, len(ids))
   119  	if lkp.Autocommit {
   120  		co = vtgatepb.CommitOrder_AUTOCOMMIT
   121  	}
   122  	var sel string
   123  	if vcursor.InTransactionAndIsDML() {
   124  		sel = lkp.selTxDml
   125  	} else {
   126  		sel = lkp.sel
   127  	}
   128  	if ids[0].IsIntegral() || lkp.BatchLookup {
   129  		// for integral types, batch query all ids and then map them back to the input order
   130  		vars, err := sqltypes.BuildBindVariable(ids)
   131  		if err != nil {
   132  			return nil, fmt.Errorf("lookup.Map: %v", err)
   133  		}
   134  		bindVars := map[string]*querypb.BindVariable{
   135  			lkp.FromColumns[0]: vars,
   136  		}
   137  		result, err := vcursor.Execute(ctx, "VindexLookup", sel, bindVars, false /* rollbackOnError */, co)
   138  		if err != nil {
   139  			return nil, fmt.Errorf("lookup.Map: %v", err)
   140  		}
   141  		resultMap := make(map[string][][]sqltypes.Value)
   142  		for _, row := range result.Rows {
   143  			resultMap[row[0].ToString()] = append(resultMap[row[0].ToString()], []sqltypes.Value{row[1]})
   144  		}
   145  
   146  		for _, id := range ids {
   147  			results = append(results, &sqltypes.Result{
   148  				Rows: resultMap[id.ToString()],
   149  			})
   150  		}
   151  	} else {
   152  		// for non integral and binary type, fallback to send query per id
   153  		for _, id := range ids {
   154  			vars, err := sqltypes.BuildBindVariable([]any{id})
   155  			if err != nil {
   156  				return nil, fmt.Errorf("lookup.Map: %v", err)
   157  			}
   158  			bindVars := map[string]*querypb.BindVariable{
   159  				lkp.FromColumns[0]: vars,
   160  			}
   161  			var result *sqltypes.Result
   162  			result, err = vcursor.Execute(ctx, "VindexLookup", sel, bindVars, false /* rollbackOnError */, co)
   163  			if err != nil {
   164  				return nil, fmt.Errorf("lookup.Map: %v", err)
   165  			}
   166  			rows := make([][]sqltypes.Value, 0, len(result.Rows))
   167  			for _, row := range result.Rows {
   168  				rows = append(rows, []sqltypes.Value{row[1]})
   169  			}
   170  			results = append(results, &sqltypes.Result{
   171  				Rows: rows,
   172  			})
   173  		}
   174  	}
   175  	return results, nil
   176  }
   177  
   178  // Verify returns true if ids map to values.
   179  func (lkp *lookupInternal) Verify(ctx context.Context, vcursor VCursor, ids, values []sqltypes.Value) ([]bool, error) {
   180  	co := vtgatepb.CommitOrder_NORMAL
   181  	if lkp.Autocommit {
   182  		co = vtgatepb.CommitOrder_AUTOCOMMIT
   183  	}
   184  	return lkp.VerifyCustom(ctx, vcursor, ids, values, co)
   185  }
   186  
   187  func (lkp *lookupInternal) VerifyCustom(ctx context.Context, vcursor VCursor, ids, values []sqltypes.Value, co vtgatepb.CommitOrder) ([]bool, error) {
   188  	out := make([]bool, len(ids))
   189  	for i, id := range ids {
   190  		bindVars := map[string]*querypb.BindVariable{
   191  			lkp.FromColumns[0]: sqltypes.ValueBindVariable(id),
   192  			lkp.To:             sqltypes.ValueBindVariable(values[i]),
   193  		}
   194  		result, err := vcursor.Execute(ctx, "VindexVerify", lkp.ver, bindVars, false /* rollbackOnError */, co)
   195  		if err != nil {
   196  			return nil, fmt.Errorf("lookup.Verify: %v", err)
   197  		}
   198  		out[i] = (len(result.Rows) != 0)
   199  	}
   200  	return out, nil
   201  }
   202  
   203  type sorter struct {
   204  	rowsColValues [][]sqltypes.Value
   205  	toValues      []sqltypes.Value
   206  }
   207  
   208  func (v *sorter) Len() int {
   209  	return len(v.toValues)
   210  }
   211  
   212  func (v *sorter) Less(i, j int) bool {
   213  	leftRow := v.rowsColValues[i]
   214  	rightRow := v.rowsColValues[j]
   215  	for cell, left := range leftRow {
   216  		right := rightRow[cell]
   217  		lBytes, _ := left.ToBytes()
   218  		rBytes, _ := right.ToBytes()
   219  		compare := bytes.Compare(lBytes, rBytes)
   220  		if compare < 0 {
   221  			return true
   222  		}
   223  		if compare > 0 {
   224  			return false
   225  		}
   226  	}
   227  	iBytes, _ := v.toValues[i].ToBytes()
   228  	jBytes, _ := v.toValues[j].ToBytes()
   229  	return bytes.Compare(iBytes, jBytes) < 0
   230  }
   231  
   232  func (v *sorter) Swap(i, j int) {
   233  	v.toValues[i], v.toValues[j] = v.toValues[j], v.toValues[i]
   234  	v.rowsColValues[i], v.rowsColValues[j] = v.rowsColValues[j], v.rowsColValues[i]
   235  }
   236  
   237  // Create creates an association between rowsColValues and toValues by inserting rows in the vindex table.
   238  // rowsColValues contains all the rows that are being inserted.
   239  // For each row, we store the value of each column defined in the vindex.
   240  // toValues contains the keyspace_id of each row being inserted.
   241  // Given a vindex with two columns and the following insert:
   242  //
   243  // INSERT INTO table_a (colum_a, column_b, column_c) VALUES (value_a0, value_b0, value_c0), (value_a1, value_b1, value_c1);
   244  // If we assume that the primary vindex is on column_c. The call to create will look like this:
   245  // Create(vcursor, [[value_a0, value_b0,], [value_a1, value_b1]], [binary(value_c0), binary(value_c1)])
   246  // Notice that toValues contains the computed binary value of the keyspace_id.
   247  func (lkp *lookupInternal) Create(ctx context.Context, vcursor VCursor, rowsColValues [][]sqltypes.Value, toValues []sqltypes.Value, ignoreMode bool) error {
   248  	if lkp.Autocommit {
   249  		return lkp.createCustom(ctx, vcursor, rowsColValues, toValues, ignoreMode, vtgatepb.CommitOrder_AUTOCOMMIT)
   250  	}
   251  	return lkp.createCustom(ctx, vcursor, rowsColValues, toValues, ignoreMode, vtgatepb.CommitOrder_NORMAL)
   252  }
   253  
   254  func (lkp *lookupInternal) createCustom(ctx context.Context, vcursor VCursor, rowsColValues [][]sqltypes.Value, toValues []sqltypes.Value, ignoreMode bool, co vtgatepb.CommitOrder) error {
   255  	// Trim rows with null values
   256  	trimmedRowsCols := make([][]sqltypes.Value, 0, len(rowsColValues))
   257  	trimmedToValues := make([]sqltypes.Value, 0, len(toValues))
   258  nextRow:
   259  	for i, row := range rowsColValues {
   260  		for j, col := range row {
   261  			if col.IsNull() {
   262  				if !lkp.IgnoreNulls {
   263  					return fmt.Errorf("lookup.Create: input has null values: row: %d, col: %d", i, j)
   264  				}
   265  				continue nextRow
   266  			}
   267  		}
   268  		trimmedRowsCols = append(trimmedRowsCols, row)
   269  		trimmedToValues = append(trimmedToValues, toValues[i])
   270  	}
   271  	if len(trimmedRowsCols) == 0 {
   272  		return nil
   273  	}
   274  	// We only need to check the first row. Number of cols per row
   275  	// is guaranteed by the engine to be uniform.
   276  	if len(trimmedRowsCols[0]) != len(lkp.FromColumns) {
   277  		return fmt.Errorf("lookup.Create: column vindex count does not match the columns in the lookup: %d vs %v", len(trimmedRowsCols[0]), lkp.FromColumns)
   278  	}
   279  	sort.Sort(&sorter{rowsColValues: trimmedRowsCols, toValues: trimmedToValues})
   280  
   281  	insStmt := "insert"
   282  	if lkp.MultiShardAutocommit {
   283  		insStmt = "insert /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */"
   284  	}
   285  	buf := new(bytes.Buffer)
   286  	if ignoreMode {
   287  		fmt.Fprintf(buf, "%s ignore into %s(", insStmt, lkp.Table)
   288  	} else {
   289  		fmt.Fprintf(buf, "%s into %s(", insStmt, lkp.Table)
   290  	}
   291  	for _, col := range lkp.FromColumns {
   292  		fmt.Fprintf(buf, "%s, ", col)
   293  	}
   294  	fmt.Fprintf(buf, "%s) values(", lkp.To)
   295  
   296  	bindVars := make(map[string]*querypb.BindVariable, 2*len(trimmedRowsCols))
   297  	for rowIdx := range trimmedToValues {
   298  		colIds := trimmedRowsCols[rowIdx]
   299  		if rowIdx != 0 {
   300  			buf.WriteString(", (")
   301  		}
   302  		for colIdx, colID := range colIds {
   303  			fromStr := lkp.FromColumns[colIdx] + "_" + strconv.Itoa(rowIdx)
   304  			bindVars[fromStr] = sqltypes.ValueBindVariable(colID)
   305  			buf.WriteString(":" + fromStr + ", ")
   306  		}
   307  		toStr := lkp.To + "_" + strconv.Itoa(rowIdx)
   308  		buf.WriteString(":" + toStr + ")")
   309  		bindVars[toStr] = sqltypes.ValueBindVariable(trimmedToValues[rowIdx])
   310  	}
   311  
   312  	if lkp.Upsert {
   313  		fmt.Fprintf(buf, " on duplicate key update ")
   314  		for _, col := range lkp.FromColumns {
   315  			fmt.Fprintf(buf, "%s=values(%s), ", col, col)
   316  		}
   317  		fmt.Fprintf(buf, "%s=values(%s)", lkp.To, lkp.To)
   318  	}
   319  
   320  	if _, err := vcursor.Execute(ctx, "VindexCreate", buf.String(), bindVars, true /* rollbackOnError */, co); err != nil {
   321  		return fmt.Errorf("lookup.Create: %v", err)
   322  	}
   323  	return nil
   324  }
   325  
   326  // Delete deletes the association between ids and value.
   327  // rowsColValues contains all the rows that are being deleted.
   328  // For each row, we store the value of each column defined in the vindex.
   329  // value cointains the keyspace_id of the vindex entry being deleted.
   330  //
   331  // Given the following information in a vindex table with two columns:
   332  //
   333  //	+------------------+-----------+--------+
   334  //	| hex(keyspace_id) | a         | b      |
   335  //	+------------------+-----------+--------+
   336  //	| 52CB7B1B31B2222E | valuea    | valueb |
   337  //	+------------------+-----------+--------+
   338  //
   339  // A call to Delete would look like this:
   340  // Delete(vcursor, [[valuea, valueb]], 52CB7B1B31B2222E)
   341  func (lkp *lookupInternal) Delete(ctx context.Context, vcursor VCursor, rowsColValues [][]sqltypes.Value, value sqltypes.Value, co vtgatepb.CommitOrder) error {
   342  	// In autocommit mode, it's not safe to delete. So, it's a no-op.
   343  	if lkp.Autocommit {
   344  		return nil
   345  	}
   346  	if len(rowsColValues) == 0 {
   347  		// This code is unreachable. It's just a failsafe.
   348  		return nil
   349  	}
   350  	// We only need to check the first row. Number of cols per row
   351  	// is guaranteed by the engine to be uniform.
   352  	if len(rowsColValues[0]) != len(lkp.FromColumns) {
   353  		return fmt.Errorf("lookup.Delete: column vindex count does not match the columns in the lookup: %d vs %v", len(rowsColValues[0]), lkp.FromColumns)
   354  	}
   355  	for _, column := range rowsColValues {
   356  		bindVars := make(map[string]*querypb.BindVariable, len(rowsColValues))
   357  		for colIdx, columnValue := range column {
   358  			bindVars[lkp.FromColumns[colIdx]] = sqltypes.ValueBindVariable(columnValue)
   359  		}
   360  		bindVars[lkp.To] = sqltypes.ValueBindVariable(value)
   361  		_, err := vcursor.Execute(ctx, "VindexDelete", lkp.del, bindVars, true /* rollbackOnError */, co)
   362  		if err != nil {
   363  			return fmt.Errorf("lookup.Delete: %v", err)
   364  		}
   365  	}
   366  	return nil
   367  }
   368  
   369  // Update implements the update functionality.
   370  func (lkp *lookupInternal) Update(ctx context.Context, vcursor VCursor, oldValues []sqltypes.Value, ksid []byte, toValue sqltypes.Value, newValues []sqltypes.Value) error {
   371  	if err := lkp.Delete(ctx, vcursor, [][]sqltypes.Value{oldValues}, toValue, vtgatepb.CommitOrder_NORMAL); err != nil {
   372  		return err
   373  	}
   374  	return lkp.Create(ctx, vcursor, [][]sqltypes.Value{newValues}, []sqltypes.Value{toValue}, false /* ignoreMode */)
   375  }
   376  
   377  func (lkp *lookupInternal) initDelStmt() string {
   378  	var delBuffer bytes.Buffer
   379  	fmt.Fprintf(&delBuffer, "delete from %s where ", lkp.Table)
   380  	for colIdx, column := range lkp.FromColumns {
   381  		if colIdx != 0 {
   382  			delBuffer.WriteString(" and ")
   383  		}
   384  		delBuffer.WriteString(column + " = :" + column)
   385  	}
   386  	delBuffer.WriteString(" and " + lkp.To + " = :" + lkp.To)
   387  	return delBuffer.String()
   388  }
   389  
   390  func (lkp *lookupInternal) query() (selQuery string, arguments []string) {
   391  	return lkp.sel, lkp.FromColumns
   392  }
   393  
   394  type commonConfig struct {
   395  	autocommit           bool
   396  	multiShardAutocommit bool
   397  }
   398  
   399  func parseCommonConfig(m map[string]string) (*commonConfig, error) {
   400  	var c commonConfig
   401  	var err error
   402  	if c.autocommit, err = boolFromMap(m, "autocommit"); err != nil {
   403  		return nil, err
   404  	}
   405  	if c.multiShardAutocommit, err = boolFromMap(m, "multi_shard_autocommit"); err != nil {
   406  		return nil, err
   407  	}
   408  	return &c, nil
   409  }
   410  
   411  func boolFromMap(m map[string]string, key string) (bool, error) {
   412  	val, ok := m[key]
   413  	if !ok {
   414  		return false, nil
   415  	}
   416  	switch val {
   417  	case "true":
   418  		return true, nil
   419  	case "false":
   420  		return false, nil
   421  	default:
   422  		return false, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "%s value must be 'true' or 'false': '%s'", key, val)
   423  	}
   424  }