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 }