From df36e71b661254a044c08735d0497d25798f6918 Mon Sep 17 00:00:00 2001 From: Andrew Mason Date: Mon, 25 Sep 2023 22:36:25 -0400 Subject: [PATCH 1/5] actually test vtcombo (#14095) Signed-off-by: Andrew Mason --- go/flags/endtoend/flags_test.go | 16 ++++++++++------ go/flags/endtoend/vtcombo.txt | 1 + 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/go/flags/endtoend/flags_test.go b/go/flags/endtoend/flags_test.go index ee24fd6a36d..713d305fc61 100644 --- a/go/flags/endtoend/flags_test.go +++ b/go/flags/endtoend/flags_test.go @@ -44,6 +44,9 @@ var ( //go:embed vtaclcheck.txt vtaclcheckTxt string + //go:embed vtcombo.txt + vtcomboTxt string + //go:embed vtexplain.txt vtexplainTxt string @@ -87,19 +90,20 @@ var ( "mysqlctl": mysqlctlTxt, "mysqlctld": mysqlctldTxt, "vtaclcheck": vtaclcheckTxt, - "vtexplain": vtexplainTxt, - "vtgate": vtgateTxt, - "vttablet": vttabletTxt, - "vttlstest": vttlstestTxt, + "vtbackup": vtbackupTxt, + "vtcombo": vtcomboTxt, "vtctld": vtctldTxt, "vtctlclient": vtctlclientTxt, "vtctldclient": vtctldclientTxt, + "vtexplain": vtexplainTxt, + "vtgate": vtgateTxt, "vtorc": vtorcTxt, + "vttablet": vttabletTxt, "vttestserver": vttestserverTxt, - "zkctld": zkctldTxt, - "vtbackup": vtbackupTxt, + "vttlstest": vttlstestTxt, "zk": zkTxt, "zkctl": zkctlTxt, + "zkctld": zkctldTxt, } ) diff --git a/go/flags/endtoend/vtcombo.txt b/go/flags/endtoend/vtcombo.txt index 89d972f2f6b..ffa2b84970a 100644 --- a/go/flags/endtoend/vtcombo.txt +++ b/go/flags/endtoend/vtcombo.txt @@ -139,6 +139,7 @@ Flags: --gc_check_interval duration Interval between garbage collection checks (default 1h0m0s) --gc_purge_check_interval duration Interval between purge discovery checks (default 1m0s) --gh-ost-path string override default gh-ost binary full path + --grpc-send-session-in-streaming If set, will send the session as last packet in streaming api to support transactions in streaming --grpc-use-effective-groups If set, and SSL is not used, will set the immediate caller's security groups from the effective caller id's groups. --grpc-use-static-authentication-callerid If set, will set the immediate caller id to the username authenticated by the static auth plugin. --grpc_auth_mode string Which auth plugin implementation to use (eg: static) From d743e1e03c51602c62beb12b80bbcd0e1ad7c3a0 Mon Sep 17 00:00:00 2001 From: Manan Gupta <35839558+GuptaManan100@users.noreply.github.com> Date: Tue, 26 Sep 2023 14:50:33 +0530 Subject: [PATCH 2/5] Fix Fk verification and update queries to accommodate for bindVariables being NULL (#14061) Signed-off-by: Manan Gupta --- go/vt/vtgate/planbuilder/operators/update.go | 98 +++++---- .../testdata/foreignkey_cases.json | 207 +++++++++++++++++- go/vt/vtgate/semantics/semantic_state.go | 4 + 3 files changed, 257 insertions(+), 52 deletions(-) diff --git a/go/vt/vtgate/planbuilder/operators/update.go b/go/vt/vtgate/planbuilder/operators/update.go index 8a16b97117e..2a101438cb2 100644 --- a/go/vt/vtgate/planbuilder/operators/update.go +++ b/go/vt/vtgate/planbuilder/operators/update.go @@ -338,7 +338,7 @@ func createFkChildForUpdate(ctx *plancontext.PlanningContext, fk vindexes.ChildF case sqlparser.Cascade: childOp, err = buildChildUpdOpForCascade(ctx, fk, updStmt, childWhereExpr, updatedTable) case sqlparser.SetNull: - childOp, err = buildChildUpdOpForSetNull(ctx, fk, updStmt, childWhereExpr, valTuple) + childOp, err = buildChildUpdOpForSetNull(ctx, fk, updStmt, childWhereExpr) case sqlparser.SetDefault: return nil, vterrors.VT09016() } @@ -398,8 +398,8 @@ func buildChildUpdOpForCascade(ctx *plancontext.PlanningContext, fk vindexes.Chi // // `UPDATE SET // WHERE IN () -// [AND NOT IN ()]` -func buildChildUpdOpForSetNull(ctx *plancontext.PlanningContext, fk vindexes.ChildFKInfo, updStmt *sqlparser.Update, childWhereExpr sqlparser.Expr, valTuple sqlparser.ValTuple) (ops.Operator, error) { +// [AND ({ IS NULL OR}... NOT IN ())]` +func buildChildUpdOpForSetNull(ctx *plancontext.PlanningContext, fk vindexes.ChildFKInfo, updStmt *sqlparser.Update, childWhereExpr sqlparser.Expr) (ops.Operator, error) { // For the SET NULL type constraint, we need to set all the child columns to NULL. var childUpdateExprs sqlparser.UpdateExprs for _, column := range fk.ChildColumns { @@ -411,24 +411,18 @@ func buildChildUpdOpForSetNull(ctx *plancontext.PlanningContext, fk vindexes.Chi // SET NULL cascade should be avoided for the case where the parent columns remains unchanged on the update. // We need to add a condition to the where clause to handle this case. - // The additional condition looks like [AND NOT IN ()]. + // The additional condition looks like [AND ({ IS NULL OR}... NOT IN ())]. // If any of the parent columns is being set to NULL, then we don't need this condition. - var updateValues sqlparser.ValTuple - colSetToNull := false - for _, updateExpr := range updStmt.Exprs { - colIdx := fk.ParentColumns.FindColumn(updateExpr.Name.Name) - if colIdx >= 0 { - if sqlparser.IsNull(updateExpr.Expr) { - colSetToNull = true - break - } - updateValues = append(updateValues, updateExpr.Expr) - } - } - if !colSetToNull { + // However, we don't necessarily know on Plan time if the Expr being updated to is NULL or not. Specifically, bindVariables in Prepared statements can be NULL on runtime. + // Therefore, in the condition we create, we also need to make it resilient to NULL values. Therefore we check if each individual value is NULL or not and OR it with the main condition. + // For example, if we are setting `update parent cola = :v1 and colb = :v2`, then on the child, the where condition would look something like this - + // `:v1 IS NULL OR :v2 IS NULL OR (child_cola, child_colb) NOT IN ((:v1,:v2))` + // So, if either of :v1 or :v2 is NULL, then the entire condition is true (which is the same as not having the condition when :v1 or :v2 is NULL). + compExpr := nullSafeNotInComparison(updStmt.Exprs, fk) + if compExpr != nil { childWhereExpr = &sqlparser.AndExpr{ Left: childWhereExpr, - Right: sqlparser.NewComparisonExpr(sqlparser.NotInOp, valTuple, sqlparser.ValTuple{updateValues}, nil), + Right: compExpr, } } childUpdStmt := &sqlparser.Update{ @@ -556,13 +550,13 @@ func createFkVerifyOpForParentFKForUpdate(ctx *plancontext.PlanningContext, updS } // Each child foreign key constraint is verified by a join query of the form: -// select 1 from child_tbl join parent_tbl on where [AND NOT IN ()] limit 1 +// select 1 from child_tbl join parent_tbl on where [AND ({ IS NULL OR}... NOT IN ())] limit 1 // E.g: // Child (c1, c2) references Parent (p1, p2) // update Parent set p1 = 1 where id = 1 // verify query: // select 1 from Child join Parent on Parent.p1 = Child.c1 and Parent.p2 = Child.c2 -// where Parent.id = 1 and (parent.p1) NOT IN ((1)) limit 1 +// where Parent.id = 1 and (1 IS NULL OR (child.c1) NOT IN ((1))) limit 1 func createFkVerifyOpForChildFKForUpdate(ctx *plancontext.PlanningContext, updStmt *sqlparser.Update, cFk vindexes.ChildFKInfo) (ops.Operator, error) { // ON UPDATE RESTRICT foreign keys that require validation, should only be allowed in the case where we // are verifying all the FKs on vtgate level. @@ -598,27 +592,16 @@ func createFkVerifyOpForChildFKForUpdate(ctx *plancontext.PlanningContext, updSt // We don't want to fail the RESTRICT for the case where the parent columns remains unchanged on the update. // We need to add a condition to the where clause to handle this case. - // The additional condition looks like [AND NOT IN ()]. + // The additional condition looks like [AND ({ IS NULL OR}... NOT IN ())]. // If any of the parent columns is being set to NULL, then we don't need this condition. - var updateValues sqlparser.ValTuple - colSetToNull := false - for _, updateExpr := range updStmt.Exprs { - colIdx := cFk.ParentColumns.FindColumn(updateExpr.Name.Name) - if colIdx >= 0 { - if sqlparser.IsNull(updateExpr.Expr) { - colSetToNull = true - break - } - updateValues = append(updateValues, updateExpr.Expr) - } - } - if !colSetToNull { - // Create a ValTuple of child column names - var valTuple sqlparser.ValTuple - for _, column := range cFk.ParentColumns { - valTuple = append(valTuple, sqlparser.NewColNameWithQualifier(column.String(), parentTbl)) - } - whereCond = sqlparser.AndExpressions(whereCond, sqlparser.NewComparisonExpr(sqlparser.NotInOp, valTuple, sqlparser.ValTuple{updateValues}, nil)) + // However, we don't necessarily know on Plan time if the Expr being updated to is NULL or not. Specifically, bindVariables in Prepared statements can be NULL on runtime. + // Therefore, in the condition we create, we also need to make it resilient to NULL values. Therefore we check if each individual value is NULL or not and OR it with the main condition. + // For example, if we are setting `update child cola = :v1 and colb = :v2`, then on the parent, the where condition would look something like this - + // `:v1 IS NULL OR :v2 IS NULL OR (cola, colb) NOT IN ((:v1,:v2))` + // So, if either of :v1 or :v2 is NULL, then the entire condition is true (which is the same as not having the condition when :v1 or :v2 is NULL). + compExpr := nullSafeNotInComparison(updStmt.Exprs, cFk) + if compExpr != nil { + whereCond = sqlparser.AndExpressions(whereCond, compExpr) } return createSelectionOp(ctx, @@ -634,3 +617,38 @@ func createFkVerifyOpForChildFKForUpdate(ctx *plancontext.PlanningContext, updSt sqlparser.NewLimitWithoutOffset(1), sqlparser.ShareModeLock) } + +// nullSafeNotInComparison is used to compare the child columns in the foreign key constraint aren't the same as the updateExpressions exactly. +// This comparison has to be null safe so we create an expression which looks like the following for a query like `update child cola = :v1 and colb = :v2` - +// `:v1 IS NULL OR :v2 IS NULL OR (cola, colb) NOT IN ((:v1,:v2))` +// So, if either of :v1 or :v2 is NULL, then the entire condition is true (which is the same as not having the condition when :v1 or :v2 is NULL) +// This expression is used in cascading SET NULLs and in verifying whether an update should be restricted. +func nullSafeNotInComparison(updateExprs sqlparser.UpdateExprs, cFk vindexes.ChildFKInfo) sqlparser.Expr { + var updateValues sqlparser.ValTuple + for _, updateExpr := range updateExprs { + colIdx := cFk.ParentColumns.FindColumn(updateExpr.Name.Name) + if colIdx >= 0 { + if sqlparser.IsNull(updateExpr.Expr) { + return nil + } + updateValues = append(updateValues, updateExpr.Expr) + } + } + // Create a ValTuple of child column names + var valTuple sqlparser.ValTuple + for _, column := range cFk.ChildColumns { + valTuple = append(valTuple, sqlparser.NewColNameWithQualifier(column.String(), cFk.Table.GetTableName())) + } + var finalExpr sqlparser.Expr = sqlparser.NewComparisonExpr(sqlparser.NotInOp, valTuple, sqlparser.ValTuple{updateValues}, nil) + for _, value := range updateValues { + finalExpr = &sqlparser.OrExpr{ + Left: &sqlparser.IsExpr{ + Left: value, + Right: sqlparser.IsNullOp, + }, + Right: finalExpr, + } + } + + return finalExpr +} diff --git a/go/vt/vtgate/planbuilder/testdata/foreignkey_cases.json b/go/vt/vtgate/planbuilder/testdata/foreignkey_cases.json index 0b88ac1e9b2..3c03a4740f6 100644 --- a/go/vt/vtgate/planbuilder/testdata/foreignkey_cases.json +++ b/go/vt/vtgate/planbuilder/testdata/foreignkey_cases.json @@ -328,7 +328,7 @@ "Cols": [ 0 ], - "Query": "update u_tbl3 set col3 = null where (col3) in ::fkc_vals and (col3) not in (('bar'))", + "Query": "update u_tbl3 set col3 = null where (col3) in ::fkc_vals and (u_tbl3.col3) not in (('bar'))", "Table": "u_tbl3" }, { @@ -439,7 +439,7 @@ "Cols": [ 0 ], - "Query": "update tbl4 set t4col4 = null where (t4col4) in ::fkc_vals and (t4col4) not in (('foo'))", + "Query": "update tbl4 set t4col4 = null where (t4col4) in ::fkc_vals and (tbl4.t4col4) not in (('foo'))", "Table": "tbl4" }, { @@ -693,7 +693,7 @@ "Cols": [ 0 ], - "Query": "update u_tbl3 set col3 = null where (col3) in ::fkc_vals1 and (col3) not in (('foo'))", + "Query": "update u_tbl3 set col3 = null where (col3) in ::fkc_vals1 and (u_tbl3.col3) not in (('foo'))", "Table": "u_tbl3" }, { @@ -727,7 +727,7 @@ "Sharded": false }, "FieldQuery": "select col9 from u_tbl9 where 1 != 1", - "Query": "select col9 from u_tbl9 where (col9) in ::fkc_vals2 and (col9) not in (('foo')) for update", + "Query": "select col9 from u_tbl9 where (col9) in ::fkc_vals2 and (u_tbl9.col9) not in (('foo')) for update", "Table": "u_tbl9" }, { @@ -755,7 +755,7 @@ "Sharded": false }, "TargetTabletType": "PRIMARY", - "Query": "update u_tbl9 set col9 = null where (col9) in ::fkc_vals2 and (col9) not in (('foo'))", + "Query": "update u_tbl9 set col9 = null where (col9) in ::fkc_vals2 and (u_tbl9.col9) not in (('foo'))", "Table": "u_tbl9" } ] @@ -832,7 +832,7 @@ "Cols": [ 0 ], - "Query": "update u_tbl3 set col3 = null where (col3) in ::fkc_vals and (col3) not in ((2))", + "Query": "update u_tbl3 set col3 = null where (col3) in ::fkc_vals and (u_tbl3.col3) not in ((2))", "Table": "u_tbl3" }, { @@ -909,7 +909,7 @@ "Cols": [ 0 ], - "Query": "update u_tbl3 set col3 = null where (col3) in ::fkc_vals1 and (col3) not in ((2))", + "Query": "update u_tbl3 set col3 = null where (col3) in ::fkc_vals1 and (u_tbl3.col3) not in ((2))", "Table": "u_tbl3" }, { @@ -943,7 +943,7 @@ "Sharded": false }, "FieldQuery": "select col9 from u_tbl9 where 1 != 1", - "Query": "select col9 from u_tbl9 where (col9) in ::fkc_vals2 and (col9) not in ((2)) for update", + "Query": "select col9 from u_tbl9 where (col9) in ::fkc_vals2 and (u_tbl9.col9) not in ((2)) for update", "Table": "u_tbl9" }, { @@ -971,7 +971,7 @@ "Sharded": false }, "TargetTabletType": "PRIMARY", - "Query": "update u_tbl9 set col9 = null where (col9) in ::fkc_vals2 and (col9) not in ((2))", + "Query": "update u_tbl9 set col9 = null where (col9) in ::fkc_vals2 and (u_tbl9.col9) not in ((2))", "Table": "u_tbl9" } ] @@ -1220,7 +1220,7 @@ "Sharded": false }, "FieldQuery": "select 1 from u_tbl4, u_tbl9 where 1 != 1", - "Query": "select 1 from u_tbl4, u_tbl9 where (u_tbl4.col4) in ::fkc_vals and (u_tbl4.col4) not in (('foo')) and u_tbl4.col4 = u_tbl9.col9 limit 1 lock in share mode", + "Query": "select 1 from u_tbl4, u_tbl9 where (u_tbl4.col4) in ::fkc_vals and (u_tbl9.col9) not in (('foo')) and u_tbl4.col4 = u_tbl9.col9 limit 1 lock in share mode", "Table": "u_tbl4, u_tbl9" }, { @@ -1259,6 +1259,95 @@ ] } }, + { + "comment": "Update that cascades and requires parent fk and restrict child fk verification - bindVariable", + "query": "update u_tbl7 set col7 = :v1", + "plan": { + "QueryType": "UPDATE", + "Original": "update u_tbl7 set col7 = :v1", + "Instructions": { + "OperatorType": "FkCascade", + "Inputs": [ + { + "InputName": "Selection", + "OperatorType": "Route", + "Variant": "Unsharded", + "Keyspace": { + "Name": "unsharded_fk_allow", + "Sharded": false + }, + "FieldQuery": "select col7 from u_tbl7 where 1 != 1", + "Query": "select col7 from u_tbl7 for update", + "Table": "u_tbl7" + }, + { + "InputName": "CascadeChild-1", + "OperatorType": "FKVerify", + "BvName": "fkc_vals", + "Cols": [ + 0 + ], + "Inputs": [ + { + "InputName": "VerifyParent-1", + "OperatorType": "Route", + "Variant": "Unsharded", + "Keyspace": { + "Name": "unsharded_fk_allow", + "Sharded": false + }, + "FieldQuery": "select 1 from u_tbl4 left join u_tbl3 on u_tbl3.col3 = :v1 where 1 != 1", + "Query": "select 1 from u_tbl4 left join u_tbl3 on u_tbl3.col3 = :v1 where (u_tbl4.col4) in ::fkc_vals and u_tbl3.col3 is null limit 1 lock in share mode", + "Table": "u_tbl3, u_tbl4" + }, + { + "InputName": "VerifyChild-2", + "OperatorType": "Route", + "Variant": "Unsharded", + "Keyspace": { + "Name": "unsharded_fk_allow", + "Sharded": false + }, + "FieldQuery": "select 1 from u_tbl4, u_tbl9 where 1 != 1", + "Query": "select 1 from u_tbl4, u_tbl9 where (u_tbl4.col4) in ::fkc_vals and (:v1 is null or (u_tbl9.col9) not in ((:v1))) and u_tbl4.col4 = u_tbl9.col9 limit 1 lock in share mode", + "Table": "u_tbl4, u_tbl9" + }, + { + "InputName": "PostVerify", + "OperatorType": "Update", + "Variant": "Unsharded", + "Keyspace": { + "Name": "unsharded_fk_allow", + "Sharded": false + }, + "TargetTabletType": "PRIMARY", + "Query": "update /*+ SET_VAR(foreign_key_checks=OFF) */ u_tbl4 set col4 = :v1 where (u_tbl4.col4) in ::fkc_vals", + "Table": "u_tbl4" + } + ] + }, + { + "InputName": "Parent", + "OperatorType": "Update", + "Variant": "Unsharded", + "Keyspace": { + "Name": "unsharded_fk_allow", + "Sharded": false + }, + "TargetTabletType": "PRIMARY", + "Query": "update u_tbl7 set col7 = :v1", + "Table": "u_tbl7" + } + ] + }, + "TablesUsed": [ + "unsharded_fk_allow.u_tbl3", + "unsharded_fk_allow.u_tbl4", + "unsharded_fk_allow.u_tbl7", + "unsharded_fk_allow.u_tbl9" + ] + } + }, { "comment": "Insert with on duplicate key update - foreign keys disallowed", "query": "insert into u_tbl1 (id, col1) values (1, 3) on duplicate key update col1 = 5", @@ -1335,7 +1424,7 @@ "Sharded": false }, "FieldQuery": "select cola, colb from u_multicol_tbl2 where 1 != 1", - "Query": "select cola, colb from u_multicol_tbl2 where (cola, colb) in ::fkc_vals and (cola, colb) not in ((1, 2)) for update", + "Query": "select cola, colb from u_multicol_tbl2 where (cola, colb) in ::fkc_vals and (u_multicol_tbl2.cola, u_multicol_tbl2.colb) not in ((1, 2)) for update", "Table": "u_multicol_tbl2" }, { @@ -1364,7 +1453,7 @@ "Sharded": false }, "TargetTabletType": "PRIMARY", - "Query": "update u_multicol_tbl2 set cola = null, colb = null where (cola, colb) in ::fkc_vals and (cola, colb) not in ((1, 2))", + "Query": "update u_multicol_tbl2 set cola = null, colb = null where (cola, colb) in ::fkc_vals and (u_multicol_tbl2.cola, u_multicol_tbl2.colb) not in ((1, 2))", "Table": "u_multicol_tbl2" } ] @@ -1390,6 +1479,100 @@ ] } }, + { + "comment": "update on a multicol foreign key that set nulls and then cascades - bindVariables", + "query": "update u_multicol_tbl1 set cola = :v1, colb = :v2 where id = :v3", + "plan": { + "QueryType": "UPDATE", + "Original": "update u_multicol_tbl1 set cola = :v1, colb = :v2 where id = :v3", + "Instructions": { + "OperatorType": "FkCascade", + "Inputs": [ + { + "InputName": "Selection", + "OperatorType": "Route", + "Variant": "Unsharded", + "Keyspace": { + "Name": "unsharded_fk_allow", + "Sharded": false + }, + "FieldQuery": "select cola, colb from u_multicol_tbl1 where 1 != 1", + "Query": "select cola, colb from u_multicol_tbl1 where id = :v3 for update", + "Table": "u_multicol_tbl1" + }, + { + "InputName": "CascadeChild-1", + "OperatorType": "FkCascade", + "BvName": "fkc_vals", + "Cols": [ + 0, + 1 + ], + "Inputs": [ + { + "InputName": "Selection", + "OperatorType": "Route", + "Variant": "Unsharded", + "Keyspace": { + "Name": "unsharded_fk_allow", + "Sharded": false + }, + "FieldQuery": "select cola, colb from u_multicol_tbl2 where 1 != 1", + "Query": "select cola, colb from u_multicol_tbl2 where (cola, colb) in ::fkc_vals and (:v2 is null or (:v1 is null or (u_multicol_tbl2.cola, u_multicol_tbl2.colb) not in ((:v1, :v2)))) for update", + "Table": "u_multicol_tbl2" + }, + { + "InputName": "CascadeChild-1", + "OperatorType": "Update", + "Variant": "Unsharded", + "Keyspace": { + "Name": "unsharded_fk_allow", + "Sharded": false + }, + "TargetTabletType": "PRIMARY", + "BvName": "fkc_vals1", + "Cols": [ + 0, + 1 + ], + "Query": "update /*+ SET_VAR(foreign_key_checks=OFF) */ u_multicol_tbl3 set cola = null, colb = null where (cola, colb) in ::fkc_vals1", + "Table": "u_multicol_tbl3" + }, + { + "InputName": "Parent", + "OperatorType": "Update", + "Variant": "Unsharded", + "Keyspace": { + "Name": "unsharded_fk_allow", + "Sharded": false + }, + "TargetTabletType": "PRIMARY", + "Query": "update u_multicol_tbl2 set cola = null, colb = null where (cola, colb) in ::fkc_vals and (:v2 is null or (:v1 is null or (u_multicol_tbl2.cola, u_multicol_tbl2.colb) not in ((:v1, :v2))))", + "Table": "u_multicol_tbl2" + } + ] + }, + { + "InputName": "Parent", + "OperatorType": "Update", + "Variant": "Unsharded", + "Keyspace": { + "Name": "unsharded_fk_allow", + "Sharded": false + }, + "TargetTabletType": "PRIMARY", + "Query": "update u_multicol_tbl1 set cola = :v1, colb = :v2 where id = :v3", + "Table": "u_multicol_tbl1" + } + ] + }, + "TablesUsed": [ + "unsharded_fk_allow.u_multicol_tbl1", + "unsharded_fk_allow.u_multicol_tbl2", + "unsharded_fk_allow.u_multicol_tbl3" + ] + } + }, { "comment": "Cascaded delete run from prepared statement", "query": "execute prep_delete using @foo", diff --git a/go/vt/vtgate/semantics/semantic_state.go b/go/vt/vtgate/semantics/semantic_state.go index 6f3a4962961..61bd8c372f0 100644 --- a/go/vt/vtgate/semantics/semantic_state.go +++ b/go/vt/vtgate/semantics/semantic_state.go @@ -488,6 +488,10 @@ func (st *SemTable) SingleUnshardedKeyspace() (*vindexes.Keyspace, []*vindexes.T // The expression in the select list is not equal to the one in the ORDER BY, // but they point to the same column and would be considered equal by this method func (st *SemTable) EqualsExpr(a, b sqlparser.Expr) bool { + // If there is no SemTable, then we cannot compare the expressions. + if st == nil { + return false + } return st.ASTEquals().Expr(a, b) } From 8278f9bbdee1d53ccfb2f96dfb14e06f41f399d6 Mon Sep 17 00:00:00 2001 From: Andrew Mason Date: Wed, 27 Sep 2023 02:35:41 -0400 Subject: [PATCH 3/5] switch casing in onlineddl subcommand help text (#14091) Signed-off-by: Andrew Mason --- go/cmd/vtctldclient/command/onlineddl.go | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/go/cmd/vtctldclient/command/onlineddl.go b/go/cmd/vtctldclient/command/onlineddl.go index 660f41f60b3..dbe927de2bf 100644 --- a/go/cmd/vtctldclient/command/onlineddl.go +++ b/go/cmd/vtctldclient/command/onlineddl.go @@ -48,7 +48,7 @@ var ( } OnlineDDLCancel = &cobra.Command{ Use: "cancel ", - Short: "cancel one or all migrations, terminating any running ones as needed.", + Short: "Cancel one or all migrations, terminating any running ones as needed.", Example: "OnlineDDL cancel test_keyspace 82fa54ac_e83e_11ea_96b7_f875a4d24e90", DisableFlagsInUseLine: true, Args: cobra.ExactArgs(2), @@ -64,7 +64,7 @@ var ( } OnlineDDLComplete = &cobra.Command{ Use: "complete ", - Short: "complete one or all migrations executed with --postpone-completion", + Short: "Complete one or all migrations executed with --postpone-completion", Example: "OnlineDDL complete test_keyspace 82fa54ac_e83e_11ea_96b7_f875a4d24e90", DisableFlagsInUseLine: true, Args: cobra.ExactArgs(2), @@ -72,7 +72,7 @@ var ( } OnlineDDLLaunch = &cobra.Command{ Use: "launch ", - Short: "launch one or all migrations executed with --postpone-launch", + Short: "Launch one or all migrations executed with --postpone-launch", Example: "OnlineDDL launch test_keyspace 82fa54ac_e83e_11ea_96b7_f875a4d24e90", DisableFlagsInUseLine: true, Args: cobra.ExactArgs(2), From 6231bc47c1c8759af5e8cb43ac35b6516c52d921 Mon Sep 17 00:00:00 2001 From: Andrew Mason Date: Wed, 27 Sep 2023 03:12:36 -0400 Subject: [PATCH 4/5] [cli] cobra zookeeper (#14094) Signed-off-by: Andrew Mason --- go/cmd/zk/command/add_auth.go | 36 + go/cmd/zk/command/cat.go | 103 ++ go/cmd/zk/command/chmod.go | 91 ++ go/cmd/zk/command/cp.go | 43 + go/cmd/zk/command/edit.go | 101 ++ go/cmd/zk/command/ls.go | 153 +++ go/cmd/zk/command/rm.go | 97 ++ go/cmd/zk/command/root.go | 66 ++ go/cmd/zk/command/stat.go | 88 ++ go/cmd/zk/command/touch.go | 93 ++ go/cmd/zk/command/unzip.go | 81 ++ go/cmd/zk/command/wait.go | 78 ++ go/cmd/zk/command/watch.go | 86 ++ go/cmd/zk/command/zip.go | 116 +++ go/cmd/zk/docgen/main.go | 37 + go/cmd/zk/internal/zkfilepath/zkfilepath.go | 75 ++ go/cmd/zk/internal/zkfs/zkfs.go | 174 ++++ go/cmd/zk/zkcmd.go | 982 +------------------- go/cmd/zkctl/command/init.go | 32 + go/cmd/zkctl/command/root.go | 63 ++ go/cmd/zkctl/command/shutdown.go | 32 + go/cmd/zkctl/command/start.go | 32 + go/cmd/zkctl/command/teardown.go | 32 + go/cmd/zkctl/docgen/main.go | 37 + go/cmd/zkctl/zkctl.go | 62 +- go/cmd/zkctld/cli/zkctld.go | 100 ++ go/cmd/zkctld/docgen/main.go | 37 + go/cmd/zkctld/zkctld.go | 70 +- go/flags/endtoend/zk.txt | 37 +- go/flags/endtoend/zkctl.txt | 19 +- go/flags/endtoend/zkctld.txt | 33 +- 31 files changed, 1951 insertions(+), 1135 deletions(-) create mode 100644 go/cmd/zk/command/add_auth.go create mode 100644 go/cmd/zk/command/cat.go create mode 100644 go/cmd/zk/command/chmod.go create mode 100644 go/cmd/zk/command/cp.go create mode 100644 go/cmd/zk/command/edit.go create mode 100644 go/cmd/zk/command/ls.go create mode 100644 go/cmd/zk/command/rm.go create mode 100644 go/cmd/zk/command/root.go create mode 100644 go/cmd/zk/command/stat.go create mode 100644 go/cmd/zk/command/touch.go create mode 100644 go/cmd/zk/command/unzip.go create mode 100644 go/cmd/zk/command/wait.go create mode 100644 go/cmd/zk/command/watch.go create mode 100644 go/cmd/zk/command/zip.go create mode 100644 go/cmd/zk/docgen/main.go create mode 100644 go/cmd/zk/internal/zkfilepath/zkfilepath.go create mode 100644 go/cmd/zk/internal/zkfs/zkfs.go create mode 100644 go/cmd/zkctl/command/init.go create mode 100644 go/cmd/zkctl/command/root.go create mode 100644 go/cmd/zkctl/command/shutdown.go create mode 100644 go/cmd/zkctl/command/start.go create mode 100644 go/cmd/zkctl/command/teardown.go create mode 100644 go/cmd/zkctl/docgen/main.go create mode 100644 go/cmd/zkctld/cli/zkctld.go create mode 100644 go/cmd/zkctld/docgen/main.go diff --git a/go/cmd/zk/command/add_auth.go b/go/cmd/zk/command/add_auth.go new file mode 100644 index 00000000000..566c463f4a8 --- /dev/null +++ b/go/cmd/zk/command/add_auth.go @@ -0,0 +1,36 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreedto in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package command + +import ( + "github.com/spf13/cobra" +) + +var AddAuth = &cobra.Command{ + Use: "addAuth ", + Args: cobra.ExactArgs(2), + RunE: commandAddAuth, +} + +func commandAddAuth(cmd *cobra.Command, args []string) error { + scheme, auth := cmd.Flags().Arg(0), cmd.Flags().Arg(1) + return fs.Conn.AddAuth(cmd.Context(), scheme, []byte(auth)) +} + +func init() { + Root.AddCommand(AddAuth) +} diff --git a/go/cmd/zk/command/cat.go b/go/cmd/zk/command/cat.go new file mode 100644 index 00000000000..1d5460f7006 --- /dev/null +++ b/go/cmd/zk/command/cat.go @@ -0,0 +1,103 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreedto in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package command + +import ( + "fmt" + "os" + + "github.com/spf13/cobra" + "github.com/z-division/go-zookeeper/zk" + "golang.org/x/term" + + "vitess.io/vitess/go/cmd/zk/internal/zkfilepath" + "vitess.io/vitess/go/vt/log" + "vitess.io/vitess/go/vt/topo" + "vitess.io/vitess/go/vt/topo/zk2topo" +) + +var ( + catArgs = struct { + LongListing bool + Force bool + DecodeProto bool + }{} + + Cat = &cobra.Command{ + Use: "cat [ ...]", + Example: `zk cat /zk/path + +# List filename before file data +zk cat -l /zk/path1 /zk/path2`, + Args: cobra.MinimumNArgs(1), + RunE: commandCat, + } +) + +func commandCat(cmd *cobra.Command, args []string) error { + resolved, err := zk2topo.ResolveWildcards(cmd.Context(), fs.Conn, cmd.Flags().Args()) + if err != nil { + return fmt.Errorf("cat: invalid wildcards: %w", err) + } + if len(resolved) == 0 { + // the wildcards didn't result in anything, we're done + return nil + } + + hasError := false + for _, arg := range resolved { + zkPath := zkfilepath.Clean(arg) + data, _, err := fs.Conn.Get(cmd.Context(), zkPath) + if err != nil { + hasError = true + if !catArgs.Force || err != zk.ErrNoNode { + log.Warningf("cat: cannot access %v: %v", zkPath, err) + } + continue + } + + if catArgs.LongListing { + fmt.Printf("%v:\n", zkPath) + } + decoded := "" + if catArgs.DecodeProto { + decoded, err = topo.DecodeContent(zkPath, data, false) + if err != nil { + log.Warningf("cat: cannot proto decode %v: %v", zkPath, err) + decoded = string(data) + } + } else { + decoded = string(data) + } + fmt.Print(decoded) + if len(decoded) > 0 && decoded[len(decoded)-1] != '\n' && (term.IsTerminal(int(os.Stdout.Fd())) || catArgs.LongListing) { + fmt.Print("\n") + } + } + if hasError { + return fmt.Errorf("cat: some paths had errors") + } + return nil +} + +func init() { + Cat.Flags().BoolVarP(&catArgs.LongListing, "longListing", "l", false, "long listing") + Cat.Flags().BoolVarP(&catArgs.Force, "force", "f", false, "no warning on nonexistent node") + Cat.Flags().BoolVarP(&catArgs.DecodeProto, "decodeProto", "p", false, "decode proto files and display them as text") + + Root.AddCommand(Cat) +} diff --git a/go/cmd/zk/command/chmod.go b/go/cmd/zk/command/chmod.go new file mode 100644 index 00000000000..39125d618c4 --- /dev/null +++ b/go/cmd/zk/command/chmod.go @@ -0,0 +1,91 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreedto in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package command + +import ( + "fmt" + + "github.com/spf13/cobra" + + "vitess.io/vitess/go/cmd/zk/internal/zkfilepath" + "vitess.io/vitess/go/cmd/zk/internal/zkfs" + "vitess.io/vitess/go/vt/log" + "vitess.io/vitess/go/vt/topo/zk2topo" +) + +var Chmod = &cobra.Command{ + Use: "chmod ", + Example: `zk chmod n-mode /zk/path +zk chmod n+mode /zk/path`, + Args: cobra.MinimumNArgs(2), + RunE: commandChmod, +} + +func commandChmod(cmd *cobra.Command, args []string) error { + mode := cmd.Flags().Arg(0) + if mode[0] != 'n' { + return fmt.Errorf("chmod: invalid mode") + } + + addPerms := false + if mode[1] == '+' { + addPerms = true + } else if mode[1] != '-' { + return fmt.Errorf("chmod: invalid mode") + } + + permMask := zkfs.ParsePermMode(mode[2:]) + + resolved, err := zk2topo.ResolveWildcards(cmd.Context(), fs.Conn, cmd.Flags().Args()[1:]) + if err != nil { + return fmt.Errorf("chmod: invalid wildcards: %w", err) + } + if len(resolved) == 0 { + // the wildcards didn't result in anything, we're done + return nil + } + + hasError := false + for _, arg := range resolved { + zkPath := zkfilepath.Clean(arg) + aclv, _, err := fs.Conn.GetACL(cmd.Context(), zkPath) + if err != nil { + hasError = true + log.Warningf("chmod: cannot set access %v: %v", zkPath, err) + continue + } + if addPerms { + aclv[0].Perms |= permMask + } else { + aclv[0].Perms &= ^permMask + } + err = fs.Conn.SetACL(cmd.Context(), zkPath, aclv, -1) + if err != nil { + hasError = true + log.Warningf("chmod: cannot set access %v: %v", zkPath, err) + continue + } + } + if hasError { + return fmt.Errorf("chmod: some paths had errors") + } + return nil +} + +func init() { + Root.AddCommand(Chmod) +} diff --git a/go/cmd/zk/command/cp.go b/go/cmd/zk/command/cp.go new file mode 100644 index 00000000000..e89486413ea --- /dev/null +++ b/go/cmd/zk/command/cp.go @@ -0,0 +1,43 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreedto in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package command + +import "github.com/spf13/cobra" + +var Cp = &cobra.Command{ + Use: "cp ", + Example: `zk cp /zk/path . +zk cp ./config /zk/path/config + +# Trailing slash indicates directory +zk cp ./config /zk/path/`, + Args: cobra.MinimumNArgs(2), + RunE: commandCp, +} + +func commandCp(cmd *cobra.Command, args []string) error { + switch cmd.Flags().NArg() { + case 2: + return fs.CopyContext(cmd.Context(), cmd.Flags().Arg(0), cmd.Flags().Arg(1)) + default: + return fs.MultiCopyContext(cmd.Context(), cmd.Flags().Args()) + } +} + +func init() { + Root.AddCommand(Cp) +} diff --git a/go/cmd/zk/command/edit.go b/go/cmd/zk/command/edit.go new file mode 100644 index 00000000000..ec4b74c4b62 --- /dev/null +++ b/go/cmd/zk/command/edit.go @@ -0,0 +1,101 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreedto in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package command + +import ( + "bytes" + "fmt" + "os" + "os/exec" + "path" + "time" + + "github.com/spf13/cobra" + "github.com/z-division/go-zookeeper/zk" + + "vitess.io/vitess/go/cmd/zk/internal/zkfilepath" + "vitess.io/vitess/go/vt/log" +) + +var ( + editArgs = struct { + Force bool + }{} + + Edit = &cobra.Command{ + Use: "edit ", + Short: "Create a local copy, edit, and write changes back to cell.", + Args: cobra.ExactArgs(1), + RunE: commandEdit, + } +) + +func commandEdit(cmd *cobra.Command, args []string) error { + arg := cmd.Flags().Arg(0) + zkPath := zkfilepath.Clean(arg) + data, stat, err := fs.Conn.Get(cmd.Context(), zkPath) + if err != nil { + if !editArgs.Force || err != zk.ErrNoNode { + log.Warningf("edit: cannot access %v: %v", zkPath, err) + } + return fmt.Errorf("edit: cannot access %v: %v", zkPath, err) + } + + name := path.Base(zkPath) + tmpPath := fmt.Sprintf("/tmp/zk-edit-%v-%v", name, time.Now().UnixNano()) + f, err := os.Create(tmpPath) + if err == nil { + _, err = f.Write(data) + f.Close() + } + if err != nil { + return fmt.Errorf("edit: cannot write file %v", err) + } + + editor := exec.Command(os.Getenv("EDITOR"), tmpPath) + editor.Stdin = os.Stdin + editor.Stdout = os.Stdout + editor.Stderr = os.Stderr + err = editor.Run() + if err != nil { + os.Remove(tmpPath) + return fmt.Errorf("edit: cannot start $EDITOR: %v", err) + } + + fileData, err := os.ReadFile(tmpPath) + if err != nil { + os.Remove(tmpPath) + return fmt.Errorf("edit: cannot read file %v", err) + } + + if !bytes.Equal(fileData, data) { + // data changed - update if we can + _, err = fs.Conn.Set(cmd.Context(), zkPath, fileData, stat.Version) + if err != nil { + os.Remove(tmpPath) + return fmt.Errorf("edit: cannot write zk file %v", err) + } + } + os.Remove(tmpPath) + return nil +} + +func init() { + Edit.Flags().BoolVarP(&editArgs.Force, "force", "f", false, "no warning on nonexistent node") + + Root.AddCommand(Edit) +} diff --git a/go/cmd/zk/command/ls.go b/go/cmd/zk/command/ls.go new file mode 100644 index 00000000000..83c1d31363b --- /dev/null +++ b/go/cmd/zk/command/ls.go @@ -0,0 +1,153 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreedto in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package command + +import ( + "fmt" + "path" + "sort" + "sync" + + "github.com/spf13/cobra" + "github.com/z-division/go-zookeeper/zk" + + "vitess.io/vitess/go/cmd/zk/internal/zkfilepath" + "vitess.io/vitess/go/vt/log" + "vitess.io/vitess/go/vt/topo/zk2topo" +) + +var ( + lsArgs = struct { + LongListing bool + DirectoryListing bool + Force bool + RecursiveListing bool + }{} + + Ls = &cobra.Command{ + Use: "ls ", + Example: `zk ls /zk +zk ls -l /zk + +# List directory node itself) +zk ls -ld /zk + +# Recursive (expensive) +zk ls -R /zk`, + Args: cobra.MinimumNArgs(1), + RunE: commandLs, + } +) + +func commandLs(cmd *cobra.Command, args []string) error { + resolved, err := zk2topo.ResolveWildcards(cmd.Context(), fs.Conn, cmd.Flags().Args()) + if err != nil { + return fmt.Errorf("ls: invalid wildcards: %v", err) + } + if len(resolved) == 0 { + // the wildcards didn't result in anything, we're + // done. + return nil + } + + hasError := false + needsHeader := len(resolved) > 1 && !lsArgs.DirectoryListing + for _, arg := range resolved { + zkPath := zkfilepath.Clean(arg) + var children []string + var err error + isDir := true + if lsArgs.DirectoryListing { + children = []string{""} + isDir = false + } else if lsArgs.RecursiveListing { + children, err = zk2topo.ChildrenRecursive(cmd.Context(), fs.Conn, zkPath) + } else { + children, _, err = fs.Conn.Children(cmd.Context(), zkPath) + // Assume this is a file node if it has no children. + if len(children) == 0 { + children = []string{""} + isDir = false + } + } + if err != nil { + hasError = true + if !lsArgs.Force || err != zk.ErrNoNode { + log.Warningf("ls: cannot access %v: %v", zkPath, err) + } + } + + // Show the full path when it helps. + showFullPath := false + if lsArgs.RecursiveListing { + showFullPath = true + } else if lsArgs.LongListing && (lsArgs.DirectoryListing || !isDir) { + showFullPath = true + } + if needsHeader { + fmt.Printf("%v:\n", zkPath) + } + if len(children) > 0 { + if lsArgs.LongListing && isDir { + fmt.Printf("total: %v\n", len(children)) + } + sort.Strings(children) + stats := make([]*zk.Stat, len(children)) + wg := sync.WaitGroup{} + f := func(i int) { + localPath := path.Join(zkPath, children[i]) + _, stat, err := fs.Conn.Exists(cmd.Context(), localPath) + if err != nil { + if !lsArgs.Force || err != zk.ErrNoNode { + log.Warningf("ls: cannot access: %v: %v", localPath, err) + } + } else { + stats[i] = stat + } + wg.Done() + } + for i := range children { + wg.Add(1) + go f(i) + } + wg.Wait() + + for i, child := range children { + localPath := path.Join(zkPath, child) + if stat := stats[i]; stat != nil { + fmt.Println(zkfilepath.Format(stat, localPath, showFullPath, lsArgs.LongListing)) + } + } + } + if needsHeader { + fmt.Println() + } + } + if hasError { + return fmt.Errorf("ls: some paths had errors") + } + return nil +} + +func init() { + Ls.Flags().BoolVarP(&lsArgs.LongListing, "longlisting", "l", false, "long listing") + Ls.Flags().BoolVarP(&lsArgs.DirectoryListing, "directorylisting", "d", false, "list directory instead of contents") + Ls.Flags().BoolVarP(&lsArgs.Force, "force", "f", false, "no warning on nonexistent node") + Ls.Flags().BoolVarP(&lsArgs.RecursiveListing, "recursivelisting", "R", false, "recursive listing") + + Root.AddCommand(Ls) +} diff --git a/go/cmd/zk/command/rm.go b/go/cmd/zk/command/rm.go new file mode 100644 index 00000000000..5e5b5f4c494 --- /dev/null +++ b/go/cmd/zk/command/rm.go @@ -0,0 +1,97 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreedto in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package command + +import ( + "fmt" + "strings" + + "github.com/spf13/cobra" + "github.com/z-division/go-zookeeper/zk" + + "vitess.io/vitess/go/cmd/zk/internal/zkfilepath" + "vitess.io/vitess/go/vt/log" + "vitess.io/vitess/go/vt/topo/zk2topo" +) + +var ( + rmArgs = struct { + Force bool + RecursiveDelete bool + }{} + + Rm = &cobra.Command{ + Use: "rm ", + Example: `zk rm /zk/path + +# Recursive. +zk rm -R /zk/path + +# No error on nonexistent node. +zk rm -f /zk/path`, + Args: cobra.MinimumNArgs(1), + RunE: commandRm, + } +) + +func commandRm(cmd *cobra.Command, args []string) error { + if rmArgs.RecursiveDelete { + for _, arg := range cmd.Flags().Args() { + zkPath := zkfilepath.Clean(arg) + if strings.Count(zkPath, "/") < 2 { + return fmt.Errorf("rm: overly general path: %v", zkPath) + } + } + } + + resolved, err := zk2topo.ResolveWildcards(cmd.Context(), fs.Conn, cmd.Flags().Args()) + if err != nil { + return fmt.Errorf("rm: invalid wildcards: %v", err) + } + if len(resolved) == 0 { + // the wildcards didn't result in anything, we're done + return nil + } + + hasError := false + for _, arg := range resolved { + zkPath := zkfilepath.Clean(arg) + var err error + if rmArgs.RecursiveDelete { + err = zk2topo.DeleteRecursive(cmd.Context(), fs.Conn, zkPath, -1) + } else { + err = fs.Conn.Delete(cmd.Context(), zkPath, -1) + } + if err != nil && (!rmArgs.Force || err != zk.ErrNoNode) { + hasError = true + log.Warningf("rm: cannot delete %v: %v", zkPath, err) + } + } + if hasError { + // to be consistent with the command line 'rm -f', return + // 0 if using 'zk rm -f' and the file doesn't exist. + return fmt.Errorf("rm: some paths had errors") + } + return nil +} + +func init() { + Rm.Flags().BoolVarP(&rmArgs.Force, "force", "f", false, "no warning on nonexistent node") + Rm.Flags().BoolVarP(&rmArgs.RecursiveDelete, "recursivedelete", "r", false, "recursive delete") + + Root.AddCommand(Rm) +} diff --git a/go/cmd/zk/command/root.go b/go/cmd/zk/command/root.go new file mode 100644 index 00000000000..f3f02e7d4f2 --- /dev/null +++ b/go/cmd/zk/command/root.go @@ -0,0 +1,66 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreedto in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package command + +import ( + "github.com/spf13/cobra" + + "vitess.io/vitess/go/acl" + "vitess.io/vitess/go/cmd/zk/internal/zkfs" + "vitess.io/vitess/go/vt/log" + "vitess.io/vitess/go/vt/logutil" + "vitess.io/vitess/go/vt/topo/zk2topo" +) + +var ( + fs *zkfs.FS + server string + + Root = &cobra.Command{ + Use: "zk", + Short: "zk is a tool for wrangling zookeeper.", + Long: `zk is a tool for wrangling zookeeper. + +It tries to mimic unix file system commands wherever possible, but +there are some slight differences in flag handling. + +The zk tool looks for the address of the cluster in /etc/zookeeper/zk_client.conf, +or the file specified in the ZK_CLIENT_CONFIG environment variable. + +The local cell may be overridden with the ZK_CLIENT_LOCAL_CELL environment +variable.`, + PersistentPreRun: func(cmd *cobra.Command, args []string) { + logutil.PurgeLogs() + + // Connect to the server. + fs = &zkfs.FS{ + Conn: zk2topo.Connect(server), + } + }, + PersistentPostRun: func(cmd *cobra.Command, args []string) { + logutil.Flush() + }, + } +) + +func init() { + Root.Flags().StringVar(&server, "server", server, "server(s) to connect to") + + log.RegisterFlags(Root.Flags()) + logutil.RegisterFlags(Root.Flags()) + acl.RegisterFlags(Root.Flags()) +} diff --git a/go/cmd/zk/command/stat.go b/go/cmd/zk/command/stat.go new file mode 100644 index 00000000000..713a68a3d4e --- /dev/null +++ b/go/cmd/zk/command/stat.go @@ -0,0 +1,88 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreedto in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package command + +import ( + "fmt" + + "github.com/spf13/cobra" + "github.com/z-division/go-zookeeper/zk" + + "vitess.io/vitess/go/cmd/zk/internal/zkfilepath" + "vitess.io/vitess/go/cmd/zk/internal/zkfs" + "vitess.io/vitess/go/vt/log" + "vitess.io/vitess/go/vt/topo/zk2topo" +) + +var ( + statArgs = struct { + Force bool + }{} + Stat = &cobra.Command{ + Use: "stat ", + Args: cobra.MinimumNArgs(1), + RunE: commandStat, + } +) + +func commandStat(cmd *cobra.Command, args []string) error { + resolved, err := zk2topo.ResolveWildcards(cmd.Context(), fs.Conn, cmd.Flags().Args()) + if err != nil { + return fmt.Errorf("stat: invalid wildcards: %v", err) + } + if len(resolved) == 0 { + // the wildcards didn't result in anything, we're done + return nil + } + + hasError := false + for _, arg := range resolved { + zkPath := zkfilepath.Clean(arg) + acls, stat, err := fs.Conn.GetACL(cmd.Context(), zkPath) + if stat == nil { + err = fmt.Errorf("no such node") + } + if err != nil { + hasError = true + if !statArgs.Force || err != zk.ErrNoNode { + log.Warningf("stat: cannot access %v: %v", zkPath, err) + } + continue + } + fmt.Printf("Path: %s\n", zkPath) + fmt.Printf("Created: %s\n", zk2topo.Time(stat.Ctime).Format(zkfilepath.TimeFmtMicro)) + fmt.Printf("Modified: %s\n", zk2topo.Time(stat.Mtime).Format(zkfilepath.TimeFmtMicro)) + fmt.Printf("Size: %v\n", stat.DataLength) + fmt.Printf("Children: %v\n", stat.NumChildren) + fmt.Printf("Version: %v\n", stat.Version) + fmt.Printf("Ephemeral: %v\n", stat.EphemeralOwner) + fmt.Printf("ACL:\n") + for _, acl := range acls { + fmt.Printf(" %v:%v %v\n", acl.Scheme, acl.ID, zkfs.FormatACL(acl)) + } + } + if hasError { + return fmt.Errorf("stat: some paths had errors") + } + return nil +} + +func init() { + Stat.Flags().BoolVarP(&statArgs.Force, "force", "f", false, "no warning on nonexistent node") + + Root.AddCommand(Stat) +} diff --git a/go/cmd/zk/command/touch.go b/go/cmd/zk/command/touch.go new file mode 100644 index 00000000000..76c390cf169 --- /dev/null +++ b/go/cmd/zk/command/touch.go @@ -0,0 +1,93 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreedto in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package command + +import ( + "fmt" + + "github.com/spf13/cobra" + "github.com/z-division/go-zookeeper/zk" + + "vitess.io/vitess/go/cmd/zk/internal/zkfilepath" + "vitess.io/vitess/go/vt/topo/zk2topo" +) + +var ( + touchArgs = struct { + CreateParents bool + TouchOnly bool + }{} + + Touch = &cobra.Command{ + Use: "touch ", + Short: "Change node access time.", + Long: `Change node access time. + +NOTE: There is no mkdir - just touch a node. +The disntinction between file and directory is not relevant in zookeeper.`, + Example: `zk touch /zk/path + +# Don't create, just touch timestamp. +zk touch -c /zk/path + +# Create all parts necessary (think mkdir -p). +zk touch -p /zk/path`, + Args: cobra.ExactArgs(1), + RunE: commandTouch, + } +) + +func commandTouch(cmd *cobra.Command, args []string) error { + zkPath := zkfilepath.Clean(cmd.Flags().Arg(0)) + var ( + version int32 = -1 + create = false + ) + + data, stat, err := fs.Conn.Get(cmd.Context(), zkPath) + switch { + case err == nil: + version = stat.Version + case err == zk.ErrNoNode: + create = true + default: + return fmt.Errorf("touch: cannot access %v: %v", zkPath, err) + } + + switch { + case !create: + _, err = fs.Conn.Set(cmd.Context(), zkPath, data, version) + case touchArgs.TouchOnly: + return fmt.Errorf("touch: no such path %v", zkPath) + case touchArgs.CreateParents: + _, err = zk2topo.CreateRecursive(cmd.Context(), fs.Conn, zkPath, data, 0, zk.WorldACL(zk.PermAll), 10) + default: + _, err = fs.Conn.Create(cmd.Context(), zkPath, data, 0, zk.WorldACL(zk.PermAll)) + } + + if err != nil { + return fmt.Errorf("touch: cannot modify %v: %v", zkPath, err) + } + return nil +} + +func init() { + Touch.Flags().BoolVarP(&touchArgs.CreateParents, "createparent", "p", false, "create parents") + Touch.Flags().BoolVarP(&touchArgs.TouchOnly, "touchonly", "c", false, "touch only - don't create") + + Root.AddCommand(Touch) +} diff --git a/go/cmd/zk/command/unzip.go b/go/cmd/zk/command/unzip.go new file mode 100644 index 00000000000..f4c800e0533 --- /dev/null +++ b/go/cmd/zk/command/unzip.go @@ -0,0 +1,81 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreedto in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package command + +import ( + "archive/zip" + "fmt" + "io" + "path" + "strings" + + "github.com/spf13/cobra" + "github.com/z-division/go-zookeeper/zk" + + "vitess.io/vitess/go/vt/topo/zk2topo" +) + +var Unzip = &cobra.Command{ + Use: "unzip ", + Example: `zk unzip zktree.zip / +zk unzip zktree.zip /zk/prefix`, + Args: cobra.ExactArgs(1), + RunE: commandUnzip, +} + +func commandUnzip(cmd *cobra.Command, args []string) error { + srcPath, dstPath := cmd.Flags().Arg(0), cmd.Flags().Arg(1) + + if !strings.HasSuffix(srcPath, ".zip") { + return fmt.Errorf("zip: need to specify src .zip path: %v", srcPath) + } + + zipReader, err := zip.OpenReader(srcPath) + if err != nil { + return fmt.Errorf("zip: error %v", err) + } + defer zipReader.Close() + + for _, zf := range zipReader.File { + rc, err := zf.Open() + if err != nil { + return fmt.Errorf("unzip: error %v", err) + } + data, err := io.ReadAll(rc) + if err != nil { + return fmt.Errorf("unzip: failed reading archive: %v", err) + } + zkPath := zf.Name + if dstPath != "/" { + zkPath = path.Join(dstPath, zkPath) + } + _, err = zk2topo.CreateRecursive(cmd.Context(), fs.Conn, zkPath, data, 0, zk.WorldACL(zk.PermAll), 10) + if err != nil && err != zk.ErrNodeExists { + return fmt.Errorf("unzip: zk create failed: %v", err) + } + _, err = fs.Conn.Set(cmd.Context(), zkPath, data, -1) + if err != nil { + return fmt.Errorf("unzip: zk set failed: %v", err) + } + rc.Close() + } + return nil +} + +func init() { + Root.AddCommand(Unzip) +} diff --git a/go/cmd/zk/command/wait.go b/go/cmd/zk/command/wait.go new file mode 100644 index 00000000000..864f6e83626 --- /dev/null +++ b/go/cmd/zk/command/wait.go @@ -0,0 +1,78 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreedto in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package command + +import ( + "fmt" + + "github.com/spf13/cobra" + "github.com/z-division/go-zookeeper/zk" + + "vitess.io/vitess/go/cmd/zk/internal/zkfilepath" +) + +var ( + waitArgs = struct { + ExitIfExists bool + }{} + + Wait = &cobra.Command{ + Use: "wait ", + Short: "Sets a watch on the node and then waits for an event to fire.", + Example: ` # Wait for node change or creation. +zk wait /zk/path + +# Trailing slash waits on children. +zk wait /zk/path/children/`, + Args: cobra.ExactArgs(1), + RunE: commandWait, + } +) + +func commandWait(cmd *cobra.Command, args []string) error { + zkPath := cmd.Flags().Arg(0) + isDir := zkPath[len(zkPath)-1] == '/' + zkPath = zkfilepath.Clean(zkPath) + + var wait <-chan zk.Event + var err error + if isDir { + _, _, wait, err = fs.Conn.ChildrenW(cmd.Context(), zkPath) + } else { + _, _, wait, err = fs.Conn.GetW(cmd.Context(), zkPath) + } + if err != nil { + if err == zk.ErrNoNode { + _, _, wait, _ = fs.Conn.ExistsW(cmd.Context(), zkPath) + } else { + return fmt.Errorf("wait: error %v: %v", zkPath, err) + } + } else { + if waitArgs.ExitIfExists { + return fmt.Errorf("already exists: %v", zkPath) + } + } + event := <-wait + fmt.Printf("event: %v\n", event) + return nil +} + +func init() { + Wait.Flags().BoolVarP(&waitArgs.ExitIfExists, "exit", "e", false, "exit if the path already exists") + + Root.AddCommand(Wait) +} diff --git a/go/cmd/zk/command/watch.go b/go/cmd/zk/command/watch.go new file mode 100644 index 00000000000..eb28cc29ca2 --- /dev/null +++ b/go/cmd/zk/command/watch.go @@ -0,0 +1,86 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreedto in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package command + +import ( + "fmt" + + "github.com/spf13/cobra" + "github.com/z-division/go-zookeeper/zk" + + "vitess.io/vitess/go/cmd/zk/internal/zkfilepath" + "vitess.io/vitess/go/vt/log" +) + +var Watch = &cobra.Command{ + Use: "watch ", + Short: "Watches for changes to nodes and prints events as they occur.", + Example: `watch /zk/path`, + Args: cobra.MinimumNArgs(1), + RunE: commandWatch, +} + +func commandWatch(cmd *cobra.Command, args []string) error { + eventChan := make(chan zk.Event, 16) + for _, arg := range cmd.Flags().Args() { + zkPath := zkfilepath.Clean(arg) + _, _, watch, err := fs.Conn.GetW(cmd.Context(), zkPath) + if err != nil { + return fmt.Errorf("watch error: %v", err) + } + go func() { + eventChan <- <-watch + }() + } + + for { + select { + case <-cmd.Context().Done(): + return nil + case event := <-eventChan: + log.Infof("watch: event %v: %v", event.Path, event) + if event.Type == zk.EventNodeDataChanged { + data, stat, watch, err := fs.Conn.GetW(cmd.Context(), event.Path) + if err != nil { + return fmt.Errorf("ERROR: failed to watch %v", err) + } + log.Infof("watch: %v %v\n", event.Path, stat) + println(data) + go func() { + eventChan <- <-watch + }() + } else if event.State == zk.StateDisconnected { + return nil + } else if event.Type == zk.EventNodeDeleted { + log.Infof("watch: %v deleted\n", event.Path) + } else { + // Most likely a session event - try t + _, _, watch, err := fs.Conn.GetW(cmd.Context(), event.Path) + if err != nil { + return fmt.Errorf("ERROR: failed to watch %v", err) + } + go func() { + eventChan <- <-watch + }() + } + } + } +} + +func init() { + Root.AddCommand(Watch) +} diff --git a/go/cmd/zk/command/zip.go b/go/cmd/zk/command/zip.go new file mode 100644 index 00000000000..b765f5bb00e --- /dev/null +++ b/go/cmd/zk/command/zip.go @@ -0,0 +1,116 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreedto in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package command + +import ( + "archive/zip" + "fmt" + "os" + "path" + "strings" + "sync" + + "github.com/spf13/cobra" + + "vitess.io/vitess/go/cmd/zk/internal/zkfilepath" + "vitess.io/vitess/go/cmd/zk/internal/zkfs" + "vitess.io/vitess/go/vt/topo/zk2topo" +) + +var Zip = &cobra.Command{ + Use: "zip [ ...] ", + Short: "Store a zk tree in a zip archive.", + Long: `Store a zk tree in a zip archive. + +Note this won't be immediately useful to the local filesystem since znodes can have data and children; +that is, even "directories" can contain data.`, + Args: cobra.MinimumNArgs(2), + RunE: commandZip, +} + +func commandZip(cmd *cobra.Command, args []string) error { + posargs := cmd.Flags().Args() + dstPath := posargs[len(posargs)-1] + paths := posargs[:len(posargs)-1] + if !strings.HasSuffix(dstPath, ".zip") { + return fmt.Errorf("zip: need to specify destination .zip path: %v", dstPath) + } + zipFile, err := os.Create(dstPath) + if err != nil { + return fmt.Errorf("zip: error %v", err) + } + + wg := sync.WaitGroup{} + items := make(chan *zkfs.Item, 64) + for _, arg := range paths { + zkPath := zkfilepath.Clean(arg) + children, err := zk2topo.ChildrenRecursive(cmd.Context(), fs.Conn, zkPath) + if err != nil { + return fmt.Errorf("zip: error %v", err) + } + for _, child := range children { + toAdd := path.Join(zkPath, child) + wg.Add(1) + go func() { + data, stat, err := fs.Conn.Get(cmd.Context(), toAdd) + items <- &zkfs.Item{ + Path: toAdd, + Data: data, + Stat: stat, + Err: err, + } + wg.Done() + }() + } + } + go func() { + wg.Wait() + close(items) + }() + + zipWriter := zip.NewWriter(zipFile) + for item := range items { + path, data, stat, err := item.Path, item.Data, item.Stat, item.Err + if err != nil { + return fmt.Errorf("zip: get failed: %v", err) + } + // Skip ephemerals - not sure why you would archive them. + if stat.EphemeralOwner > 0 { + continue + } + fi := &zip.FileHeader{Name: path, Method: zip.Deflate} + fi.Modified = zk2topo.Time(stat.Mtime) + f, err := zipWriter.CreateHeader(fi) + if err != nil { + return fmt.Errorf("zip: create failed: %v", err) + } + _, err = f.Write(data) + if err != nil { + return fmt.Errorf("zip: create failed: %v", err) + } + } + err = zipWriter.Close() + if err != nil { + return fmt.Errorf("zip: close failed: %v", err) + } + zipFile.Close() + return nil +} + +func init() { + Root.AddCommand(Zip) +} diff --git a/go/cmd/zk/docgen/main.go b/go/cmd/zk/docgen/main.go new file mode 100644 index 00000000000..b8a7bde3d14 --- /dev/null +++ b/go/cmd/zk/docgen/main.go @@ -0,0 +1,37 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package main + +import ( + "github.com/spf13/cobra" + + "vitess.io/vitess/go/cmd/internal/docgen" + "vitess.io/vitess/go/cmd/zk/command" +) + +func main() { + var dir string + cmd := cobra.Command{ + Use: "docgen [-d ]", + RunE: func(cmd *cobra.Command, args []string) error { + return docgen.GenerateMarkdownTree(command.Root, dir) + }, + } + + cmd.Flags().StringVarP(&dir, "dir", "d", "doc", "output directory to write documentation") + _ = cmd.Execute() +} diff --git a/go/cmd/zk/internal/zkfilepath/zkfilepath.go b/go/cmd/zk/internal/zkfilepath/zkfilepath.go new file mode 100644 index 00000000000..7febc7a9677 --- /dev/null +++ b/go/cmd/zk/internal/zkfilepath/zkfilepath.go @@ -0,0 +1,75 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreedto in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +// Package zkfilepath provides filepath utilities specialized to zookeeper. +package zkfilepath + +import ( + "fmt" + "path" + "strings" + + "github.com/z-division/go-zookeeper/zk" + + "vitess.io/vitess/go/vt/topo/zk2topo" +) + +const ( + TimeFmt = "2006-01-02 15:04:05" + TimeFmtMicro = "2006-01-02 15:04:05.000000" +) + +// Clean returns the shortest path name of a zookeeper path after trimming +// trailing slashes. +func Clean(zkPath string) string { + if zkPath != "/" { + zkPath = strings.TrimSuffix(zkPath, "/") + } + + return path.Clean(zkPath) +} + +// Format returns a path formatted to a canonical string. +func Format(stat *zk.Stat, zkPath string, showFullPath bool, longListing bool) string { + var name, perms string + + if !showFullPath { + name = path.Base(zkPath) + } else { + name = zkPath + } + + if longListing { + if stat.NumChildren > 0 { + // FIXME(msolomon) do permissions check? + perms = "drwxrwxrwx" + if stat.DataLength > 0 { + // give a visual indication that this node has data as well as children + perms = "nrw-rw-rw-" + } + } else if stat.EphemeralOwner != 0 { + perms = "erw-rw-rw-" + } else { + perms = "-rw-rw-rw-" + } + // always print the Local version of the time. zookeeper's + // go / C library would return a local time anyway, but + // might as well be sure. + return fmt.Sprintf("%v %v %v % 8v % 20v %v\n", perms, "zk", "zk", stat.DataLength, zk2topo.Time(stat.Mtime).Local().Format(TimeFmt), name) + } else { + return fmt.Sprintf("%v\n", name) + } +} diff --git a/go/cmd/zk/internal/zkfs/zkfs.go b/go/cmd/zk/internal/zkfs/zkfs.go new file mode 100644 index 00000000000..9bab19ec1e4 --- /dev/null +++ b/go/cmd/zk/internal/zkfs/zkfs.go @@ -0,0 +1,174 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreedto in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +// Package zkfs provides utilities for working with zookeepr in a filesystem-like manner. +package zkfs + +import ( + "context" + "fmt" + "io" + "os" + "path" + "strings" + "syscall" + + "github.com/z-division/go-zookeeper/zk" + + "vitess.io/vitess/go/cmd/zk/internal/zkfilepath" + "vitess.io/vitess/go/vt/topo/zk2topo" +) + +// FS wraps a zk2topo connection to provide FS utility methods. +type FS struct { + Conn *zk2topo.ZkConn +} + +// CopyContext copies the contents of src to dst. +func (fs *FS) CopyContext(ctx context.Context, src, dst string) error { + dstIsDir := dst[len(dst)-1] == '/' + src = zkfilepath.Clean(src) + dst = zkfilepath.Clean(dst) + + if !IsFile(src) && !IsFile(dst) { + return fmt.Errorf("cp: neither src nor dst is a /zk file") + } + + data, err := fs.ReadContext(ctx, src) + if err != nil { + return fmt.Errorf("cp: cannot read %v: %v", src, err) + } + + // If we are copying to a local directory - say '.', make the filename + // the same as the source. + if !IsFile(dst) { + fileInfo, err := os.Stat(dst) + if err != nil { + if err.(*os.PathError).Err != syscall.ENOENT { + return fmt.Errorf("cp: cannot stat %v: %v", dst, err) + } + } else if fileInfo.IsDir() { + dst = path.Join(dst, path.Base(src)) + } + } else if dstIsDir { + // If we are copying into zk, interpret trailing slash as treating the + // dst as a directory. + dst = path.Join(dst, path.Base(src)) + } + if err := fs.WriteContext(ctx, dst, data); err != nil { + return fmt.Errorf("cp: cannot write %v: %v", dst, err) + } + return nil +} + +// MultiCopyContext copies the contents of multiple sources to a single dst directory. +func (fs *FS) MultiCopyContext(ctx context.Context, args []string) error { + dstPath := args[len(args)-1] + if dstPath[len(dstPath)-1] != '/' { + // In multifile context, dstPath must be a directory. + dstPath += "/" + } + + for _, srcPath := range args[:len(args)-1] { + if err := fs.CopyContext(ctx, srcPath, dstPath); err != nil { + return err + } + } + return nil +} + +// ReadContext reads the data stored at path. +func (fs *FS) ReadContext(ctx context.Context, path string) (data []byte, err error) { + if !IsFile(path) { + data, _, err = fs.Conn.Get(ctx, path) + return data, err + } + + file, err := os.Open(path) + if err != nil { + return nil, err + } + + data, err = io.ReadAll(file) + return data, err +} + +// WriteContext writes the given data to path. +func (fs *FS) WriteContext(ctx context.Context, path string, data []byte) (err error) { + if IsFile(path) { + _, err = fs.Conn.Set(ctx, path, data, -1) + if err == zk.ErrNoNode { + _, err = zk2topo.CreateRecursive(ctx, fs.Conn, path, data, 0, zk.WorldACL(zk.PermAll), 10) + } + return err + } + return os.WriteFile(path, []byte(data), 0666) +} + +var ( + charPermMap map[string]int32 + permCharMap map[int32]string +) + +func init() { + charPermMap = map[string]int32{ + "r": zk.PermRead, + "w": zk.PermWrite, + "d": zk.PermDelete, + "c": zk.PermCreate, + "a": zk.PermAdmin, + } + permCharMap = make(map[int32]string) + for c, p := range charPermMap { + permCharMap[p] = c + } +} + +// FormatACL returns a string representation of a zookeeper ACL permission. +func FormatACL(acl zk.ACL) string { + s := "" + + for _, perm := range []int32{zk.PermRead, zk.PermWrite, zk.PermDelete, zk.PermCreate, zk.PermAdmin} { + if acl.Perms&perm != 0 { + s += permCharMap[perm] + } else { + s += "-" + } + } + return s +} + +// IsFile returns true if the path is a zk type of file. +func IsFile(path string) bool { + return strings.HasPrefix(path, "/zk") +} + +// ParsePermMode parses the mode string as a perm mask. +func ParsePermMode(mode string) (mask int32) { + for _, c := range mode[2:] { + mask |= charPermMap[string(c)] + } + + return mask +} + +// Item represents an item in a zookeeper filesystem. +type Item struct { + Path string + Data []byte + Stat *zk.Stat + Err error +} diff --git a/go/cmd/zk/zkcmd.go b/go/cmd/zk/zkcmd.go index 8d456f6b081..f03ac41c6ef 100644 --- a/go/cmd/zk/zkcmd.go +++ b/go/cmd/zk/zkcmd.go @@ -17,156 +17,17 @@ limitations under the License. package main import ( - "archive/zip" - "bytes" "context" - "fmt" - "io" "os" - "os/exec" "os/signal" - "path" - "sort" - "strings" - "sync" - "syscall" - "time" - "github.com/spf13/pflag" - "github.com/z-division/go-zookeeper/zk" - "golang.org/x/term" - - "vitess.io/vitess/go/acl" + "vitess.io/vitess/go/cmd/zk/command" "vitess.io/vitess/go/exit" "vitess.io/vitess/go/vt/log" - "vitess.io/vitess/go/vt/logutil" - "vitess.io/vitess/go/vt/topo" - "vitess.io/vitess/go/vt/topo/zk2topo" -) - -var doc = ` -zk is a tool for wrangling the zookeeper - -It tries to mimic unix file system commands wherever possible, but -there are some slight differences in flag handling. - -zk -h - provide help on overriding cell selection - -zk addAuth digest user:pass - -zk cat /zk/path -zk cat -l /zk/path1 /zk/path2 (list filename before file data) - -zk chmod n-mode /zk/path -zk chmod n+mode /zk/path - -zk cp /zk/path . -zk cp ./config /zk/path/config -zk cp ./config /zk/path/ (trailing slash indicates directory) - -zk edit /zk/path (create a local copy, edit and write changes back to cell) - -zk ls /zk -zk ls -l /zk -zk ls -ld /zk (list directory node itself) -zk ls -R /zk (recursive, expensive) - -zk stat /zk/path - -zk touch /zk/path -zk touch -c /zk/path (don't create, just touch timestamp) -zk touch -p /zk/path (create all parts necessary, think mkdir -p) -NOTE: there is no mkdir - just touch a node. The distinction -between file and directory is just not relevant in zookeeper. - -zk rm /zk/path -zk rm -r /zk/path (recursive) -zk rm -f /zk/path (no error on nonexistent node) - -zk wait /zk/path (wait for node change or creation) -zk wait /zk/path/children/ (trailing slash waits on children) - -zk watch /zk/path (print changes) - -zk unzip zktree.zip / -zk unzip zktree.zip /zk/prefix - -zk zip /zk/root zktree.zip -NOTE: zip file can't be dumped to the file system since znodes -can have data and children. - -The zk tool looks for the address of the cluster in /etc/zookeeper/zk_client.conf, -or the file specified in the ZK_CLIENT_CONFIG environment variable. - -The local cell may be overridden with the ZK_CLIENT_LOCAL_CELL environment -variable. -` - -const ( - timeFmt = "2006-01-02 15:04:05" - timeFmtMicro = "2006-01-02 15:04:05.000000" ) -type cmdFunc func(ctx context.Context, subFlags *pflag.FlagSet, args []string) error - -var cmdMap map[string]cmdFunc -var zconn *zk2topo.ZkConn -var server string - -func init() { - cmdMap = map[string]cmdFunc{ - "addAuth": cmdAddAuth, - "cat": cmdCat, - "chmod": cmdChmod, - "cp": cmdCp, - "edit": cmdEdit, - "ls": cmdLs, - "rm": cmdRm, - "stat": cmdStat, - "touch": cmdTouch, - "unzip": cmdUnzip, - "wait": cmdWait, - "watch": cmdWatch, - "zip": cmdZip, - } -} - func main() { defer exit.Recover() - defer logutil.Flush() - pflag.StringVar(&server, "server", server, "server(s) to connect to") - // handling case of --help & -h - var help bool - pflag.BoolVarP(&help, "help", "h", false, "display usage and exit") - log.RegisterFlags(pflag.CommandLine) - logutil.RegisterFlags(pflag.CommandLine) - acl.RegisterFlags(pflag.CommandLine) - pflag.CommandLine.Usage = func() { - fmt.Fprint(os.Stderr, doc) - pflag.Usage() - } - - pflag.Parse() - logutil.PurgeLogs() - - if help || pflag.Arg(0) == "help" { - pflag.Usage() - os.Exit(0) - } - - // if no zk command is provided after --server then we need to print doc & usage both - args := pflag.Args() - if len(args) == 0 { - pflag.CommandLine.Usage() - exit.Return(1) - } - cmdName := args[0] - args = args[1:] - cmd, ok := cmdMap[cmdName] - if !ok { - log.Exitf("Unknown command %v", cmdName) - } - subFlags := pflag.NewFlagSet(cmdName, pflag.ContinueOnError) // Create a context for the command, cancel it if we get a signal. ctx, cancel := context.WithCancel(context.Background()) @@ -177,848 +38,9 @@ func main() { cancel() }() - // Connect to the server. - zconn = zk2topo.Connect(server) - // Run the command. - if err := cmd(ctx, subFlags, args); err != nil { + if err := command.Root.ExecuteContext(ctx); err != nil { log.Error(err) exit.Return(1) } } - -func fixZkPath(zkPath string) string { - if zkPath != "/" { - zkPath = strings.TrimSuffix(zkPath, "/") - } - return path.Clean(zkPath) -} - -func isZkFile(path string) bool { - return strings.HasPrefix(path, "/zk") -} - -func cmdWait(ctx context.Context, subFlags *pflag.FlagSet, args []string) error { - var exitIfExists bool - subFlags.BoolVarP(&exitIfExists, "exit", "e", false, "exit if the path already exists") - - if err := subFlags.Parse(args); err != nil { - return err - } - - if subFlags.NArg() != 1 { - return fmt.Errorf("wait: can only wait for one path") - } - zkPath := subFlags.Arg(0) - isDir := zkPath[len(zkPath)-1] == '/' - zkPath = fixZkPath(zkPath) - - var wait <-chan zk.Event - var err error - if isDir { - _, _, wait, err = zconn.ChildrenW(ctx, zkPath) - } else { - _, _, wait, err = zconn.GetW(ctx, zkPath) - } - if err != nil { - if err == zk.ErrNoNode { - _, _, wait, _ = zconn.ExistsW(ctx, zkPath) - } else { - return fmt.Errorf("wait: error %v: %v", zkPath, err) - } - } else { - if exitIfExists { - return fmt.Errorf("already exists: %v", zkPath) - } - } - event := <-wait - fmt.Printf("event: %v\n", event) - return nil -} - -// Watch for changes to the node. -func cmdWatch(ctx context.Context, subFlags *pflag.FlagSet, args []string) error { - if err := subFlags.Parse(args); err != nil { - return err - } - - eventChan := make(chan zk.Event, 16) - for _, arg := range subFlags.Args() { - zkPath := fixZkPath(arg) - _, _, watch, err := zconn.GetW(ctx, zkPath) - if err != nil { - return fmt.Errorf("watch error: %v", err) - } - go func() { - eventChan <- <-watch - }() - } - - for { - select { - case <-ctx.Done(): - return nil - case event := <-eventChan: - log.Infof("watch: event %v: %v", event.Path, event) - if event.Type == zk.EventNodeDataChanged { - data, stat, watch, err := zconn.GetW(ctx, event.Path) - if err != nil { - return fmt.Errorf("ERROR: failed to watch %v", err) - } - log.Infof("watch: %v %v\n", event.Path, stat) - println(data) - go func() { - eventChan <- <-watch - }() - } else if event.State == zk.StateDisconnected { - return nil - } else if event.Type == zk.EventNodeDeleted { - log.Infof("watch: %v deleted\n", event.Path) - } else { - // Most likely a session event - try t - _, _, watch, err := zconn.GetW(ctx, event.Path) - if err != nil { - return fmt.Errorf("ERROR: failed to watch %v", err) - } - go func() { - eventChan <- <-watch - }() - } - } - } -} - -func cmdLs(ctx context.Context, subFlags *pflag.FlagSet, args []string) error { - var ( - longListing bool - directoryListing bool - force bool - recursiveListing bool - ) - subFlags.BoolVarP(&longListing, "longlisting", "l", false, "long listing") - subFlags.BoolVarP(&directoryListing, "directorylisting", "d", false, "list directory instead of contents") - subFlags.BoolVarP(&force, "force", "f", false, "no warning on nonexistent node") - subFlags.BoolVarP(&recursiveListing, "recursivelisting", "R", false, "recursive listing") - - if err := subFlags.Parse(args); err != nil { - return err - } - if subFlags.NArg() == 0 { - return fmt.Errorf("ls: no path specified") - } - resolved, err := zk2topo.ResolveWildcards(ctx, zconn, subFlags.Args()) - if err != nil { - return fmt.Errorf("ls: invalid wildcards: %v", err) - } - if len(resolved) == 0 { - // the wildcards didn't result in anything, we're - // done. - return nil - } - - hasError := false - needsHeader := len(resolved) > 1 && !directoryListing - for _, arg := range resolved { - zkPath := fixZkPath(arg) - var children []string - var err error - isDir := true - if directoryListing { - children = []string{""} - isDir = false - } else if recursiveListing { - children, err = zk2topo.ChildrenRecursive(ctx, zconn, zkPath) - } else { - children, _, err = zconn.Children(ctx, zkPath) - // Assume this is a file node if it has no children. - if len(children) == 0 { - children = []string{""} - isDir = false - } - } - if err != nil { - hasError = true - if !force || err != zk.ErrNoNode { - log.Warningf("ls: cannot access %v: %v", zkPath, err) - } - } - - // Show the full path when it helps. - showFullPath := false - if recursiveListing { - showFullPath = true - } else if longListing && (directoryListing || !isDir) { - showFullPath = true - } - if needsHeader { - fmt.Printf("%v:\n", zkPath) - } - if len(children) > 0 { - if longListing && isDir { - fmt.Printf("total: %v\n", len(children)) - } - sort.Strings(children) - stats := make([]*zk.Stat, len(children)) - wg := sync.WaitGroup{} - f := func(i int) { - localPath := path.Join(zkPath, children[i]) - _, stat, err := zconn.Exists(ctx, localPath) - if err != nil { - if !force || err != zk.ErrNoNode { - log.Warningf("ls: cannot access: %v: %v", localPath, err) - } - } else { - stats[i] = stat - } - wg.Done() - } - for i := range children { - wg.Add(1) - go f(i) - } - wg.Wait() - - for i, child := range children { - localPath := path.Join(zkPath, child) - if stat := stats[i]; stat != nil { - fmtPath(stat, localPath, showFullPath, longListing) - } - } - } - if needsHeader { - fmt.Println() - } - } - if hasError { - return fmt.Errorf("ls: some paths had errors") - } - return nil -} - -func fmtPath(stat *zk.Stat, zkPath string, showFullPath bool, longListing bool) { - var name, perms string - - if !showFullPath { - name = path.Base(zkPath) - } else { - name = zkPath - } - - if longListing { - if stat.NumChildren > 0 { - // FIXME(msolomon) do permissions check? - perms = "drwxrwxrwx" - if stat.DataLength > 0 { - // give a visual indication that this node has data as well as children - perms = "nrw-rw-rw-" - } - } else if stat.EphemeralOwner != 0 { - perms = "erw-rw-rw-" - } else { - perms = "-rw-rw-rw-" - } - // always print the Local version of the time. zookeeper's - // go / C library would return a local time anyway, but - // might as well be sure. - fmt.Printf("%v %v %v % 8v % 20v %v\n", perms, "zk", "zk", stat.DataLength, zk2topo.Time(stat.Mtime).Local().Format(timeFmt), name) - } else { - fmt.Printf("%v\n", name) - } -} - -func cmdTouch(ctx context.Context, subFlags *pflag.FlagSet, args []string) error { - var ( - createParents bool - touchOnly bool - ) - - subFlags.BoolVarP(&createParents, "createparent", "p", false, "create parents") - subFlags.BoolVarP(&touchOnly, "touchonly", "c", false, "touch only - don't create") - - if err := subFlags.Parse(args); err != nil { - return err - } - if subFlags.NArg() != 1 { - return fmt.Errorf("touch: need to specify exactly one path") - } - - zkPath := fixZkPath(subFlags.Arg(0)) - - var ( - version int32 = -1 - create = false - ) - - data, stat, err := zconn.Get(ctx, zkPath) - switch { - case err == nil: - version = stat.Version - case err == zk.ErrNoNode: - create = true - default: - return fmt.Errorf("touch: cannot access %v: %v", zkPath, err) - } - - switch { - case !create: - _, err = zconn.Set(ctx, zkPath, data, version) - case touchOnly: - return fmt.Errorf("touch: no such path %v", zkPath) - case createParents: - _, err = zk2topo.CreateRecursive(ctx, zconn, zkPath, data, 0, zk.WorldACL(zk.PermAll), 10) - default: - _, err = zconn.Create(ctx, zkPath, data, 0, zk.WorldACL(zk.PermAll)) - } - - if err != nil { - return fmt.Errorf("touch: cannot modify %v: %v", zkPath, err) - } - return nil -} - -func cmdRm(ctx context.Context, subFlags *pflag.FlagSet, args []string) error { - var ( - force bool - recursiveDelete bool - ) - subFlags.BoolVarP(&force, "force", "f", false, "no warning on nonexistent node") - subFlags.BoolVarP(&recursiveDelete, "recursivedelete", "r", false, "recursive delete") - - if err := subFlags.Parse(args); err != nil { - return err - } - - if subFlags.NArg() == 0 { - return fmt.Errorf("rm: no path specified") - } - - if recursiveDelete { - for _, arg := range subFlags.Args() { - zkPath := fixZkPath(arg) - if strings.Count(zkPath, "/") < 2 { - return fmt.Errorf("rm: overly general path: %v", zkPath) - } - } - } - - resolved, err := zk2topo.ResolveWildcards(ctx, zconn, subFlags.Args()) - if err != nil { - return fmt.Errorf("rm: invalid wildcards: %v", err) - } - if len(resolved) == 0 { - // the wildcards didn't result in anything, we're done - return nil - } - - hasError := false - for _, arg := range resolved { - zkPath := fixZkPath(arg) - var err error - if recursiveDelete { - err = zk2topo.DeleteRecursive(ctx, zconn, zkPath, -1) - } else { - err = zconn.Delete(ctx, zkPath, -1) - } - if err != nil && (!force || err != zk.ErrNoNode) { - hasError = true - log.Warningf("rm: cannot delete %v: %v", zkPath, err) - } - } - if hasError { - // to be consistent with the command line 'rm -f', return - // 0 if using 'zk rm -f' and the file doesn't exist. - return fmt.Errorf("rm: some paths had errors") - } - return nil -} - -func cmdAddAuth(ctx context.Context, subFlags *pflag.FlagSet, args []string) error { - if err := subFlags.Parse(args); err != nil { - return err - } - if subFlags.NArg() < 2 { - return fmt.Errorf("addAuth: expected args ") - } - scheme, auth := subFlags.Arg(0), subFlags.Arg(1) - return zconn.AddAuth(ctx, scheme, []byte(auth)) -} - -func cmdCat(ctx context.Context, subFlags *pflag.FlagSet, args []string) error { - var ( - longListing bool - force bool - decodeProto bool - ) - subFlags.BoolVarP(&longListing, "longListing", "l", false, "long listing") - subFlags.BoolVarP(&force, "force", "f", false, "no warning on nonexistent node") - subFlags.BoolVarP(&decodeProto, "decodeProto", "p", false, "decode proto files and display them as text") - - if err := subFlags.Parse(args); err != nil { - return err - } - if subFlags.NArg() == 0 { - return fmt.Errorf("cat: no path specified") - } - resolved, err := zk2topo.ResolveWildcards(ctx, zconn, subFlags.Args()) - if err != nil { - return fmt.Errorf("cat: invalid wildcards: %v", err) - } - if len(resolved) == 0 { - // the wildcards didn't result in anything, we're done - return nil - } - - hasError := false - for _, arg := range resolved { - zkPath := fixZkPath(arg) - data, _, err := zconn.Get(ctx, zkPath) - if err != nil { - hasError = true - if !force || err != zk.ErrNoNode { - log.Warningf("cat: cannot access %v: %v", zkPath, err) - } - continue - } - - if longListing { - fmt.Printf("%v:\n", zkPath) - } - decoded := "" - if decodeProto { - decoded, err = topo.DecodeContent(zkPath, data, false) - if err != nil { - log.Warningf("cat: cannot proto decode %v: %v", zkPath, err) - decoded = string(data) - } - } else { - decoded = string(data) - } - fmt.Print(decoded) - if len(decoded) > 0 && decoded[len(decoded)-1] != '\n' && (term.IsTerminal(int(os.Stdout.Fd())) || longListing) { - fmt.Print("\n") - } - } - if hasError { - return fmt.Errorf("cat: some paths had errors") - } - return nil -} - -func cmdEdit(ctx context.Context, subFlags *pflag.FlagSet, args []string) error { - var force bool - subFlags.BoolVarP(&force, "force", "f", false, "no warning on nonexistent node") - - if err := subFlags.Parse(args); err != nil { - return err - } - if subFlags.NArg() == 0 { - return fmt.Errorf("edit: no path specified") - } - arg := subFlags.Arg(0) - zkPath := fixZkPath(arg) - data, stat, err := zconn.Get(ctx, zkPath) - if err != nil { - if !force || err != zk.ErrNoNode { - log.Warningf("edit: cannot access %v: %v", zkPath, err) - } - return fmt.Errorf("edit: cannot access %v: %v", zkPath, err) - } - - name := path.Base(zkPath) - tmpPath := fmt.Sprintf("/tmp/zk-edit-%v-%v", name, time.Now().UnixNano()) - f, err := os.Create(tmpPath) - if err == nil { - _, err = f.Write(data) - f.Close() - } - if err != nil { - return fmt.Errorf("edit: cannot write file %v", err) - } - - cmd := exec.Command(os.Getenv("EDITOR"), tmpPath) - cmd.Stdin = os.Stdin - cmd.Stdout = os.Stdout - cmd.Stderr = os.Stderr - err = cmd.Run() - if err != nil { - os.Remove(tmpPath) - return fmt.Errorf("edit: cannot start $EDITOR: %v", err) - } - - fileData, err := os.ReadFile(tmpPath) - if err != nil { - os.Remove(tmpPath) - return fmt.Errorf("edit: cannot read file %v", err) - } - - if !bytes.Equal(fileData, data) { - // data changed - update if we can - _, err = zconn.Set(ctx, zkPath, fileData, stat.Version) - if err != nil { - os.Remove(tmpPath) - return fmt.Errorf("edit: cannot write zk file %v", err) - } - } - os.Remove(tmpPath) - return nil -} - -func cmdStat(ctx context.Context, subFlags *pflag.FlagSet, args []string) error { - var force bool - subFlags.BoolVarP(&force, "force", "f", false, "no warning on nonexistent node") - - if err := subFlags.Parse(args); err != nil { - return err - } - - if subFlags.NArg() == 0 { - return fmt.Errorf("stat: no path specified") - } - - resolved, err := zk2topo.ResolveWildcards(ctx, zconn, subFlags.Args()) - if err != nil { - return fmt.Errorf("stat: invalid wildcards: %v", err) - } - if len(resolved) == 0 { - // the wildcards didn't result in anything, we're done - return nil - } - - hasError := false - for _, arg := range resolved { - zkPath := fixZkPath(arg) - acls, stat, err := zconn.GetACL(ctx, zkPath) - if stat == nil { - err = fmt.Errorf("no such node") - } - if err != nil { - hasError = true - if !force || err != zk.ErrNoNode { - log.Warningf("stat: cannot access %v: %v", zkPath, err) - } - continue - } - fmt.Printf("Path: %s\n", zkPath) - fmt.Printf("Created: %s\n", zk2topo.Time(stat.Ctime).Format(timeFmtMicro)) - fmt.Printf("Modified: %s\n", zk2topo.Time(stat.Mtime).Format(timeFmtMicro)) - fmt.Printf("Size: %v\n", stat.DataLength) - fmt.Printf("Children: %v\n", stat.NumChildren) - fmt.Printf("Version: %v\n", stat.Version) - fmt.Printf("Ephemeral: %v\n", stat.EphemeralOwner) - fmt.Printf("ACL:\n") - for _, acl := range acls { - fmt.Printf(" %v:%v %v\n", acl.Scheme, acl.ID, fmtACL(acl)) - } - } - if hasError { - return fmt.Errorf("stat: some paths had errors") - } - return nil -} - -var charPermMap map[string]int32 -var permCharMap map[int32]string - -func init() { - charPermMap = map[string]int32{ - "r": zk.PermRead, - "w": zk.PermWrite, - "d": zk.PermDelete, - "c": zk.PermCreate, - "a": zk.PermAdmin, - } - permCharMap = make(map[int32]string) - for c, p := range charPermMap { - permCharMap[p] = c - } -} - -func fmtACL(acl zk.ACL) string { - s := "" - - for _, perm := range []int32{zk.PermRead, zk.PermWrite, zk.PermDelete, zk.PermCreate, zk.PermAdmin} { - if acl.Perms&perm != 0 { - s += permCharMap[perm] - } else { - s += "-" - } - } - return s -} - -func cmdChmod(ctx context.Context, subFlags *pflag.FlagSet, args []string) error { - if err := subFlags.Parse(args); err != nil { - return err - } - if subFlags.NArg() < 2 { - return fmt.Errorf("chmod: no permission specified") - } - mode := subFlags.Arg(0) - if mode[0] != 'n' { - return fmt.Errorf("chmod: invalid mode") - } - - addPerms := false - if mode[1] == '+' { - addPerms = true - } else if mode[1] != '-' { - return fmt.Errorf("chmod: invalid mode") - } - - var permMask int32 - for _, c := range mode[2:] { - permMask |= charPermMap[string(c)] - } - - resolved, err := zk2topo.ResolveWildcards(ctx, zconn, subFlags.Args()[1:]) - if err != nil { - return fmt.Errorf("chmod: invalid wildcards: %v", err) - } - if len(resolved) == 0 { - // the wildcards didn't result in anything, we're done - return nil - } - - hasError := false - for _, arg := range resolved { - zkPath := fixZkPath(arg) - aclv, _, err := zconn.GetACL(ctx, zkPath) - if err != nil { - hasError = true - log.Warningf("chmod: cannot set access %v: %v", zkPath, err) - continue - } - if addPerms { - aclv[0].Perms |= permMask - } else { - aclv[0].Perms &= ^permMask - } - err = zconn.SetACL(ctx, zkPath, aclv, -1) - if err != nil { - hasError = true - log.Warningf("chmod: cannot set access %v: %v", zkPath, err) - continue - } - } - if hasError { - return fmt.Errorf("chmod: some paths had errors") - } - return nil -} - -func cmdCp(ctx context.Context, subFlags *pflag.FlagSet, args []string) error { - if err := subFlags.Parse(args); err != nil { - return err - } - switch { - case subFlags.NArg() < 2: - return fmt.Errorf("cp: need to specify source and destination paths") - case subFlags.NArg() == 2: - return fileCp(ctx, args[0], args[1]) - default: - return multiFileCp(ctx, args) - } -} - -func getPathData(ctx context.Context, filePath string) ([]byte, error) { - if isZkFile(filePath) { - data, _, err := zconn.Get(ctx, filePath) - return data, err - } - var err error - file, err := os.Open(filePath) - if err == nil { - data, err := io.ReadAll(file) - if err == nil { - return data, err - } - } - return nil, err -} - -func setPathData(ctx context.Context, filePath string, data []byte) error { - if isZkFile(filePath) { - _, err := zconn.Set(ctx, filePath, data, -1) - if err == zk.ErrNoNode { - _, err = zk2topo.CreateRecursive(ctx, zconn, filePath, data, 0, zk.WorldACL(zk.PermAll), 10) - } - return err - } - return os.WriteFile(filePath, []byte(data), 0666) -} - -func fileCp(ctx context.Context, srcPath, dstPath string) error { - dstIsDir := dstPath[len(dstPath)-1] == '/' - srcPath = fixZkPath(srcPath) - dstPath = fixZkPath(dstPath) - - if !isZkFile(srcPath) && !isZkFile(dstPath) { - return fmt.Errorf("cp: neither src nor dst is a /zk file: exitting") - } - - data, err := getPathData(ctx, srcPath) - if err != nil { - return fmt.Errorf("cp: cannot read %v: %v", srcPath, err) - } - - // If we are copying to a local directory - say '.', make the filename - // the same as the source. - if !isZkFile(dstPath) { - fileInfo, err := os.Stat(dstPath) - if err != nil { - if err.(*os.PathError).Err != syscall.ENOENT { - return fmt.Errorf("cp: cannot stat %v: %v", dstPath, err) - } - } else if fileInfo.IsDir() { - dstPath = path.Join(dstPath, path.Base(srcPath)) - } - } else if dstIsDir { - // If we are copying into zk, interpret trailing slash as treating the - // dstPath as a directory. - dstPath = path.Join(dstPath, path.Base(srcPath)) - } - if err := setPathData(ctx, dstPath, data); err != nil { - return fmt.Errorf("cp: cannot write %v: %v", dstPath, err) - } - return nil -} - -func multiFileCp(ctx context.Context, args []string) error { - dstPath := args[len(args)-1] - if dstPath[len(dstPath)-1] != '/' { - // In multifile context, dstPath must be a directory. - dstPath += "/" - } - - for _, srcPath := range args[:len(args)-1] { - if err := fileCp(ctx, srcPath, dstPath); err != nil { - return err - } - } - return nil -} - -type zkItem struct { - path string - data []byte - stat *zk.Stat - err error -} - -// Store a zk tree in a zip archive. This won't be immediately useful to -// zip tools since even "directories" can contain data. -func cmdZip(ctx context.Context, subFlags *pflag.FlagSet, args []string) error { - if err := subFlags.Parse(args); err != nil { - return err - } - if subFlags.NArg() < 2 { - return fmt.Errorf("zip: need to specify source and destination paths") - } - - dstPath := subFlags.Arg(subFlags.NArg() - 1) - paths := subFlags.Args()[:len(args)-1] - if !strings.HasSuffix(dstPath, ".zip") { - return fmt.Errorf("zip: need to specify destination .zip path: %v", dstPath) - } - zipFile, err := os.Create(dstPath) - if err != nil { - return fmt.Errorf("zip: error %v", err) - } - - wg := sync.WaitGroup{} - items := make(chan *zkItem, 64) - for _, arg := range paths { - zkPath := fixZkPath(arg) - children, err := zk2topo.ChildrenRecursive(ctx, zconn, zkPath) - if err != nil { - return fmt.Errorf("zip: error %v", err) - } - for _, child := range children { - toAdd := path.Join(zkPath, child) - wg.Add(1) - go func() { - data, stat, err := zconn.Get(ctx, toAdd) - items <- &zkItem{toAdd, data, stat, err} - wg.Done() - }() - } - } - go func() { - wg.Wait() - close(items) - }() - - zipWriter := zip.NewWriter(zipFile) - for item := range items { - path, data, stat, err := item.path, item.data, item.stat, item.err - if err != nil { - return fmt.Errorf("zip: get failed: %v", err) - } - // Skip ephemerals - not sure why you would archive them. - if stat.EphemeralOwner > 0 { - continue - } - fi := &zip.FileHeader{Name: path, Method: zip.Deflate} - fi.Modified = zk2topo.Time(stat.Mtime) - f, err := zipWriter.CreateHeader(fi) - if err != nil { - return fmt.Errorf("zip: create failed: %v", err) - } - _, err = f.Write(data) - if err != nil { - return fmt.Errorf("zip: create failed: %v", err) - } - } - err = zipWriter.Close() - if err != nil { - return fmt.Errorf("zip: close failed: %v", err) - } - zipFile.Close() - return nil -} - -func cmdUnzip(ctx context.Context, subFlags *pflag.FlagSet, args []string) error { - if err := subFlags.Parse(args); err != nil { - return err - } - if subFlags.NArg() != 2 { - return fmt.Errorf("zip: need to specify source and destination paths") - } - - srcPath, dstPath := subFlags.Arg(0), subFlags.Arg(1) - - if !strings.HasSuffix(srcPath, ".zip") { - return fmt.Errorf("zip: need to specify src .zip path: %v", srcPath) - } - - zipReader, err := zip.OpenReader(srcPath) - if err != nil { - return fmt.Errorf("zip: error %v", err) - } - defer zipReader.Close() - - for _, zf := range zipReader.File { - rc, err := zf.Open() - if err != nil { - return fmt.Errorf("unzip: error %v", err) - } - data, err := io.ReadAll(rc) - if err != nil { - return fmt.Errorf("unzip: failed reading archive: %v", err) - } - zkPath := zf.Name - if dstPath != "/" { - zkPath = path.Join(dstPath, zkPath) - } - _, err = zk2topo.CreateRecursive(ctx, zconn, zkPath, data, 0, zk.WorldACL(zk.PermAll), 10) - if err != nil && err != zk.ErrNodeExists { - return fmt.Errorf("unzip: zk create failed: %v", err) - } - _, err = zconn.Set(ctx, zkPath, data, -1) - if err != nil { - return fmt.Errorf("unzip: zk set failed: %v", err) - } - rc.Close() - } - return nil -} diff --git a/go/cmd/zkctl/command/init.go b/go/cmd/zkctl/command/init.go new file mode 100644 index 00000000000..518b4a6239d --- /dev/null +++ b/go/cmd/zkctl/command/init.go @@ -0,0 +1,32 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package command + +import "github.com/spf13/cobra" + +var Init = &cobra.Command{ + Use: "init", + Short: "Generates a new config and then starts zookeeper.", + Args: cobra.ExactArgs(0), + RunE: func(cmd *cobra.Command, args []string) error { + return zkd.Init() + }, +} + +func init() { + Root.AddCommand(Init) +} diff --git a/go/cmd/zkctl/command/root.go b/go/cmd/zkctl/command/root.go new file mode 100644 index 00000000000..3399ed8c4cb --- /dev/null +++ b/go/cmd/zkctl/command/root.go @@ -0,0 +1,63 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package command + +import ( + "github.com/spf13/cobra" + + "vitess.io/vitess/go/vt/logutil" + "vitess.io/vitess/go/vt/servenv" + "vitess.io/vitess/go/vt/zkctl" +) + +var ( + zkCfg = "6@:3801:3802:3803" + myID uint + zkExtra []string + + zkd *zkctl.Zkd + + Root = &cobra.Command{ + Use: "zkctl", + Short: "Initializes and controls zookeeper with Vitess-specific configuration.", + PersistentPreRunE: func(cmd *cobra.Command, args []string) error { + if err := servenv.CobraPreRunE(cmd, args); err != nil { + return err + } + + zkConfig := zkctl.MakeZkConfigFromString(zkCfg, uint32(myID)) + zkConfig.Extra = zkExtra + zkd = zkctl.NewZkd(zkConfig) + + return nil + }, + PersistentPostRun: func(cmd *cobra.Command, args []string) { + logutil.Flush() + }, + } +) + +func init() { + Root.PersistentFlags().StringVar(&zkCfg, "zk.cfg", zkCfg, + "zkid@server1:leaderPort1:electionPort1:clientPort1,...)") + Root.PersistentFlags().UintVar(&myID, "zk.myid", myID, + "which server do you want to be? only needed when running multiple instance on one box, otherwise myid is implied by hostname") + Root.PersistentFlags().StringArrayVar(&zkExtra, "zk.extra", zkExtra, + "extra config line(s) to append verbatim to config (flag can be specified more than once)") + + servenv.MovePersistentFlagsToCobraCommand(Root) +} diff --git a/go/cmd/zkctl/command/shutdown.go b/go/cmd/zkctl/command/shutdown.go new file mode 100644 index 00000000000..b3166bbd36b --- /dev/null +++ b/go/cmd/zkctl/command/shutdown.go @@ -0,0 +1,32 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package command + +import "github.com/spf13/cobra" + +var Shutdown = &cobra.Command{ + Use: "shutdown", + Short: "Terminates a zookeeper server but keeps its data dir intact.", + Args: cobra.ExactArgs(0), + RunE: func(cmd *cobra.Command, args []string) error { + return zkd.Shutdown() + }, +} + +func init() { + Root.AddCommand(Shutdown) +} diff --git a/go/cmd/zkctl/command/start.go b/go/cmd/zkctl/command/start.go new file mode 100644 index 00000000000..1ed31d0ed54 --- /dev/null +++ b/go/cmd/zkctl/command/start.go @@ -0,0 +1,32 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package command + +import "github.com/spf13/cobra" + +var Start = &cobra.Command{ + Use: "start", + Short: "Runs an already initialized zookeeper server.", + Args: cobra.ExactArgs(0), + RunE: func(cmd *cobra.Command, args []string) error { + return zkd.Start() + }, +} + +func init() { + Root.AddCommand(Start) +} diff --git a/go/cmd/zkctl/command/teardown.go b/go/cmd/zkctl/command/teardown.go new file mode 100644 index 00000000000..14fe7278835 --- /dev/null +++ b/go/cmd/zkctl/command/teardown.go @@ -0,0 +1,32 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package command + +import "github.com/spf13/cobra" + +var Teardown = &cobra.Command{ + Use: "teardown", + Short: "Shuts down the zookeeper server and removes its data dir.", + Args: cobra.ExactArgs(0), + RunE: func(cmd *cobra.Command, args []string) error { + return zkd.Teardown() + }, +} + +func init() { + Root.AddCommand(Teardown) +} diff --git a/go/cmd/zkctl/docgen/main.go b/go/cmd/zkctl/docgen/main.go new file mode 100644 index 00000000000..c35da8930e4 --- /dev/null +++ b/go/cmd/zkctl/docgen/main.go @@ -0,0 +1,37 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package main + +import ( + "github.com/spf13/cobra" + + "vitess.io/vitess/go/cmd/internal/docgen" + "vitess.io/vitess/go/cmd/zkctl/command" +) + +func main() { + var dir string + cmd := cobra.Command{ + Use: "docgen [-d ]", + RunE: func(cmd *cobra.Command, args []string) error { + return docgen.GenerateMarkdownTree(command.Root, dir) + }, + } + + cmd.Flags().StringVarP(&dir, "dir", "d", "doc", "output directory to write documentation") + _ = cmd.Execute() +} diff --git a/go/cmd/zkctl/zkctl.go b/go/cmd/zkctl/zkctl.go index 566631e8cc2..b00e3eb4812 100644 --- a/go/cmd/zkctl/zkctl.go +++ b/go/cmd/zkctl/zkctl.go @@ -14,75 +14,19 @@ See the License for the specific language governing permissions and limitations under the License. */ -// zkctl initializes and controls ZooKeeper with Vitess-specific configuration. package main import ( - "github.com/spf13/pflag" - + "vitess.io/vitess/go/cmd/zkctl/command" "vitess.io/vitess/go/exit" "vitess.io/vitess/go/vt/log" - "vitess.io/vitess/go/vt/logutil" - "vitess.io/vitess/go/vt/servenv" - "vitess.io/vitess/go/vt/zkctl" ) -var usage = ` -Commands: - - init | start | shutdown | teardown -` - -var ( - zkCfg = "6@:3801:3802:3803" - myID uint - zkExtra []string -) - -func registerZkctlFlags(fs *pflag.FlagSet) { - fs.StringVar(&zkCfg, "zk.cfg", zkCfg, - "zkid@server1:leaderPort1:electionPort1:clientPort1,...)") - fs.UintVar(&myID, "zk.myid", myID, - "which server do you want to be? only needed when running multiple instance on one box, otherwise myid is implied by hostname") - fs.StringArrayVar(&zkExtra, "zk.extra", zkExtra, - "extra config line(s) to append verbatim to config (flag can be specified more than once)") -} - -func init() { - servenv.OnParse(registerZkctlFlags) -} - func main() { defer exit.Recover() - defer logutil.Flush() - - fs := pflag.NewFlagSet("zkctl", pflag.ExitOnError) - log.RegisterFlags(fs) - logutil.RegisterFlags(fs) - args := servenv.ParseFlagsWithArgs("zkctl") - zkConfig := zkctl.MakeZkConfigFromString(zkCfg, uint32(myID)) - zkConfig.Extra = zkExtra - zkd := zkctl.NewZkd(zkConfig) - - action := args[0] - var err error - switch action { - case "init": - err = zkd.Init() - case "shutdown": - err = zkd.Shutdown() - case "start": - err = zkd.Start() - case "teardown": - err = zkd.Teardown() - default: - log.Errorf("invalid action: %v", action) - log.Errorf(usage) - exit.Return(1) - } - if err != nil { - log.Errorf("failed %v: %v", action, err) + if err := command.Root.Execute(); err != nil { + log.Error(err) exit.Return(1) } } diff --git a/go/cmd/zkctld/cli/zkctld.go b/go/cmd/zkctld/cli/zkctld.go new file mode 100644 index 00000000000..101f1013722 --- /dev/null +++ b/go/cmd/zkctld/cli/zkctld.go @@ -0,0 +1,100 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package cli + +import ( + "fmt" + "os" + "os/signal" + "syscall" + + "github.com/spf13/cobra" + "github.com/spf13/pflag" + + "vitess.io/vitess/go/acl" + "vitess.io/vitess/go/vt/log" + "vitess.io/vitess/go/vt/logutil" + "vitess.io/vitess/go/vt/servenv" + "vitess.io/vitess/go/vt/zkctl" +) + +var ( + zkCfg = "6@:3801:3802:3803" + myID uint + zkExtra []string + + Main = &cobra.Command{ + Use: "zkctld", + Short: "zkctld is a daemon that starts or initializes ZooKeeper with Vitess-specific configuration. It will stay running as long as the underlying ZooKeeper server, and will pass along SIGTERM.", + Args: cobra.NoArgs, + PersistentPreRunE: servenv.CobraPreRunE, + PostRun: func(cmd *cobra.Command, args []string) { + logutil.Flush() + }, + RunE: run, + } +) + +func init() { + servenv.OnParse(registerFlags) +} + +func registerFlags(fs *pflag.FlagSet) { + fs.StringVar(&zkCfg, "zk.cfg", zkCfg, + "zkid@server1:leaderPort1:electionPort1:clientPort1,...)") + fs.UintVar(&myID, "zk.myid", myID, + "which server do you want to be? only needed when running multiple instance on one box, otherwise myid is implied by hostname") + fs.StringArrayVar(&zkExtra, "zk.extra", zkExtra, + "extra config line(s) to append verbatim to config (flag can be specified more than once)") + acl.RegisterFlags(fs) +} + +func run(cmd *cobra.Command, args []string) error { + servenv.Init() + zkConfig := zkctl.MakeZkConfigFromString(zkCfg, uint32(myID)) + zkConfig.Extra = zkExtra + zkd := zkctl.NewZkd(zkConfig) + + if zkd.Inited() { + log.Infof("already initialized, starting without init...") + if err := zkd.Start(); err != nil { + return fmt.Errorf("failed start: %v", err) + } + } else { + log.Infof("initializing...") + if err := zkd.Init(); err != nil { + return fmt.Errorf("failed init: %v", err) + } + } + + log.Infof("waiting for signal or server shutdown...") + sig := make(chan os.Signal, 1) + signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM) + select { + case <-zkd.Done(): + log.Infof("server shut down on its own") + case <-sig: + log.Infof("signal received, shutting down server") + + // Action to perform if there is an error + if err := zkd.Shutdown(); err != nil { + return fmt.Errorf("error during shutdown:%v", err) + } + } + + return nil +} diff --git a/go/cmd/zkctld/docgen/main.go b/go/cmd/zkctld/docgen/main.go new file mode 100644 index 00000000000..9915d641352 --- /dev/null +++ b/go/cmd/zkctld/docgen/main.go @@ -0,0 +1,37 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package main + +import ( + "github.com/spf13/cobra" + + "vitess.io/vitess/go/cmd/internal/docgen" + "vitess.io/vitess/go/cmd/vttablet/cli" +) + +func main() { + var dir string + cmd := cobra.Command{ + Use: "docgen [-d ]", + RunE: func(cmd *cobra.Command, args []string) error { + return docgen.GenerateMarkdownTree(cli.Main, dir) + }, + } + + cmd.Flags().StringVarP(&dir, "dir", "d", "doc", "output directory to write documentation") + _ = cmd.Execute() +} diff --git a/go/cmd/zkctld/zkctld.go b/go/cmd/zkctld/zkctld.go index 0d1ee413a66..211b63325eb 100644 --- a/go/cmd/zkctld/zkctld.go +++ b/go/cmd/zkctld/zkctld.go @@ -20,77 +20,15 @@ limitations under the License. package main import ( - "os" - "os/signal" - "syscall" - - "github.com/spf13/pflag" - - "vitess.io/vitess/go/acl" + "vitess.io/vitess/go/cmd/zkctld/cli" "vitess.io/vitess/go/exit" "vitess.io/vitess/go/vt/log" - "vitess.io/vitess/go/vt/logutil" - "vitess.io/vitess/go/vt/servenv" - "vitess.io/vitess/go/vt/zkctl" -) - -var ( - zkCfg = "6@:3801:3802:3803" - myID uint - zkExtra []string ) -func init() { - servenv.OnParse(registerFlags) -} - -func registerFlags(fs *pflag.FlagSet) { - fs.StringVar(&zkCfg, "zk.cfg", zkCfg, - "zkid@server1:leaderPort1:electionPort1:clientPort1,...)") - fs.UintVar(&myID, "zk.myid", myID, - "which server do you want to be? only needed when running multiple instance on one box, otherwise myid is implied by hostname") - fs.StringArrayVar(&zkExtra, "zk.extra", zkExtra, - "extra config line(s) to append verbatim to config (flag can be specified more than once)") - acl.RegisterFlags(fs) -} - func main() { defer exit.Recover() - defer logutil.Flush() - - servenv.ParseFlags("zkctld") - servenv.Init() - zkConfig := zkctl.MakeZkConfigFromString(zkCfg, uint32(myID)) - zkConfig.Extra = zkExtra - zkd := zkctl.NewZkd(zkConfig) - - if zkd.Inited() { - log.Infof("already initialized, starting without init...") - if err := zkd.Start(); err != nil { - log.Errorf("failed start: %v", err) - exit.Return(255) - } - } else { - log.Infof("initializing...") - if err := zkd.Init(); err != nil { - log.Errorf("failed init: %v", err) - exit.Return(255) - } - } - - log.Infof("waiting for signal or server shutdown...") - sig := make(chan os.Signal, 1) - signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM) - select { - case <-zkd.Done(): - log.Infof("server shut down on its own") - case <-sig: - log.Infof("signal received, shutting down server") - - // Action to perform if there is an error - if err := zkd.Shutdown(); err != nil { - log.Errorf("error during shutdown:%v", err) - exit.Return(1) - } + if err := cli.Main.Execute(); err != nil { + log.Error(err) + exit.Return(1) } } diff --git a/go/flags/endtoend/zk.txt b/go/flags/endtoend/zk.txt index 443bf0b9ca2..add1b6b6803 100644 --- a/go/flags/endtoend/zk.txt +++ b/go/flags/endtoend/zk.txt @@ -1,8 +1,41 @@ -Usage of zk: - -h, --help display usage and exit +zk is a tool for wrangling zookeeper. + +It tries to mimic unix file system commands wherever possible, but +there are some slight differences in flag handling. + +The zk tool looks for the address of the cluster in /etc/zookeeper/zk_client.conf, +or the file specified in the ZK_CLIENT_CONFIG environment variable. + +The local cell may be overridden with the ZK_CLIENT_LOCAL_CELL environment +variable. + +Usage: + zk [command] + +Available Commands: + addAuth + cat + chmod + completion Generate the autocompletion script for the specified shell + cp + edit Create a local copy, edit, and write changes back to cell. + help Help about any command + ls + rm + stat + touch Change node access time. + unzip + wait Sets a watch on the node and then waits for an event to fire. + watch Watches for changes to nodes and prints events as they occur. + zip Store a zk tree in a zip archive. + +Flags: + -h, --help help for zk --keep_logs duration keep logs for this long (using ctime) (zero to keep forever) --keep_logs_by_mtime duration keep logs for this long (using mtime) (zero to keep forever) --log_rotate_max_size uint size in bytes at which logs are rotated (glog.MaxSize) (default 1887436800) --purge_logs_interval duration how often try to remove old logs (default 1h0m0s) --security_policy string the name of a registered security policy to use for controlling access to URLs - empty means allow all for anyone (built-in policies: deny-all, read-only) --server string server(s) to connect to + +Use "zk [command] --help" for more information about a command. diff --git a/go/flags/endtoend/zkctl.txt b/go/flags/endtoend/zkctl.txt index 6b0473d1cb2..d1aea061ea5 100644 --- a/go/flags/endtoend/zkctl.txt +++ b/go/flags/endtoend/zkctl.txt @@ -1,4 +1,17 @@ -Usage of zkctl: +Initializes and controls zookeeper with Vitess-specific configuration. + +Usage: + zkctl [command] + +Available Commands: + completion Generate the autocompletion script for the specified shell + help Help about any command + init Generates a new config and then starts zookeeper. + shutdown Terminates a zookeeper server but keeps its data dir intact. + start Runs an already initialized zookeeper server. + teardown Shuts down the zookeeper server and removes its data dir. + +Flags: --alsologtostderr log to standard error as well as files --config-file string Full path of the config file (with extension) to use. If set, --config-path, --config-type, and --config-name are ignored. --config-file-not-found-handling ConfigFileNotFoundHandling Behavior when a config file is not found. (Options: error, exit, ignore, warn) (default warn) @@ -6,7 +19,7 @@ Usage of zkctl: --config-path strings Paths to search for config files in. (default [{{ .Workdir }}]) --config-persistence-min-interval duration minimum interval between persisting dynamic config changes back to disk (if no change has occurred, nothing is done). (default 1s) --config-type string Config file type (omit to infer config type from file extension). - -h, --help display usage and exit + -h, --help help for zkctl --keep_logs duration keep logs for this long (using ctime) (zero to keep forever) --keep_logs_by_mtime duration keep logs for this long (using mtime) (zero to keep forever) --log_backtrace_at traceLocation when logging hits line file:N, emit a stack trace (default :0) @@ -23,3 +36,5 @@ Usage of zkctl: --zk.cfg string zkid@server1:leaderPort1:electionPort1:clientPort1,...) (default "6@:3801:3802:3803") --zk.extra stringArray extra config line(s) to append verbatim to config (flag can be specified more than once) --zk.myid uint which server do you want to be? only needed when running multiple instance on one box, otherwise myid is implied by hostname + +Use "zkctl [command] --help" for more information about a command. diff --git a/go/flags/endtoend/zkctld.txt b/go/flags/endtoend/zkctld.txt index e957f7a3b3c..d808bd7ce67 100644 --- a/go/flags/endtoend/zkctld.txt +++ b/go/flags/endtoend/zkctld.txt @@ -1,26 +1,7 @@ -Usage of zkctld: - --alsologtostderr log to standard error as well as files - --config-file string Full path of the config file (with extension) to use. If set, --config-path, --config-type, and --config-name are ignored. - --config-file-not-found-handling ConfigFileNotFoundHandling Behavior when a config file is not found. (Options: error, exit, ignore, warn) (default warn) - --config-name string Name of the config file (without extension) to search for. (default "vtconfig") - --config-path strings Paths to search for config files in. (default [{{ .Workdir }}]) - --config-persistence-min-interval duration minimum interval between persisting dynamic config changes back to disk (if no change has occurred, nothing is done). (default 1s) - --config-type string Config file type (omit to infer config type from file extension). - -h, --help display usage and exit - --keep_logs duration keep logs for this long (using ctime) (zero to keep forever) - --keep_logs_by_mtime duration keep logs for this long (using mtime) (zero to keep forever) - --log_backtrace_at traceLocation when logging hits line file:N, emit a stack trace (default :0) - --log_dir string If non-empty, write log files in this directory - --log_err_stacks log stack traces for errors - --log_rotate_max_size uint size in bytes at which logs are rotated (glog.MaxSize) (default 1887436800) - --logtostderr log to standard error instead of files - --pprof strings enable profiling - --purge_logs_interval duration how often try to remove old logs (default 1h0m0s) - --security_policy string the name of a registered security policy to use for controlling access to URLs - empty means allow all for anyone (built-in policies: deny-all, read-only) - --stderrthreshold severity logs at or above this threshold go to stderr (default 1) - --v Level log level for V logs - -v, --version print binary version - --vmodule moduleSpec comma-separated list of pattern=N settings for file-filtered logging - --zk.cfg string zkid@server1:leaderPort1:electionPort1:clientPort1,...) (default "6@:3801:3802:3803") - --zk.extra stringArray extra config line(s) to append verbatim to config (flag can be specified more than once) - --zk.myid uint which server do you want to be? only needed when running multiple instance on one box, otherwise myid is implied by hostname +zkctld is a daemon that starts or initializes ZooKeeper with Vitess-specific configuration. It will stay running as long as the underlying ZooKeeper server, and will pass along SIGTERM. + +Usage: + zkctld [flags] + +Flags: + -h, --help help for zkctld From b793adf139af8ecbc1e44bc540617ad155ffc505 Mon Sep 17 00:00:00 2001 From: Harshit Gangal Date: Wed, 27 Sep 2023 14:32:28 +0530 Subject: [PATCH 5/5] Reduce network pressure on multi row insert (#14064) Signed-off-by: Harshit Gangal --- go/test/endtoend/cluster/cluster_util.go | 2 +- go/test/endtoend/utils/utils.go | 2 +- .../queries/benchmark/benchmark_test.go | 80 +++++++ .../vtgate/queries/benchmark/main_test.go | 124 ++++++++++ .../queries/benchmark/sharded_schema.sql | 16 ++ .../vtgate/queries/benchmark/vschema.json | 18 ++ go/vt/vtgate/autocommit_test.go | 6 - go/vt/vtgate/engine/cached_size.go | 13 +- go/vt/vtgate/engine/insert.go | 20 +- go/vt/vtgate/engine/insert_test.go | 212 ++++++++++-------- go/vt/vtgate/executor_dml_test.go | 63 ++---- go/vt/vtgate/executor_vschema_ddl_test.go | 55 +++-- .../planbuilder/operator_transformers.go | 18 +- go/vt/vtgate/vtgate_test.go | 10 +- 14 files changed, 445 insertions(+), 194 deletions(-) create mode 100644 go/test/endtoend/vtgate/queries/benchmark/benchmark_test.go create mode 100644 go/test/endtoend/vtgate/queries/benchmark/main_test.go create mode 100644 go/test/endtoend/vtgate/queries/benchmark/sharded_schema.sql create mode 100644 go/test/endtoend/vtgate/queries/benchmark/vschema.json diff --git a/go/test/endtoend/cluster/cluster_util.go b/go/test/endtoend/cluster/cluster_util.go index 1af9504389b..3d442bbb576 100644 --- a/go/test/endtoend/cluster/cluster_util.go +++ b/go/test/endtoend/cluster/cluster_util.go @@ -127,7 +127,7 @@ func VerifyRowsInTablet(t *testing.T, vttablet *Vttablet, ksName string, expecte } // PanicHandler handles the panic in the testcase. -func PanicHandler(t *testing.T) { +func PanicHandler(t testing.TB) { err := recover() if t == nil { return diff --git a/go/test/endtoend/utils/utils.go b/go/test/endtoend/utils/utils.go index c0137b27066..4fc7d7cfecf 100644 --- a/go/test/endtoend/utils/utils.go +++ b/go/test/endtoend/utils/utils.go @@ -169,7 +169,7 @@ func ExecCompareMySQL(t *testing.T, vtConn, mysqlConn *mysql.Conn, query string) // ExecAllowError executes the given query without failing the test if it produces // an error. The error is returned to the client, along with the result set. -func ExecAllowError(t *testing.T, conn *mysql.Conn, query string) (*sqltypes.Result, error) { +func ExecAllowError(t testing.TB, conn *mysql.Conn, query string) (*sqltypes.Result, error) { t.Helper() return conn.ExecuteFetch(query, 1000, true) } diff --git a/go/test/endtoend/vtgate/queries/benchmark/benchmark_test.go b/go/test/endtoend/vtgate/queries/benchmark/benchmark_test.go new file mode 100644 index 00000000000..3fd7edd14de --- /dev/null +++ b/go/test/endtoend/vtgate/queries/benchmark/benchmark_test.go @@ -0,0 +1,80 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package dml + +import ( + "fmt" + "math/rand" + "strconv" + "strings" + "testing" + + "vitess.io/vitess/go/test/endtoend/utils" +) + +type testQuery struct { + tableName string + cols []string + intTyp []bool +} + +func (tq *testQuery) getInsertQuery(rows int) string { + var allRows []string + for i := 0; i < rows; i++ { + var row []string + for _, isInt := range tq.intTyp { + if isInt { + row = append(row, strconv.Itoa(i)) + continue + } + row = append(row, "'"+getRandomString(50)+"'") + } + allRows = append(allRows, "("+strings.Join(row, ",")+")") + } + return fmt.Sprintf("insert into %s(%s) values %s", tq.tableName, strings.Join(tq.cols, ","), strings.Join(allRows, ",")) +} + +func getRandomString(size int) string { + var str strings.Builder + + for i := 0; i < size; i++ { + str.WriteByte(byte(rand.Intn(27) + 97)) + } + return str.String() +} + +func BenchmarkShardedTblNoLookup(b *testing.B) { + conn, closer := start(b) + defer closer() + + cols := []string{"id", "c1", "c2", "c3", "c4", "c5", "c6", "c7", "c8", "c9", "c10", "c11", "c12"} + intType := make([]bool, len(cols)) + intType[0] = true + tq := &testQuery{ + tableName: "tbl_no_lkp_vdx", + cols: cols, + intTyp: intType, + } + for _, rows := range []int{1, 10, 100, 500, 1000, 5000, 10000} { + insStmt := tq.getInsertQuery(rows) + b.Run(fmt.Sprintf("16-shards-%d-rows", rows), func(b *testing.B) { + for i := 0; i < b.N; i++ { + _ = utils.Exec(b, conn, insStmt) + } + }) + } +} diff --git a/go/test/endtoend/vtgate/queries/benchmark/main_test.go b/go/test/endtoend/vtgate/queries/benchmark/main_test.go new file mode 100644 index 00000000000..6978d0b9428 --- /dev/null +++ b/go/test/endtoend/vtgate/queries/benchmark/main_test.go @@ -0,0 +1,124 @@ +/* +Copyright 2023 The Vitess Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package dml + +import ( + "context" + _ "embed" + "flag" + "os" + "testing" + + "github.com/stretchr/testify/require" + + "vitess.io/vitess/go/mysql" + "vitess.io/vitess/go/test/endtoend/cluster" + "vitess.io/vitess/go/test/endtoend/utils" +) + +var ( + clusterInstance *cluster.LocalProcessCluster + vtParams mysql.ConnParams + mysqlParams mysql.ConnParams + sKs = "sks" + uKs = "uks" + cell = "test" + + //go:embed sharded_schema.sql + sSchemaSQL string + + //go:embed vschema.json + sVSchema string +) + +var ( + shards4 = []string{ + "-40", "40-80", "80-c0", "c0-", + } + + shards8 = []string{ + "-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-", + } + + shards16 = []string{ + "-10", "10-20", "20-30", "30-40", "40-50", "50-60", "60-70", "70-80", "80-90", "90-a0", "a0-b0", "b0-c0", "c0-d0", "d0-e0", "e0-f0", "f0-", + } + + shards32 = []string{ + "-05", "05-10", "10-15", "15-20", "20-25", "25-30", "30-35", "35-40", "40-45", "45-50", "50-55", "55-60", "60-65", "65-70", "70-75", "75-80", + "80-85", "85-90", "90-95", "95-a0", "a0-a5", "a5-b0", "b0-b5", "b5-c0", "c0-c5", "c5-d0", "d0-d5", "d5-e0", "e0-e5", "e5-f0", "f0-f5", "f5-", + } +) + +func TestMain(m *testing.M) { + defer cluster.PanicHandler(nil) + flag.Parse() + + exitCode := func() int { + clusterInstance = cluster.NewCluster(cell, "localhost") + defer clusterInstance.Teardown() + + // Start topo server + err := clusterInstance.StartTopo() + if err != nil { + return 1 + } + + // Start sharded keyspace + sKeyspace := &cluster.Keyspace{ + Name: sKs, + SchemaSQL: sSchemaSQL, + VSchema: sVSchema, + } + + err = clusterInstance.StartKeyspace(*sKeyspace, shards4, 0, false) + if err != nil { + return 1 + } + + // Start vtgate + err = clusterInstance.StartVtgate() + if err != nil { + return 1 + } + + vtParams = clusterInstance.GetVTParams(sKs) + + return m.Run() + }() + os.Exit(exitCode) +} + +func start(b *testing.B) (*mysql.Conn, func()) { + conn, err := mysql.Connect(context.Background(), &vtParams) + require.NoError(b, err) + + deleteAll := func() { + tables := []string{"tbl_no_lkp_vdx"} + for _, table := range tables { + _, _ = utils.ExecAllowError(b, conn, "delete from "+table) + } + } + + deleteAll() + + return conn, func() { + deleteAll() + conn.Close() + cluster.PanicHandler(b) + } +} diff --git a/go/test/endtoend/vtgate/queries/benchmark/sharded_schema.sql b/go/test/endtoend/vtgate/queries/benchmark/sharded_schema.sql new file mode 100644 index 00000000000..850b6ffc15a --- /dev/null +++ b/go/test/endtoend/vtgate/queries/benchmark/sharded_schema.sql @@ -0,0 +1,16 @@ +create table tbl_no_lkp_vdx +( + id bigint, + c1 varchar(50), + c2 varchar(50), + c3 varchar(50), + c4 varchar(50), + c5 varchar(50), + c6 varchar(50), + c7 varchar(50), + c8 varchar(50), + c9 varchar(50), + c10 varchar(50), + c11 varchar(50), + c12 varchar(50) +) Engine = InnoDB; \ No newline at end of file diff --git a/go/test/endtoend/vtgate/queries/benchmark/vschema.json b/go/test/endtoend/vtgate/queries/benchmark/vschema.json new file mode 100644 index 00000000000..4970e8b7437 --- /dev/null +++ b/go/test/endtoend/vtgate/queries/benchmark/vschema.json @@ -0,0 +1,18 @@ +{ + "sharded": true, + "vindexes": { + "xxhash": { + "type": "xxhash" + } + }, + "tables": { + "tbl_no_lkp_vdx": { + "column_vindexes": [ + { + "column": "id", + "name": "xxhash" + } + ] + } + } +} \ No newline at end of file diff --git a/go/vt/vtgate/autocommit_test.go b/go/vt/vtgate/autocommit_test.go index 0d55bbf2875..fa63695bfbd 100644 --- a/go/vt/vtgate/autocommit_test.go +++ b/go/vt/vtgate/autocommit_test.go @@ -274,7 +274,6 @@ func TestAutocommitInsertLookup(t *testing.T) { BindVariables: map[string]*querypb.BindVariable{ "_Id_0": sqltypes.Int64BindVariable(1), "_name_0": sqltypes.StringBindVariable("myname"), - "__seq0": sqltypes.Int64BindVariable(1), }, }}) testCommitCount(t, "sbc1", sbc1, 1) @@ -292,7 +291,6 @@ func TestAutocommitInsertMultishardAutoCommit(t *testing.T) { Sql: "insert /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */ into user_extra(user_id, v) values (:_user_id_0, 2)", BindVariables: map[string]*querypb.BindVariable{ "_user_id_0": sqltypes.Int64BindVariable(1), - "_user_id_1": sqltypes.Int64BindVariable(3), }, }}) testCommitCount(t, "sbc1", sbc1, 0) @@ -300,7 +298,6 @@ func TestAutocommitInsertMultishardAutoCommit(t *testing.T) { assertQueries(t, sbc2, []*querypb.BoundQuery{{ Sql: "insert /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */ into user_extra(user_id, v) values (:_user_id_1, 4)", BindVariables: map[string]*querypb.BindVariable{ - "_user_id_0": sqltypes.Int64BindVariable(1), "_user_id_1": sqltypes.Int64BindVariable(3), }, }}) @@ -321,7 +318,6 @@ func TestAutocommitInsertMultishardAutoCommit(t *testing.T) { assertQueries(t, sbc2, []*querypb.BoundQuery{{ Sql: "insert /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */ into user_extra(user_id, v) values (:_user_id_1, 4)", BindVariables: map[string]*querypb.BindVariable{ - "_user_id_0": sqltypes.Int64BindVariable(1), "_user_id_1": sqltypes.Int64BindVariable(3), }, }}) @@ -339,7 +335,6 @@ func TestAutocommitInsertMultishard(t *testing.T) { Sql: "insert into user_extra(user_id, v) values (:_user_id_0, 2)", BindVariables: map[string]*querypb.BindVariable{ "_user_id_0": sqltypes.Int64BindVariable(1), - "_user_id_1": sqltypes.Int64BindVariable(3), }, }}) testCommitCount(t, "sbc1", sbc1, 1) @@ -347,7 +342,6 @@ func TestAutocommitInsertMultishard(t *testing.T) { assertQueries(t, sbc2, []*querypb.BoundQuery{{ Sql: "insert into user_extra(user_id, v) values (:_user_id_1, 4)", BindVariables: map[string]*querypb.BindVariable{ - "_user_id_0": sqltypes.Int64BindVariable(1), "_user_id_1": sqltypes.Int64BindVariable(3), }, }}) diff --git a/go/vt/vtgate/engine/cached_size.go b/go/vt/vtgate/engine/cached_size.go index dcaefd270ed..b2d8913f6be 100644 --- a/go/vt/vtgate/engine/cached_size.go +++ b/go/vt/vtgate/engine/cached_size.go @@ -417,11 +417,18 @@ func (cached *Insert) CachedSize(alloc bool) int64 { size += cached.Generate.CachedSize(true) // field Prefix string size += hack.RuntimeAllocSize(int64(len(cached.Prefix))) - // field Mid []string + // field Mid vitess.io/vitess/go/vt/sqlparser.Values { - size += hack.RuntimeAllocSize(int64(cap(cached.Mid)) * int64(16)) + size += hack.RuntimeAllocSize(int64(cap(cached.Mid)) * int64(24)) for _, elem := range cached.Mid { - size += hack.RuntimeAllocSize(int64(len(elem))) + { + size += hack.RuntimeAllocSize(int64(cap(elem)) * int64(16)) + for _, elem := range elem { + if cc, ok := elem.(cachedObject); ok { + size += cc.CachedSize(true) + } + } + } } } // field Suffix string diff --git a/go/vt/vtgate/engine/insert.go b/go/vt/vtgate/engine/insert.go index 394ccb8ecce..212bcd6620c 100644 --- a/go/vt/vtgate/engine/insert.go +++ b/go/vt/vtgate/engine/insert.go @@ -25,6 +25,7 @@ import ( "sync" "time" + "vitess.io/vitess/go/slice" "vitess.io/vitess/go/sqltypes" "vitess.io/vitess/go/vt/key" querypb "vitess.io/vitess/go/vt/proto/query" @@ -75,7 +76,7 @@ type ( // Prefix, Mid and Suffix are for sharded insert plans. Prefix string - Mid []string + Mid sqlparser.Values Suffix string // Option to override the standard behavior and allow a multi-shard insert @@ -132,7 +133,7 @@ func NewInsert( vindexValues [][][]evalengine.Expr, table *vindexes.Table, prefix string, - mid []string, + mid sqlparser.Values, suffix string, ) *Insert { ins := &Insert{ @@ -758,17 +759,23 @@ func (ins *Insert) getInsertShardedRoute( queries := make([]*querypb.BoundQuery, len(rss)) for i := range rss { + shardBindVars := map[string]*querypb.BindVariable{} var mids []string for _, indexValue := range indexesPerRss[i] { index, _ := strconv.ParseInt(string(indexValue.Value), 0, 64) if keyspaceIDs[index] != nil { - mids = append(mids, ins.Mid[index]) + mids = append(mids, sqlparser.String(ins.Mid[index])) + for _, expr := range ins.Mid[index] { + if arg, ok := expr.(*sqlparser.Argument); ok { + shardBindVars[arg.Name] = bindVars[arg.Name] + } + } } } rewritten := ins.Prefix + strings.Join(mids, ",") + ins.Suffix queries[i] = &querypb.BoundQuery{ Sql: rewritten, - BindVariables: bindVars, + BindVariables: shardBindVars, } } @@ -985,7 +992,10 @@ func (ins *Insert) description() PrimitiveDescription { other["VindexOffsetFromSelect"] = valuesOffsets } if len(ins.Mid) > 0 { - shardQuery := fmt.Sprintf("%s%s%s", ins.Prefix, strings.Join(ins.Mid, ", "), ins.Suffix) + mids := slice.Map(ins.Mid, func(from sqlparser.ValTuple) string { + return sqlparser.String(from) + }) + shardQuery := fmt.Sprintf("%s%s%s", ins.Prefix, strings.Join(mids, ", "), ins.Suffix) if shardQuery != ins.Query { other["ShardedQuery"] = shardQuery } diff --git a/go/vt/vtgate/engine/insert_test.go b/go/vt/vtgate/engine/insert_test.go index b651efe2b03..014654f37d6 100644 --- a/go/vt/vtgate/engine/insert_test.go +++ b/go/vt/vtgate/engine/insert_test.go @@ -21,17 +21,15 @@ import ( "errors" "testing" - topodatapb "vitess.io/vitess/go/vt/proto/topodata" - - "vitess.io/vitess/go/vt/vtgate/evalengine" - "github.com/stretchr/testify/require" "vitess.io/vitess/go/sqltypes" - "vitess.io/vitess/go/vt/vtgate/vindexes" - querypb "vitess.io/vitess/go/vt/proto/query" + topodatapb "vitess.io/vitess/go/vt/proto/topodata" vschemapb "vitess.io/vitess/go/vt/proto/vschema" + "vitess.io/vitess/go/vt/sqlparser" + "vitess.io/vitess/go/vt/vtgate/evalengine" + "vitess.io/vitess/go/vt/vtgate/vindexes" ) func TestInsertUnsharded(t *testing.T) { @@ -212,7 +210,9 @@ func TestInsertShardedSimple(t *testing.T) { }}, ks.Tables["t1"], "prefix", - []string{" mid1"}, + sqlparser.Values{ + {&sqlparser.Argument{Name: "_id_0", Type: sqltypes.Int64}}, + }, " suffix", ) vc := newDMLTestVCursor("-20", "20-") @@ -227,7 +227,7 @@ func TestInsertShardedSimple(t *testing.T) { `ResolveDestinations sharded [value:"0"] Destinations:DestinationKeyspaceID(166b40b44aba4bd6)`, // Row 2 will go to -20, rows 1 & 3 will go to 20- `ExecuteMultiShard ` + - `sharded.20-: prefix mid1 suffix {_id_0: type:INT64 value:"1"} ` + + `sharded.20-: prefix(:_id_0 /* INT64 */) suffix {_id_0: type:INT64 value:"1"} ` + `true true`, }) @@ -247,7 +247,11 @@ func TestInsertShardedSimple(t *testing.T) { }}, ks.Tables["t1"], "prefix", - []string{" mid1", " mid2", " mid3"}, + sqlparser.Values{ + {&sqlparser.Argument{Name: "_id_0", Type: sqltypes.Int64}}, + {&sqlparser.Argument{Name: "_id_1", Type: sqltypes.Int64}}, + {&sqlparser.Argument{Name: "_id_2", Type: sqltypes.Int64}}, + }, " suffix", ) vc = newDMLTestVCursor("-20", "20-") @@ -262,8 +266,8 @@ func TestInsertShardedSimple(t *testing.T) { `ResolveDestinations sharded [value:"0" value:"1" value:"2"] Destinations:DestinationKeyspaceID(166b40b44aba4bd6),DestinationKeyspaceID(06e7ea22ce92708f),DestinationKeyspaceID(4eb190c9a2fa169c)`, // Row 2 will go to -20, rows 1 & 3 will go to 20- `ExecuteMultiShard ` + - `sharded.20-: prefix mid1, mid3 suffix {_id_0: type:INT64 value:"1" _id_1: type:INT64 value:"2" _id_2: type:INT64 value:"3"} ` + - `sharded.-20: prefix mid2 suffix {_id_0: type:INT64 value:"1" _id_1: type:INT64 value:"2" _id_2: type:INT64 value:"3"} ` + + `sharded.20-: prefix(:_id_0 /* INT64 */),(:_id_2 /* INT64 */) suffix {_id_0: type:INT64 value:"1" _id_2: type:INT64 value:"3"} ` + + `sharded.-20: prefix(:_id_1 /* INT64 */) suffix {_id_1: type:INT64 value:"2"} ` + `true false`, }) @@ -284,7 +288,11 @@ func TestInsertShardedSimple(t *testing.T) { ks.Tables["t1"], "prefix", - []string{" mid1", " mid2", " mid3"}, + sqlparser.Values{ + {&sqlparser.Argument{Name: "_id_0", Type: sqltypes.Int64}}, + {&sqlparser.Argument{Name: "_id_1", Type: sqltypes.Int64}}, + {&sqlparser.Argument{Name: "_id_2", Type: sqltypes.Int64}}, + }, " suffix", ) ins.MultiShardAutocommit = true @@ -301,8 +309,8 @@ func TestInsertShardedSimple(t *testing.T) { `ResolveDestinations sharded [value:"0" value:"1" value:"2"] Destinations:DestinationKeyspaceID(166b40b44aba4bd6),DestinationKeyspaceID(06e7ea22ce92708f),DestinationKeyspaceID(4eb190c9a2fa169c)`, // Row 2 will go to -20, rows 1 & 3 will go to 20- `ExecuteMultiShard ` + - `sharded.20-: prefix mid1, mid3 suffix {_id_0: type:INT64 value:"1" _id_1: type:INT64 value:"2" _id_2: type:INT64 value:"3"} ` + - `sharded.-20: prefix mid2 suffix {_id_0: type:INT64 value:"1" _id_1: type:INT64 value:"2" _id_2: type:INT64 value:"3"} ` + + `sharded.20-: prefix(:_id_0 /* INT64 */),(:_id_2 /* INT64 */) suffix {_id_0: type:INT64 value:"1" _id_2: type:INT64 value:"3"} ` + + `sharded.-20: prefix(:_id_1 /* INT64 */) suffix {_id_1: type:INT64 value:"2"} ` + `true true`, }) } @@ -349,7 +357,9 @@ func TestInsertShardedFail(t *testing.T) { ks.Tables["t1"], "prefix", - []string{" mid1", " mid2", " mid3"}, + sqlparser.Values{ + {&sqlparser.Argument{Name: "_id_0", Type: sqltypes.Int64}}, + }, " suffix", ) @@ -399,7 +409,11 @@ func TestInsertShardedGenerate(t *testing.T) { }}, ks.Tables["t1"], "prefix", - []string{" mid1", " mid2", " mid3"}, + sqlparser.Values{ + {&sqlparser.Argument{Name: "__seq0", Type: sqltypes.Int64}}, + {&sqlparser.Argument{Name: "__seq1", Type: sqltypes.Int64}}, + {&sqlparser.Argument{Name: "__seq2", Type: sqltypes.Int64}}, + }, " suffix", ) @@ -412,7 +426,7 @@ func TestInsertShardedGenerate(t *testing.T) { Values: evalengine.NewTupleExpr( evalengine.NewLiteralInt(1), evalengine.NullExpr, - evalengine.NewLiteralInt(2), + evalengine.NewLiteralInt(3), ), } @@ -440,12 +454,10 @@ func TestInsertShardedGenerate(t *testing.T) { `ResolveDestinations sharded [value:"0" value:"1" value:"2"] Destinations:DestinationKeyspaceID(166b40b44aba4bd6),DestinationKeyspaceID(06e7ea22ce92708f),DestinationKeyspaceID(4eb190c9a2fa169c)`, // Row 2 will go to -20, rows 1 & 3 will go to 20- `ExecuteMultiShard ` + - `sharded.20-: prefix mid1, mid3 suffix ` + - `{__seq0: type:INT64 value:"1" __seq1: type:INT64 value:"2" __seq2: type:INT64 value:"2" ` + - `_id_0: type:INT64 value:"1" _id_1: type:INT64 value:"2" _id_2: type:INT64 value:"3"} ` + - `sharded.-20: prefix mid2 suffix ` + - `{__seq0: type:INT64 value:"1" __seq1: type:INT64 value:"2" __seq2: type:INT64 value:"2" ` + - `_id_0: type:INT64 value:"1" _id_1: type:INT64 value:"2" _id_2: type:INT64 value:"3"} ` + + `sharded.20-: prefix(:__seq0 /* INT64 */),(:__seq2 /* INT64 */) suffix ` + + `{__seq0: type:INT64 value:"1" __seq2: type:INT64 value:"3"} ` + + `sharded.-20: prefix(:__seq1 /* INT64 */) suffix ` + + `{__seq1: type:INT64 value:"2"} ` + `true false`, }) @@ -535,7 +547,11 @@ func TestInsertShardedOwned(t *testing.T) { }}, ks.Tables["t1"], "prefix", - []string{" mid1", " mid2", " mid3"}, + sqlparser.Values{ + {&sqlparser.Argument{Name: "_id_0", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c1_0", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c2_0", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c3_0", Type: sqltypes.Int64}}, + {&sqlparser.Argument{Name: "_id_1", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c1_1", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c2_1", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c3_1", Type: sqltypes.Int64}}, + {&sqlparser.Argument{Name: "_id_2", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c1_2", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c2_2", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c3_2", Type: sqltypes.Int64}}, + }, " suffix", ) @@ -557,16 +573,15 @@ func TestInsertShardedOwned(t *testing.T) { // Based on shardForKsid, values returned will be 20-, -20, 20-. `ResolveDestinations sharded [value:"0" value:"1" value:"2"] Destinations:DestinationKeyspaceID(166b40b44aba4bd6),DestinationKeyspaceID(06e7ea22ce92708f),DestinationKeyspaceID(4eb190c9a2fa169c)`, `ExecuteMultiShard ` + - `sharded.20-: prefix mid1, mid3 suffix ` + - `{_c1_0: type:INT64 value:"4" _c1_1: type:INT64 value:"5" _c1_2: type:INT64 value:"6" ` + - `_c2_0: type:INT64 value:"7" _c2_1: type:INT64 value:"8" _c2_2: type:INT64 value:"9" ` + - `_c3_0: type:INT64 value:"10" _c3_1: type:INT64 value:"11" _c3_2: type:INT64 value:"12" ` + - `_id_0: type:INT64 value:"1" _id_1: type:INT64 value:"2" _id_2: type:INT64 value:"3"} ` + - `sharded.-20: prefix mid2 suffix ` + - `{_c1_0: type:INT64 value:"4" _c1_1: type:INT64 value:"5" _c1_2: type:INT64 value:"6" ` + - `_c2_0: type:INT64 value:"7" _c2_1: type:INT64 value:"8" _c2_2: type:INT64 value:"9" ` + - `_c3_0: type:INT64 value:"10" _c3_1: type:INT64 value:"11" _c3_2: type:INT64 value:"12" ` + - `_id_0: type:INT64 value:"1" _id_1: type:INT64 value:"2" _id_2: type:INT64 value:"3"} ` + + `sharded.20-: prefix(:_id_0 /* INT64 */, :_c1_0 /* INT64 */, :_c2_0 /* INT64 */, :_c3_0 /* INT64 */)` + + `,(:_id_2 /* INT64 */, :_c1_2 /* INT64 */, :_c2_2 /* INT64 */, :_c3_2 /* INT64 */) suffix ` + + `{_c1_0: type:INT64 value:"4" _c1_2: type:INT64 value:"6" ` + + `_c2_0: type:INT64 value:"7" _c2_2: type:INT64 value:"9" ` + + `_c3_0: type:INT64 value:"10" _c3_2: type:INT64 value:"12" ` + + `_id_0: type:INT64 value:"1" _id_2: type:INT64 value:"3"} ` + + `sharded.-20: prefix(:_id_1 /* INT64 */, :_c1_1 /* INT64 */, :_c2_1 /* INT64 */, :_c3_1 /* INT64 */) suffix ` + + `{_c1_1: type:INT64 value:"5" _c2_1: type:INT64 value:"8" _c3_1: type:INT64 value:"11" ` + + `_id_1: type:INT64 value:"2"} ` + `true false`, }) } @@ -626,7 +641,9 @@ func TestInsertShardedOwnedWithNull(t *testing.T) { }}, ks.Tables["t1"], "prefix", - []string{" mid1", " mid2", " mid3"}, + sqlparser.Values{ + {&sqlparser.Argument{Name: "_id_0", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c3_0", Type: sqltypes.Null}}, + }, " suffix", ) @@ -639,7 +656,7 @@ func TestInsertShardedOwnedWithNull(t *testing.T) { } vc.ExpectLog(t, []string{ `ResolveDestinations sharded [value:"0"] Destinations:DestinationKeyspaceID(166b40b44aba4bd6)`, - `ExecuteMultiShard sharded.20-: prefix mid1 suffix ` + + `ExecuteMultiShard sharded.20-: prefix(:_id_0 /* INT64 */, :_c3_0 /* NULL_TYPE */) suffix ` + `{_c3_0: _id_0: type:INT64 value:"1"} true true`, }) } @@ -709,7 +726,10 @@ func TestInsertShardedGeo(t *testing.T) { }}, ks.Tables["t1"], "prefix", - []string{" mid1", " mid2"}, + sqlparser.Values{ + {&sqlparser.Argument{Name: "_region_0", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_id_0", Type: sqltypes.Int64}}, + {&sqlparser.Argument{Name: "_region_1", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_id_1", Type: sqltypes.Int64}}, + }, " suffix", ) @@ -725,12 +745,10 @@ func TestInsertShardedGeo(t *testing.T) { `id_0: type:INT64 value:"1" id_1: type:INT64 value:"1" ` + `keyspace_id_0: type:VARBINARY value:"\x01\x16k@\xb4J\xbaK\xd6" keyspace_id_1: type:VARBINARY value:"\xff\x16k@\xb4J\xbaK\xd6" true`, `ResolveDestinations sharded [value:"0" value:"1"] Destinations:DestinationKeyspaceID(01166b40b44aba4bd6),DestinationKeyspaceID(ff166b40b44aba4bd6)`, - `ExecuteMultiShard sharded.20-: prefix mid1 suffix ` + - `{_id_0: type:INT64 value:"1" _id_1: type:INT64 value:"1" ` + - `_region_0: type:INT64 value:"1" _region_1: type:INT64 value:"255"} ` + - `sharded.-20: prefix mid2 suffix ` + - `{_id_0: type:INT64 value:"1" _id_1: type:INT64 value:"1" ` + - `_region_0: type:INT64 value:"1" _region_1: type:INT64 value:"255"} ` + + `ExecuteMultiShard sharded.20-: prefix(:_region_0 /* INT64 */, :_id_0 /* INT64 */) suffix ` + + `{_id_0: type:INT64 value:"1" _region_0: type:INT64 value:"1"} ` + + `sharded.-20: prefix(:_region_1 /* INT64 */, :_id_1 /* INT64 */) suffix ` + + `{_id_1: type:INT64 value:"1" _region_1: type:INT64 value:"255"} ` + `true false`, }) } @@ -830,7 +848,12 @@ func TestInsertShardedIgnoreOwned(t *testing.T) { }}, ks.Tables["t1"], "prefix", - []string{" mid1", " mid2", " mid3", " mid4"}, + sqlparser.Values{ + {&sqlparser.Argument{Name: "_id_0", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c1_0", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c2_0", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c3_0", Type: sqltypes.Int64}}, + {&sqlparser.Argument{Name: "_id_1", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c1_1", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c2_1", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c3_1", Type: sqltypes.Int64}}, + {&sqlparser.Argument{Name: "_id_2", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c1_2", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c2_2", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c3_2", Type: sqltypes.Int64}}, + {&sqlparser.Argument{Name: "_id_3", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c1_3", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c2_3", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c3_3", Type: sqltypes.Int64}}, + }, " suffix", ) @@ -894,16 +917,10 @@ func TestInsertShardedIgnoreOwned(t *testing.T) { `ResolveDestinations sharded [value:"0" value:"3"] Destinations:DestinationKeyspaceID(00),DestinationKeyspaceID(00)`, // Bind vars for rows 2 & 3 may be missing because they were not sent. `ExecuteMultiShard ` + - `sharded.20-: prefix mid1 suffix ` + - `{_c1_0: type:INT64 value:"5" _c1_3: type:INT64 value:"8" ` + - `_c2_0: type:INT64 value:"9" _c2_3: type:INT64 value:"12" ` + - `_c3_0: type:INT64 value:"13" _c3_3: type:INT64 value:"16" ` + - `_id_0: type:INT64 value:"1" _id_3: type:INT64 value:"4"} ` + - `sharded.-20: prefix mid4 suffix ` + - `{_c1_0: type:INT64 value:"5" _c1_3: type:INT64 value:"8" ` + - `_c2_0: type:INT64 value:"9" _c2_3: type:INT64 value:"12" ` + - `_c3_0: type:INT64 value:"13" _c3_3: type:INT64 value:"16" ` + - `_id_0: type:INT64 value:"1" _id_3: type:INT64 value:"4"} ` + + `sharded.20-: prefix(:_id_0 /* INT64 */, :_c1_0 /* INT64 */, :_c2_0 /* INT64 */, :_c3_0 /* INT64 */) suffix ` + + `{_c1_0: type:INT64 value:"5" _c2_0: type:INT64 value:"9" _c3_0: type:INT64 value:"13" _id_0: type:INT64 value:"1"} ` + + `sharded.-20: prefix(:_id_3 /* INT64 */, :_c1_3 /* INT64 */, :_c2_3 /* INT64 */, :_c3_3 /* INT64 */) suffix ` + + `{_c1_3: type:INT64 value:"8" _c2_3: type:INT64 value:"12" _c3_3: type:INT64 value:"16" _id_3: type:INT64 value:"4"} ` + `true false`, }) } @@ -964,7 +981,9 @@ func TestInsertShardedIgnoreOwnedWithNull(t *testing.T) { }}, ks.Tables["t1"], "prefix", - []string{" mid1", " mid2", " mid3", " mid4"}, + sqlparser.Values{ + {&sqlparser.Argument{Name: "_id_0", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c3_0", Type: sqltypes.Int64}}, + }, " suffix", ) @@ -990,7 +1009,7 @@ func TestInsertShardedIgnoreOwnedWithNull(t *testing.T) { vc.ExpectLog(t, []string{ `Execute select from from lkp1 where from = :from and toc = :toc from: toc: type:VARBINARY value:"\x16k@\xb4J\xbaK\xd6" false`, `ResolveDestinations sharded [value:"0"] Destinations:DestinationKeyspaceID(166b40b44aba4bd6)`, - `ExecuteMultiShard sharded.-20: prefix mid1 suffix ` + + `ExecuteMultiShard sharded.-20: prefix(:_id_0 /* INT64 */, :_c3_0 /* INT64 */) suffix ` + `{_c3_0: _id_0: type:INT64 value:"1"} true true`, }) } @@ -1078,7 +1097,11 @@ func TestInsertShardedUnownedVerify(t *testing.T) { }}, ks.Tables["t1"], "prefix", - []string{" mid1", " mid2", " mid3"}, + sqlparser.Values{ + {&sqlparser.Argument{Name: "_id_0", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c1_0", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c2_0", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c3_0", Type: sqltypes.Int64}}, + {&sqlparser.Argument{Name: "_id_1", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c1_1", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c2_1", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c3_1", Type: sqltypes.Int64}}, + {&sqlparser.Argument{Name: "_id_2", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c1_2", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c2_2", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c3_2", Type: sqltypes.Int64}}, + }, " suffix", ) @@ -1117,16 +1140,15 @@ func TestInsertShardedUnownedVerify(t *testing.T) { // Based on shardForKsid, values returned will be 20-, -20, 20-. `ResolveDestinations sharded [value:"0" value:"1" value:"2"] Destinations:DestinationKeyspaceID(166b40b44aba4bd6),DestinationKeyspaceID(06e7ea22ce92708f),DestinationKeyspaceID(4eb190c9a2fa169c)`, `ExecuteMultiShard ` + - `sharded.20-: prefix mid1, mid3 suffix ` + - `{_c1_0: type:INT64 value:"4" _c1_1: type:INT64 value:"5" _c1_2: type:INT64 value:"6" ` + - `_c2_0: type:INT64 value:"7" _c2_1: type:INT64 value:"8" _c2_2: type:INT64 value:"9" ` + - `_c3_0: type:INT64 value:"10" _c3_1: type:INT64 value:"11" _c3_2: type:INT64 value:"12" ` + - `_id_0: type:INT64 value:"1" _id_1: type:INT64 value:"2" _id_2: type:INT64 value:"3"} ` + - `sharded.-20: prefix mid2 suffix ` + - `{_c1_0: type:INT64 value:"4" _c1_1: type:INT64 value:"5" _c1_2: type:INT64 value:"6" ` + - `_c2_0: type:INT64 value:"7" _c2_1: type:INT64 value:"8" _c2_2: type:INT64 value:"9" ` + - `_c3_0: type:INT64 value:"10" _c3_1: type:INT64 value:"11" _c3_2: type:INT64 value:"12" ` + - `_id_0: type:INT64 value:"1" _id_1: type:INT64 value:"2" _id_2: type:INT64 value:"3"} ` + + `sharded.20-: prefix(:_id_0 /* INT64 */, :_c1_0 /* INT64 */, :_c2_0 /* INT64 */, :_c3_0 /* INT64 */),` + + `(:_id_2 /* INT64 */, :_c1_2 /* INT64 */, :_c2_2 /* INT64 */, :_c3_2 /* INT64 */) suffix ` + + `{_c1_0: type:INT64 value:"4" _c1_2: type:INT64 value:"6" ` + + `_c2_0: type:INT64 value:"7" _c2_2: type:INT64 value:"9" ` + + `_c3_0: type:INT64 value:"10" _c3_2: type:INT64 value:"12" ` + + `_id_0: type:INT64 value:"1" _id_2: type:INT64 value:"3"} ` + + `sharded.-20: prefix(:_id_1 /* INT64 */, :_c1_1 /* INT64 */, :_c2_1 /* INT64 */, :_c3_1 /* INT64 */) suffix ` + + `{_c1_1: type:INT64 value:"5" _c2_1: type:INT64 value:"8" ` + + `_c3_1: type:INT64 value:"11" _id_1: type:INT64 value:"2"} ` + `true false`, }) } @@ -1189,7 +1211,11 @@ func TestInsertShardedIgnoreUnownedVerify(t *testing.T) { }}, ks.Tables["t1"], "prefix", - []string{" mid1", " mid2", " mid3"}, + sqlparser.Values{ + {&sqlparser.Argument{Name: "_id_0", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c3_0", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "v1", Type: sqltypes.VarChar}}, + {&sqlparser.Argument{Name: "_id_1", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c3_1", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "v2", Type: sqltypes.VarChar}}, + {&sqlparser.Argument{Name: "_id_2", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c3_2", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "v3", Type: sqltypes.VarChar}}, + }, " suffix", ) @@ -1210,7 +1236,9 @@ func TestInsertShardedIgnoreUnownedVerify(t *testing.T) { {}, nonemptyResult, } - _, err := ins.TryExecute(context.Background(), vc, map[string]*querypb.BindVariable{}, false) + _, err := ins.TryExecute(context.Background(), vc, map[string]*querypb.BindVariable{ + "v1": sqltypes.StringBindVariable("a"), "v2": sqltypes.StringBindVariable("b"), "v3": sqltypes.StringBindVariable("c"), + }, false) if err != nil { t.Fatal(err) } @@ -1223,12 +1251,10 @@ func TestInsertShardedIgnoreUnownedVerify(t *testing.T) { // Based on shardForKsid, values returned will be 20-, -20. `ResolveDestinations sharded [value:"0" value:"2"] Destinations:DestinationKeyspaceID(166b40b44aba4bd6),DestinationKeyspaceID(4eb190c9a2fa169c)`, `ExecuteMultiShard ` + - `sharded.20-: prefix mid1 suffix ` + - `{_c3_0: type:INT64 value:"10" _c3_2: type:INT64 value:"12" ` + - `_id_0: type:INT64 value:"1" _id_2: type:INT64 value:"3"} ` + - `sharded.-20: prefix mid3 suffix ` + - `{_c3_0: type:INT64 value:"10" _c3_2: type:INT64 value:"12" ` + - `_id_0: type:INT64 value:"1" _id_2: type:INT64 value:"3"} ` + + `sharded.20-: prefix(:_id_0 /* INT64 */, :_c3_0 /* INT64 */, :v1 /* VARCHAR */) suffix ` + + `{_c3_0: type:INT64 value:"10" _id_0: type:INT64 value:"1" v1: type:VARCHAR value:"a"} ` + + `sharded.-20: prefix(:_id_2 /* INT64 */, :_c3_2 /* INT64 */, :v3 /* VARCHAR */) suffix ` + + `{_c3_2: type:INT64 value:"12" _id_2: type:INT64 value:"3" v3: type:VARCHAR value:"c"} ` + `true false`, }) } @@ -1287,7 +1313,9 @@ func TestInsertShardedIgnoreUnownedVerifyFail(t *testing.T) { }}, ks.Tables["t1"], "prefix", - []string{" mid1", " mid2", " mid3"}, + sqlparser.Values{ + {&sqlparser.Argument{Name: "_id_0", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c3_0", Type: sqltypes.Int64}}, + }, " suffix", ) @@ -1380,7 +1408,11 @@ func TestInsertShardedUnownedReverseMap(t *testing.T) { }}, ks.Tables["t1"], "prefix", - []string{" mid1", " mid2", " mid3"}, + sqlparser.Values{ + {&sqlparser.Argument{Name: "_id_0", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c1_0", Type: sqltypes.Null}, &sqlparser.Argument{Name: "_c2_0", Type: sqltypes.Null}, &sqlparser.Argument{Name: "_c3_0", Type: sqltypes.Null}}, + {&sqlparser.Argument{Name: "_id_1", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c1_1", Type: sqltypes.Null}, &sqlparser.Argument{Name: "_c2_1", Type: sqltypes.Null}, &sqlparser.Argument{Name: "_c3_1", Type: sqltypes.Null}}, + {&sqlparser.Argument{Name: "_id_2", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c1_2", Type: sqltypes.Null}, &sqlparser.Argument{Name: "_c2_2", Type: sqltypes.Null}, &sqlparser.Argument{Name: "_c3_2", Type: sqltypes.Null}}, + }, " suffix", ) @@ -1405,18 +1437,16 @@ func TestInsertShardedUnownedReverseMap(t *testing.T) { } vc.ExpectLog(t, []string{ `ResolveDestinations sharded [value:"0" value:"1" value:"2"] Destinations:DestinationKeyspaceID(166b40b44aba4bd6),DestinationKeyspaceID(06e7ea22ce92708f),DestinationKeyspaceID(4eb190c9a2fa169c)`, - `ExecuteMultiShard ` + - `sharded.20-: prefix mid1, mid3 suffix ` + - `{_c1_0: type:UINT64 value:"1" _c1_1: type:UINT64 value:"2" _c1_2: type:UINT64 value:"3" ` + - `_c2_0: _c2_1: _c2_2: ` + - `_c3_0: type:UINT64 value:"1" _c3_1: type:UINT64 value:"2" _c3_2: type:UINT64 value:"3" ` + - `_id_0: type:INT64 value:"1" _id_1: type:INT64 value:"2" _id_2: type:INT64 value:"3"} ` + - `sharded.-20: prefix mid2 suffix ` + - `{_c1_0: type:UINT64 value:"1" _c1_1: type:UINT64 value:"2" _c1_2: type:UINT64 value:"3" ` + - `_c2_0: _c2_1: _c2_2: ` + - `_c3_0: type:UINT64 value:"1" _c3_1: type:UINT64 value:"2" _c3_2: type:UINT64 value:"3" ` + - `_id_0: type:INT64 value:"1" _id_1: type:INT64 value:"2" _id_2: type:INT64 value:"3"} ` + - `true false`, + `ExecuteMultiShard sharded.20-: ` + + `prefix(:_id_0 /* INT64 */, :_c1_0 /* NULL_TYPE */, :_c2_0 /* NULL_TYPE */, :_c3_0 /* NULL_TYPE */),` + + `(:_id_2 /* INT64 */, :_c1_2 /* NULL_TYPE */, :_c2_2 /* NULL_TYPE */, :_c3_2 /* NULL_TYPE */) suffix ` + + `{_c1_0: type:UINT64 value:"1" _c1_2: type:UINT64 value:"3" ` + + `_c2_0: _c2_2: ` + + `_c3_0: type:UINT64 value:"1" _c3_2: type:UINT64 value:"3" ` + + `_id_0: type:INT64 value:"1" _id_2: type:INT64 value:"3"} ` + + `sharded.-20: ` + + `prefix(:_id_1 /* INT64 */, :_c1_1 /* NULL_TYPE */, :_c2_1 /* NULL_TYPE */, :_c3_1 /* NULL_TYPE */) suffix ` + + `{_c1_1: type:UINT64 value:"2" _c2_1: _c3_1: type:UINT64 value:"2" _id_1: type:INT64 value:"2"} true false`, }) } @@ -1474,7 +1504,9 @@ func TestInsertShardedUnownedReverseMapSuccess(t *testing.T) { }}, ks.Tables["t1"], "prefix", - []string{" mid1", " mid2", " mid3"}, + sqlparser.Values{ + {&sqlparser.Argument{Name: "_id_0", Type: sqltypes.Int64}, &sqlparser.Argument{Name: "_c3_0", Type: sqltypes.Null}}, + }, " suffix", ) diff --git a/go/vt/vtgate/executor_dml_test.go b/go/vt/vtgate/executor_dml_test.go index 59fbe314346..c53d24eb6fb 100644 --- a/go/vt/vtgate/executor_dml_test.go +++ b/go/vt/vtgate/executor_dml_test.go @@ -1273,7 +1273,6 @@ func TestInsertSharded(t *testing.T) { BindVariables: map[string]*querypb.BindVariable{ "_Id_0": sqltypes.Int64BindVariable(1), "_name_0": sqltypes.StringBindVariable("myname"), - "__seq0": sqltypes.Int64BindVariable(1), }, }} assertQueries(t, sbc1, wantQueries) @@ -1299,7 +1298,6 @@ func TestInsertSharded(t *testing.T) { Sql: "insert into `user`(id, v, `name`) values (:_Id_0, 2, :_name_0)", BindVariables: map[string]*querypb.BindVariable{ "_Id_0": sqltypes.Int64BindVariable(3), - "__seq0": sqltypes.Int64BindVariable(3), "_name_0": sqltypes.StringBindVariable("myname2"), }, }} @@ -1345,10 +1343,7 @@ func TestInsertSharded(t *testing.T) { BindVariables: map[string]*querypb.BindVariable{ "_Id_0": sqltypes.Int64BindVariable(1), "_name_0": sqltypes.BytesBindVariable([]byte("myname")), - "__seq0": sqltypes.Int64BindVariable(1), - "vtg1": sqltypes.Int64BindVariable(1), "vtg2": sqltypes.Int64BindVariable(2), - "vtg3": sqltypes.StringBindVariable("myname"), }, }} assertQueries(t, sbc1, wantQueries) @@ -1448,7 +1443,6 @@ func TestInsertShardedAutocommitLookup(t *testing.T) { "_Id_0": sqltypes.Int64BindVariable(1), "_music_0": sqltypes.StringBindVariable("star"), "_name_0": sqltypes.StringBindVariable("myname"), - "__seq0": sqltypes.Int64BindVariable(1), }, }} assertQueries(t, sbc1, wantQueries) @@ -1510,27 +1504,18 @@ func TestInsertShardedIgnore(t *testing.T) { BindVariables: map[string]*querypb.BindVariable{ "_pv_0": sqltypes.Int64BindVariable(1), "_pv_4": sqltypes.Int64BindVariable(5), - "_pv_5": sqltypes.Int64BindVariable(6), "_owned_0": sqltypes.Int64BindVariable(1), "_owned_4": sqltypes.Int64BindVariable(5), - "_owned_5": sqltypes.Int64BindVariable(6), "_verify_0": sqltypes.Int64BindVariable(1), "_verify_4": sqltypes.Int64BindVariable(1), - "_verify_5": sqltypes.Int64BindVariable(3), }, }} assertQueries(t, sbc1, wantQueries) wantQueries = []*querypb.BoundQuery{{ Sql: "insert ignore into insert_ignore_test(pv, owned, verify) values (:_pv_5, :_owned_5, :_verify_5)", BindVariables: map[string]*querypb.BindVariable{ - "_pv_0": sqltypes.Int64BindVariable(1), - "_pv_4": sqltypes.Int64BindVariable(5), "_pv_5": sqltypes.Int64BindVariable(6), - "_owned_0": sqltypes.Int64BindVariable(1), - "_owned_4": sqltypes.Int64BindVariable(5), "_owned_5": sqltypes.Int64BindVariable(6), - "_verify_0": sqltypes.Int64BindVariable(1), - "_verify_4": sqltypes.Int64BindVariable(1), "_verify_5": sqltypes.Int64BindVariable(3), }, }} @@ -1700,7 +1685,6 @@ func TestInsertComments(t *testing.T) { BindVariables: map[string]*querypb.BindVariable{ "_Id_0": sqltypes.Int64BindVariable(1), "_name_0": sqltypes.StringBindVariable("myname"), - "__seq0": sqltypes.Int64BindVariable(1), }, }} assertQueries(t, sbc1, wantQueries) @@ -1734,7 +1718,6 @@ func TestInsertGeneratorSharded(t *testing.T) { Sql: "insert into `user`(v, `name`, id) values (2, :_name_0, :_Id_0)", BindVariables: map[string]*querypb.BindVariable{ "_Id_0": sqltypes.Int64BindVariable(1), - "__seq0": sqltypes.Int64BindVariable(1), "_name_0": sqltypes.StringBindVariable("myname"), }, }} @@ -1856,7 +1839,6 @@ func TestInsertLookupOwned(t *testing.T) { BindVariables: map[string]*querypb.BindVariable{ "_user_id_0": sqltypes.Int64BindVariable(2), "_id_0": sqltypes.Int64BindVariable(3), - "__seq0": sqltypes.Int64BindVariable(3), }, }} assertQueries(t, sbc, wantQueries) @@ -1890,7 +1872,6 @@ func TestInsertLookupOwnedGenerator(t *testing.T) { BindVariables: map[string]*querypb.BindVariable{ "_user_id_0": sqltypes.Int64BindVariable(2), "_id_0": sqltypes.Int64BindVariable(4), - "__seq0": sqltypes.Int64BindVariable(4), }, }} assertQueries(t, sbc, wantQueries) @@ -2020,7 +2001,6 @@ func TestInsertPartialFail2(t *testing.T) { Sql: "insert into `user`(id, v, `name`) values (:_Id_0, 2, :_name_0)", BindVariables: map[string]*querypb.BindVariable{ "_Id_0": sqltypes.Int64BindVariable(1), - "__seq0": sqltypes.Int64BindVariable(1), "_name_0": sqltypes.StringBindVariable("myname"), }, }, { @@ -2043,22 +2023,14 @@ func TestMultiInsertSharded(t *testing.T) { BindVariables: map[string]*querypb.BindVariable{ "_Id_0": sqltypes.Int64BindVariable(1), "_name_0": sqltypes.StringBindVariable("myname1"), - "__seq0": sqltypes.Int64BindVariable(1), - "_Id_1": sqltypes.Int64BindVariable(3), - "_name_1": sqltypes.StringBindVariable("myname3"), - "__seq1": sqltypes.Int64BindVariable(3), }, }} wantQueries2 := []*querypb.BoundQuery{{ Sql: "insert into `user`(id, v, `name`) values (:_Id_1, 3, :_name_1)", BindVariables: map[string]*querypb.BindVariable{ - "_Id_0": sqltypes.Int64BindVariable(1), - "_name_0": sqltypes.StringBindVariable("myname1"), - "__seq0": sqltypes.Int64BindVariable(1), "_Id_1": sqltypes.Int64BindVariable(3), "_name_1": sqltypes.StringBindVariable("myname3"), - "__seq1": sqltypes.Int64BindVariable(3), }, }} assertQueries(t, sbc1, wantQueries1) @@ -2084,10 +2056,8 @@ func TestMultiInsertSharded(t *testing.T) { Sql: "insert into `user`(id, v, `name`) values (:_Id_0, 1, :_name_0),(:_Id_1, 2, :_name_1)", BindVariables: map[string]*querypb.BindVariable{ "_Id_0": sqltypes.Int64BindVariable(1), - "__seq0": sqltypes.Int64BindVariable(1), "_name_0": sqltypes.StringBindVariable("myname1"), "_Id_1": sqltypes.Int64BindVariable(2), - "__seq1": sqltypes.Int64BindVariable(2), "_name_1": sqltypes.StringBindVariable("myname2"), }, }} @@ -2117,12 +2087,18 @@ func TestMultiInsertSharded(t *testing.T) { "_id_0": sqltypes.Int64BindVariable(2), "_name_0": sqltypes.StringBindVariable("myname"), "_lastname_0": sqltypes.StringBindVariable("mylastname"), + }, + }} + assertQueries(t, sbc1, wantQueries) + wantQueries = []*querypb.BoundQuery{{ + Sql: "insert into user2(id, `name`, lastname) values (:_id_1, :_name_1, :_lastname_1)", + BindVariables: map[string]*querypb.BindVariable{ "_id_1": sqltypes.Int64BindVariable(3), "_name_1": sqltypes.StringBindVariable("myname2"), "_lastname_1": sqltypes.StringBindVariable("mylastname2"), }, }} - assertQueries(t, sbc1, wantQueries) + assertQueries(t, sbc2, wantQueries) wantQueries = []*querypb.BoundQuery{{ Sql: "insert into name_lastname_keyspace_id_map(`name`, lastname, keyspace_id) values (:name_0, :lastname_0, :keyspace_id_0), (:name_1, :lastname_1, :keyspace_id_1)", BindVariables: map[string]*querypb.BindVariable{ @@ -2155,12 +2131,9 @@ func TestMultiInsertGenerator(t *testing.T) { wantQueries := []*querypb.BoundQuery{{ Sql: "insert into music(user_id, `name`, id) values (:_user_id_0, 'myname1', :_id_0),(:_user_id_1, 'myname2', :_id_1)", BindVariables: map[string]*querypb.BindVariable{ - "u": sqltypes.Int64BindVariable(2), "_id_0": sqltypes.Int64BindVariable(1), - "__seq0": sqltypes.Int64BindVariable(1), "_user_id_0": sqltypes.Int64BindVariable(2), "_id_1": sqltypes.Int64BindVariable(2), - "__seq1": sqltypes.Int64BindVariable(2), "_user_id_1": sqltypes.Int64BindVariable(2), }, }} @@ -2203,15 +2176,11 @@ func TestMultiInsertGeneratorSparse(t *testing.T) { wantQueries := []*querypb.BoundQuery{{ Sql: "insert into music(id, user_id, `name`) values (:_id_0, :_user_id_0, 'myname1'),(:_id_1, :_user_id_1, 'myname2'),(:_id_2, :_user_id_2, 'myname3')", BindVariables: map[string]*querypb.BindVariable{ - "u": sqltypes.Int64BindVariable(2), "_id_0": sqltypes.Int64BindVariable(1), - "__seq0": sqltypes.Int64BindVariable(1), "_user_id_0": sqltypes.Int64BindVariable(2), "_id_1": sqltypes.Int64BindVariable(2), - "__seq1": sqltypes.Int64BindVariable(2), "_user_id_1": sqltypes.Int64BindVariable(2), "_id_2": sqltypes.Int64BindVariable(2), - "__seq2": sqltypes.Int64BindVariable(2), "_user_id_2": sqltypes.Int64BindVariable(2), }, }} @@ -2735,12 +2704,8 @@ func TestPartialVindexInsertQueryFailure(t *testing.T) { }, { Sql: "insert into t1_lkp_idx(unq_col, keyspace_id) values (:_unq_col_0, :keyspace_id_0)", BindVariables: map[string]*querypb.BindVariable{ - "unq_col_0": sqltypes.Int64BindVariable(1), "keyspace_id_0": sqltypes.BytesBindVariable([]byte("\x16k@\xb4J\xbaK\xd6")), - "unq_col_1": sqltypes.Int64BindVariable(3), - "keyspace_id_1": sqltypes.BytesBindVariable([]byte("\x06\xe7\xea\"Βp\x8f")), "_unq_col_0": sqltypes.Int64BindVariable(1), - "_unq_col_1": sqltypes.Int64BindVariable(3), }, }, { Sql: "rollback to x", @@ -2757,6 +2722,10 @@ func TestPartialVindexInsertQueryFailure(t *testing.T) { // only parameter in expected query changes wantQ[1].Sql = "insert into t1_lkp_idx(unq_col, keyspace_id) values (:_unq_col_1, :keyspace_id_1)" + wantQ[1].BindVariables = map[string]*querypb.BindVariable{ + "keyspace_id_1": sqltypes.BytesBindVariable([]byte("\x06\xe7\xea\"Βp\x8f")), + "_unq_col_1": sqltypes.Int64BindVariable(3), + } assertQueriesWithSavepoint(t, sbc2, wantQ) testQueryLog(t, executor, logChan, "TestExecute", "BEGIN", "begin", 0) @@ -2780,12 +2749,8 @@ func TestPartialVindexInsertQueryFailureAutoCommit(t *testing.T) { wantQ := []*querypb.BoundQuery{{ Sql: "insert into t1_lkp_idx(unq_col, keyspace_id) values (:_unq_col_0, :keyspace_id_0)", BindVariables: map[string]*querypb.BindVariable{ - "unq_col_0": sqltypes.Int64BindVariable(1), "keyspace_id_0": sqltypes.BytesBindVariable([]byte("\x16k@\xb4J\xbaK\xd6")), - "unq_col_1": sqltypes.Int64BindVariable(3), - "keyspace_id_1": sqltypes.BytesBindVariable([]byte("\x06\xe7\xea\"Βp\x8f")), "_unq_col_0": sqltypes.Int64BindVariable(1), - "_unq_col_1": sqltypes.Int64BindVariable(3), }, }} @@ -2799,6 +2764,10 @@ func TestPartialVindexInsertQueryFailureAutoCommit(t *testing.T) { // only parameter in expected query changes wantQ[0].Sql = "insert into t1_lkp_idx(unq_col, keyspace_id) values (:_unq_col_1, :keyspace_id_1)" + wantQ[0].BindVariables = map[string]*querypb.BindVariable{ + "keyspace_id_1": sqltypes.BytesBindVariable([]byte("\x06\xe7\xea\"Βp\x8f")), + "_unq_col_1": sqltypes.Int64BindVariable(3), + } assertQueriesWithSavepoint(t, sbc2, wantQ) testQueryLog(t, executor, logChan, "VindexCreate", "INSERT", "insert into t1_lkp_idx(unq_col, keyspace_id) values (:unq_col_0, :keyspace_id_0), (:unq_col_1, :keyspace_id_1)", 2) @@ -2827,7 +2796,6 @@ func TestMultiInternalSavepoint(t *testing.T) { Sql: "insert into user_extra(user_id) values (:_user_id_0)", BindVariables: map[string]*querypb.BindVariable{ "_user_id_0": sqltypes.Int64BindVariable(1), - "_user_id_1": sqltypes.Int64BindVariable(4), }, }} assertQueriesWithSavepoint(t, sbc1, wantQ) @@ -2846,7 +2814,6 @@ func TestMultiInternalSavepoint(t *testing.T) { Sql: "insert into user_extra(user_id) values (:_user_id_0)", BindVariables: map[string]*querypb.BindVariable{ "_user_id_0": sqltypes.Int64BindVariable(3), - "_user_id_1": sqltypes.Int64BindVariable(6), }, }} assertQueriesWithSavepoint(t, sbc2, wantQ) diff --git a/go/vt/vtgate/executor_vschema_ddl_test.go b/go/vt/vtgate/executor_vschema_ddl_test.go index 4218dcf93e0..1c2813a33c4 100644 --- a/go/vt/vtgate/executor_vschema_ddl_test.go +++ b/go/vt/vtgate/executor_vschema_ddl_test.go @@ -17,8 +17,9 @@ limitations under the License. package vtgate import ( + "context" "reflect" - "sort" + "slices" "testing" "time" @@ -78,13 +79,19 @@ func waitForVschemaTables(t *testing.T, ks string, tables []string, executor *Ex // Wait up to 100ms until the vindex manager gets notified of the update for i := 0; i < 10; i++ { vschema := executor.vm.GetCurrentSrvVschema() - gotTables := []string{} + var gotTables []string for t := range vschema.Keyspaces[ks].Tables { gotTables = append(gotTables, t) } - sort.Strings(tables) - sort.Strings(gotTables) - if reflect.DeepEqual(tables, gotTables) { + + foundAll := true + for _, expTbl := range tables { + if !slices.Contains(gotTables, expTbl) { + foundAll = false + break + } + } + if foundAll { return vschema } time.Sleep(10 * time.Millisecond) @@ -399,30 +406,32 @@ func TestExecutorDropSequenceDDL(t *testing.T) { t.Fatalf("test_seq should not exist in original vschema") } - var vschemaTables []string - for t := range vschema.Keyspaces[ks].Tables { - vschemaTables = append(vschemaTables, t) - } - session := NewSafeSession(&vtgatepb.Session{TargetString: ks}) // add test sequence stmt := "alter vschema add sequence test_seq" _, err := executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) - _ = waitForVschemaTables(t, ks, append(vschemaTables, []string{"test_seq"}...), executor) + _ = waitForVschemaTables(t, ks, []string{"test_seq"}, executor) vschema = executor.vm.GetCurrentSrvVschema() table := vschema.Keyspaces[ks].Tables["test_seq"] wantType := "sequence" - if table.Type != wantType { - t.Errorf("want table type sequence got %v", table) - } + require.Equal(t, wantType, table.Type) + + // note the last vschema updated time. + ts := executor.VSchema().GetCreated() // drop existing test sequence stmt = "alter vschema drop sequence test_seq" _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) + ctxWithTimeout, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + if !waitForNewerVSchema(ctxWithTimeout, executor, ts) { + t.Fatalf("vschema did not drop the sequene 'test_seq'") + } + // Should fail dropping a non-existing test sequence session = NewSafeSession(&vtgatepb.Session{TargetString: ks}) @@ -449,21 +458,33 @@ func TestExecutorDropAutoIncDDL(t *testing.T) { _, err := executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) + _ = waitForVschemaTables(t, ks, []string{"test_table"}, executor) + ts := executor.VSchema().GetCreated() + stmt = "alter vschema on test_table add auto_increment id using `db-name`.`test_seq`" _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) + ctxWithTimeout, cancel := context.WithTimeout(ctx, 5*time.Second) + defer cancel() + if !waitForNewerVSchema(ctxWithTimeout, executor, ts) { + t.Fatalf("vschema did not update with auto_increment for 'test_table'") + } + ts = executor.VSchema().GetCreated() wantAutoInc := &vschemapb.AutoIncrement{Column: "id", Sequence: "`db-name`.test_seq"} gotAutoInc := executor.vm.GetCurrentSrvVschema().Keyspaces[ks].Tables["test_table"].AutoIncrement - if !reflect.DeepEqual(wantAutoInc, gotAutoInc) { - t.Errorf("want autoinc %v, got autoinc %v", wantAutoInc, gotAutoInc) - } + utils.MustMatch(t, wantAutoInc, gotAutoInc) stmt = "alter vschema on test_table drop auto_increment" _, err = executor.Execute(ctx, nil, "TestExecute", session, stmt, nil) require.NoError(t, err) + ctxWithTimeout, cancel2 := context.WithTimeout(ctx, 5*time.Second) + defer cancel2() + if !waitForNewerVSchema(ctxWithTimeout, executor, ts) { + t.Fatalf("vschema did not drop the auto_increment for 'test_table'") + } if executor.vm.GetCurrentSrvVschema().Keyspaces[ks].Tables["test_table"].AutoIncrement != nil { t.Errorf("auto increment should be nil after drop") } diff --git a/go/vt/vtgate/planbuilder/operator_transformers.go b/go/vt/vtgate/planbuilder/operator_transformers.go index 04ce68564c0..f2896f72525 100644 --- a/go/vt/vtgate/planbuilder/operator_transformers.go +++ b/go/vt/vtgate/planbuilder/operator_transformers.go @@ -542,8 +542,8 @@ func autoIncGenerate(gen *operators.Generate) *engine.Generate { } } -func generateInsertShardedQuery(ins *sqlparser.Insert) (prefix string, mid []string, suffix string) { - valueTuples, isValues := ins.Rows.(sqlparser.Values) +func generateInsertShardedQuery(ins *sqlparser.Insert) (prefix string, mids sqlparser.Values, suffix string) { + mids, isValues := ins.Rows.(sqlparser.Values) prefixFormat := "insert %v%sinto %v%v " if isValues { // the mid values are filled differently @@ -560,20 +560,6 @@ func generateInsertShardedQuery(ins *sqlparser.Insert) (prefix string, mid []str suffixBuf := sqlparser.NewTrackedBuffer(dmlFormatter) suffixBuf.Myprintf("%v", ins.OnDup) suffix = suffixBuf.String() - - if !isValues { - // this is a insert query using select to insert the rows. - return - } - - midBuf := sqlparser.NewTrackedBuffer(dmlFormatter) - mid = make([]string, len(valueTuples)) - for rowNum, val := range valueTuples { - midBuf.Myprintf("%v", val) - mid[rowNum] = midBuf.String() - midBuf.Reset() - } - return } diff --git a/go/vt/vtgate/vtgate_test.go b/go/vt/vtgate/vtgate_test.go index 6f21158d7bb..c113ed16308 100644 --- a/go/vt/vtgate/vtgate_test.go +++ b/go/vt/vtgate/vtgate_test.go @@ -626,14 +626,14 @@ func TestMultiInternalSavepointVtGate(t *testing.T) { Sql: "insert into sp_tbl(user_id) values (:_user_id_0)", BindVariables: map[string]*querypb.BindVariable{ "_user_id_0": sqltypes.Int64BindVariable(1), - "_user_id_1": sqltypes.Int64BindVariable(3), - "vtg1": sqltypes.Int64BindVariable(1), - "vtg2": sqltypes.Int64BindVariable(3), }, }} assertQueriesWithSavepoint(t, sbc1, wantQ) wantQ[1].Sql = "insert into sp_tbl(user_id) values (:_user_id_1)" + wantQ[1].BindVariables = map[string]*querypb.BindVariable{ + "_user_id_1": sqltypes.Int64BindVariable(3), + } assertQueriesWithSavepoint(t, sbc2, wantQ) assert.Len(t, sbc3.Queries, 0) // internal savepoint should be removed. @@ -650,10 +650,7 @@ func TestMultiInternalSavepointVtGate(t *testing.T) { }, { Sql: "insert into sp_tbl(user_id) values (:_user_id_1)", BindVariables: map[string]*querypb.BindVariable{ - "_user_id_0": sqltypes.Int64BindVariable(2), "_user_id_1": sqltypes.Int64BindVariable(4), - "vtg1": sqltypes.Int64BindVariable(2), - "vtg2": sqltypes.Int64BindVariable(4), }, }} assertQueriesWithSavepoint(t, sbc3, wantQ) @@ -669,7 +666,6 @@ func TestMultiInternalSavepointVtGate(t *testing.T) { Sql: "insert into sp_tbl(user_id) values (:_user_id_0)", BindVariables: map[string]*querypb.BindVariable{ "_user_id_0": sqltypes.Int64BindVariable(5), - "vtg1": sqltypes.Int64BindVariable(5), }, }} assertQueriesWithSavepoint(t, sbc2, wantQ)