YQ kqprun added recipe (#9453)

Author: GrigoriyPA

ydb/tests/tools/kqprun/kqprun.cpp

@@ -317,15 +317,9 @@ void RunArgumentQueries(const TExecutionOptions& executionOptions, NKqpRun::TKqp
 void RunAsDaemon() {
     NColorizer::TColors colors = NColorizer::AutoColors(Cout);
 
-    Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Started reading commands" << colors.Default() << Endl;
+    Cout << colors.Yellow() << TInstant::Now().ToIsoStringLocal() << " Initialization finished" << colors.Default() << Endl;
     while (true) {
-        TString command;
-        Cin >> command;
-
-        if (command == "exit") {
-            break;
-        }
-        Cerr << colors.Red() << TInstant::Now().ToIsoStringLocal() << " Invalid command '" << command << "'" << colors.Default() << Endl;
+        pause();
     }
 }
 

ydb/tests/tools/kqprun/recipe/__main__.py

@@ -0,0 +1,75 @@
+import argparse
+import logging
+import os
+
+from library.python.testing.recipe import declare_recipe, set_env
+from library.recipes import common as recipes_common
+from yatest.common.network import PortManager
+from ydb.tests.library.common import yatest_common
+
+
+PID_FILENAME = "kqprun_daemon.pid"
+KQPRUN_PATH = os.getenv("KQPRUN_EXECUTABLE") or "ydb/tests/tools/kqprun/kqprun"
+INITIALIZATION_IMEOUT_RATIO = 2
+
+
+def is_kqprun_daemon_ready() -> bool:
+    with open(yatest_common.output_path("kqprun_daemon.out.log"), "r") as outFile:
+        return "Initialization finished" in outFile.read()
+
+
+def build_start_comand(argv: list[str], grpc_port: int) -> tuple[int, list[str]]:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--query", action="append", type=str, default=[])
+    parser.add_argument("--config", action='store', type=str, default="ydb/tests/tools/kqprun/kqprun/configuration/app_config.conf")
+    parser.add_argument("--timeout-ms", action='store', type=int, default=30000)
+    parsed, _ = parser.parse_known_args(argv)
+
+    cmd = [
+        yatest_common.binary_path(KQPRUN_PATH),
+        "--log-file", yatest_common.output_path("kqprun_daemon.ydb.log"),
+        "--app-config", yatest_common.source_path(parsed.config),
+        "--grpc", str(grpc_port),
+        "--timeout", str(parsed.timeout_ms)
+    ]
+
+    if parsed.query:
+        cmd.append("--execution-case")
+        cmd.append("query")
+
+    for query in parsed.query:
+        cmd.append("--script-query")
+        cmd.append(yatest_common.source_path(query))
+
+    return (parsed.timeout_ms, cmd)
+
+
+def start(argv: list[str]):
+    logging.debug("Starting kqprun daemon")
+
+    portManager = PortManager()
+    grpc_port = portManager.get_port()
+    timeout_ms, cmd = build_start_comand(argv, grpc_port)
+
+    recipes_common.start_daemon(
+        command=cmd,
+        environment=None,
+        is_alive_check=is_kqprun_daemon_ready,
+        pid_file_name=PID_FILENAME,
+        timeout=INITIALIZATION_IMEOUT_RATIO * (timeout_ms // 1000),
+        daemon_name="kqprun_daemon"
+    )
+
+    set_env("KQPRUN_ENDPOINT", f"grpc://localhost:{grpc_port}")
+    logging.debug(f"kqprun daemon has been started on port: {grpc_port}")
+
+
+def stop(argv: list[str]):
+    logging.debug("Stop kqprun daemon")
+    with open(PID_FILENAME, "r") as pidFile:
+        pid = int(pidFile.read())
+        recipes_common.stop_daemon(pid)
+
+
+if __name__ == "__main__":
+    declare_recipe(start, stop)

ydb/tests/tools/kqprun/recipe/ya.make

@@ -0,0 +1,12 @@
+PY3_PROGRAM(kqprun_recipe)
+
+PY_SRCS(__main__.py)
+
+PEERDIR(
+    library/python/testing/recipe
+    library/python/testing/yatest_common
+    library/recipes/common
+    ydb/tests/library
+)
+
+END()

ydb/tests/tools/kqprun/tests/cfg/config.conf

@@ -0,0 +1,3 @@
+LogConfig {
+  DefaultLevel: 5
+}

ydb/tests/tools/kqprun/tests/cfg/create_tables.sql

@@ -0,0 +1,4 @@
+CREATE TABLE test_table (
+    Key Int32,
+    PRIMARY KEY (Key)
+);

ydb/tests/tools/kqprun/tests/cfg/fill_tables.sql

@@ -0,0 +1,2 @@
+INSERT INTO test_table
+SELECT 42 AS Key;

ydb/tests/tools/kqprun/tests/test_kqprun_recipe.py

@@ -0,0 +1,18 @@
+import os
+
+from ydb.tests.oss.ydb_sdk_import import ydb
+
+
+class TestKqprunRecipe(object):
+    def test_query_execution(self):
+        with ydb.Driver(
+            endpoint=os.getenv("KQPRUN_ENDPOINT"),
+            database="/Root"
+        ) as driver:
+            driver.wait(timeout=5, fail_fast=True)
+
+            with ydb.QuerySessionPool(driver) as pool:
+                result_sets = pool.execute_with_retries("SELECT * FROM test_table")
+                rows = result_sets[0].rows
+                assert len(rows) == 1
+                assert rows[0].Key == 42

ydb/tests/tools/kqprun/tests/ya.make

@@ -0,0 +1,29 @@
+PY3TEST()
+
+DATA(
+    arcadia/ydb/tests/tools/kqprun/tests/cfg
+)
+
+TEST_SRCS(
+    test_kqprun_recipe.py
+)
+
+PEERDIR(
+    ydb/tests/oss/ydb_sdk_import
+)
+
+DEPENDS(
+    ydb/tests/tools/kqprun
+    ydb/tests/tools/kqprun/recipe
+)
+
+USE_RECIPE(
+    ydb/tests/tools/kqprun/recipe/kqprun_recipe 
+        --config ydb/tests/tools/kqprun/tests/cfg/config.conf
+        --query ydb/tests/tools/kqprun/tests/cfg/create_tables.sql
+        --query ydb/tests/tools/kqprun/tests/cfg/fill_tables.sql
+)
+
+SIZE(MEDIUM)
+
+END()

ydb/tests/tools/kqprun/ya.make

@@ -1,4 +1,4 @@
-PROGRAM()
+PROGRAM(kqprun)
 
 SRCS(
     kqprun.cpp
@@ -24,3 +24,11 @@ PEERDIR(
 YQL_LAST_ABI_VERSION()
 
 END()
+
+RECURSE(
+    recipe
+)
+
+RECURSE_FOR_TESTS(
+    tests
+)