feat : support set tiflash mode ddl action

dbms/src/Debug/DBGInvoker.cpp

@@ -120,6 +120,8 @@ DBGInvoker::DBGInvoker()
     regSchemafulFunc("query_mapped", dbgFuncQueryMapped);
     regSchemalessFunc("get_tiflash_replica_count", dbgFuncGetTiflashReplicaCount);
     regSchemalessFunc("get_partition_tables_tiflash_replica_count", dbgFuncGetPartitionTablesTiflashReplicaCount);
+    regSchemalessFunc("get_tiflash_mode", dbgFuncGetTiflashMode);
+    regSchemalessFunc("get_partition_tables_tiflash_mode", dbgFuncGetPartitionTablesTiflashMode);
 
     regSchemalessFunc("search_log_for_key", dbgFuncSearchLogForKey);
     regSchemalessFunc("tidb_dag", dbgFuncTiDBQueryFromNaturalDag);

dbms/src/Debug/dbgFuncSchemaName.cpp

@@ -181,4 +181,56 @@ void dbgFuncGetPartitionTablesTiflashReplicaCount(Context & context, const ASTs
     output(fmt_buf.toString());
 }
 
+void dbgFuncGetTiflashMode(Context & context, const ASTs & args, DBGInvoker::Printer output)
+{
+    if (args.empty() || args.size() != 2)
+        throw Exception("Args not matched, should be: database-name[, table-name]", ErrorCodes::BAD_ARGUMENTS);
+
+    const String & database_name = typeid_cast<const ASTIdentifier &>(*args[0]).name;
+    FmtBuffer fmt_buf;
+
+    const String & table_name = typeid_cast<const ASTIdentifier &>(*args[1]).name;
+    auto mapped = mappedTable(context, database_name, table_name);
+    auto storage = context.getTable(mapped->first, mapped->second);
+    auto managed_storage = std::dynamic_pointer_cast<IManageableStorage>(storage);
+    if (!managed_storage)
+        throw Exception(database_name + "." + table_name + " is not ManageableStorage", ErrorCodes::BAD_ARGUMENTS);
+
+    fmt_buf.append((TiFlashModeToString(managed_storage->getTableInfo().tiflash_mode)));
+
+    output(fmt_buf.toString());
+}
+
+void dbgFuncGetPartitionTablesTiflashMode(Context & context, const ASTs & args, DBGInvoker::Printer output)
+{
+    if (args.empty() || args.size() != 2)
+        throw Exception("Args not matched, should be: database-name[, table-name]", ErrorCodes::BAD_ARGUMENTS);
+
+    const String & database_name = typeid_cast<const ASTIdentifier &>(*args[0]).name;
+    FmtBuffer fmt_buf;
+
+    const String & table_name = typeid_cast<const ASTIdentifier &>(*args[1]).name;
+    auto mapped = mappedTable(context, database_name, table_name);
+    auto storage = context.getTable(mapped->first, mapped->second);
+    auto managed_storage = std::dynamic_pointer_cast<IManageableStorage>(storage);
+    if (!managed_storage)
+        throw Exception(database_name + "." + table_name + " is not ManageableStorage", ErrorCodes::BAD_ARGUMENTS);
+
+    auto table_info = managed_storage->getTableInfo();
+
+    if (!table_info.isLogicalPartitionTable())
+        throw Exception(database_name + "." + table_name + " is not logical partition table", ErrorCodes::BAD_ARGUMENTS);
+
+    SchemaNameMapper name_mapper;
+    for (const auto & part_def : table_info.partition.definitions)
+    {
+        auto paritition_table_info = table_info.producePartitionTableInfo(part_def.id, name_mapper);
+        auto partition_storage = context.getTMTContext().getStorages().get(paritition_table_info->id);
+        fmt_buf.append((TiFlashModeToString(partition_storage->getTableInfo().tiflash_mode)));
+        fmt_buf.append("/");
+    }
+
+    output(fmt_buf.toString());
+}
+
 } // namespace DB

dbms/src/Debug/dbgFuncSchemaName.h

@@ -50,4 +50,14 @@ void dbgFuncGetTiflashReplicaCount(Context & context, const ASTs & args, DBGInvo
 //   ./storage-client.sh "DBGInvoke get_partition_tables_tiflash_replica_count(db_name, table_name)"
 void dbgFuncGetPartitionTablesTiflashReplicaCount(Context & context, const ASTs & args, DBGInvoker::Printer output);
 
+// Get table's tiflash mode with mapped table name
+// Usage:
+//   ./storage-client.sh "DBGInvoke get_tiflash_mode(db_name, table_name)"
+void dbgFuncGetTiflashMode(Context & context, const ASTs & args, DBGInvoker::Printer output);
+
+// Get the logical table's partition tables' tiflash replica counts with mapped table name
+// Usage:
+//   ./storage-client.sh "DBGInvoke get_partition_tables_tiflash_mode(db_name, table_name)"
+void dbgFuncGetPartitionTablesTiflashMode(Context & context, const ASTs & args, DBGInvoker::Printer output);
+
 } // namespace DB

dbms/src/Storages/Transaction/TiDB.cpp

@@ -25,6 +25,7 @@
 #include <Storages/Transaction/Collator.h>
 #include <Storages/Transaction/TiDB.h>
 #include <TiDB/Schema/SchemaNameMapper.h>
+#include <common/logger_useful.h>
 
 #include <cmath>
 
@@ -772,6 +773,37 @@ catch (const Poco::Exception & e)
         DB::Exception(e));
 }
 
+String TiFlashModeToString(TiFlashMode tiflash_mode)
+{
+    switch (tiflash_mode)
+    {
+    case TiFlashMode::Normal:
+        return "";
+    case TiFlashMode::Fast:
+        return "fast";
+    default:
+        LOG_FMT_WARNING(&Poco::Logger::get("TiDB"), "TiFlashModeToString with invalid tiflash mode {}", tiflash_mode);
+        return "";
+    }
+}
+
+TiFlashMode parseTiFlashMode(String mode_str)
+{
+    if (mode_str.empty())
+    {
+        return TiFlashMode::Normal;
+    }
+    else if (mode_str == "fast")
+    {
+        return TiFlashMode::Fast;
+    }
+    else
+    {
+        throw DB::Exception(
+            std::string(__PRETTY_FUNCTION__)
+            + " ParseTiFlashMode Failed. mode " + mode_str + " is unvalid, please set mode as fast/normal");
+    }
+}
 ///////////////////////
 ////// TableInfo //////
 ///////////////////////
@@ -840,6 +872,8 @@ try
 
     json->set("tiflash_replica", replica_info.getJSONObject());
 
+    json->set("tiflash_mode", std::string(TiFlashModeToString(tiflash_mode)));
+
     json->stringify(buf);
 
     return buf.str();
@@ -926,6 +960,14 @@ try
             replica_info.deserialize(replica_obj);
         }
     }
+    if (obj->has("tiflash_mode"))
+    {
+        auto mode = obj->getValue<String>("tiflash_mode");
+        if (!mode.empty())
+        {
+            tiflash_mode = parseTiFlashMode(mode);
+        }
+    }
     if (is_common_handle && index_infos.size() != 1)
     {
         throw DB::Exception(

dbms/src/Storages/Transaction/TiDB.h

@@ -333,6 +333,12 @@ struct IndexInfo
     bool is_global;
 };
 
+enum class TiFlashMode
+{
+    Normal,
+    Fast,
+};
+
 struct TableInfo
 {
     TableInfo() = default;
@@ -382,6 +388,8 @@ struct TableInfo
     // The TiFlash replica info persisted by TiDB
     TiFlashReplicaInfo replica_info;
 
+    TiFlashMode tiflash_mode = TiFlashMode::Normal;
+
     ::TiDB::StorageEngine engine_type = ::TiDB::StorageEngine::UNSPECIFIED;
 
     ColumnID getColumnID(const String & name) const;
@@ -413,4 +421,7 @@ String genJsonNull();
 tipb::FieldType columnInfoToFieldType(const ColumnInfo & ci);
 ColumnInfo fieldTypeToColumnInfo(const tipb::FieldType & field_type);
 
+String TiFlashModeToString(TiFlashMode tiflash_mode);
+TiFlashMode parseTiFlashMode(String mode_str);
+
 } // namespace TiDB