Compile PY protospecs with cmake when YDB recipe is in use by a test

compile_protos.py now returns a correct result code
A custom target is added to the ydb CMakelist which invokes compile_protos.py (using an .sh wrapper to suppress warnings)
Those tests which require YDB recipe must have an epilogue.cmake added to their directory, adding a dependency to the py protos compile target

8 files changed, 75 insertions, 6 deletions

diff --git a/ydb/CMakeLists.txt b/ydb/CMakeLists.txt
index 92e0648cae..066036e84b 100644
--- a/ydb/CMakeLists.txt
+++ b/ydb/CMakeLists.txt
@@ -12,3 +12,12 @@ add_subdirectory(library)
 add_subdirectory(public)
 add_subdirectory(services)
 add_subdirectory(tests)
+add_custom_target(ydb_py_protos)
+
+add_custom_command(
+    TARGET ydb_py_protos POST_BUILD
+    COMMAND ${CMAKE_SOURCE_DIR}/ydb/tests/oss/launch/compile_protos.sh ${CMAKE_SOURCE_DIR} ydb library/cpp/actors
+    WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}
+    COMMENT "Compiling PY protos..."
+)
+
diff --git a/ydb/apps/ydb/ut/CMakeLists.linux-aarch64.txt b/ydb/apps/ydb/ut/CMakeLists.linux-aarch64.txt
index f6bfcce26e..7b93fec089 100644
--- a/ydb/apps/ydb/ut/CMakeLists.linux-aarch64.txt
+++ b/ydb/apps/ydb/ut/CMakeLists.linux-aarch64.txt
@@ -68,3 +68,5 @@ target_allocator(ydb-apps-ydb-ut
   cpp-malloc-jemalloc
 )
 vcs_info(ydb-apps-ydb-ut)
+add_dependencies(ydb-apps-ydb-ut ydb_py_protos)
+
diff --git a/ydb/apps/ydb/ut/CMakeLists.linux-x86_64.txt b/ydb/apps/ydb/ut/CMakeLists.linux-x86_64.txt
index 4ef3e88724..16b271598b 100644
--- a/ydb/apps/ydb/ut/CMakeLists.linux-x86_64.txt
+++ b/ydb/apps/ydb/ut/CMakeLists.linux-x86_64.txt
@@ -70,3 +70,5 @@ target_allocator(ydb-apps-ydb-ut
   libs-tcmalloc-no_percpu_cache
 )
 vcs_info(ydb-apps-ydb-ut)
+add_dependencies(ydb-apps-ydb-ut ydb_py_protos)
+
diff --git a/ydb/apps/ydb/ut/epilogue.cmake b/ydb/apps/ydb/ut/epilogue.cmake
new file mode 100644
index 0000000000..c14b6679de
--- /dev/null
+++ b/ydb/apps/ydb/ut/epilogue.cmake
@@ -0,0 +1 @@
+add_dependencies(ydb-apps-ydb-ut ydb_py_protos)
diff --git a/ydb/tests/oss/launch/compile_protos.py b/ydb/tests/oss/launch/compile_protos.py
index 2c21ad7613..416e2047a8 100644
--- a/ydb/tests/oss/launch/compile_protos.py
+++ b/ydb/tests/oss/launch/compile_protos.py
@@ -1,10 +1,55 @@
+import os
+import sys
 from argparse import ArgumentParser
-from grpc_tools import command
+from grpc_tools import protoc
+
+
+def build_ydb_protos(ydb_repo_root, proto_dir='ydb', strict_mode=False):
+    proto_files = []
+    repo_root = os.path.abspath(ydb_repo_root)
+    files_root = os.path.join(repo_root, proto_dir)
+    # print("Repo root:", repo_root, file=sys.stderr)
+    # print("Proto files root:", files_root, file=sys.stderr)
+    for root, _, files in os.walk(files_root):
+        for filename in files:
+            if filename.endswith('.proto'):
+                proto_files.append(os.path.abspath(os.path.join(root,
+                                                                filename)))
+
+    include_paths = [
+        'contrib/libs/googleapis-common-protos',
+        'contrib/libs/protobuf/src',
+    ]
+
+    command_fix = [
+        'grpc_tools.protoc',
+        '--proto_path={}'.format(repo_root),
+        '--python_out={}'.format(repo_root),
+        '--pyi_out={}'.format(repo_root),
+        '--grpc_python_out={}'.format(repo_root),
+    ]
+
+    for ipath in include_paths:
+        command_fix = command_fix + [
+            '--proto_path={}'.format(os.path.join(repo_root, ipath)),
+        ]
+
+    for proto_file in proto_files:
+        # print(proto_file, file=sys.stderr)
+        command = command_fix + [proto_file]
+        if protoc.main(command) != 0:
+            if strict_mode:
+                raise Exception('error: {} failed'.format(command))
+            else:
+                sys.stderr.write('warning: {} failed'.format(command))
 
 
 if __name__ == '__main__':
     parser = ArgumentParser()
-    parser.add_argument('--source-root', required=True, help='YDB source directory')
+    parser.add_argument('directory', help='Directory inside YDB repo to scan for .proto files', default='ydb', nargs='+')
+    parser.add_argument('--source-root', required=True, help='YDB repo source directory')
+    parser.add_argument('--strict-mode', required=False, default=False, help='Strict mode (fail with result code 1 on errors)')
     args = parser.parse_args()
-
-    command.build_package_protos(args.source_root)
+    for dr in args.directory:
+        # print("Dir:", dr, file=sys.stderr)
+        build_ydb_protos(args.source_root, dr, args.strict_mode)
diff --git a/ydb/tests/oss/launch/compile_protos.sh b/ydb/tests/oss/launch/compile_protos.sh
new file mode 100755
index 0000000000..fa972823b9
--- /dev/null
+++ b/ydb/tests/oss/launch/compile_protos.sh
@@ -0,0 +1,6 @@
+#!/usr/bin/env bash
+set -o pipefail
+echo Compiling protospecs for Python...
+# Usage: compile_protos.sh root_ydb_repo_dir subdir_to_scan_for_protospecs
+SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+python3 $SCRIPT_DIR/compile_protos.py --source-root $1 --strict-mode 1 "${@:2}" 2>&1 >/dev/null | grep -v "No syntax specified for the proto file"
diff --git a/ydb/tests/oss/launch/prepare.sh b/ydb/tests/oss/launch/prepare.sh
index b084d294a8..286257e371 100755
--- a/ydb/tests/oss/launch/prepare.sh
+++ b/ydb/tests/oss/launch/prepare.sh
@@ -1,6 +1,6 @@
 #!/bin/bash
 
-python ${source_root}/ydb/tests/oss/launch/compile_protos.py --source-root ${source_root} 2>/dev/null
+${source_root}/ydb/tests/oss/launch/compile_protos.sh ${source_root} ydb library/cpp/actors
 
 #testresults=${source_root}/ydb/tests/functional/test-results
 #
diff --git a/ydb/tests/oss/launch/ya.make b/ydb/tests/oss/launch/ya.make
index 61ba373fb9..20d1a700dd 100644
--- a/ydb/tests/oss/launch/ya.make
+++ b/ydb/tests/oss/launch/ya.make
@@ -1,4 +1,4 @@
-PY23_LIBRARY()
+PY3_LIBRARY()
 
 PY_SRCS(
     __init__.py
@@ -7,4 +7,8 @@ PY_SRCS(
     launch.py
 )
 
+FILES(
+    compile_protos.sh
+)
+
 END()