, fixed truncated flag

Fixed truncated flag

2 files changed, 40 insertions, 1 deletions

diff --git a/ydb/core/kqp/run_script_actor/kqp_run_script_actor.cpp b/ydb/core/kqp/run_script_actor/kqp_run_script_actor.cpp
index 65231f36e3..d3c53e3148 100644
--- a/ydb/core/kqp/run_script_actor/kqp_run_script_actor.cpp
+++ b/ydb/core/kqp/run_script_actor/kqp_run_script_actor.cpp
@@ -259,7 +259,7 @@ private:
             std::vector<TString> serializedRows;
 
             for (const auto& row : ev->Get()->Record.GetResultSet().rows()) {
-                if (QueryServiceConfig.GetScriptResultRowsLimit() && rowCount > QueryServiceConfig.GetScriptResultRowsLimit()) {
+                if (QueryServiceConfig.GetScriptResultRowsLimit() && rowCount + 1 > QueryServiceConfig.GetScriptResultRowsLimit()) {
                     Truncated[resultSetIndex] = true;
                     break;
                 }
diff --git a/ydb/core/kqp/ut/service/kqp_qs_scripts_ut.cpp b/ydb/core/kqp/ut/service/kqp_qs_scripts_ut.cpp
index 8aaafae20f..718988eba9 100644
--- a/ydb/core/kqp/ut/service/kqp_qs_scripts_ut.cpp
+++ b/ydb/core/kqp/ut/service/kqp_qs_scripts_ut.cpp
@@ -656,6 +656,45 @@ Y_UNIT_TEST_SUITE(KqpQueryServiceScripts) {
             settings.FetchToken(results.GetNextFetchToken());
         }
     }
+
+    Y_UNIT_TEST(TestTruncatedByRows) {
+        constexpr size_t ROWS_LIMIT = 2000;
+
+        NKikimrConfig::TAppConfig appCfg;
+        appCfg.MutableQueryServiceConfig()->set_scriptresultrowslimit(ROWS_LIMIT);
+        appCfg.MutableTableServiceConfig()->MutableQueryLimits()->set_resultrowslimit(ROWS_LIMIT);
+
+        auto kikimr = DefaultKikimrRunner({}, appCfg);
+        auto db = kikimr.GetQueryClient();
+
+        auto scriptExecutionOperationTruncated = CreateScriptExecutionOperation(ROWS_LIMIT + 1, db, kikimr.GetDriver());
+        TFetchScriptResultsResult resultsTruncated = db.FetchScriptResults(scriptExecutionOperationTruncated.Id(), 0).ExtractValueSync();
+        UNIT_ASSERT_C(resultsTruncated.IsSuccess(), resultsTruncated.GetIssues().ToString());
+        UNIT_ASSERT(resultsTruncated.GetResultSet().Truncated());
+
+        auto scriptExecutionOperationNotTruncated = CreateScriptExecutionOperation(ROWS_LIMIT, db, kikimr.GetDriver());
+        TFetchScriptResultsResult resultsNotTruncated = db.FetchScriptResults(scriptExecutionOperationNotTruncated.Id(), 0).ExtractValueSync();
+        UNIT_ASSERT_C(resultsNotTruncated.IsSuccess(), resultsNotTruncated.GetIssues().ToString());
+        UNIT_ASSERT(!resultsNotTruncated.GetResultSet().Truncated());
+    }
+
+    Y_UNIT_TEST(TestTruncatedBySize) {
+        constexpr size_t NUMER_ROWS = 500;
+
+        NKikimrConfig::TAppConfig appCfg;
+        appCfg.MutableQueryServiceConfig()->set_scriptresultsizelimit(NUMER_ROWS / 2);
+        appCfg.MutableTableServiceConfig()->MutableQueryLimits()->set_resultrowslimit(NUMER_ROWS);
+
+        auto kikimr = DefaultKikimrRunner({}, appCfg);
+        auto db = kikimr.GetQueryClient();
+
+        auto scriptExecutionOperation = CreateScriptExecutionOperation(NUMER_ROWS, db, kikimr.GetDriver());
+
+        TFetchScriptResultsResult results = db.FetchScriptResults(scriptExecutionOperation.Id(), 0).ExtractValueSync();
+        UNIT_ASSERT_C(results.IsSuccess(), results.GetIssues().ToString());
+        
+        UNIT_ASSERT(results.GetResultSet().Truncated());
+    }
 }
 
 } // namespace NKqp