Move idx_test in to ydb

idx_test is a tool to perform consistency check for index table under
various load scenario.

12 files changed, 516 insertions, 0 deletions

diff --git a/ydb/CMakeLists.txt b/ydb/CMakeLists.txt
index f8a673b1a9..92e0648cae 100644
--- a/ydb/CMakeLists.txt
+++ b/ydb/CMakeLists.txt
@@ -11,3 +11,4 @@ add_subdirectory(core)
 add_subdirectory(library)
 add_subdirectory(public)
 add_subdirectory(services)
+add_subdirectory(tests)
diff --git a/ydb/tests/CMakeLists.txt b/ydb/tests/CMakeLists.txt
new file mode 100644
index 0000000000..ba1c90c695
--- /dev/null
+++ b/ydb/tests/CMakeLists.txt
@@ -0,0 +1,9 @@
+
+# This file was generated by the build system used internally in the Yandex monorepo.
+# Only simple modifications are allowed (adding source-files to targets, adding simple properties
+# like target_include_directories). These modifications will be ported to original
+# ya.make files by maintainers. Any complex modifications which can't be ported back to the
+# original buildsystem will not be accepted.
+
+
+add_subdirectory(tools)
diff --git a/ydb/tests/tools/CMakeLists.txt b/ydb/tests/tools/CMakeLists.txt
new file mode 100644
index 0000000000..3ec7cd1baa
--- /dev/null
+++ b/ydb/tests/tools/CMakeLists.txt
@@ -0,0 +1,9 @@
+
+# This file was generated by the build system used internally in the Yandex monorepo.
+# Only simple modifications are allowed (adding source-files to targets, adding simple properties
+# like target_include_directories). These modifications will be ported to original
+# ya.make files by maintainers. Any complex modifications which can't be ported back to the
+# original buildsystem will not be accepted.
+
+
+add_subdirectory(idx_test)
diff --git a/ydb/tests/tools/idx_test/CMakeLists.darwin-x86_64.txt b/ydb/tests/tools/idx_test/CMakeLists.darwin-x86_64.txt
new file mode 100644
index 0000000000..85f13fe714
--- /dev/null
+++ b/ydb/tests/tools/idx_test/CMakeLists.darwin-x86_64.txt
@@ -0,0 +1,49 @@
+
+# This file was generated by the build system used internally in the Yandex monorepo.
+# Only simple modifications are allowed (adding source-files to targets, adding simple properties
+# like target_include_directories). These modifications will be ported to original
+# ya.make files by maintainers. Any complex modifications which can't be ported back to the
+# original buildsystem will not be accepted.
+
+
+get_built_tool_path(
+  TOOL_rescompiler_bin
+  TOOL_rescompiler_dependency
+  tools/rescompiler/bin
+  rescompiler
+)
+
+add_executable(idx_test)
+target_link_libraries(idx_test PUBLIC
+  contrib-libs-cxxsupp
+  yutil
+  library-cpp-cpuid_check
+  library-cpp-getopt
+  library-cpp-resource
+  public-lib-idx_test
+  public-lib-yson_value
+  cpp-client-ydb_scheme
+  cpp-client-ydb_table
+)
+target_link_options(idx_test PRIVATE
+  -Wl,-platform_version,macos,11.0,11.0
+  -fPIC
+  -fPIC
+  -framework
+  CoreFoundation
+)
+target_sources(idx_test PRIVATE
+  ${CMAKE_SOURCE_DIR}/ydb/tests/tools/idx_test/main.cpp
+  ${CMAKE_BINARY_DIR}/ydb/tests/tools/idx_test/5f46adc80fd9bb1406279a9088a76e67.cpp
+)
+resources(idx_test
+  ${CMAKE_BINARY_DIR}/ydb/tests/tools/idx_test/5f46adc80fd9bb1406279a9088a76e67.cpp
+  INPUTS
+  ${CMAKE_SOURCE_DIR}/ydb/tests/tools/idx_test/sql/create_table1.sql
+  KEYS
+  create_table1
+)
+target_allocator(idx_test
+  system_allocator
+)
+vcs_info(idx_test)
diff --git a/ydb/tests/tools/idx_test/CMakeLists.linux-aarch64.txt b/ydb/tests/tools/idx_test/CMakeLists.linux-aarch64.txt
new file mode 100644
index 0000000000..166ecb0c6e
--- /dev/null
+++ b/ydb/tests/tools/idx_test/CMakeLists.linux-aarch64.txt
@@ -0,0 +1,52 @@
+
+# This file was generated by the build system used internally in the Yandex monorepo.
+# Only simple modifications are allowed (adding source-files to targets, adding simple properties
+# like target_include_directories). These modifications will be ported to original
+# ya.make files by maintainers. Any complex modifications which can't be ported back to the
+# original buildsystem will not be accepted.
+
+
+get_built_tool_path(
+  TOOL_rescompiler_bin
+  TOOL_rescompiler_dependency
+  tools/rescompiler/bin
+  rescompiler
+)
+
+add_executable(idx_test)
+target_link_libraries(idx_test PUBLIC
+  contrib-libs-linux-headers
+  contrib-libs-cxxsupp
+  yutil
+  library-cpp-getopt
+  library-cpp-resource
+  public-lib-idx_test
+  public-lib-yson_value
+  cpp-client-ydb_scheme
+  cpp-client-ydb_table
+)
+target_link_options(idx_test PRIVATE
+  -ldl
+  -lrt
+  -Wl,--no-as-needed
+  -fPIC
+  -fPIC
+  -lpthread
+  -lrt
+  -ldl
+)
+target_sources(idx_test PRIVATE
+  ${CMAKE_SOURCE_DIR}/ydb/tests/tools/idx_test/main.cpp
+  ${CMAKE_BINARY_DIR}/ydb/tests/tools/idx_test/5f46adc80fd9bb1406279a9088a76e67.cpp
+)
+resources(idx_test
+  ${CMAKE_BINARY_DIR}/ydb/tests/tools/idx_test/5f46adc80fd9bb1406279a9088a76e67.cpp
+  INPUTS
+  ${CMAKE_SOURCE_DIR}/ydb/tests/tools/idx_test/sql/create_table1.sql
+  KEYS
+  create_table1
+)
+target_allocator(idx_test
+  cpp-malloc-jemalloc
+)
+vcs_info(idx_test)
diff --git a/ydb/tests/tools/idx_test/CMakeLists.linux-x86_64.txt b/ydb/tests/tools/idx_test/CMakeLists.linux-x86_64.txt
new file mode 100644
index 0000000000..2bd7d6d6b6
--- /dev/null
+++ b/ydb/tests/tools/idx_test/CMakeLists.linux-x86_64.txt
@@ -0,0 +1,54 @@
+
+# This file was generated by the build system used internally in the Yandex monorepo.
+# Only simple modifications are allowed (adding source-files to targets, adding simple properties
+# like target_include_directories). These modifications will be ported to original
+# ya.make files by maintainers. Any complex modifications which can't be ported back to the
+# original buildsystem will not be accepted.
+
+
+get_built_tool_path(
+  TOOL_rescompiler_bin
+  TOOL_rescompiler_dependency
+  tools/rescompiler/bin
+  rescompiler
+)
+
+add_executable(idx_test)
+target_link_libraries(idx_test PUBLIC
+  contrib-libs-linux-headers
+  contrib-libs-cxxsupp
+  yutil
+  library-cpp-cpuid_check
+  library-cpp-getopt
+  library-cpp-resource
+  public-lib-idx_test
+  public-lib-yson_value
+  cpp-client-ydb_scheme
+  cpp-client-ydb_table
+)
+target_link_options(idx_test PRIVATE
+  -ldl
+  -lrt
+  -Wl,--no-as-needed
+  -fPIC
+  -fPIC
+  -lpthread
+  -lrt
+  -ldl
+)
+target_sources(idx_test PRIVATE
+  ${CMAKE_SOURCE_DIR}/ydb/tests/tools/idx_test/main.cpp
+  ${CMAKE_BINARY_DIR}/ydb/tests/tools/idx_test/5f46adc80fd9bb1406279a9088a76e67.cpp
+)
+resources(idx_test
+  ${CMAKE_BINARY_DIR}/ydb/tests/tools/idx_test/5f46adc80fd9bb1406279a9088a76e67.cpp
+  INPUTS
+  ${CMAKE_SOURCE_DIR}/ydb/tests/tools/idx_test/sql/create_table1.sql
+  KEYS
+  create_table1
+)
+target_allocator(idx_test
+  cpp-malloc-tcmalloc
+  libs-tcmalloc-no_percpu_cache
+)
+vcs_info(idx_test)
diff --git a/ydb/tests/tools/idx_test/CMakeLists.txt b/ydb/tests/tools/idx_test/CMakeLists.txt
new file mode 100644
index 0000000000..f8b31df0c1
--- /dev/null
+++ b/ydb/tests/tools/idx_test/CMakeLists.txt
@@ -0,0 +1,17 @@
+
+# This file was generated by the build system used internally in the Yandex monorepo.
+# Only simple modifications are allowed (adding source-files to targets, adding simple properties
+# like target_include_directories). These modifications will be ported to original
+# ya.make files by maintainers. Any complex modifications which can't be ported back to the
+# original buildsystem will not be accepted.
+
+
+if (CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR STREQUAL "aarch64" AND NOT HAVE_CUDA)
+  include(CMakeLists.linux-aarch64.txt)
+elseif (CMAKE_SYSTEM_NAME STREQUAL "Darwin" AND CMAKE_SYSTEM_PROCESSOR STREQUAL "x86_64")
+  include(CMakeLists.darwin-x86_64.txt)
+elseif (WIN32 AND CMAKE_SYSTEM_PROCESSOR STREQUAL "AMD64" AND NOT HAVE_CUDA)
+  include(CMakeLists.windows-x86_64.txt)
+elseif (CMAKE_SYSTEM_NAME STREQUAL "Linux" AND CMAKE_SYSTEM_PROCESSOR STREQUAL "x86_64" AND NOT HAVE_CUDA)
+  include(CMakeLists.linux-x86_64.txt)
+endif()
diff --git a/ydb/tests/tools/idx_test/CMakeLists.windows-x86_64.txt b/ydb/tests/tools/idx_test/CMakeLists.windows-x86_64.txt
new file mode 100644
index 0000000000..49248143b6
--- /dev/null
+++ b/ydb/tests/tools/idx_test/CMakeLists.windows-x86_64.txt
@@ -0,0 +1,42 @@
+
+# This file was generated by the build system used internally in the Yandex monorepo.
+# Only simple modifications are allowed (adding source-files to targets, adding simple properties
+# like target_include_directories). These modifications will be ported to original
+# ya.make files by maintainers. Any complex modifications which can't be ported back to the
+# original buildsystem will not be accepted.
+
+
+get_built_tool_path(
+  TOOL_rescompiler_bin
+  TOOL_rescompiler_dependency
+  tools/rescompiler/bin
+  rescompiler
+)
+
+add_executable(idx_test)
+target_link_libraries(idx_test PUBLIC
+  contrib-libs-cxxsupp
+  yutil
+  library-cpp-cpuid_check
+  library-cpp-getopt
+  library-cpp-resource
+  public-lib-idx_test
+  public-lib-yson_value
+  cpp-client-ydb_scheme
+  cpp-client-ydb_table
+)
+target_sources(idx_test PRIVATE
+  ${CMAKE_SOURCE_DIR}/ydb/tests/tools/idx_test/main.cpp
+  ${CMAKE_BINARY_DIR}/ydb/tests/tools/idx_test/5f46adc80fd9bb1406279a9088a76e67.cpp
+)
+resources(idx_test
+  ${CMAKE_BINARY_DIR}/ydb/tests/tools/idx_test/5f46adc80fd9bb1406279a9088a76e67.cpp
+  INPUTS
+  ${CMAKE_SOURCE_DIR}/ydb/tests/tools/idx_test/sql/create_table1.sql
+  KEYS
+  create_table1
+)
+target_allocator(idx_test
+  system_allocator
+)
+vcs_info(idx_test)
diff --git a/ydb/tests/tools/idx_test/main.cpp b/ydb/tests/tools/idx_test/main.cpp
new file mode 100644
index 0000000000..03ca36b845
--- /dev/null
+++ b/ydb/tests/tools/idx_test/main.cpp
@@ -0,0 +1,254 @@
+#include <util/system/env.h>
+#include <util/string/printf.h>
+#include <util/stream/file.h>
+
+#include <library/cpp/getopt/last_getopt.h>
+#include <library/cpp/json/json_writer.h>
+
+#include <ydb/public/lib/idx_test/idx_test.h>
+#include <ydb/public/sdk/cpp/client/ydb_scheme/scheme.h>
+#include <ydb/public/sdk/cpp/client/ydb_table/table.h>
+
+using namespace NYdb;
+using namespace NYdb::NTable;
+using namespace NYdb::NScheme;
+using namespace NIdxTest;
+using namespace NLastGetopt;
+
+class TYdbErrorException : public yexception {
+    mutable TString StrBuf_;
+public:
+    TYdbErrorException(const NYdb::TStatus& status)
+        : Status(status) {}
+
+    NYdb::TStatus Status;
+    const char* what() const noexcept override {
+        TStringStream ss;
+        ss << "Status: " << Status.GetStatus() << "\n";
+        ss << "Issues: " << Status.GetIssues().ToString() << "\n";
+        StrBuf_ = ss.Str();
+        return StrBuf_.c_str();
+    }
+};
+
+static void ThrowOnError(const TStatus& status) {
+    if (!status.IsSuccess()) {
+        throw TYdbErrorException(status);
+    }
+}
+
+static TMap<TString, IWorkLoader::ELoadCommand> Cmds {
+    {"upsert", IWorkLoader::LC_UPSERT},
+    {"insert", IWorkLoader::LC_INSERT},
+    {"update", IWorkLoader::LC_UPDATE},
+    {"update_on", IWorkLoader::LC_UPDATE_ON},
+    {"replace", IWorkLoader::LC_REPLACE},
+    {"delete", IWorkLoader::LC_DELETE},
+    {"delete_on", IWorkLoader::LC_DELETE_ON},
+    {"select", IWorkLoader::LC_SELECT},
+    {"upsert_if_uniq", IWorkLoader::LC_UPSERT_IF_UNIQ},
+    {"add_index", IWorkLoader::LC_ALTER_ADD_INDEX},
+    {"add_covering_index", IWorkLoader::LC_ALTER_ADD_INDEX_WITH_DATA_COLUMN},
+};
+
+static TString GetWorkLoadList() {
+    TString res;
+    auto i = Cmds.size();
+    for (const auto& c : Cmds) {
+        res += c.first;
+        if (--i)
+            res += " ";
+    }
+    return res;
+}
+
+static ui32 ParseCmd(int argc, char** argv) {
+    ui32 res = 0;
+    for (int i = 0; i < argc; i++) {
+        auto it = Cmds.find(TString(argv[i]));
+        if (it != Cmds.end()) {
+            res |= it->second;
+        }
+    }
+    return res;
+}
+
+static void ExecuteDDL(TTableClient client, const TString& sql) {
+    Cerr << "Try to execute DDL query: " << sql << Endl;
+    ThrowOnError(client.RetryOperationSync([sql](TSession session) {
+        return session.ExecuteSchemeQuery(sql).GetValueSync();
+    }));
+}
+
+static void CreatePath(TSchemeClient scheme, const TString& database, const TString prefix) {
+    size_t prevPos = 0;
+    size_t pos = 0;
+    TString curPath = database + "/";
+    curPath.reserve(database.size() + prefix.size() + 1);
+    for (;;) {
+        prevPos = pos;
+        pos = prefix.find('/', pos+1);
+        if (pos == prefix.npos)
+            break;
+        curPath += prefix.substr(prevPos, pos - prevPos);
+        ThrowOnError(scheme.MakeDirectory(curPath, TMakeDirectorySettings()
+            .ClientTimeout(TDuration::Seconds(10))).GetValueSync());
+    }
+}
+
+static void ExecuteCreateUniformTable(TTableClient client, const TString& tablePath, ui32 shardsCount) {
+    Cerr << "Try to create table " << tablePath << " shards count: " << shardsCount << Endl;
+
+    auto builder = TTableBuilder()
+        .AddNullableColumn("key", EPrimitiveType::Uint32)
+        .AddNullableColumn("index1", EPrimitiveType::Uint32)
+        .AddNullableColumn("value", EPrimitiveType::String)
+        .SetPrimaryKeyColumn("key")
+        .AddSecondaryIndex("index_name", "index1");
+
+    auto desc = builder.Build();
+
+    ThrowOnError(client.RetryOperationSync([tablePath, desc, shardsCount](TSession session) {
+        auto d = desc;
+        return session.CreateTable(tablePath,
+            std::move(d),
+            TCreateTableSettings()
+                .PartitioningPolicy(
+                    TPartitioningPolicy()
+                        .UniformPartitions(shardsCount)
+                )
+            ).GetValueSync();
+    }));
+}
+/*
+static void UploadData(TDriver driver, const TString& tablePath, ui32 rows) {
+    auto param = TUploaderParams{1};
+    // Uploader will try to create table again, but table already exixts. It is not a error
+    auto uploader = NIdxTest::CreateUploader(driver, tablePath, param);
+    TTableClient client(driver);
+    TMaybe<TTableDescription> desc;
+    ThrowOnError(client.RetryOperationSync([&desc, tablePath](TSession session) {
+        auto result = session.DescribeTable(tablePath).GetValueSync();
+        if (result.IsSuccess()) {
+            desc = result.GetTableDescription();
+        }
+        return result;
+    }));
+    uploader->Run(CreateDataProvider(rows, 1, desc.GetRef()));
+}
+*/
+static void RunOperations(TDriver driver, const TString& tablePath, ui32 op, const IWorkLoader::TRunSettings& settings, const TString& json) {
+    auto tracker = CreateStderrProgressTracker(500, "OPERATIONS LOADER");
+    auto workLoader = CreateWorkLoader(driver, std::move(tracker));
+    auto report = workLoader->Run(tablePath, op, settings);
+    if (json) {
+        TOFStream jStream{json};
+        NJson::WriteJson(&jStream, &report, /*formatOutput*/ true);
+        jStream.Finish();
+    }
+}
+
+int main(int argc, char** argv) {
+
+    TString endpoint;
+    TString database;
+    TString prefix;
+    TString table = "TestTable1";
+    TString json;
+    bool createTable = false;
+    ui32 rows = 1000;
+    ui32 opsPerTx = 1;
+    ui32 createUniformTable = 0;
+    ui32 infly = 100;
+    bool skipCheck = false;
+
+    TOpts opts = TOpts::Default();
+
+    opts.AddLongOption('e', "endpoint", "YDB endpoint").Required().RequiredArgument("HOST:PORT")
+        .StoreResult(&endpoint);
+    opts.AddLongOption('d', "database", "YDB database").Required().RequiredArgument("PATH")
+        .StoreResult(&database);
+    opts.AddLongOption('p', "prefix", "Base prefix for tables").Optional().RequiredArgument("PATH")
+        .StoreResult(&prefix);
+    opts.AddLongOption('t', "table", "Table name to use").Optional().RequiredArgument("NAME")
+        .StoreResult(&table);
+    opts.AddLongOption('i', "infly", "Number of requests infly").Optional().RequiredArgument("NUM")
+        .StoreResult(&infly);
+    opts.AddLongOption('j', "json", "Path to json stat result").Optional().RequiredArgument("PATH")
+        .StoreResult(&json);
+
+    opts.AddLongOption("create", "Create table before run").Optional()
+        .StoreTrue(&createTable);
+    opts.AddLongOption("skip-check", "Skip data and index table comparation").Optional()
+        .StoreTrue(&skipCheck);
+    opts.AddLongOption("create_uniform", "Create table with uniform partitioning using given number of shards")
+        .Optional().RequiredArgument("SHARD_COUNT").StoreResult(&createUniformTable);
+    opts.AddLongOption("rows", "Number of string to touch").Optional()
+        .StoreResult(&rows);
+    opts.AddLongOption("ops", "Number of operations per transaction").Optional()
+        .StoreResult(&opsPerTx);
+
+    opts.SetFreeArgsMin(0);
+    opts.SetFreeArgTitle(0, "<WORKLOAD>", GetWorkLoadList());
+
+    TOptsParseResult res(&opts, argc, argv);
+    size_t freeArgsPos = res.GetFreeArgsPos();
+
+    argv += freeArgsPos;
+
+    ui32 op = ParseCmd(argc - freeArgsPos, argv);
+
+    auto driver = TDriver(
+        TDriverConfig()
+            .SetEndpoint(endpoint)
+            .SetDatabase(database)
+            .SetAuthToken(GetEnv("YDB_TOKEN"))
+            .SetGRpcKeepAliveTimeout(TDuration::Seconds(20))
+            .SetGRpcKeepAlivePermitWithoutCalls(true));
+
+    auto client = TTableClient(driver);
+    if (prefix) {
+        if (prefix[prefix.size() - 1] != '/') {
+            prefix += "/";
+        }
+    }
+    const auto tablePath = database + "/" + prefix + table;
+
+    try {
+        if (createTable && createUniformTable) {
+            Cerr << "Only one option create or create_uniform allowed" << Endl;
+            return 1;
+        }
+        if (createTable) {
+            auto sql = Sprintf(NResource::Find("create_table1").c_str(), tablePath.c_str());
+            ExecuteDDL(client, sql);
+            Cerr << "Table created" << Endl;
+        }
+        if (createUniformTable) {
+            CreatePath(TSchemeClient(driver), database, prefix);
+            ExecuteCreateUniformTable(client, tablePath, createUniformTable);
+            Cerr << "Table created" << Endl;
+        }
+
+        if (op) {
+            RunOperations(driver, tablePath, op, IWorkLoader::TRunSettings{rows, infly, opsPerTx}, json);
+        }
+
+        if (!skipCheck) {
+            auto tracker = CreateStderrProgressTracker(500, "CHECKER");
+            auto checker = CreateChecker(driver, std::move(tracker));
+            checker->Run(tablePath);
+        }
+    } catch (const std::exception& ex) {
+        driver.Stop(true);
+        Cerr << "Test failed: " << ex.what() << Endl;
+        Cerr << "..." << Endl;
+        auto ydbEx = dynamic_cast<const TYdbErrorException*>(&ex);
+        if (ydbEx) {
+            Cerr << "Ydb error: " << ydbEx->Status.GetStatus() << " " << ydbEx->Status.GetIssues().ToString() << Endl;
+        }
+        return 1;
+    }
+    driver.Stop(true);
+    return 0;
+}
diff --git a/ydb/tests/tools/idx_test/sql/create_table1.sql b/ydb/tests/tools/idx_test/sql/create_table1.sql
new file mode 100644
index 0000000000..edc25fc55b
--- /dev/null
+++ b/ydb/tests/tools/idx_test/sql/create_table1.sql
@@ -0,0 +1,8 @@
+--!syntax_v1
+CREATE TABLE `%s` (
+    Key Int32,
+    Index1 Int32,
+    Value String,
+    PRIMARY KEY(Key),
+    INDEX NewIndex1 GLOBAL ON (Index1)
+);
diff --git a/ydb/tests/tools/idx_test/ya.make b/ydb/tests/tools/idx_test/ya.make
new file mode 100644
index 0000000000..b3dbed0435
--- /dev/null
+++ b/ydb/tests/tools/idx_test/ya.make
@@ -0,0 +1,20 @@
+PROGRAM()
+
+SRCS(
+    main.cpp
+)
+
+PEERDIR(
+    library/cpp/getopt
+    library/cpp/resource
+    ydb/public/lib/idx_test
+    ydb/public/lib/yson_value
+    ydb/public/sdk/cpp/client/ydb_scheme
+    ydb/public/sdk/cpp/client/ydb_table
+)
+
+RESOURCE(
+    ./sql/create_table1.sql create_table1
+)
+
+END()
diff --git a/ydb/tests/tools/ya.make b/ydb/tests/tools/ya.make
index 5665022cb8..45ba0fa0b8 100644
--- a/ydb/tests/tools/ya.make
+++ b/ydb/tests/tools/ya.make
@@ -1,4 +1,5 @@
 RECURSE(
+    idx_test
     s3_recipe
     ydb_serializable
     ydb_serializable/replay