Allow to decrease dq query timeout from SQL (#12285)

4 files changed, 18 insertions, 2 deletions

diff --git a/ydb/library/yql/providers/dq/actors/resource_allocator.cpp b/ydb/library/yql/providers/dq/actors/resource_allocator.cpp
index 34c412a082..8c143ec3da 100644
--- a/ydb/library/yql/providers/dq/actors/resource_allocator.cpp
+++ b/ydb/library/yql/providers/dq/actors/resource_allocator.cpp
@@ -92,7 +92,7 @@ private:
     static TDuration GetAllocationTimeout(int workerCount, const TDqConfiguration::TPtr settings) {
         ui64 timeout = workerCount == 1
             ? settings->_LiteralTimeout.Get().GetOrElse(TDqSettings::TDefault::LiteralTimeout)
-            : settings->_TableTimeout.Get().GetOrElse(TDqSettings::TDefault::TableTimeout);
+            : settings->GetQueryTimeout();
         if (timeout & (1ULL<<63)) {
             return TDuration(); // no timeout
         }
diff --git a/ydb/library/yql/providers/dq/common/yql_dq_settings.cpp b/ydb/library/yql/providers/dq/common/yql_dq_settings.cpp
index 282fa1dffd..92252eeb11 100644
--- a/ydb/library/yql/providers/dq/common/yql_dq_settings.cpp
+++ b/ydb/library/yql/providers/dq/common/yql_dq_settings.cpp
@@ -31,6 +31,8 @@ TDqConfiguration::TDqConfiguration() {
 
     REGISTER_SETTING(*this, _LiteralTimeout);
     REGISTER_SETTING(*this, _TableTimeout);
+    REGISTER_SETTING(*this, QueryTimeout);
+
     REGISTER_SETTING(*this, _LongWorkersAllocationWarnTimeout);
     REGISTER_SETTING(*this, _LongWorkersAllocationFailTimeout);
 
diff --git a/ydb/library/yql/providers/dq/common/yql_dq_settings.h b/ydb/library/yql/providers/dq/common/yql_dq_settings.h
index 8ea2477e6a..852df36bd6 100644
--- a/ydb/library/yql/providers/dq/common/yql_dq_settings.h
+++ b/ydb/library/yql/providers/dq/common/yql_dq_settings.h
@@ -15,6 +15,7 @@
 namespace NYql {
 
 struct TDqSettings {
+    friend struct TDqConfiguration;
 
     enum class ETaskRunnerStats {
         Disable,
@@ -89,7 +90,10 @@ struct TDqSettings {
     NCommon::TConfSetting<ui64, false> ChunkSizeLimit;
     NCommon::TConfSetting<NSize::TSize, false> MemoryLimit;
     NCommon::TConfSetting<ui64, false> _LiteralTimeout;
+private:
     NCommon::TConfSetting<ui64, false> _TableTimeout;
+    NCommon::TConfSetting<ui64, false> QueryTimeout; // less or equal than _TableTimeout
+public:
     NCommon::TConfSetting<ui64, false> _LongWorkersAllocationWarnTimeout;
     NCommon::TConfSetting<ui64, false> _LongWorkersAllocationFailTimeout;
     NCommon::TConfSetting<bool, false> EnableInsert;
@@ -168,6 +172,7 @@ struct TDqSettings {
         SAVE_SETTING(MemoryLimit);
         SAVE_SETTING(_LiteralTimeout);
         SAVE_SETTING(_TableTimeout);
+        SAVE_SETTING(QueryTimeout);
         SAVE_SETTING(_LongWorkersAllocationWarnTimeout);
         SAVE_SETTING(_LongWorkersAllocationFailTimeout);
         SAVE_SETTING(_AllResultsBytesLimit);
@@ -220,6 +225,15 @@ struct TDqSettings {
         }
     }
 
+    ui64 GetQueryTimeout() const {
+        auto upper = _TableTimeout.Get().GetOrElse(TDefault::TableTimeout);
+        if (QueryTimeout.Get().Defined()) {
+            return Min(*QueryTimeout.Get(), upper);
+        }
+
+        return upper;
+    }
+
     bool IsSpillingEngineEnabled() const {
         return SpillingEngine.Get().GetOrElse(TDqSettings::TDefault::SpillingEngine) != ESpillingEngine::Disable;
     }
diff --git a/ydb/library/yql/providers/dq/provider/exec/yql_dq_exectransformer.cpp b/ydb/library/yql/providers/dq/provider/exec/yql_dq_exectransformer.cpp
index 942d09bac4..f1daa3b760 100644
--- a/ydb/library/yql/providers/dq/provider/exec/yql_dq_exectransformer.cpp
+++ b/ydb/library/yql/providers/dq/provider/exec/yql_dq_exectransformer.cpp
@@ -1267,7 +1267,7 @@ private:
         YQL_CLOG(TRACE, ProviderDq) << "HandlePull " << NCommon::ExprToPrettyString(ctx, *input);
 
         TInstant startTime = TInstant::Now();
-        ui64 executionTimeout = State->Settings->_TableTimeout.Get().GetOrElse(TDqSettings::TDefault::TableTimeout);
+        ui64 executionTimeout = State->Settings->GetQueryTimeout();
         auto pull = TPull(input);
 
         THashMap<TString, TString> pullSettings;