Fix panic for unknown columns in foreign key managed mode (#15025)

Signed-off-by: Manan Gupta <manan@planetscale.com>

go/vt/vtgate/planbuilder/plan_test.go

@@ -183,8 +183,10 @@ func setFks(t *testing.T, vschema *vindexes.VSchema) {
 		// FK from tblrefDef referencing tbl20 that is shard scoped of SET-Default types.
 		_ = vschema.AddForeignKey("sharded_fk_allow", "tblrefDef", createFkDefinition([]string{"ref"}, "tbl20", []string{"col2"}, sqlparser.SetDefault, sqlparser.SetDefault))
 
+		// FK from tbl_auth referencing tbl20 that is shard scoped of CASCADE types.
+		_ = vschema.AddForeignKey("sharded_fk_allow", "tbl_auth", createFkDefinition([]string{"id"}, "tbl20", []string{"col2"}, sqlparser.Cascade, sqlparser.Cascade))
 		addPKs(t, vschema, "sharded_fk_allow", []string{"tbl1", "tbl2", "tbl3", "tbl4", "tbl5", "tbl6", "tbl7", "tbl9", "tbl10",
-			"multicol_tbl1", "multicol_tbl2", "tblrefDef", "tbl20"})
+			"multicol_tbl1", "multicol_tbl2", "tbl_auth", "tblrefDef", "tbl20"})
 	}
 	if vschema.Keyspaces["unsharded_fk_allow"] != nil {
 		// u_tbl2(col2)  -> u_tbl1(col1)  Cascade.

go/vt/vtgate/planbuilder/testdata/foreignkey_cases.json

@@ -3043,5 +3043,10 @@
         "unsharded_fk_allow.u_tbl3"
       ]
     }
+  },
+  {
+    "comment": "Unknown update column in foreign keys",
+    "query": "update tbl_auth set unknown_col = 'verified' where id = 1",
+    "plan": "column 'unknown_col' not found in table 'tbl_auth'"
   }
 ]

go/vt/vtgate/planbuilder/testdata/vschemas/schema.json

@@ -723,6 +723,20 @@
             }
           ]
         },
+        "tbl_auth": {
+          "columns": [
+            {
+              "name": "id"
+            }
+          ],
+          "column_vindexes": [
+            {
+              "column": "id",
+              "name": "hash_vin"
+            }
+          ],
+          "column_list_authoritative": true
+        },
         "tblrefDef": {
           "column_vindexes": [
             {

go/vt/vtgate/semantics/analyzer.go

@@ -352,25 +352,24 @@ func (a *analyzer) getInvolvedForeignKeys(statement sqlparser.Statement, fkCheck
 		}
 		// If only a certain set of columns are being updated, then there might be some child foreign keys that don't need any consideration since their columns aren't being updated.
 		// So, we filter these child foreign keys out. We can't filter any parent foreign keys because the statement will INSERT a row too, which requires validating all the parent foreign keys.
-		updatedChildFks, _, childFkToUpdExprs := a.filterForeignKeysUsingUpdateExpressions(allChildFks, nil, sqlparser.UpdateExprs(stmt.OnDup))
-		return updatedChildFks, allParentFKs, childFkToUpdExprs, nil
+		updatedChildFks, _, childFkToUpdExprs, err := a.filterForeignKeysUsingUpdateExpressions(allChildFks, nil, sqlparser.UpdateExprs(stmt.OnDup))
+		return updatedChildFks, allParentFKs, childFkToUpdExprs, err
 	case *sqlparser.Update:
 		// For UPDATE queries we get all the parent and child foreign keys, but we can filter some of them out if the columns that they consist off aren't being updated or are set to NULLs.
 		allChildFks, allParentFks, err := a.getAllManagedForeignKeys()
 		if err != nil {
 			return nil, nil, nil, err
 		}
-		childFks, parentFks, childFkToUpdExprs := a.filterForeignKeysUsingUpdateExpressions(allChildFks, allParentFks, stmt.Exprs)
-		return childFks, parentFks, childFkToUpdExprs, nil
+		return a.filterForeignKeysUsingUpdateExpressions(allChildFks, allParentFks, stmt.Exprs)
 	default:
 		return nil, nil, nil, nil
 	}
 }
 
 // filterForeignKeysUsingUpdateExpressions filters the child and parent foreign key constraints that don't require any validations/cascades given the updated expressions.
-func (a *analyzer) filterForeignKeysUsingUpdateExpressions(allChildFks map[TableSet][]vindexes.ChildFKInfo, allParentFks map[TableSet][]vindexes.ParentFKInfo, updExprs sqlparser.UpdateExprs) (map[TableSet][]vindexes.ChildFKInfo, map[TableSet][]vindexes.ParentFKInfo, map[string]sqlparser.UpdateExprs) {
+func (a *analyzer) filterForeignKeysUsingUpdateExpressions(allChildFks map[TableSet][]vindexes.ChildFKInfo, allParentFks map[TableSet][]vindexes.ParentFKInfo, updExprs sqlparser.UpdateExprs) (map[TableSet][]vindexes.ChildFKInfo, map[TableSet][]vindexes.ParentFKInfo, map[string]sqlparser.UpdateExprs, error) {
 	if len(allChildFks) == 0 && len(allParentFks) == 0 {
-		return nil, nil, nil
+		return nil, nil, nil, nil
 	}
 
 	pFksRequired := make(map[TableSet][]bool, len(allParentFks))
@@ -392,7 +391,10 @@ func (a *analyzer) filterForeignKeysUsingUpdateExpressions(allChildFks map[Table
 	for _, updateExpr := range updExprs {
 		deps := a.binder.direct.dependencies(updateExpr.Name)
 		if deps.NumberOfTables() != 1 {
-			panic("expected to have single table dependency")
+			// If we don't get exactly one table for the given update expression, we would have definitely run into an error
+			// during the binder phase that we would have stored. We should return that error, since we can't safely proceed with
+			// foreign key related changes without having all the information.
+			return nil, nil, nil, a.getError()
 		}
 		updExprToTableSet[updateExpr.Name] = deps
 		// Get all the child and parent foreign keys for the given table that the update expression belongs to.
@@ -459,7 +461,18 @@ func (a *analyzer) filterForeignKeysUsingUpdateExpressions(allChildFks map[Table
 		}
 		cFksNeedsHandling[ts] = cFKNeeded
 	}
-	return cFksNeedsHandling, pFksNeedsHandling, childFKToUpdExprs
+	return cFksNeedsHandling, pFksNeedsHandling, childFKToUpdExprs, nil
+}
+
+// getError gets the error stored in the analyzer during previous phases.
+func (a *analyzer) getError() error {
+	if a.projErr != nil {
+		return a.projErr
+	}
+	if a.unshardedErr != nil {
+		return a.unshardedErr
+	}
+	return a.err
 }
 
 // getAllManagedForeignKeys gets all the foreign keys for the query we are analyzing that Vitess is responsible for managing.

go/vt/vtgate/semantics/analyzer_fk_test.go

@@ -235,6 +235,7 @@ func TestFilterForeignKeysUsingUpdateExpressions(t *testing.T) {
 				cold: SingleTableSet(1),
 			},
 		},
+		unshardedErr: fmt.Errorf("ambiguous test error"),
 		tables: &tableCollector{
 			Tables: []TableInfo{
 				tbl["t4"],
@@ -261,6 +262,7 @@ func TestFilterForeignKeysUsingUpdateExpressions(t *testing.T) {
 		updExprs        sqlparser.UpdateExprs
 		childFksWanted  map[TableSet][]vindexes.ChildFKInfo
 		parentFksWanted map[TableSet][]vindexes.ParentFKInfo
+		errWanted       string
 	}{
 		{
 			name:         "Child Foreign Keys Filtering",
@@ -301,13 +303,30 @@ func TestFilterForeignKeysUsingUpdateExpressions(t *testing.T) {
 					pkInfo(parentTbl, []string{"pcolc", "pcolx"}, []string{"colc", "colx"}),
 				},
 			},
+		}, {
+			name:     "Unknown column",
+			analyzer: a,
+			allParentFks: map[TableSet][]vindexes.ParentFKInfo{
+				SingleTableSet(0): tbl["t4"].(*RealTable).Table.ParentForeignKeys,
+				SingleTableSet(1): tbl["t5"].(*RealTable).Table.ParentForeignKeys,
+			},
+			allChildFks: nil,
+			updExprs: sqlparser.UpdateExprs{
+				&sqlparser.UpdateExpr{Name: sqlparser.NewColName("unknownCol"), Expr: sqlparser.NewIntLiteral("1")},
+			},
+			errWanted: "ambiguous test error",
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			childFks, parentFks, _ := tt.analyzer.filterForeignKeysUsingUpdateExpressions(tt.allChildFks, tt.allParentFks, tt.updExprs)
+			childFks, parentFks, _, err := tt.analyzer.filterForeignKeysUsingUpdateExpressions(tt.allChildFks, tt.allParentFks, tt.updExprs)
 			require.EqualValues(t, tt.childFksWanted, childFks)
 			require.EqualValues(t, tt.parentFksWanted, parentFks)
+			if tt.errWanted != "" {
+				require.EqualError(t, err, tt.errWanted)
+			} else {
+				require.NoError(t, err)
+			}
 		})
 	}
 }