onlineddl: ALTER VITESS_MIGRATION CANCEL CONTEXT 'ctx' (#20146)

Signed-off-by: Shlomi Noach <2607934+shlomi-noach@users.noreply.github.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: shlomi-noach

go/test/endtoend/onlineddl/scheduler/onlineddl_scheduler_test.go

@@ -2042,6 +2042,25 @@ func testScheduler(t *testing.T) {
 		})
 		testTableCompletionTimes(t, t1uuid, v1uuid)
 	})
+	t.Run("cancel migrations by context", func(t *testing.T) {
+		// Submit two migrations with the same explicit context, both postponed so they stay running.
+		t1uuid = testOnlineDDLStatement(t, &testOnlineDDLStatementParams{ddlStatement: trivialAlterT1Statement, ddlStrategy: ddlStrategy + " --allow-concurrent --postpone-completion", executeStrategy: "vtctl", migrationContext: "ctx-cancel-by-context", skipWait: true})
+		t2uuid = testOnlineDDLStatement(t, &testOnlineDDLStatementParams{ddlStatement: trivialAlterT2Statement, ddlStrategy: ddlStrategy + " --allow-concurrent --postpone-completion", executeStrategy: "vtctl", migrationContext: "ctx-cancel-by-context", skipWait: true})
+		onlineddl.WaitForMigrationStatus(t, &vtParams, shards, t1uuid, normalWaitTime, schema.OnlineDDLStatusRunning)
+		onlineddl.WaitForMigrationStatus(t, &vtParams, shards, t2uuid, normalWaitTime, schema.OnlineDDLStatusRunning)
+
+		// A non-matching context cancels nothing.
+		onlineddl.CheckCancelContextMigrations(t, &vtParams, "ctx-cancel-by-context-other", 0)
+		onlineddl.CheckMigrationStatus(t, &vtParams, shards, t1uuid, schema.OnlineDDLStatusRunning)
+		onlineddl.CheckMigrationStatus(t, &vtParams, shards, t2uuid, schema.OnlineDDLStatusRunning)
+
+		// Cancel by context: both migrations must be cancelled.
+		onlineddl.CheckCancelContextMigrations(t, &vtParams, "ctx-cancel-by-context", 2)
+		onlineddl.WaitForMigrationStatus(t, &vtParams, shards, t1uuid, normalWaitTime, schema.OnlineDDLStatusCancelled, schema.OnlineDDLStatusFailed)
+		onlineddl.WaitForMigrationStatus(t, &vtParams, shards, t2uuid, normalWaitTime, schema.OnlineDDLStatusCancelled, schema.OnlineDDLStatusFailed)
+		onlineddl.CheckMigrationStatus(t, &vtParams, shards, t1uuid, schema.OnlineDDLStatusCancelled)
+		onlineddl.CheckMigrationStatus(t, &vtParams, shards, t2uuid, schema.OnlineDDLStatusCancelled)
+	})
 }
 
 func testSingleton(t *testing.T) {

go/test/endtoend/onlineddl/vtgate_util.go

@@ -257,6 +257,17 @@ func CheckCancelAllMigrations(t *testing.T, vtParams *mysql.ConnParams, expectCo
 	}
 }
 
+// CheckCancelContextMigrations cancels all pending migrations with a given context and expect number of affected rows
+// A negative value for expectCount indicates "don't care, no need to check"
+func CheckCancelContextMigrations(t *testing.T, vtParams *mysql.ConnParams, migrationContext string, expectCount int) {
+	cancelQuery := fmt.Sprintf("alter vitess_migration cancel context '%s'", migrationContext)
+	r := VtgateExecQuery(t, vtParams, cancelQuery, "")
+
+	if expectCount >= 0 {
+		assert.Equal(t, expectCount, int(r.RowsAffected))
+	}
+}
+
 // CheckCleanupAllMigrations cleans up all applicable migrations and expect number of affected rows
 // A negative value for expectCount indicates "don't care, no need to check"
 func CheckCleanupAllMigrations(t *testing.T, vtParams *mysql.ConnParams, expectCount int) uint64 {

go/vt/sqlparser/ast.go

@@ -523,6 +523,7 @@ type (
 	AlterMigration struct {
 		Type      AlterMigrationType
 		UUID      string
+		Context   string
 		Expire    string
 		Ratio     *Literal
 		Threshold string

go/vt/sqlparser/ast_equals.go

@@ -337,7 +337,11 @@ func (node *AlterMigration) Format(buf *TrackedBuffer) {
 	case CancelMigrationType:
 		alterType = "cancel"
 	case CancelAllMigrationType:
-		alterType = "cancel all"
+		if node.Context != "" {
+			alterType = "cancel"
+		} else {
+			alterType = "cancel all"
+		}
 	case ThrottleMigrationType:
 		alterType = "throttle"
 	case ThrottleAllMigrationType:
@@ -355,16 +359,19 @@ func (node *AlterMigration) Format(buf *TrackedBuffer) {
 	}
 	buf.astPrintf(node, " %#s", alterType)
 	if node.Threshold != "" {
-		buf.astPrintf(node, " '%#s'", node.Threshold)
+		buf.astPrintf(node, " %#s", encodeSQLString(node.Threshold))
 	}
 	if node.Expire != "" {
-		buf.astPrintf(node, " expire '%#s'", node.Expire)
+		buf.astPrintf(node, " expire %#s", encodeSQLString(node.Expire))
 	}
 	if node.Ratio != nil {
 		buf.astPrintf(node, " ratio %v", node.Ratio)
 	}
+	if node.Context != "" {
+		buf.astPrintf(node, " context %#s", encodeSQLString(node.Context))
+	}
 	if node.Shards != "" {
-		buf.astPrintf(node, " vitess_shards '%#s'", node.Shards)
+		buf.astPrintf(node, " vitess_shards %#s", encodeSQLString(node.Shards))
 	}
 }
 

go/vt/sqlparser/ast_format.go

@@ -191,6 +191,7 @@ var keywords = []keyword{
 	{"constraint_catalog", CONSTRAINT_CATALOG},
 	{"constraint_name", CONSTRAINT_NAME},
 	{"constraint_schema", CONSTRAINT_SCHEMA},
+	{"context", CONTEXT},
 	{"continue", CONTINUE},
 	{"convert", CONVERT},
 	{"copy", COPY},

go/vt/sqlparser/ast_format_fast.go

@@ -2689,6 +2689,8 @@ var validSQL = []struct {
 }, {
 	input:  "alter vitess_migration '9748c3b7_7fdb_11eb_ac2c_f875a4d24e90' FORCE_CUTOVER",
 	output: "alter vitess_migration '9748c3b7_7fdb_11eb_ac2c_f875a4d24e90' force_cutover",
+}, {
+	input: "alter vitess_migration cancel context 'some-context'",
 }, {
 	input: "alter vitess_migration cancel all",
 }, {
@@ -6488,6 +6490,9 @@ var invalidSQL = []struct {
 	input  string
 	output string
 }{{
+	input:  "alter vitess_migration cancel context ''",
+	output: "migration context cannot be empty at position 41",
+}, {
 	input:  "select : from t",
 	output: "syntax error at position 9 near ':'",
 }, {