Check TPC-C path and print HR errors (#17333) (#21211)

7 files changed, 273 insertions, 0 deletions

diff --git a/ydb/library/workload/tpcc/constants.h b/ydb/library/workload/tpcc/constants.h
index 8948a4516c7..95e4f8115e0 100644
--- a/ydb/library/workload/tpcc/constants.h
+++ b/ydb/library/workload/tpcc/constants.h
@@ -68,6 +68,18 @@ constexpr const char* TABLE_STOCK = "stock";
 constexpr const char* TABLE_ORDER_LINE = "order_line";
 constexpr const char* TABLE_HISTORY = "history";
 
+constexpr std::array<const char*, 9> TPCC_TABLES = {
+    TABLE_CUSTOMER,
+    TABLE_WAREHOUSE,
+    TABLE_DISTRICT,
+    TABLE_NEW_ORDER,
+    TABLE_OORDER,
+    TABLE_ITEM,
+    TABLE_STOCK,
+    TABLE_ORDER_LINE,
+    TABLE_HISTORY
+};
+
 // Index/View names
 constexpr const char* INDEX_CUSTOMER_NAME = "idx_customer_name";
 constexpr const char* INDEX_ORDER = "idx_order";
diff --git a/ydb/library/workload/tpcc/import.cpp b/ydb/library/workload/tpcc/import.cpp
index 711e4651430..fd374daa48a 100644
--- a/ydb/library/workload/tpcc/import.cpp
+++ b/ydb/library/workload/tpcc/import.cpp
@@ -5,6 +5,7 @@
 #include "import_tui.h"
 #include "log.h"
 #include "log_backend.h"
+#include "path_checker.h"
 #include "util.h"
 
 #include <ydb/public/lib/ydb_cli/commands/ydb_command.h>
@@ -838,6 +839,8 @@ public:
             std::exit(1);
         }
 
+        CheckPathForImport(ConnectionConfig, Config.Path);
+
         Config.SetDisplay();
         CalculateApproximateDataSize();
 
diff --git a/ydb/library/workload/tpcc/init.cpp b/ydb/library/workload/tpcc/init.cpp
index bc926fa1803..258ca18001a 100644
--- a/ydb/library/workload/tpcc/init.cpp
+++ b/ydb/library/workload/tpcc/init.cpp
@@ -3,6 +3,7 @@
 #include "constants.h"
 #include "data_splitter.h"
 #include "log.h"
+#include "path_checker.h"
 
 #include <ydb/public/lib/ydb_cli/commands/ydb_command.h>
 #include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/query/client.h>
@@ -78,6 +79,8 @@ void InitSync(const NConsoleClient::TClientCommand::TConfig& connectionConfig, c
     auto driver = NConsoleClient::TYdbCommand::CreateDriver(connectionConfigCopy);
     NQuery::TQueryClient client(driver);
 
+    CheckPathForInit(connectionConfig, runConfig.Path);
+
     TDataSplitter dataSplitter(runConfig.WarehouseCount);
 
     // Get split keys for different table types
diff --git a/ydb/library/workload/tpcc/path_checker.cpp b/ydb/library/workload/tpcc/path_checker.cpp
new file mode 100644
index 00000000000..ef0d9dcf8bf
--- /dev/null
+++ b/ydb/library/workload/tpcc/path_checker.cpp
@@ -0,0 +1,228 @@
+#include "path_checker.h"
+
+#include "constants.h"
+
+#include <ydb/public/lib/ydb_cli/commands/ydb_command.h>
+#include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/scheme/scheme.h>
+#include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/driver/driver.h>
+#include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/query/client.h>
+
+#include <util/generic/hash_set.h>
+
+namespace NYdb::NTPCC {
+
+namespace {
+
+void ExitIfError(const TStatus& status, const TString& what) {
+    if (status.GetStatus() == EStatus::SUCCESS) {
+        return;
+    }
+
+    TStringStream ss;
+    ss << what << ": " << ToString(status.GetStatus());
+    const auto& issues = status.GetIssues();
+    if (issues) {
+        ss << ", issues: ";
+        issues.PrintTo(ss, true);
+    }
+
+    Cerr << ss.Str() << Endl;
+    std::exit(1);
+}
+
+THashSet<std::string> ListPath(TDriver& driver, const TString& path) {
+    NScheme::TSchemeClient schemeClient(driver);
+    auto listResult = schemeClient.ListDirectory(path, {}).GetValueSync();
+    ExitIfError(listResult, TStringBuilder() << "Failed to list path '" << path << "'");
+
+    THashSet<std::string> entries;
+
+    const auto& children = listResult.GetChildren();
+    for (const auto& child: children) {
+        entries.insert(child.Name);
+    }
+
+    return entries;
+}
+
+void CheckTablesExist(TDriver& driver, const TString& path, const char* what) {
+    auto entries = ListPath(driver, path);
+
+    for (const char* table: TPCC_TABLES) {
+        if (!entries.contains(table)) {
+            Cerr << "TPC-C table '" << table << "' is missing in " << path << ". " << what << Endl;
+            std::exit(1);
+        }
+    }
+}
+
+void CheckNoTablesExist(TDriver& driver, const TString& path, const char* what) {
+    auto entries = ListPath(driver, path);
+
+    for (const char* table: TPCC_TABLES) {
+        if (entries.contains(table)) {
+            Cerr << "TPC-C table '" << table << "' already exists in " << path << ". " << what << Endl;
+            std::exit(1);
+        }
+    }
+}
+
+NTable::TTableDescription DescribeTable(TDriver& driver, const TString& path, const TString& tableName) {
+    TString fullPath = path + "/" + tableName;
+    NTable::TTableClient client(driver);
+    NTable::TCreateSessionResult sessionResult = client.GetSession(NTable::TCreateSessionSettings()).GetValueSync();
+    ExitIfError(sessionResult, "Failed to obtain a session to check the tables");
+    auto session = sessionResult.GetSession();
+    auto result = session.DescribeTable(fullPath).GetValueSync();
+    ExitIfError(result, TStringBuilder() << "Failed to describe the table '" << fullPath << "'");
+
+    return result.GetTableDescription();
+}
+
+void CheckNoIndexExists(TDriver& driver, const TString& path, const TString& tableName) {
+    TString fullPath = path + "/" + tableName;
+    auto desc = DescribeTable(driver, path, tableName);
+    const std::vector<NTable::TIndexDescription>& indexes = desc.GetIndexDescriptions();
+    if (!indexes.empty()) {
+        Cerr << "Table '" << fullPath
+             << "' already has index/indices. Are you importing to the existing data?" << Endl;
+        std::exit(1);
+    }
+}
+
+void CheckIndexExistsAndReady(TDriver& driver, const TString& path, const TString& tableName) {
+    TString fullPath = path + "/" + tableName;
+    auto desc = DescribeTable(driver, path, tableName);
+    const std::vector<NTable::TIndexDescription>& indexes = desc.GetIndexDescriptions();
+    if (indexes.empty()) {
+        Cerr << "Table '" << fullPath << "' has no index" << Endl;
+        std::exit(1);
+    }
+
+    if (indexes.size() > 1) {
+        Cerr << "Table '" << fullPath << "' has more than one index: " << indexes.size() << Endl;
+        std::exit(1);
+    }
+
+    const auto& index = indexes[0];
+    const auto& indexName = index.GetIndexName();
+    auto indexType = index.GetIndexType();
+    if (indexType != NTable::EIndexType::GlobalSync) {
+        Cerr << "Index '" << indexName << "' has unexpected type: " << int(indexType) << Endl;
+        std::exit(1);
+    }
+
+    if (tableName == TABLE_CUSTOMER) {
+        if (indexName != INDEX_CUSTOMER_NAME) {
+            Cerr << "Table '" << fullPath << "' is expected to have index '" << INDEX_CUSTOMER_NAME
+                << ", but index '" << indexName << "' found" << Endl;
+            std::exit(1);
+        }
+    } else if (tableName == TABLE_OORDER) {
+        if (indexName != INDEX_ORDER) {
+            Cerr << "Table '" << fullPath << "' is expected to have index '" << INDEX_ORDER
+                << ", but index '" << indexName << "' found" << Endl;
+            std::exit(1);
+        }
+    } else {
+        Cerr << "Checking index for unexpected table '" << fullPath << "'" << Endl;
+        std::exit(1);
+    }
+}
+
+int GetWarehouseCount(TDriver& driver, const TString& path) {
+    // just check the warehouse table and assume that the rest is consistent
+
+    TString query = std::format(R"(
+        PRAGMA TablePathPrefix("{}");
+
+        SELECT COUNT(*) FROM `{}`;
+    )", path.c_str(), TABLE_WAREHOUSE);
+
+    NQuery::TQueryClient queryClient(driver);
+    auto result = queryClient.ExecuteQuery(query, NQuery::TTxControl::NoTx()).GetValueSync();
+    ExitIfError(result, TStringBuilder() << "Failed to count warehouses in " << path << "/" << TABLE_WAREHOUSE);
+
+    TResultSetParser parser(result.GetResultSet(0));
+    if (!parser.TryNextRow()) {
+        return 0;
+    }
+
+    try {
+        return parser.ColumnParser("column0").GetUint64();
+    } catch (std::exception& e) {
+        Cerr << "Failed to count warehouses in " << path << "/" << TABLE_WAREHOUSE << ": " << e.what() << Endl;
+        std::exit(1);
+    }
+}
+
+} // anonymous
+
+void CheckPathForInit(
+    const NConsoleClient::TClientCommand::TConfig& connectionConfig,
+    const TString& path) noexcept
+{
+    auto connectionConfigCopy = connectionConfig;
+    TDriver driver = NConsoleClient::TYdbCommand::CreateDriver(connectionConfigCopy);
+
+    CheckNoTablesExist(driver, path, "Already inited or forgot to clean?");
+}
+
+void CheckPathForImport(
+    const NConsoleClient::TClientCommand::TConfig& connectionConfig,
+    const TString& path) noexcept
+{
+    auto connectionConfigCopy = connectionConfig;
+    TDriver driver = NConsoleClient::TYdbCommand::CreateDriver(connectionConfigCopy);
+
+    // 1. Check tables exist
+
+    CheckTablesExist(driver, path, "Run 'init' first.");
+
+    // 2. Check no wh populated
+
+    int whCount = GetWarehouseCount(driver, path);
+    if (whCount != 0) {
+        Cerr << path << " already has " << whCount << " warehouses. Are you importing to the existing data?" << Endl;
+        std::exit(1);
+    }
+
+    // 3. Check no indices
+
+    CheckNoIndexExists(driver, path, TABLE_CUSTOMER);
+    CheckNoIndexExists(driver, path, TABLE_OORDER);
+}
+
+void CheckPathForRun(
+    const NConsoleClient::TClientCommand::TConfig& connectionConfig,
+    const TString& path,
+    int expectedWhCount) noexcept
+{
+    auto connectionConfigCopy = connectionConfig;
+    TDriver driver = NConsoleClient::TYdbCommand::CreateDriver(connectionConfigCopy);
+
+    // 1. Check tables exist
+
+    CheckTablesExist(driver, path, "Import the data first.");
+
+    // 2. Check indices exist
+
+    CheckIndexExistsAndReady(driver, path, TABLE_CUSTOMER);
+    CheckIndexExistsAndReady(driver, path, TABLE_OORDER);
+
+    // 3. Check the number of warehouses
+
+    int whCount = GetWarehouseCount(driver, path);
+    if (whCount != expectedWhCount) {
+        if (whCount == 0) {
+            Cerr << "Empty warehouse table (and maybe missing other TPC-C data), run import first" << Endl;
+            std::exit(1);
+        }
+        if (whCount < expectedWhCount) {
+            Cerr << "Expected data for " << expectedWhCount << " warehouses, but found for " << whCount << Endl;
+            std::exit(1);
+        }
+    }
+}
+
+} // namespace NYdb::NTPCC
diff --git a/ydb/library/workload/tpcc/path_checker.h b/ydb/library/workload/tpcc/path_checker.h
new file mode 100644
index 00000000000..8e80ccf6e29
--- /dev/null
+++ b/ydb/library/workload/tpcc/path_checker.h
@@ -0,0 +1,23 @@
+#pragma once
+
+#include <ydb/public/lib/ydb_cli/common/command.h>
+
+namespace NYdb::NTPCC {
+
+// do a very basic data check before doing anything else,
+// in case of any issues print error and exit
+
+void CheckPathForInit(
+    const NConsoleClient::TClientCommand::TConfig& connectionConfig,
+    const TString& path) noexcept;
+
+void CheckPathForImport(
+    const NConsoleClient::TClientCommand::TConfig& connectionConfig,
+    const TString& path) noexcept;
+
+void CheckPathForRun(
+    const NConsoleClient::TClientCommand::TConfig& connectionConfig,
+    const TString& path,
+    int expectedWhCount) noexcept;
+
+} // namespace NYdb::NTPCC
diff --git a/ydb/library/workload/tpcc/runner.cpp b/ydb/library/workload/tpcc/runner.cpp
index c107b442ad1..ef176f22da3 100644
--- a/ydb/library/workload/tpcc/runner.cpp
+++ b/ydb/library/workload/tpcc/runner.cpp
@@ -3,6 +3,7 @@
 #include "constants.h"
 #include "log.h"
 #include "log_backend.h"
+#include "path_checker.h"
 #include "runner_display_data.h"
 #include "runner_tui.h"
 #include "task_queue.h"
@@ -135,6 +136,8 @@ TPCCRunner::TPCCRunner(const NConsoleClient::TClientCommand::TConfig& connection
         std::exit(1);
     }
 
+    CheckPathForRun(connectionConfig, Config.Path, Config.WarehouseCount);
+
     const size_t terminalsCount = Config.WarehouseCount * TERMINALS_PER_WAREHOUSE;
 
     size_t threadCount = 0;
diff --git a/ydb/library/workload/tpcc/ya.make b/ydb/library/workload/tpcc/ya.make
index 28742f3968c..492f0a3aeb0 100644
--- a/ydb/library/workload/tpcc/ya.make
+++ b/ydb/library/workload/tpcc/ya.make
@@ -10,6 +10,7 @@ SRCS(
     import_tui.cpp
     log_backend.cpp
     logs_scroller.cpp
+    path_checker.cpp
     runner.cpp
     runner_tui.cpp
     scroller.cpp