From eb88e211bb43cbf215be4b3aaf83b4a6eea3bd61 Mon Sep 17 00:00:00 2001
From: Alexandr Gorshenin <alexandr268@ydb.tech>
Date: Wed, 13 Aug 2025 11:46:08 +0100
Subject: [PATCH] Added option to read raw protobuf message in QueryStream

---
 query/src/main/java/tech/ydb/query/QueryStream.java      | 8 ++++++++
 query/src/main/java/tech/ydb/query/impl/SessionImpl.java | 3 +--
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/query/src/main/java/tech/ydb/query/QueryStream.java b/query/src/main/java/tech/ydb/query/QueryStream.java
index a4ee3ba7..53fcfeef 100644
--- a/query/src/main/java/tech/ydb/query/QueryStream.java
+++ b/query/src/main/java/tech/ydb/query/QueryStream.java
@@ -6,6 +6,7 @@
 
 import tech.ydb.core.Issue;
 import tech.ydb.core.Result;
+import tech.ydb.proto.ValueProtos;
 import tech.ydb.query.result.QueryInfo;
 import tech.ydb.query.result.QueryResultPart;
 
@@ -15,9 +16,16 @@
  */
 @ExperimentalApi("QueryService is experimental and API may change without notice")
 public interface QueryStream {
+
+    @FunctionalInterface
     interface PartsHandler {
         default void onIssues(Issue[] issues) { }
+
         void onNextPart(QueryResultPart part);
+
+        default void onNextRawPart(long index, ValueProtos.ResultSet rs) {
+            onNextPart(new QueryResultPart(index, rs));
+        }
     }
 
     CompletableFuture<Result<QueryInfo>> execute(PartsHandler handler);
diff --git a/query/src/main/java/tech/ydb/query/impl/SessionImpl.java b/query/src/main/java/tech/ydb/query/impl/SessionImpl.java
index b533b2fd..ef4d4c6b 100644
--- a/query/src/main/java/tech/ydb/query/impl/SessionImpl.java
+++ b/query/src/main/java/tech/ydb/query/impl/SessionImpl.java
@@ -29,7 +29,6 @@
 import tech.ydb.query.QueryStream;
 import tech.ydb.query.QueryTransaction;
 import tech.ydb.query.result.QueryInfo;
-import tech.ydb.query.result.QueryResultPart;
 import tech.ydb.query.result.QueryStats;
 import tech.ydb.query.settings.AttachSessionSettings;
 import tech.ydb.query.settings.BeginTransactionSettings;
@@ -304,7 +303,7 @@ public CompletableFuture<Result<QueryInfo>> execute(PartsHandler handler) {
                 if (msg.hasResultSet()) {
                     long index = msg.getResultSetIndex();
                     if (handler != null) {
-                        handler.onNextPart(new QueryResultPart(index, msg.getResultSet()));
+                        handler.onNextRawPart(index, msg.getResultSet());
                     } else {
                         logger.trace("{} lost result set part with index {}", SessionImpl.this, index);
                     }