vitess.io/vitess@v0.16.2/go/vt/vtgate/planbuilder/update_planner.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 planbuilder
    18  
    19  import (
    20  	"fmt"
    21  
    22  	"vitess.io/vitess/go/vt/sqlparser"
    23  	"vitess.io/vitess/go/vt/vterrors"
    24  	"vitess.io/vitess/go/vt/vtgate/engine"
    25  	"vitess.io/vitess/go/vt/vtgate/evalengine"
    26  	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
    27  
    28  	"vitess.io/vitess/go/vt/vtgate/semantics"
    29  	"vitess.io/vitess/go/vt/vtgate/vindexes"
    30  )
    31  
    32  // buildUpdatePlan returns a stmtPlanner that builds the instructions for an UPDATE statement.
    33  func buildUpdatePlan(string) stmtPlanner {
    34  	return func(stmt sqlparser.Statement, reservedVars *sqlparser.ReservedVars, vschema plancontext.VSchema) (*planResult, error) {
    35  		upd := stmt.(*sqlparser.Update)
    36  		if upd.With != nil {
    37  			return nil, vterrors.VT12001("WITH expression in UPDATE statement")
    38  		}
    39  		dml, tables, ksidVindex, err := buildDMLPlan(vschema, "update", stmt, reservedVars, upd.TableExprs, upd.Where, upd.OrderBy, upd.Limit, upd.Comments, upd.Exprs)
    40  		if err != nil {
    41  			return nil, err
    42  		}
    43  		eupd := &engine.Update{DML: dml}
    44  
    45  		if dml.Opcode == engine.Unsharded {
    46  			return newPlanResult(eupd, tables...), nil
    47  		}
    48  		eupdTable, err := eupd.GetSingleTable()
    49  		if err != nil {
    50  			return nil, err
    51  		}
    52  		cvv, ovq, err := buildChangedVindexesValues(upd, eupdTable, ksidVindex.Columns)
    53  		if err != nil {
    54  			return nil, err
    55  		}
    56  		eupd.ChangedVindexValues = cvv
    57  		eupd.OwnedVindexQuery = ovq
    58  		if len(eupd.ChangedVindexValues) != 0 {
    59  			eupd.KsidVindex = ksidVindex.Vindex
    60  			eupd.KsidLength = len(ksidVindex.Columns)
    61  		}
    62  		return newPlanResult(eupd, tables...), nil
    63  	}
    64  }
    65  
    66  // buildChangedVindexesValues adds to the plan all the lookup vindexes that are changing.
    67  // Updates can only be performed to secondary lookup vindexes with no complex expressions
    68  // in the set clause.
    69  func buildChangedVindexesValues(update *sqlparser.Update, table *vindexes.Table, ksidCols []sqlparser.IdentifierCI) (map[string]*engine.VindexValues, string, error) {
    70  	changedVindexes := make(map[string]*engine.VindexValues)
    71  	buf, offset := initialQuery(ksidCols, table)
    72  	for i, vindex := range table.ColumnVindexes {
    73  		vindexValueMap := make(map[string]evalengine.Expr)
    74  		first := true
    75  		for _, vcol := range vindex.Columns {
    76  			// Searching in order of columns in colvindex.
    77  			found := false
    78  			for _, assignment := range update.Exprs {
    79  				if !vcol.Equal(assignment.Name.Name) {
    80  					continue
    81  				}
    82  				if found {
    83  					return nil, "", vterrors.VT03015(assignment.Name.Name)
    84  				}
    85  				found = true
    86  				pv, err := extractValueFromUpdate(assignment)
    87  				if err != nil {
    88  					return nil, "", err
    89  				}
    90  				vindexValueMap[vcol.String()] = pv
    91  				if first {
    92  					buf.Myprintf(", %v", assignment)
    93  					first = false
    94  				} else {
    95  					buf.Myprintf(" and %v", assignment)
    96  				}
    97  			}
    98  		}
    99  		if len(vindexValueMap) == 0 {
   100  			// Vindex not changing, continue
   101  			continue
   102  		}
   103  
   104  		if update.Limit != nil && len(update.OrderBy) == 0 {
   105  			return nil, "", vterrors.VT12001(fmt.Sprintf("need to provide ORDER BY clause when using LIMIT; invalid update on vindex: %v", vindex.Name))
   106  		}
   107  		if i == 0 {
   108  			return nil, "", vterrors.VT12001(fmt.Sprintf("you cannot update primary vindex columns; invalid update on vindex: %v", vindex.Name))
   109  		}
   110  		if _, ok := vindex.Vindex.(vindexes.Lookup); !ok {
   111  			return nil, "", vterrors.VT12001(fmt.Sprintf("you can only update lookup vindexes; invalid update on vindex: %v", vindex.Name))
   112  		}
   113  		changedVindexes[vindex.Name] = &engine.VindexValues{
   114  			PvMap:  vindexValueMap,
   115  			Offset: offset,
   116  		}
   117  		offset++
   118  	}
   119  	if len(changedVindexes) == 0 {
   120  		return nil, "", nil
   121  	}
   122  	// generate rest of the owned vindex query.
   123  	aTblExpr, ok := update.TableExprs[0].(*sqlparser.AliasedTableExpr)
   124  	if !ok {
   125  		return nil, "", vterrors.VT12001("UPDATE on complex table expression")
   126  	}
   127  	tblExpr := &sqlparser.AliasedTableExpr{Expr: sqlparser.TableName{Name: table.Name}, As: aTblExpr.As}
   128  	buf.Myprintf(" from %v%v%v%v for update", tblExpr, update.Where, update.OrderBy, update.Limit)
   129  	return changedVindexes, buf.String(), nil
   130  }
   131  
   132  func initialQuery(ksidCols []sqlparser.IdentifierCI, table *vindexes.Table) (*sqlparser.TrackedBuffer, int) {
   133  	buf := sqlparser.NewTrackedBuffer(nil)
   134  	offset := 0
   135  	for _, col := range ksidCols {
   136  		if offset == 0 {
   137  			buf.Myprintf("select %v", col)
   138  		} else {
   139  			buf.Myprintf(", %v", col)
   140  		}
   141  		offset++
   142  	}
   143  	for _, cv := range table.Owned {
   144  		for _, column := range cv.Columns {
   145  			buf.Myprintf(", %v", column)
   146  			offset++
   147  		}
   148  	}
   149  	return buf, offset
   150  }
   151  
   152  // extractValueFromUpdate given an UpdateExpr attempts to extracts the Value
   153  // it's holding. At the moment it only supports: StrVal, HexVal, IntVal, ValArg.
   154  // If a complex expression is provided (e.g set name = name + 1), the update will be rejected.
   155  func extractValueFromUpdate(upd *sqlparser.UpdateExpr) (evalengine.Expr, error) {
   156  	pv, err := evalengine.Translate(upd.Expr, semantics.EmptySemTable())
   157  	if err != nil || sqlparser.IsSimpleTuple(upd.Expr) {
   158  		err := vterrors.VT12001(fmt.Sprintf("only values are supported: invalid update on column: `%s` with expr: [%s]", upd.Name.Name.String(), sqlparser.String(upd.Expr)))
   159  		return nil, err
   160  	}
   161  	return pv, nil
   162  }