Skip to content

YQ kqprun added recipe #9453

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Sep 19, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
)
Loading