Integration tests for YDB external datasource (over YQ)

ydb/tests/fq/generic/conftest.py

@@ -6,6 +6,7 @@
 from ydb.tests.tools.fq_runner.kikimr_utils import YQv2Extension
 from ydb.tests.tools.fq_runner.kikimr_utils import TokenAccessorExtension
 from ydb.tests.tools.fq_runner.kikimr_utils import MDBExtension
+from ydb.tests.tools.fq_runner.kikimr_utils import YdbMvpExtension
 from ydb.tests.tools.fq_runner.kikimr_utils import start_kikimr
 
 from utils.settings import Settings
@@ -22,6 +23,7 @@ def kikimr(request: pytest.FixtureRequest, settings: Settings, yq_version: str):
         ConnectorExtension(settings.connector.grpc_host, settings.connector.grpc_port, False),
         TokenAccessorExtension(settings.token_accessor_mock.endpoint, settings.token_accessor_mock.hmac_secret_file),
         MDBExtension(settings.mdb_mock.endpoint),
+        YdbMvpExtension(settings.ydb_mvp_mock.endpoint),
         YQv2Extension(yq_version),
     ]
     with start_kikimr(request, kikimr_extensions) as kikimr:

ydb/tests/fq/generic/docker-compose.yml

@@ -3,7 +3,7 @@ services:
   postgresql:
     image: "postgres:15-bullseye@sha256:3411b9f2e5239cd7867f34fcf22fe964230f7d447a71d63c283e3593d3f84085"
     # to be able to run tests by different users on the same machine we set prefix to ${USER}
-    container_name: ${USER}_ydb_tests_fq_generic_postgresql
+    container_name: tests-fq-generic-postgresql
     environment:
       POSTGRES_DB: db
       POSTGRES_USER: user
@@ -13,16 +13,31 @@ services:
     command: -p 6432
   clickhouse:
     image: "clickhouse/clickhouse-server:23-alpine@sha256:b078c1cd294632afa2aeba3530e7ba2e568513da23304354f455a25fab575c06"
-    container_name: ${USER}_ydb_tests_fq_generic_clickhouse
+    container_name: tests-fq-generic-clickhouse
     environment:
       CLICKHOUSE_DB: db
       CLICKHOUSE_USER: user
       CLICKHOUSE_DEFAULT_ACCESS_MANAGEMENT: 1
       CLICKHOUSE_PASSWORD: password
     volumes:
       - ./clickhouse:/docker-entrypoint-initdb.d
-  connector:
-    image: "ghcr.io/ydb-platform/fq-connector-go:v0.1.1-rc.2@sha256:e5c2d86bce9cb43420eed0ed534afe760fb90ad41229dbbf34af28023b219af3"
-    container_name: ${USER}_ydb_tests_fq_generic_connector
+  ydb:
+    image: ghcr.io/ydb-platform/local-ydb:latest@sha256:9045e00afec1923dc3277564c7b2f829087c2115f45f18e1d38b80bb89f98be6
+    container_name: tests-fq-generic-ydb
+    hostname: tests-fq-generic-ydb
+    environment:
+      YDB_DEFAULT_LOG_LEVEL: DEBUG
+      POSTGRES_USER: user
+      POSTGRES_PASSWORD: password
+    volumes:
+       - ./ydb/init_ydb:/init_ydb
+       - ./ydb/01_basic.sh:/01_basic.sh
+  fq-connector-go:
+    image: "ghcr.io/ydb-platform/fq-connector-go:v0.2.20@sha256:a1771f348dc8be6219865e332f788429907cdfec3677b3e98f0bc6f7dd542dc6"
+    container_name: tests-fq-generic-fq-connector-go
     ports:
-      - '50051'
+      - '2130'
+    command: >
+      sh -c "
+        echo \"$$(dig fq-tests-ydb-ydb +short) fq-tests-ydb-ydb\" >> /etc/hosts; cat /etc/hosts;
+        /opt/ydb/bin/fq-connector-go server -c /opt/ydb/cfg/fq-connector-go.yaml"

ydb/tests/fq/generic/test_join.py

@@ -15,6 +15,7 @@ def test_simple(self, fq_client: FederatedQueryClient, settings: Settings):
         table_name = 'join_table'
         ch_conn_name = f'ch_conn_{table_name}'
         pg_conn_name = f'pg_conn_{table_name}'
+        ydb_conn_name = f'ydb_conn_{table_name}'
         query_name = f'query_{table_name}'
 
         fq_client.create_postgresql_connection(
@@ -33,11 +34,18 @@ def test_simple(self, fq_client: FederatedQueryClient, settings: Settings):
             password=settings.clickhouse.password,
         )
 
+        fq_client.create_ydb_connection(
+            name=ydb_conn_name,
+            database_id=settings.ydb.dbname,
+        )
+
         sql = fR'''
-            SELECT pg.data AS data_pg, ch.data AS data_ch
+            SELECT pg.data AS data_pg, ch.data AS data_ch, ydb.data AS data_ydb
             FROM {pg_conn_name}.{table_name} AS pg
             JOIN {ch_conn_name}.{table_name} AS ch
-            ON pg.id = ch.id;
+            ON pg.id = ch.id
+            JOIN {ydb_conn_name}.{table_name} AS ydb
+            ON pg.id = ydb.id;
             '''
 
         query_id = fq_client.create_query(query_name, sql, type=fq.QueryContent.QueryType.ANALYTICS).result.query_id
@@ -46,13 +54,17 @@ def test_simple(self, fq_client: FederatedQueryClient, settings: Settings):
         data = fq_client.get_result_data(query_id)
         result_set = data.result.result_set
         logging.debug(str(result_set))
-        assert len(result_set.columns) == 2
+        assert len(result_set.columns) == 3
         assert result_set.columns[0].name == "data_pg"
         assert result_set.columns[1].name == "data_ch"
+        assert result_set.columns[2].name == "data_ydb"
         assert len(result_set.rows) == 3
         assert result_set.rows[0].items[0].bytes_value == b'pg10'
         assert result_set.rows[0].items[1].bytes_value == b'ch10'
+        assert result_set.rows[0].items[2].bytes_value == b'ydb10', result_set
         assert result_set.rows[1].items[0].bytes_value == b'pg20'
         assert result_set.rows[1].items[1].bytes_value == b'ch20'
+        assert result_set.rows[1].items[2].bytes_value == b'ydb20'
         assert result_set.rows[2].items[0].bytes_value == b'pg30'
         assert result_set.rows[2].items[1].bytes_value == b'ch30'
+        assert result_set.rows[2].items[2].bytes_value == b'ydb30'

ydb/tests/fq/generic/test_ydb.py

@@ -0,0 +1,44 @@
+import logging
+import pytest
+
+import ydb.public.api.protos.draft.fq_pb2 as fq
+import ydb.public.api.protos.ydb_value_pb2 as ydb
+from ydb.tests.tools.fq_runner.kikimr_utils import yq_v2
+
+from ydb.tests.tools.fq_runner.fq_client import FederatedQueryClient
+from utils.settings import Settings
+
+
+class TestYdb:
+    @yq_v2
+    @pytest.mark.parametrize("fq_client", [{"folder_id": "my_folder"}], indirect=True)
+    def test_simple(self, fq_client: FederatedQueryClient, settings: Settings):
+        table_name = 'simple_table'
+        conn_name = f'conn_{table_name}'
+        query_name = f'query_{table_name}'
+
+        fq_client.create_ydb_connection(
+            name=conn_name,
+            database_id=settings.ydb.dbname,
+        )
+
+        sql = fR'''
+            SELECT *
+            FROM {conn_name}.{table_name};
+            '''
+
+        query_id = fq_client.create_query(query_name, sql, type=fq.QueryContent.QueryType.ANALYTICS).result.query_id
+        fq_client.wait_query_status(query_id, fq.QueryMeta.COMPLETED)
+
+        data = fq_client.get_result_data(query_id)
+        result_set = data.result.result_set
+        logging.debug(str(result_set))
+        assert len(result_set.columns) == 1
+        assert result_set.columns[0].name == "number"
+        assert result_set.columns[0].type == ydb.Type(
+            optional_type=ydb.OptionalType(item=ydb.Type(type_id=ydb.Type.INT32))
+        )
+        assert len(result_set.rows) == 3
+        assert result_set.rows[0].items[0].int32_value == 1
+        assert result_set.rows[1].items[0].int32_value == 2
+        assert result_set.rows[2].items[0].int32_value == 3

ydb/tests/fq/generic/utils/settings.py

@@ -4,7 +4,7 @@
 
 import yatest.common
 
-from ydb.tests.tools.docker_compose_helpers.endpoint_determiner import EndpointDeterminer
+from ydb.tests.fq.generic.utils.endpoint_determiner import EndpointDeterminer
 
 
 @dataclass
@@ -22,6 +22,12 @@ class MdbMock:
 
     mdb_mock: MdbMock
 
+    @dataclass
+    class YdbMvpMock:
+        endpoint: str
+
+    ydb_mvp_mock: YdbMvpMock
+
     @dataclass
     class TokenAccessorMock:
         endpoint: str
@@ -46,6 +52,14 @@ class PostgreSQL:
 
     postgresql: PostgreSQL
 
+    @dataclass
+    class Ydb:
+        dbname: str
+        username: str
+        password: str
+
+    ydb: Ydb
+
     @classmethod
     def from_env(cls) -> 'Settings':
         docker_compose_file = yatest.common.source_path('ydb/tests/fq/generic/docker-compose.yml')
@@ -54,11 +68,14 @@ def from_env(cls) -> 'Settings':
         s = cls(
             connector=cls.Connector(
                 grpc_host='localhost',
-                grpc_port=endpoint_determiner.get_port('connector', 50051),
+                grpc_port=endpoint_determiner.get_port('fq-connector-go', 2130),
             ),
             mdb_mock=cls.MdbMock(
                 endpoint=environ['MDB_MOCK_ENDPOINT'],
             ),
+            ydb_mvp_mock=cls.YdbMvpMock(
+                endpoint=environ['YDB_MVP_MOCK_ENDPOINT'],
+            ),
             token_accessor_mock=cls.TokenAccessorMock(
                 endpoint=environ['TOKEN_ACCESSOR_MOCK_ENDPOINT'],
                 hmac_secret_file=environ['TOKEN_ACCESSOR_HMAC_SECRET_FILE'],
@@ -74,6 +91,7 @@ def from_env(cls) -> 'Settings':
                 username='user',
                 password='password',
             ),
+            ydb=cls.Ydb(dbname='local', username='user', password='password'),
         )
 
         return s

ydb/tests/fq/generic/utils/ya.make

@@ -3,7 +3,12 @@ PY3_LIBRARY()
 STYLE_PYTHON()
 
 PY_SRCS(
+    endpoint_determiner.py
     settings.py
 )
 
+PEERDIR(
+    library/python/testing/yatest_common
+)
+
 END()

ydb/tests/fq/generic/ya.make

@@ -5,31 +5,49 @@ PY3TEST()
 STYLE_PYTHON()
 NO_CHECK_IMPORTS()
 
-TAG(
-    ya:external
-    ya:force_sandbox
-    ya:fat
-)
-
-REQUIREMENTS(
-    container:4467981730
-    cpu:all
-    dns:dns64
-)
-
+INCLUDE(${ARCADIA_ROOT}/ydb/tests/tools/fq_runner/ydb_runner_with_datastreams.inc)
 INCLUDE(${ARCADIA_ROOT}/ydb/tests/tools/mdb_mock/recipe.inc)
 INCLUDE(${ARCADIA_ROOT}/ydb/tests/tools/token_accessor_mock/recipe.inc)
-INCLUDE(${ARCADIA_ROOT}/ydb/tests/tools/fq_runner/ydb_runner_with_datastreams.inc)
-INCLUDE(${ARCADIA_ROOT}/library/recipes/docker_compose/recipe.inc)
+INCLUDE(${ARCADIA_ROOT}/ydb/tests/tools/ydb_mvp_mock/recipe.inc)
 
-# Including of docker_compose/recipe.inc automatically converts these tests into LARGE,
-# which makes it impossible to run them during precommit checks on Github CI.
-# Next several lines forces these tests to be MEDIUM. To see discussion, visit YDBOPS-8928.
+IF (AUTOCHECK) 
+    # Temporarily disable these tests due to infrastructure incompatibility
+    SKIP_TEST("DEVTOOLSUPPORT-44637")
+
+    # Split tests to chunks only when they're running on different machines with distbuild,
+    # otherwise this directive will slow down local test execution.
+    # Look through DEVTOOLSSUPPORT-39642 for more information.
+    FORK_SUBTESTS()
+
+    # TAG and REQUIREMENTS are copied from: https://docs.yandex-team.ru/devtools/test/environment#docker-compose
+    TAG(
+        ya:external
+        ya:force_sandbox
+        ya:fat
+    )
+
+    REQUIREMENTS(
+        cpu:all
+        container:4467981730
+        dns:dns64
+    )
+ENDIF()
+
+INCLUDE(${ARCADIA_ROOT}/library/recipes/docker_compose/recipe.inc)
 
 IF (OPENSOURCE)
+    # Including of docker_compose/recipe.inc automatically converts these tests into LARGE, 
+    # which makes it impossible to run them during precommit checks on Github CI. 
+    # Next several lines forces these tests to be MEDIUM. To see discussion, visit YDBOPS-8928.
     SIZE(MEDIUM)
     SET(TEST_TAGS_VALUE)
     SET(TEST_REQUIREMENTS_VALUE)
+
+    # This requirement forces tests to be launched consequently,
+    # otherwise CI system would be overloaded due to simultaneous launch of many Docker containers.
+    # See DEVTOOLSSUPPORT-44103, YA-1759 for details.
+    TAG(ya:not_autocheck)
+    REQUIREMENTS(cpu:all)
 ENDIF()
 
 PEERDIR(
@@ -39,7 +57,6 @@ PEERDIR(
     library/python/testing/yatest_common
     library/recipes/common
     ydb/tests/tools/fq_runner
-    ydb/tests/tools/docker_compose_helpers
     ydb/public/api/protos
 
     contrib/python/pytest
@@ -50,6 +67,7 @@ TEST_SRCS(
     test_clickhouse.py
     test_join.py
     test_postgresql.py
+    test_ydb.py
 )
 
 END()

ydb/tests/fq/generic/ydb/01_basic.sh

@@ -0,0 +1,29 @@
+#!/bin/bash
+
+set -ex
+
+/ydb -p ${PROFILE} yql -s '
+    CREATE TABLE simple_table (number Int32, PRIMARY KEY (number));
+    COMMIT;
+    INSERT INTO simple_table (number) VALUES
+      (1),
+      (2),
+      (3);
+    COMMIT;
+
+    CREATE TABLE join_table (id Int32, data STRING, PRIMARY KEY (id));
+    COMMIT;
+    INSERT INTO join_table (id, data) VALUES
+      (1, "ydb10"),
+      (2, "ydb20"),
+      (3, "ydb30");
+    COMMIT;
+  '
+
+retVal=$?
+if [ $retVal -ne 0 ]; then
+  echo $retVal
+  exit $retVal
+fi
+
+echo $(date +"%T.%6N") "SUCCESS"

ydb/tests/fq/generic/ydb/init_ydb

@@ -0,0 +1,12 @@
+#!/bin/bash
+
+set -ex
+
+export PROFILE=tests-ydb-client-$(LC_ALL=C tr -dc A-Za-z0-9 </dev/urandom | head -c 8)
+
+# Wait until server is up
+timeout 30s bash -c 'until /ydb --endpoint grpc://tests-fq-generic-ydb:2136 --database /local scheme ls; do sleep 3; done'
+
+# Run initialization scripts
+/ydb config profile create ${PROFILE} --endpoint grpc://tests-fq-generic-ydb:2136 --database /local
+/bin/bash ./01_basic.sh

ydb/tests/tools/fq_runner/fq_client.py

@@ -390,13 +390,13 @@ def create_connection(self, request, check_issues=True):
         return FederatedQueryClient.Response(response.operation.issues, result, check_issues)
 
     @retry.retry_intrusive
-    def create_ydb_connection(self, name, database, endpoint, visibility=fq.Acl.Visibility.PRIVATE,
-                              auth_method=AuthMethod.no_auth(), check_issues=True):
+    def create_ydb_connection(self, name, database_id, 
+                              secure=False, visibility=fq.Acl.Visibility.PRIVATE, auth_method=AuthMethod.service_account('sa'), check_issues=True):
         request = fq.CreateConnectionRequest()
         request.content.name = name
         ydb = request.content.setting.ydb_database
-        ydb.database = database
-        ydb.endpoint = endpoint
+        ydb.database_id = database_id
+        ydb.secure = secure
 
         ydb.auth.CopyFrom(auth_method)
         request.content.acl.visibility = visibility