YQ-2730: use docker-compose random port mapping in connector integration tests (#978)

* Use docker-compose random port mapping

* Retry creating connections for PostgreSQL

* Close ClickHouse and PostgreSQL connections after use

11 files changed, 114 insertions, 41 deletions

diff --git a/.github/config/muted_ya.txt b/.github/config/muted_ya.txt
index 3ab191db3a..57508bf6ac 100644
--- a/.github/config/muted_ya.txt
+++ b/.github/config/muted_ya.txt
@@ -19,7 +19,6 @@ ydb/core/tx/columnshard/ut_schema TColumnShardTestSchema.RebootForgetAfterFail
 ydb/core/tx/columnshard/engines/ut *
 ydb/core/tx/coordinator/ut Coordinator.RestoreTenantConfiguration
 ydb/core/tx/schemeshard/ut_split_merge TSchemeShardSplitBySizeTest.Merge1KShards
-ydb/library/yql/providers/generic/connector/tests test.py.test_select_positive_postgresql*
 ydb/library/yql/sql/pg/ut PgSqlParsingAutoparam.AutoParamValues_DifferentTypes
 ydb/library/yql/tests/sql/dq_file/part16 test.py.test[expr-as_dict_list_key-default.txt-Analyze]
 ydb/library/yql/tests/sql/dq_file/part18 test.py.test[expr-cast_type_bind-default.txt-Analyze]
diff --git a/ydb/library/yql/providers/generic/connector/tests/conftest.py b/ydb/library/yql/providers/generic/connector/tests/conftest.py
index 554fa42f78..e89287fd8c 100644
--- a/ydb/library/yql/providers/generic/connector/tests/conftest.py
+++ b/ydb/library/yql/providers/generic/connector/tests/conftest.py
@@ -21,7 +21,9 @@ def settings() -> Settings:
 
 @pytest.fixture
 def clickhouse_client(settings) -> utils.clickhouse.Client:
-    return utils.clickhouse.make_client(settings.clickhouse)
+    client = utils.clickhouse.make_client(settings.clickhouse)
+    yield client
+    client.close()
 
 
 @pytest.fixture
diff --git a/ydb/library/yql/providers/generic/connector/tests/docker-compose.yml b/ydb/library/yql/providers/generic/connector/tests/docker-compose.yml
index 0521e59630..a8bfa53fb1 100644
--- a/ydb/library/yql/providers/generic/connector/tests/docker-compose.yml
+++ b/ydb/library/yql/providers/generic/connector/tests/docker-compose.yml
@@ -2,24 +2,26 @@ version: '3.4'
 services:
   postgresql:
     image: postgres:15-bullseye@sha256:3411b9f2e5239cd7867f34fcf22fe964230f7d447a71d63c283e3593d3f84085
+    container_name: ${USER}_connector-integration-tests-postgresql
     environment:
       POSTGRES_DB: db
       POSTGRES_USER: user
       POSTGRES_PASSWORD: password
     ports:
-      - 15432:5432
+      - 5432
   clickhouse:
     image: clickhouse/clickhouse-server:23-alpine@sha256:b078c1cd294632afa2aeba3530e7ba2e568513da23304354f455a25fab575c06
+    container_name: ${USER}_connector-integration-tests-clickhouse
     environment:
       CLICKHOUSE_DB: db
       CLICKHOUSE_USER: user
       CLICKHOUSE_DEFAULT_ACCESS_MANAGEMENT: 1
       CLICKHOUSE_PASSWORD: password
     ports:
-      - 19000:9000
-      - 18123:8123
+      - 9000
+      - 8123
   fq-connector-go:
+    container_name: ${USER}_connector-integration-tests-fq-connector-go
     image: ghcr.io/ydb-platform/fq-connector-go:v0.1.1@sha256:47e24df143aee31a83d4a4cd0acc20b4cab8c03a9c63e81a6e99cb017a31f916
     ports:
-      - 50051:50051
-    network_mode: host
+      - 50051
diff --git a/ydb/library/yql/providers/generic/connector/tests/utils/clickhouse.py b/ydb/library/yql/providers/generic/connector/tests/utils/clickhouse.py
index 293c48527d..991b808c1c 100644
--- a/ydb/library/yql/providers/generic/connector/tests/utils/clickhouse.py
+++ b/ydb/library/yql/providers/generic/connector/tests/utils/clickhouse.py
@@ -18,10 +18,10 @@ def make_client(s: Settings.ClickHouse) -> Client:
         attempt += 1
         try:
             client = clickhouse_connect.get_client(
-                host=s.host, port=s.http_port, username=s.username, password=s.password
+                host=s.host_external, port=s.http_port_external, username=s.username, password=s.password
             )
         except Exception as e:
-            sys.stderr.write(f"attempt #{attempt}: {e}")
+            sys.stderr.write(f"attempt #{attempt}: {e}\n")
             time.sleep(5)
             continue
 
diff --git a/ydb/library/yql/providers/generic/connector/tests/utils/docker_compose.py b/ydb/library/yql/providers/generic/connector/tests/utils/docker_compose.py
new file mode 100644
index 0000000000..7d5b2b13d8
--- /dev/null
+++ b/ydb/library/yql/providers/generic/connector/tests/utils/docker_compose.py
@@ -0,0 +1,22 @@
+import os
+import subprocess
+
+import yatest.common
+
+
+class EndpointDeterminer:
+    docker_compose_bin: os.PathLike
+    docker_compose_yml: os.PathLike
+
+    def __init__(self, docker_compose_yml: os.PathLike):
+        self.docker_compose_bin = yatest.common.build_path('library/recipes/docker_compose/bin/docker-compose')
+        self.docker_compose_yml = docker_compose_yml
+
+    def get_port(self, service_name: str, internal_port: int) -> int:
+        cmd = [self.docker_compose_bin, '-f', self.docker_compose_yml, 'port', service_name, str(internal_port)]
+        try:
+            out = subprocess.check_output(cmd, stderr=subprocess.STDOUT)
+            external_port = int(out.split(b':')[1])
+            return external_port
+        except subprocess.CalledProcessError as e:
+            raise RuntimeError(f"docker-compose error: {e.output} (code {e.returncode})")
diff --git a/ydb/library/yql/providers/generic/connector/tests/utils/dqrun.py b/ydb/library/yql/providers/generic/connector/tests/utils/dqrun.py
index dcd1f17c3f..affacc0657 100644
--- a/ydb/library/yql/providers/generic/connector/tests/utils/dqrun.py
+++ b/ydb/library/yql/providers/generic/connector/tests/utils/dqrun.py
@@ -66,17 +66,17 @@ Generic {
 {% for cluster in generic_settings.clickhouse_clusters %}
 
 {% if cluster.protocol == EProtocol.NATIVE %}
-{% set CLICKHOUSE_PORT = settings.clickhouse.native_port %}
+{% set CLICKHOUSE_PORT = settings.clickhouse.native_port_internal %}
 {% set CLICKHOUSE_PROTOCOL = NATIVE %}
 {% elif cluster.protocol == EProtocol.HTTP %}
-{% set CLICKHOUSE_PORT = settings.clickhouse.http_port %}
+{% set CLICKHOUSE_PORT = settings.clickhouse.http_port_internal %}
 {% set CLICKHOUSE_PROTOCOL = HTTP %}
 {% endif %}
 
 {{ data_source(
     CLICKHOUSE,
     settings.clickhouse.cluster_name,
-    settings.clickhouse.host,
+    settings.clickhouse.host_internal,
     CLICKHOUSE_PORT,
     settings.clickhouse.username,
     settings.clickhouse.password,
@@ -90,8 +90,8 @@ Generic {
 {{ data_source(
     POSTGRESQL,
     settings.postgresql.cluster_name,
-    settings.postgresql.host,
-    settings.postgresql.port,
+    settings.postgresql.host_internal,
+    settings.postgresql.port_internal,
     settings.postgresql.username,
     settings.postgresql.password,
     NATIVE,
diff --git a/ydb/library/yql/providers/generic/connector/tests/utils/kqprun.py b/ydb/library/yql/providers/generic/connector/tests/utils/kqprun.py
index d37c861504..039db0e1d2 100644
--- a/ydb/library/yql/providers/generic/connector/tests/utils/kqprun.py
+++ b/ydb/library/yql/providers/generic/connector/tests/utils/kqprun.py
@@ -52,17 +52,17 @@ CREATE EXTERNAL DATA SOURCE {{data_source}} WITH (
 {% for cluster in generic_settings.clickhouse_clusters %}
 
 {% if cluster.protocol == EProtocol.NATIVE %}
-{% set CLICKHOUSE_PORT = settings.clickhouse.native_port %}
+{% set CLICKHOUSE_PORT = settings.clickhouse.native_port_internal %}
 {% set CLICKHOUSE_PROTOCOL = NATIVE %}
 {% elif cluster.protocol == EProtocol.HTTP %}
-{% set CLICKHOUSE_PORT = settings.clickhouse.http_port %}
+{% set CLICKHOUSE_PORT = settings.clickhouse.http_port_internal %}
 {% set CLICKHOUSE_PROTOCOL = HTTP %}
 {% endif %}
 
 {{ create_data_source(
     CLICKHOUSE,
     settings.clickhouse.cluster_name,
-    settings.clickhouse.host,
+    settings.clickhouse.host_internal,
     CLICKHOUSE_PORT,
     settings.clickhouse.username,
     settings.clickhouse.password,
@@ -76,8 +76,8 @@ CREATE EXTERNAL DATA SOURCE {{data_source}} WITH (
 {{ create_data_source(
     POSTGRESQL,
     settings.postgresql.cluster_name,
-    settings.postgresql.host,
-    settings.postgresql.port,
+    settings.postgresql.host_internal,
+    settings.postgresql.port_internal,
     settings.postgresql.username,
     settings.postgresql.password,
     NATIVE,
diff --git a/ydb/library/yql/providers/generic/connector/tests/utils/postgresql.py b/ydb/library/yql/providers/generic/connector/tests/utils/postgresql.py
index 9c290a4fe9..8e9389da94 100644
--- a/ydb/library/yql/providers/generic/connector/tests/utils/postgresql.py
+++ b/ydb/library/yql/providers/generic/connector/tests/utils/postgresql.py
@@ -1,5 +1,9 @@
 from contextlib import contextmanager
 import abc
+import time
+from datetime import datetime
+from typing import Tuple
+import sys
 
 import pg8000.dbapi
 
@@ -16,17 +20,39 @@ class Client:
 
     @contextmanager
     def get_cursor(self, dbname: str):
-        conn = pg8000.dbapi.Connection(
-            user=self.settings.username,
-            password=self.settings.password,
-            host=self.settings.host,
-            port=self.settings.port,
-            database=dbname,
-        )
-        conn.autocommit = True
-
-        cur = conn.cursor()
-        yield conn, cur
+        conn, cursor = self._make_cursor(dbname=dbname)
+        yield conn, cursor
+        cursor.close()
+        conn.close()
+
+    def _make_cursor(self, dbname: str) -> Tuple[pg8000.dbapi.Connection, pg8000.dbapi.Cursor]:
+        start = datetime.now()
+        attempt = 0
+
+        while (datetime.now() - start).total_seconds() < 10:
+            attempt += 1
+            try:
+                sys.stdout.write(
+                    f"Trying to connect PostgreSQL: {self.settings.host_external}:{self.settings.port_external}\n"
+                )
+                conn = pg8000.dbapi.Connection(
+                    user=self.settings.username,
+                    password=self.settings.password,
+                    host=self.settings.host_external,
+                    port=self.settings.port_external,
+                    database=dbname,
+                    timeout=1,
+                )
+                conn.autocommit = True
+
+                cur = conn.cursor()
+                return conn, cur
+            except Exception as e:
+                sys.stderr.write(f"attempt #{attempt} failed: {e} {e.args}\n")
+                time.sleep(3)
+                continue
+
+        raise Exception(f"Failed to connect PostgreSQL in {attempt} attempt(s)")
 
 
 class Type(abc.ABC):
diff --git a/ydb/library/yql/providers/generic/connector/tests/utils/settings.py b/ydb/library/yql/providers/generic/connector/tests/utils/settings.py
index bd78fdc822..97bfde8d33 100644
--- a/ydb/library/yql/providers/generic/connector/tests/utils/settings.py
+++ b/ydb/library/yql/providers/generic/connector/tests/utils/settings.py
@@ -1,8 +1,11 @@
 from dataclasses import dataclass
 from typing import Optional, Sequence
 
+import yatest.common
+
 from ydb.library.yql.providers.generic.connector.api.common.data_source_pb2 import EDataSourceKind, EProtocol
 from ydb.library.yql.providers.generic.connector.api.service.protos.connector_pb2 import EDateTimeFormat
+from ydb.library.yql.providers.generic.connector.tests.utils.docker_compose import EndpointDeterminer
 
 
 @dataclass
@@ -21,9 +24,12 @@ class Settings:
         cluster_name: str
         username: str
         password: str
-        host: str
-        http_port: int
-        native_port: int
+        host_external: str
+        host_internal: str
+        http_port_external: int
+        http_port_internal: int
+        native_port_external: int
+        native_port_internal: int
         protocol: str
 
     clickhouse: ClickHouse
@@ -34,33 +40,45 @@ class Settings:
         cluster_name: str
         username: str
         password: Optional[str]
-        host: str
-        port: int
+        host_external: str
+        host_internal: str
+        port_external: int
+        port_internal: int
 
     postgresql: PostgreSQL
 
     @classmethod
     def from_env(cls) -> 'Settings':
+        docker_compose_file = yatest.common.source_path(
+            'ydb/library/yql/providers/generic/connector/tests/docker-compose.yml'
+        )
+        endpoint_determiner = EndpointDeterminer(docker_compose_file)
+
         return cls(
             connector=cls.Connector(
                 grpc_host='localhost',
-                grpc_port=50051,
+                grpc_port=endpoint_determiner.get_port('fq-connector-go', 50051),
                 paging_bytes_per_page=4 * 1024 * 1024,
                 paging_prefetch_queue_capacity=2,
             ),
             clickhouse=cls.ClickHouse(
                 cluster_name='clickhouse_integration_test',
-                host='localhost',
-                http_port=18123,
-                native_port=19000,
+                host_external='localhost',
+                host_internal='clickhouse',
+                http_port_external=endpoint_determiner.get_port('clickhouse', 8123),
+                native_port_external=endpoint_determiner.get_port('clickhouse', 9000),
+                http_port_internal=8123,
+                native_port_internal=9000,
                 username='user',
                 password='password',
                 protocol='native',
             ),
             postgresql=cls.PostgreSQL(
                 cluster_name='postgresql_integration_test',
-                host='localhost',
-                port=15432,
+                host_external='localhost',
+                host_internal='postgresql',
+                port_external=endpoint_determiner.get_port('postgresql', 5432),
+                port_internal=5432,
                 dbname='db',
                 username='user',
                 password='password',
diff --git a/ydb/library/yql/providers/generic/connector/tests/utils/ya.make b/ydb/library/yql/providers/generic/connector/tests/utils/ya.make
index 5b53dcacd7..d61320c8d4 100644
--- a/ydb/library/yql/providers/generic/connector/tests/utils/ya.make
+++ b/ydb/library/yql/providers/generic/connector/tests/utils/ya.make
@@ -7,6 +7,7 @@ PY_SRCS(
     comparator.py
     database.py
     dqrun.py
+    docker_compose.py
     generate.py
     kqprun.py
     log.py
diff --git a/ydb/library/yql/providers/generic/connector/tests/ya.make b/ydb/library/yql/providers/generic/connector/tests/ya.make
index 0ae1a3dcb6..0100555d73 100644
--- a/ydb/library/yql/providers/generic/connector/tests/ya.make
+++ b/ydb/library/yql/providers/generic/connector/tests/ya.make
@@ -3,6 +3,8 @@ PY3TEST()
 STYLE_PYTHON()
 NO_CHECK_IMPORTS()
 
+DATA(arcadia/ydb/library/yql/providers/generic/connector/tests/docker-compose.yml)
+
 IF (AUTOCHECK) 
     # Split tests to chunks only when they're running on different machines with distbuild,
     # otherwise this directive will slow down local test execution.
@@ -61,6 +63,7 @@ PEERDIR(
 DEPENDS(
     ydb/library/yql/tools/dqrun
     ydb/tests/tools/kqprun
+    library/recipes/docker_compose/bin
 )
 
 END()