YQ-4502 Streaming: type tests and compatibility (#28799)

kardymonds · web-flow · commit 91428ef65d5b · 2025-11-20T15:19:46.000+03:00
diff --git a/ydb/tests/compatibility/streaming/test_streaming.py b/ydb/tests/compatibility/streaming/test_streaming.py
@@ -0,0 +1,208 @@
+# -*- coding: utf-8 -*-
+import logging
+import os
+import pytest
+import time
+
+from ydb.tests.library.compatibility.fixtures import MixedClusterFixture, RestartToAnotherVersionFixture, RollingUpgradeAndDowngradeFixture
+from ydb.tests.library.harness.util import LogLevels
+from ydb.tests.library.test_meta import link_test_case
+from ydb.tests.oss.ydb_sdk_import import ydb
+from ydb.tests.tools.datastreams_helpers.data_plane import write_stream, read_stream
+
+logger = logging.getLogger(__name__)
+
+
+class StreamingTestBase:
+    def setup_cluster(self):
+        logger.debug(f"setup_cluster, versions {self.versions}")
+
+        if min(self.versions) < (25, 4):
+            logger.debug("skip test, only available since 25-4")
+            pytest.skip("Only available since 25-4")
+
+        os.environ["YDB_TEST_DEFAULT_CHECKPOINTING_PERIOD_MS"] = "200"
+        os.environ["YDB_TEST_LEASE_DURATION_SEC"] = "15"
+        yield from super().setup_cluster(
+            extra_feature_flags={
+                "enable_external_data_sources": True,
+                "enable_streaming_queries": True
+            },
+            additional_log_configs={
+                'KQP_COMPUTE': LogLevels.TRACE,
+                'STREAMS_CHECKPOINT_COORDINATOR': LogLevels.TRACE,
+                'STREAMS_STORAGE_SERVICE': LogLevels.TRACE,
+                'FQ_ROW_DISPATCHER': LogLevels.TRACE,
+                'KQP_PROXY': LogLevels.DEBUG,
+                'KQP_EXECUTOR': LogLevels.DEBUG},
+        )
+
+    def create_topics(self):
+        logger.debug("create_topics")
+        self.input_topic = 'streaming_recipe/input_topic'
+        self.output_topic = 'streaming_recipe/output_topic'
+        self.consumer_name = 'consumer_name'
+        with ydb.QuerySessionPool(self.driver) as session_pool:
+            query = f"""
+                CREATE TOPIC `{self.input_topic}`;
+                CREATE TOPIC `{self.output_topic}` (CONSUMER {self.consumer_name});
+            """
+            session_pool.execute_with_retries(query)
+
+    def create_external_data_source(self):
+        logger.debug("create_external_data_source")
+        endpoint = f"localhost:{self.cluster.nodes[1].port}"
+        with ydb.QuerySessionPool(self.driver) as session_pool:
+            query = f"""
+                CREATE EXTERNAL DATA SOURCE source_name WITH (
+                    SOURCE_TYPE="Ydb",
+                    LOCATION="{endpoint}",
+                    DATABASE_NAME="{self.database_path}",
+                    SHARED_READING="false",
+                    AUTH_METHOD="NONE");
+            """
+            session_pool.execute_with_retries(query)
+
+    def create_streaming_query(self):
+        logger.debug("create_streaming_query")
+        with ydb.QuerySessionPool(self.driver) as session_pool:
+            query = f"""
+                CREATE STREAMING QUERY `my_queries/query_name` AS DO BEGIN
+                $input = (
+                    SELECT * FROM
+                        source_name.`{self.input_topic}` WITH (
+                            FORMAT = 'json_each_row',
+                            SCHEMA (time String NOT NULL, level String NOT NULL, host String NOT NULL)
+                        )
+                );
+                $filtered = (SELECT * FROM $input WHERE level == 'error');
+
+                $number_errors = (
+                    SELECT host, COUNT(*) AS error_count, CAST(HOP_START() AS String) AS ts
+                    FROM $filtered
+                    GROUP BY
+                        HoppingWindow(CAST(time AS Timestamp), 'PT600S', 'PT600S'),
+                        host
+                );
+
+                $json = (SELECT ToBytes(Unwrap(Yson::SerializeJson(Yson::From(TableRow()))))
+                    FROM $number_errors
+                );
+
+                INSERT INTO source_name.`{self.output_topic}`
+                SELECT * FROM $json;
+                END DO;
+
+            """
+            session_pool.execute_with_retries(query)
+
+    def create_simple_streaming_query(self):
+        logger.debug("create_simple_streaming_query")
+        with ydb.QuerySessionPool(self.driver) as session_pool:
+            query = f"""
+                CREATE STREAMING QUERY `my_queries/query_name` AS DO BEGIN
+                $input = (
+                    SELECT
+                        *
+                    FROM
+                        source_name.`{self.input_topic}` WITH (
+                            FORMAT = 'json_each_row',
+                            SCHEMA (time String NOT NULL, level String NOT NULL, host String NOT NULL)
+                        )
+                );
+
+                $json = (SELECT ToBytes(Unwrap(Yson::SerializeJson(Yson::From(TableRow()))))
+                    FROM $input
+                );
+
+                INSERT INTO source_name.`{self.output_topic}`
+                SELECT * FROM $json;
+                END DO;
+
+            """
+            session_pool.execute_with_retries(query)
+
+    def do_write_read(self, input, expected_output):
+        logger.debug("do_write_read")
+        endpoint = f"localhost:{self.cluster.nodes[1].port}"
+        time.sleep(2)
+        logger.debug("write data to stream")
+        write_stream(path=self.input_topic, data=input, database=self.database_path, endpoint=endpoint)
+        logger.debug("read data from stream")
+        assert sorted(read_stream(
+            path=self.output_topic,
+            messages_count=len(expected_output),
+            consumer_name=self.consumer_name,
+            database=self.database_path,
+            endpoint=endpoint)) == sorted(expected_output)
+
+    def do_test_part1(self):
+        input = [
+            '{"time": "2025-01-01T00:00:00.000000Z", "level": "error", "host": "host-1"}',
+            '{"time": "2025-01-01T00:04:00.000000Z", "level": "error", "host": "host-2"}',
+            '{"time": "2025-01-01T00:08:00.000000Z", "level": "error", "host": "host-1"}',
+            '{"time": "2025-01-01T00:12:00.000000Z", "level": "error", "host": "host-2"}',
+            '{"time": "2025-01-01T00:12:00.000000Z", "level": "error", "host": "host-1"}']
+        expected_data = sorted([
+            '{"error_count":1,"host":"host-2","ts":"2025-01-01T00:00:00Z"}',
+            '{"error_count":2,"host":"host-1","ts":"2025-01-01T00:00:00Z"}'])
+        self.do_write_read(input, expected_data)
+
+    def do_test_part2(self):
+        input = [
+            '{"time": "2025-01-01T00:15:00.000000Z", "level": "error", "host": "host-2"}',
+            '{"time": "2025-01-01T00:22:00.000000Z", "level": "error", "host": "host-1"}',
+            '{"time": "2025-01-01T00:22:00.000000Z", "level": "error", "host": "host-2"}']
+        expected_data = sorted([
+            '{"error_count":2,"host":"host-2","ts":"2025-01-01T00:10:00Z"}',
+            '{"error_count":1,"host":"host-1","ts":"2025-01-01T00:10:00Z"}'])
+        self.do_write_read(input, expected_data)
+
+
+class TestStreamingMixedCluster(StreamingTestBase, MixedClusterFixture):
+    @pytest.fixture(autouse=True, scope="function")
+    def setup(self):
+        yield from self.setup_cluster()
+
+    @link_test_case("#27924")
+    def test_mixed_cluster(self):
+        self.create_topics()
+        self.create_external_data_source()
+        self.create_streaming_query()
+        self.do_test_part1()
+        self.do_test_part2()
+
+
+class TestStreamingRestartToAnotherVersion(StreamingTestBase, RestartToAnotherVersionFixture):
+    @pytest.fixture(autouse=True, scope="function")
+    def setup(self):
+        yield from self.setup_cluster()
+
+    @link_test_case("#27924")
+    def test_restart_to_another_version(self):
+        self.create_topics()
+        self.create_external_data_source()
+        self.create_streaming_query()
+        self.do_test_part1()
+        self.change_cluster_version()
+        self.do_test_part2()
+
+
+class TestStreamingRollingUpgradeAndDowngrade(StreamingTestBase, RollingUpgradeAndDowngradeFixture):
+    @pytest.fixture(autouse=True, scope="function")
+    def setup(self):
+        yield from self.setup_cluster()
+
+    @link_test_case("#27924")
+    def test_rolling_upgrage(self):
+        self.create_topics()
+        self.create_external_data_source()
+        self.create_simple_streaming_query()
+
+        for _ in self.roll():  # every iteration is a step in rolling upgrade process
+            #
+            # 2. check written data is correct during rolling upgrade
+            #
+            input = ['{"time": "2025-01-01T00:15:00.000000Z", "level": "error", "host": "host-2"}']
+            expected_data = ['{"host":"host-2","level":"error","time":"2025-01-01T00:15:00.000000Z"}']
+            self.do_write_read(input, expected_data)
diff --git a/ydb/tests/compatibility/streaming/ya.make b/ydb/tests/compatibility/streaming/ya.make
@@ -0,0 +1,32 @@
+PY3TEST()
+INCLUDE(${ARCADIA_ROOT}/ydb/tests/ydbd_dep.inc)
+
+
+FORK_TEST_FILES()
+FORK_TESTS()
+FORK_SUBTESTS()
+SPLIT_FACTOR(10)
+
+TEST_SRCS(
+    test_streaming.py
+)
+
+SIZE(LARGE)
+REQUIREMENTS(cpu:16)
+INCLUDE(${ARCADIA_ROOT}/ydb/tests/large.inc)
+
+
+DEPENDS(
+    ydb/tests/library/compatibility/binaries
+    ydb/tests/tools/pq_read
+)
+
+PEERDIR(
+    contrib/python/boto3
+    ydb/tests/library
+    ydb/tests/library/compatibility
+    ydb/tests/library/test_meta
+    ydb/tests/tools/datastreams_helpers
+)
+
+END()
diff --git a/ydb/tests/compatibility/ya.make b/ydb/tests/compatibility/ya.make
@@ -56,4 +56,5 @@ RECURSE(
     federated_queries
     s3_backups
     olap
+    streaming
 )
diff --git a/ydb/tests/fq/streaming/test_streaming.py b/ydb/tests/fq/streaming/test_streaming.py
@@ -414,3 +414,109 @@ def test_pragma(self, kikimr):
         assert self.read_stream(1, topic_path=self.output_topic) == ['lunch time']
 
         kikimr.YdbClient.query(f"DROP STREAMING QUERY `{query_name}`")
+
+    def test_types(self, kikimr):
+        sourceName = "test_types"
+        self.init_topics(sourceName, partitions_count=1)
+
+        self.create_source(kikimr, sourceName)
+
+        query_name = "test_types1"
+
+        def test_type(self, kikimr, type, input, expected_output):
+            sql = R'''
+                CREATE STREAMING QUERY `{query_name}` AS
+                DO BEGIN
+                    $in = SELECT field_name FROM {source_name}.`{input_topic}`
+                    WITH (
+                        FORMAT="json_each_row",
+                        SCHEMA=(field_name {type_name} NOT NULL));
+                    INSERT INTO {source_name}.`{output_topic}` SELECT CAST(field_name as String) FROM $in;
+                END DO;'''
+
+            kikimr.YdbClient.query(sql.format(query_name=query_name, source_name=sourceName, type_name=type, input_topic=self.input_topic, output_topic=self.output_topic))
+            self.write_stream([f"{{\"field_name\": {input}}}"])
+            assert self.read_stream(1, topic_path=self.output_topic) == [expected_output]
+            kikimr.YdbClient.query(f"DROP STREAMING QUERY `{query_name}`")
+
+        test_type(self, kikimr, type="String", input='"lunch time"', expected_output='lunch time')
+        test_type(self, kikimr, type="Utf8", input='"Relativitätstheorie"', expected_output='Relativitätstheorie')
+        test_type(self, kikimr, type="Int8", input='42', expected_output='42')
+        test_type(self, kikimr, type="Uint64", input='777', expected_output='777')
+        test_type(self, kikimr, type="Float", input='1024.1024', expected_output='1024.1024')
+        test_type(self, kikimr, type="Double", input='-777.777', expected_output='-777.777')
+        test_type(self, kikimr, type="Bool", input='true', expected_output='true')
+        test_type(self, kikimr, type="Uuid", input='"3d6c7233-d082-4b25-83e2-10d271bbc911"', expected_output='3d6c7233-d082-4b25-83e2-10d271bbc911')
+        # Unsupported
+        # test_type(self, kikimr, type="Timestamp", input='"2025-08-25 10:49:00"', expected_output='2025-08-25T10:49:00Z')
+        # test_type(self, kikimr, type="Json", input='{"name": "value"}', expected_output='{"name": "value"}')
+        # test_type(self, kikimr, type="JsonDocument", input='{"name": "value"}', expected_output='lunch time')
+
+    def test_raw_format(self, kikimr):
+        sourceName = "test_restart_query" + ''.join(random.choices(string.ascii_letters + string.digits, k=8))
+        self.init_topics(sourceName, partitions_count=10)
+        self.create_source(kikimr, sourceName, False)
+
+        query_name = "test_raw_format_string"
+        sql = R'''
+            CREATE STREAMING QUERY `{query_name}` AS
+            DO BEGIN
+                $input = SELECT CAST(data AS Json) AS json FROM {source_name}.`{input_topic}`
+                WITH (
+                    FORMAT="raw",
+                    SCHEMA=(data String));
+                $parsed = SELECT JSON_VALUE(json, "$.time") as k, JSON_VALUE(json, "$.value") as v FROM $input;
+                INSERT INTO {source_name}.`{output_topic}` SELECT ToBytes(Unwrap(Json::SerializeJson(Yson::From(TableRow())))) FROM $parsed;
+            END DO;'''
+        path = f"/Root/{query_name}"
+        kikimr.YdbClient.query(sql.format(query_name=query_name, source_name=sourceName, input_topic=self.input_topic, output_topic=self.output_topic))
+        self.wait_completed_checkpoints(kikimr, path)
+
+        data = ['{"time": "2020-01-01T13:00:00.000000Z", "value": "lunch time"}']
+        expected_data = ['{"k":"2020-01-01T13:00:00.000000Z","v":"lunch time"}']
+        self.write_stream(data)
+
+        assert self.read_stream(len(expected_data), topic_path=self.output_topic) == expected_data
+        kikimr.YdbClient.query(f"DROP STREAMING QUERY `{query_name}`")
+
+        query_name = "test_raw_format_default"
+        sql = R'''
+            CREATE STREAMING QUERY `{query_name}` AS
+            DO BEGIN
+                $input = SELECT CAST(Data AS Json) AS json FROM {source_name}.`{input_topic}`;
+                $parsed = SELECT JSON_VALUE(json, "$.time") as k, JSON_VALUE(json, "$.value") as v FROM $input;
+                INSERT INTO {source_name}.`{output_topic}` SELECT ToBytes(Unwrap(Json::SerializeJson(Yson::From(TableRow())))) FROM $parsed;
+            END DO;'''
+        path = f"/Root/{query_name}"
+        kikimr.YdbClient.query(sql.format(query_name=query_name, source_name=sourceName, input_topic=self.input_topic, output_topic=self.output_topic))
+        self.wait_completed_checkpoints(kikimr, path)
+
+        data = ['{"time": "2020-01-01T13:00:00.000000Z", "value": "lunch time"}']
+        expected_data = ['{"k":"2020-01-01T13:00:00.000000Z","v":"lunch time"}']
+        self.write_stream(data)
+
+        assert self.read_stream(len(expected_data), topic_path=self.output_topic) == expected_data
+        kikimr.YdbClient.query(f"DROP STREAMING QUERY `{query_name}`")
+
+        query_name = "test_raw_format_json"
+        sql = R'''
+            CREATE STREAMING QUERY `{query_name}` AS
+            DO BEGIN
+                $input = SELECT data AS json FROM {source_name}.`{input_topic}`
+                WITH (
+                    FORMAT="raw",
+                    SCHEMA=(data Json));
+                $parsed = SELECT JSON_VALUE(json, "$.time") as k, JSON_VALUE(json, "$.value") as v FROM $input;
+                INSERT INTO {source_name}.`{output_topic}` SELECT ToBytes(Unwrap(Json::SerializeJson(Yson::From(TableRow())))) FROM $parsed;
+            END DO;'''
+        path = f"/Root/{query_name}"
+        kikimr.YdbClient.query(sql.format(query_name=query_name, source_name=sourceName, input_topic=self.input_topic, output_topic=self.output_topic))
+        self.wait_completed_checkpoints(kikimr, path)
+
+        data = ['{"time": "2020-01-01T13:00:00.000000Z", "value": "lunch time"}']
+        expected_data = ['{"k":"2020-01-01T13:00:00.000000Z","v":"lunch time"}']
+        self.write_stream(data)
+
+        assert self.read_stream(len(expected_data), topic_path=self.output_topic) == expected_data
+
+        kikimr.YdbClient.query(f"DROP STREAMING QUERY `{query_name}`")
diff --git a/ydb/tests/tools/datastreams_helpers/data_plane.py b/ydb/tests/tools/datastreams_helpers/data_plane.py
diff --git a/ydb/tests/tools/datastreams_helpers/test_yds_base.py b/ydb/tests/tools/datastreams_helpers/test_yds_base.py

Original file line number	Diff line number	Diff line change
`@@ -56,4 +56,5 @@ RECURSE(`
`56`	`56`	`federated_queries`
`57`	`57`	`s3_backups`
`58`	`58`	`olap`
	`59`	`+ streaming`
`59`	`60`	`)`