From d1134043da9c2447054e8c8d8007e35a993758f7 Mon Sep 17 00:00:00 2001
From: Lloyd-Pottiger <60744015+Lloyd-Pottiger@users.noreply.github.com>
Date: Tue, 27 Aug 2024 11:37:47 +0800
Subject: [PATCH] storage: Add keyspace_id to DT system tables (#9371)

ref pingcap/tiflash#9032

Signed-off-by: Wish <breezewish@outlook.com>

Co-authored-by: Wenxuan <breezewish@outlook.com>
Co-authored-by: JaySon <tshent@qq.com>
---
 .../Storages/System/StorageSystemDTSegments.cpp   | 14 ++++++++++++--
 .../src/Storages/System/StorageSystemDTTables.cpp | 15 +++++++++++++--
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/dbms/src/Storages/System/StorageSystemDTSegments.cpp b/dbms/src/Storages/System/StorageSystemDTSegments.cpp
index c10165141c2..0092d984d49 100644
--- a/dbms/src/Storages/System/StorageSystemDTSegments.cpp
+++ b/dbms/src/Storages/System/StorageSystemDTSegments.cpp
@@ -14,6 +14,7 @@
 
 #include <Columns/ColumnString.h>
 #include <DataStreams/OneBlockInputStream.h>
+#include <DataTypes/DataTypeNullable.h>
 #include <DataTypes/DataTypeString.h>
 #include <DataTypes/DataTypesNumber.h>
 #include <Databases/DatabaseTiFlash.h>
@@ -37,6 +38,7 @@ StorageSystemDTSegments::StorageSystemDTSegments(const std::string & name_)
 
         {"tidb_database", std::make_shared<DataTypeString>()},
         {"tidb_table", std::make_shared<DataTypeString>()},
+        {"keyspace_id", std::make_shared<DataTypeNullable>(std::make_shared<DataTypeUInt64>())},
         {"table_id", std::make_shared<DataTypeInt64>()},
         {"is_tombstone", std::make_shared<DataTypeUInt64>()},
 
@@ -115,11 +117,19 @@ BlockInputStreams StorageSystemDTSegments::read(
                 res_columns[j++]->insert(table_name);
 
                 String tidb_db_name;
+                KeyspaceID keyspace_id = NullspaceID;
                 if (db_tiflash)
-                    tidb_db_name = mapper.displayDatabaseName(db_tiflash->getDatabaseInfo());
+                {
+                    tidb_db_name = db_tiflash->getDatabaseInfo().name;
+                    keyspace_id = db_tiflash->getDatabaseInfo().keyspace_id;
+                }
                 res_columns[j++]->insert(tidb_db_name);
-                String tidb_table_name = mapper.displayTableName(table_info);
+                String tidb_table_name = table_info.name;
                 res_columns[j++]->insert(tidb_table_name);
+                if (keyspace_id == NullspaceID)
+                    res_columns[j++]->insert(Field());
+                else
+                    res_columns[j++]->insert(static_cast<UInt64>(keyspace_id));
                 res_columns[j++]->insert(table_id);
                 res_columns[j++]->insert(dm_storage->getTombstone());
 
diff --git a/dbms/src/Storages/System/StorageSystemDTTables.cpp b/dbms/src/Storages/System/StorageSystemDTTables.cpp
index e4dbdbc2f64..2a0d45957a1 100644
--- a/dbms/src/Storages/System/StorageSystemDTTables.cpp
+++ b/dbms/src/Storages/System/StorageSystemDTTables.cpp
@@ -14,12 +14,14 @@
 
 #include <Columns/ColumnString.h>
 #include <DataStreams/OneBlockInputStream.h>
+#include <DataTypes/DataTypeNullable.h>
 #include <DataTypes/DataTypeString.h>
 #include <DataTypes/DataTypesNumber.h>
 #include <Databases/DatabaseTiFlash.h>
 #include <Databases/IDatabase.h>
 #include <Interpreters/Context.h>
 #include <Storages/DeltaMerge/DeltaMergeStore.h>
+#include <Storages/KVStore/Types.h>
 #include <Storages/MutableSupport.h>
 #include <Storages/StorageDeltaMerge.h>
 #include <Storages/System/StorageSystemDTTables.h>
@@ -37,6 +39,7 @@ StorageSystemDTTables::StorageSystemDTTables(const std::string & name_)
 
         {"tidb_database", std::make_shared<DataTypeString>()},
         {"tidb_table", std::make_shared<DataTypeString>()},
+        {"keyspace_id", std::make_shared<DataTypeNullable>(std::make_shared<DataTypeUInt64>())},
         {"table_id", std::make_shared<DataTypeInt64>()},
         {"is_tombstone", std::make_shared<DataTypeUInt64>()},
 
@@ -146,11 +149,19 @@ BlockInputStreams StorageSystemDTTables::read(
             res_columns[j++]->insert(table_name);
 
             String tidb_db_name;
+            KeyspaceID keyspace_id = NullspaceID;
             if (db_tiflash)
-                tidb_db_name = mapper.displayDatabaseName(db_tiflash->getDatabaseInfo());
+            {
+                tidb_db_name = db_tiflash->getDatabaseInfo().name;
+                keyspace_id = db_tiflash->getDatabaseInfo().keyspace_id;
+            }
             res_columns[j++]->insert(tidb_db_name);
-            String tidb_table_name = mapper.displayTableName(table_info);
+            String tidb_table_name = table_info.name;
             res_columns[j++]->insert(tidb_table_name);
+            if (keyspace_id == NullspaceID)
+                res_columns[j++]->insert(Field());
+            else
+                res_columns[j++]->insert(static_cast<UInt64>(keyspace_id));
             res_columns[j++]->insert(table_id);
             res_columns[j++]->insert(dm_storage->getTombstone());