From a5882f6d2809bc95c6c174b56d2a1b4ee2ad3d52 Mon Sep 17 00:00:00 2001 From: Hangjie Mo Date: Tue, 22 Nov 2022 15:38:33 +0800 Subject: [PATCH 1/2] *: fix tidb_decode_key for partition table --- expression/integration_test.go | 16 ++++++++++++++++ planner/core/expression_rewriter.go | 26 +++++++++++++++++++++++--- 2 files changed, 39 insertions(+), 3 deletions(-) diff --git a/expression/integration_test.go b/expression/integration_test.go index 305ed81295a3f..a11285ac41d1c 100644 --- a/expression/integration_test.go +++ b/expression/integration_test.go @@ -3050,6 +3050,22 @@ func TestTiDBDecodeKeyFunc(t *testing.T) { sql = fmt.Sprintf("select tidb_decode_key( '%s' )", hexKey) rs = fmt.Sprintf(`{"%s":%d,"table_id":"%d"}`, tbl.Meta().GetPkName().String(), rowID, tbl.Meta().ID) tk.MustQuery(sql).Check(testkit.Rows(rs)) + + // Test partition table. + tk.MustExec("drop table if exists t;") + tk.MustExec("create table t (a int primary key clustered, b int, key bk (b)) PARTITION BY RANGE (a) (PARTITION p0 VALUES LESS THAN (1), PARTITION p1 VALUES LESS THAN (2));") + dom = domain.GetDomain(tk.Session()) + is = dom.InfoSchema() + tbl, err = is.TableByName(model.NewCIStr("test"), model.NewCIStr("t")) + require.NoError(t, err) + require.NotNil(t, tbl.Meta().Partition) + hexKey = buildTableRowKey(tbl.Meta().Partition.Definitions[0].ID, rowID) + sql = fmt.Sprintf("select tidb_decode_key( '%s' )", hexKey) + rs = fmt.Sprintf(`{"%s":%d,"partition_id":%d,"table_id":"%d"}`, tbl.Meta().GetPkName().String(), rowID, tbl.Meta().Partition.Definitions[0].ID, tbl.Meta().ID) + hexKey = tablecodec.EncodeTablePrefix(tbl.Meta().Partition.Definitions[0].ID).String() + sql = fmt.Sprintf("select tidb_decode_key( '%s' )", hexKey) + rs = fmt.Sprintf(`{"partition_id":%d,"table_id":%d}`, tbl.Meta().Partition.Definitions[0].ID, tbl.Meta().ID) + tk.MustQuery(sql).Check(testkit.Rows(rs)) } func TestTwoDecimalTruncate(t *testing.T) { diff --git a/planner/core/expression_rewriter.go b/planner/core/expression_rewriter.go index 2e22a752df310..522dedacca24a 100644 --- a/planner/core/expression_rewriter.go +++ b/planner/core/expression_rewriter.go @@ -2169,6 +2169,9 @@ func decodeKeyFromString(ctx sessionctx.Context, s string) string { return s } tbl, _ := is.TableByID(tableID) + if tbl == nil { + tbl, _, _ = is.FindTableByPartitionID(tableID) + } loc := ctx.GetSessionVars().Location() if tablecodec.IsRecordKey(key) { ret, err := decodeRecordKey(key, tableID, tbl, loc) @@ -2185,7 +2188,7 @@ func decodeKeyFromString(ctx sessionctx.Context, s string) string { } return ret } else if tablecodec.IsTableKey(key) { - ret, err := decodeTableKey(key, tableID) + ret, err := decodeTableKey(key, tableID, tbl) if err != nil { sc.AppendWarning(err) return s @@ -2203,6 +2206,10 @@ func decodeRecordKey(key []byte, tableID int64, tbl table.Table, loc *time.Locat } if handle.IsInt() { ret := make(map[string]interface{}) + if tbl != nil && tbl.Meta().Partition != nil { + ret["partition_id"] = tableID + tableID = tbl.Meta().ID + } ret["table_id"] = strconv.FormatInt(tableID, 10) // When the clustered index is enabled, we should show the PK name. if tbl != nil && tbl.Meta().HasClusteredIndex() { @@ -2239,6 +2246,10 @@ func decodeRecordKey(key []byte, tableID int64, tbl table.Table, loc *time.Locat return "", errors.Trace(err) } ret := make(map[string]interface{}) + if tbl.Meta().Partition != nil { + ret["partition_id"] = tableID + tableID = tbl.Meta().ID + } ret["table_id"] = tableID handleRet := make(map[string]interface{}) for colID := range datumMap { @@ -2308,6 +2319,10 @@ func decodeIndexKey(key []byte, tableID int64, tbl table.Table, loc *time.Locati ds = append(ds, d) } ret := make(map[string]interface{}) + if tbl.Meta().Partition != nil { + ret["partition_id"] = tableID + tableID = tbl.Meta().ID + } ret["table_id"] = tableID ret["index_id"] = indexID idxValMap := make(map[string]interface{}, len(targetIndex.Columns)) @@ -2340,8 +2355,13 @@ func decodeIndexKey(key []byte, tableID int64, tbl table.Table, loc *time.Locati return string(retStr), nil } -func decodeTableKey(_ []byte, tableID int64) (string, error) { - ret := map[string]int64{"table_id": tableID} +func decodeTableKey(_ []byte, tableID int64, tbl table.Table) (string, error) { + ret := map[string]int64{} + if tbl != nil && tbl.Meta().GetPartitionInfo() != nil { + ret["partition_id"] = tableID + tableID = tbl.Meta().ID + } + ret["table_id"] = tableID retStr, err := json.Marshal(ret) if err != nil { return "", errors.Trace(err) From 95c0e0ce4f06b3274381e2f11dd63914e05f0aba Mon Sep 17 00:00:00 2001 From: Hangjie Mo Date: Tue, 22 Nov 2022 15:58:49 +0800 Subject: [PATCH 2/2] update test --- expression/integration_test.go | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/expression/integration_test.go b/expression/integration_test.go index a11285ac41d1c..83aaea0a8675d 100644 --- a/expression/integration_test.go +++ b/expression/integration_test.go @@ -3062,10 +3062,18 @@ func TestTiDBDecodeKeyFunc(t *testing.T) { hexKey = buildTableRowKey(tbl.Meta().Partition.Definitions[0].ID, rowID) sql = fmt.Sprintf("select tidb_decode_key( '%s' )", hexKey) rs = fmt.Sprintf(`{"%s":%d,"partition_id":%d,"table_id":"%d"}`, tbl.Meta().GetPkName().String(), rowID, tbl.Meta().Partition.Definitions[0].ID, tbl.Meta().ID) + tk.MustQuery(sql).Check(testkit.Rows(rs)) + hexKey = tablecodec.EncodeTablePrefix(tbl.Meta().Partition.Definitions[0].ID).String() sql = fmt.Sprintf("select tidb_decode_key( '%s' )", hexKey) rs = fmt.Sprintf(`{"partition_id":%d,"table_id":%d}`, tbl.Meta().Partition.Definitions[0].ID, tbl.Meta().ID) tk.MustQuery(sql).Check(testkit.Rows(rs)) + + data = []types.Datum{types.NewIntDatum(100)} + hexKey = buildIndexKeyFromData(tbl.Meta().Partition.Definitions[0].ID, tbl.Indices()[0].Meta().ID, data) + sql = fmt.Sprintf("select tidb_decode_key( '%s' )", hexKey) + rs = fmt.Sprintf(`{"index_id":1,"index_vals":{"b":"100"},"partition_id":%d,"table_id":%d}`, tbl.Meta().Partition.Definitions[0].ID, tbl.Meta().ID) + tk.MustQuery(sql).Check(testkit.Rows(rs)) } func TestTwoDecimalTruncate(t *testing.T) {