Skip to content

Commit

Permalink
YQ kqprun added recipe (ydb-platform#9453)
Browse files Browse the repository at this point in the history
  • Loading branch information
GrigoriyPA authored Sep 19, 2024
1 parent a430757 commit 2a0bfb0
Show file tree
Hide file tree
Showing 9 changed files with 154 additions and 9 deletions.
10 changes: 2 additions & 8 deletions ydb/tests/tools/kqprun/kqprun.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
}

Expand Down
75 changes: 75 additions & 0 deletions ydb/tests/tools/kqprun/recipe/__main__.py
Original file line number Diff line number Diff line change
@@ -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)
12 changes: 12 additions & 0 deletions ydb/tests/tools/kqprun/recipe/ya.make
Original file line number Diff line number Diff line change
@@ -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()
3 changes: 3 additions & 0 deletions ydb/tests/tools/kqprun/tests/cfg/config.conf
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
LogConfig {
DefaultLevel: 5
}
4 changes: 4 additions & 0 deletions ydb/tests/tools/kqprun/tests/cfg/create_tables.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
CREATE TABLE test_table (
Key Int32,
PRIMARY KEY (Key)
);
2 changes: 2 additions & 0 deletions ydb/tests/tools/kqprun/tests/cfg/fill_tables.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
INSERT INTO test_table
SELECT 42 AS Key;
18 changes: 18 additions & 0 deletions ydb/tests/tools/kqprun/tests/test_kqprun_recipe.py
Original file line number Diff line number Diff line change
@@ -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
29 changes: 29 additions & 0 deletions ydb/tests/tools/kqprun/tests/ya.make
Original file line number Diff line number Diff line change
@@ -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()
10 changes: 9 additions & 1 deletion ydb/tests/tools/kqprun/ya.make
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
PROGRAM()
PROGRAM(kqprun)

SRCS(
kqprun.cpp
Expand All @@ -24,3 +24,11 @@ PEERDIR(
YQL_LAST_ABI_VERSION()

END()

RECURSE(
recipe
)

RECURSE_FOR_TESTS(
tests
)

0 comments on commit 2a0bfb0

Please sign in to comment.