abort hidden on DqEngine="disable"

12 files changed, 37 insertions, 16 deletions

diff --git a/ydb/library/yql/core/facade/yql_facade.cpp b/ydb/library/yql/core/facade/yql_facade.cpp
index 35d571b1f4..43ec044998 100644
--- a/ydb/library/yql/core/facade/yql_facade.cpp
+++ b/ydb/library/yql/core/facade/yql_facade.cpp
@@ -198,7 +198,7 @@ TProgramPtr TProgramFactory::Create(
         CreateDeterministicTimeProvider(10000000) : CreateDefaultTimeProvider();
 
     TUdfIndex::TPtr udfIndex = UdfIndex_ ? UdfIndex_->Clone() : nullptr;
-    TUdfIndexPackageSet::TPtr udfIndexPackageSet = UdfIndexPackageSet_ ? UdfIndexPackageSet_->Clone() : nullptr;
+    TUdfIndexPackageSet::TPtr udfIndexPackageSet = (UdfIndexPackageSet_ && hiddenMode == EHiddenMode::Disable) ? UdfIndexPackageSet_->Clone() : nullptr;
     IModuleResolver::TPtr moduleResolver = Modules_ ? Modules_->CreateMutableChild() : nullptr;
     auto udfResolver = udfIndex ? NCommon::CreateUdfResolverWithIndex(udfIndex, UdfResolver_, FileStorage_) : UdfResolver_;
 
@@ -1007,7 +1007,7 @@ TFuture<IGraphTransformer::TStatus> TProgram::AsyncTransformWithFallback(bool ap
             FallbackCounter ++;
             // don't execute recapture again
             ExprCtx_->Step.Done(TExprStep::Recapture);
-            BeforeFallback();
+            AbortHidden_();
             return AsyncTransformWithFallback(false);
         }
         if (status == IGraphTransformer::TStatus::Error && (TypeCtx_->DqFallbackPolicy == "never" || TypeCtx_->ForceDq)) {
@@ -1317,7 +1317,8 @@ TTypeAnnotationContextPtr TProgram::BuildTypeAnnotationContext(const TString& us
             RandomProvider_,
             typeAnnotationContext,
             ProgressWriter_,
-            OperationOptions_
+            OperationOptions_,
+            AbortHidden_
         );
         if (HiddenMode_ != EHiddenMode::Disable && !dp.SupportsHidden) {
             continue;
diff --git a/ydb/library/yql/core/facade/yql_facade.h b/ydb/library/yql/core/facade/yql_facade.h
index cf1fc92f57..5d64fc336a 100644
--- a/ydb/library/yql/core/facade/yql_facade.h
+++ b/ydb/library/yql/core/facade/yql_facade.h
@@ -296,8 +296,8 @@ public:
 
     IPlanBuilder& GetPlanBuilder();
 
-    void SetBeforeFallback(std::function<void()>&& func) {
-        BeforeFallback = std::move(func);
+    void SetAbortHidden(std::function<void()>&& func) {
+        AbortHidden_ = std::move(func);
     }
 
 private:
@@ -397,8 +397,8 @@ private:
     const bool EnableRangeComputeFor_;
     const IArrowResolver::TPtr ArrowResolver_;
     i64 FallbackCounter = 0;
-    std::function<void()> BeforeFallback = [](){};
     const EHiddenMode HiddenMode_ = EHiddenMode::Disable;
+    THiddenQueryAborter AbortHidden_ = [](){};
 };
 
 } // namspace NYql
diff --git a/ydb/library/yql/core/yql_data_provider.h b/ydb/library/yql/core/yql_data_provider.h
index ad2667e3e1..ab4fefae5f 100644
--- a/ydb/library/yql/core/yql_data_provider.h
+++ b/ydb/library/yql/core/yql_data_provider.h
@@ -231,6 +231,7 @@ struct TDataProviderInfo {
     std::function<TMaybe<TString>(const TString& url)> TokenResolver;
 };
 
+using THiddenQueryAborter = std::function<void()>; // aborts hidden query, which is running within a separate TProgram
 using TDataProviderInitializer = std::function<TDataProviderInfo(
     const TString& userName,
     const TString& sessionId,
@@ -239,6 +240,7 @@ using TDataProviderInitializer = std::function<TDataProviderInfo(
     TIntrusivePtr<IRandomProvider> randomProvider,
     TIntrusivePtr<TTypeAnnotationContext> typeCtx,
     const TOperationProgressWriter& progressWriter,
-    const TYqlOperationOptions& operationOptions)>;
+    const TYqlOperationOptions& operationOptions,
+    THiddenQueryAborter hiddenAborter)>;
 
 } // namespace NYql
diff --git a/ydb/library/yql/providers/clickhouse/provider/yql_clickhouse_provider.cpp b/ydb/library/yql/providers/clickhouse/provider/yql_clickhouse_provider.cpp
index 3f30984933..07fc270978 100644
--- a/ydb/library/yql/providers/clickhouse/provider/yql_clickhouse_provider.cpp
+++ b/ydb/library/yql/providers/clickhouse/provider/yql_clickhouse_provider.cpp
@@ -16,7 +16,8 @@ TDataProviderInitializer GetClickHouseDataProviderInitializer(
         TIntrusivePtr<IRandomProvider> randomProvider,
         TIntrusivePtr<TTypeAnnotationContext> typeCtx,
         const TOperationProgressWriter& progressWriter,
-        const TYqlOperationOptions& operationOptions)
+        const TYqlOperationOptions& operationOptions,
+        THiddenQueryAborter)
     {
         Y_UNUSED(sessionId);
         Y_UNUSED(userName);
diff --git a/ydb/library/yql/providers/dq/provider/yql_dq_provider.cpp b/ydb/library/yql/providers/dq/provider/yql_dq_provider.cpp
index 5521d96e15..3c297af994 100644
--- a/ydb/library/yql/providers/dq/provider/yql_dq_provider.cpp
+++ b/ydb/library/yql/providers/dq/provider/yql_dq_provider.cpp
@@ -31,7 +31,8 @@ TDataProviderInitializer GetDqDataProviderInitializer(
         TIntrusivePtr<IRandomProvider> randomProvider,
         TIntrusivePtr<TTypeAnnotationContext> typeCtx,
         const TOperationProgressWriter& progressWriter,
-        const TYqlOperationOptions& operationOptions
+        const TYqlOperationOptions& operationOptions,
+        THiddenQueryAborter hiddenAborter
     ) {
         Y_UNUSED(userName);
 
@@ -54,7 +55,8 @@ TDataProviderInitializer GetDqDataProviderInitializer(
             fileStorage,
             dqGateway ? dqGateway->GetVanillaJobPath() : "",
             dqGateway ? dqGateway->GetVanillaJobMd5() : "",
-            externalUser
+            externalUser,
+            std::move(hiddenAborter)
         );
 
         TDataProviderInfo info;
diff --git a/ydb/library/yql/providers/dq/provider/yql_dq_recapture.cpp b/ydb/library/yql/providers/dq/provider/yql_dq_recapture.cpp
index 86791f7d6c..c792fbe607 100644
--- a/ydb/library/yql/providers/dq/provider/yql_dq_recapture.cpp
+++ b/ydb/library/yql/providers/dq/provider/yql_dq_recapture.cpp
@@ -52,6 +52,8 @@ public:
 
         if (State_->ExternalUser && !State_->Settings->_EnablePorto.Get().GetOrElse(TDqSettings::TDefault::EnablePorto)) {
             Statistics_["DqExternalUser"]++;
+            YQL_CLOG(DEBUG, ProviderDq) << "abort hidden";
+            State_->AbortHidden();
             return TStatus::Ok;
         }
 
@@ -69,6 +71,8 @@ public:
             }
 
             if (State_->TypeCtx->PureResultDataSource != DqProviderName || !State_->Settings->AnalyzeQuery.Get().GetOrElse(false)) {
+                YQL_CLOG(DEBUG, ProviderDq) << "abort hidden";
+                State_->AbortHidden();
                 return TStatus::Ok;
             }
 
@@ -93,6 +97,8 @@ public:
             }
 
             if (!good) {
+                YQL_CLOG(DEBUG, ProviderDq) << "abort hidden";
+                State_->AbortHidden();
                 YQL_CLOG(DEBUG, ProviderDq) << "good: " << good << " hasJoin: " << hasJoin << " dataSize: " << dataSize;
                 return TStatus::Ok;
             }
diff --git a/ydb/library/yql/providers/dq/provider/yql_dq_state.h b/ydb/library/yql/providers/dq/provider/yql_dq_state.h
index a3fcaf1dbd..42e149a8aa 100644
--- a/ydb/library/yql/providers/dq/provider/yql_dq_state.h
+++ b/ydb/library/yql/providers/dq/provider/yql_dq_state.h
@@ -36,6 +36,8 @@ struct TDqState: public TThrRefBase {
     THashMap<ui32, TOperationStatistics> Statistics;
     std::atomic<ui32> MetricId = 1;
 
+    std::function<void()> AbortHidden = [](){};
+
     TDqState(
         const IDqGateway::TPtr& dqGateway,
         const TGatewaysConfig* gatewaysConfig,
@@ -51,7 +53,8 @@ struct TDqState: public TThrRefBase {
         const TFileStoragePtr& fileStorage,
         const TString& vanillaJobPath,
         const TString& vanillaJobMd5,
-        bool externalUser)
+        bool externalUser,
+        THiddenQueryAborter&& hiddenAborter)
         : DqGateway(dqGateway)
         , GatewaysConfig(gatewaysConfig)
         , FunctionRegistry(functionRegistry)
@@ -67,6 +70,7 @@ struct TDqState: public TThrRefBase {
         , VanillaJobPath(vanillaJobPath)
         , VanillaJobMd5(vanillaJobMd5)
         , ExternalUser(externalUser)
+        , AbortHidden(std::move(hiddenAborter))
     { }
 };
 
diff --git a/ydb/library/yql/providers/function/provider/dq_function_provider.cpp b/ydb/library/yql/providers/function/provider/dq_function_provider.cpp
index 08bb1dc080..b0152f1de1 100644
--- a/ydb/library/yql/providers/function/provider/dq_function_provider.cpp
+++ b/ydb/library/yql/providers/function/provider/dq_function_provider.cpp
@@ -23,7 +23,8 @@ TDataProviderInitializer GetDqFunctionDataProviderInitializer(
                 TIntrusivePtr<IRandomProvider> randomProvider,
                 TIntrusivePtr<TTypeAnnotationContext> typeCtx,
                 const TOperationProgressWriter& progressWriter,
-                const TYqlOperationOptions& operationOptions)
+                const TYqlOperationOptions& operationOptions,
+                THiddenQueryAborter)
         {
 
             Y_UNUSED(userName);
diff --git a/ydb/library/yql/providers/pq/provider/yql_pq_provider.cpp b/ydb/library/yql/providers/pq/provider/yql_pq_provider.cpp
index d7f6cc3828..953f8ed9f7 100644
--- a/ydb/library/yql/providers/pq/provider/yql_pq_provider.cpp
+++ b/ydb/library/yql/providers/pq/provider/yql_pq_provider.cpp
@@ -21,7 +21,8 @@ TDataProviderInitializer GetPqDataProviderInitializer(
                TIntrusivePtr<IRandomProvider> randomProvider,
                TIntrusivePtr<TTypeAnnotationContext> typeCtx,
                const TOperationProgressWriter& progressWriter,
-               const TYqlOperationOptions& operationOptions)
+               const TYqlOperationOptions& operationOptions,
+               THiddenQueryAborter)
         {
             Y_UNUSED(userName);
             Y_UNUSED(functionRegistry);
diff --git a/ydb/library/yql/providers/s3/provider/yql_s3_provider.cpp b/ydb/library/yql/providers/s3/provider/yql_s3_provider.cpp
index 4b11f78ad7..9b1d35e98e 100644
--- a/ydb/library/yql/providers/s3/provider/yql_s3_provider.cpp
+++ b/ydb/library/yql/providers/s3/provider/yql_s3_provider.cpp
@@ -13,7 +13,8 @@ TDataProviderInitializer GetS3DataProviderInitializer(IHTTPGateway::TPtr gateway
         TIntrusivePtr<IRandomProvider> randomProvider,
         TIntrusivePtr<TTypeAnnotationContext> typeCtx,
         const TOperationProgressWriter& progressWriter,
-        const TYqlOperationOptions& operationOptions)
+        const TYqlOperationOptions& operationOptions,
+        THiddenQueryAborter)
     {
         Y_UNUSED(sessionId);
         Y_UNUSED(userName);
diff --git a/ydb/library/yql/providers/solomon/provider/yql_solomon_provider.cpp b/ydb/library/yql/providers/solomon/provider/yql_solomon_provider.cpp
index d84f1502a6..8f62c5354a 100644
--- a/ydb/library/yql/providers/solomon/provider/yql_solomon_provider.cpp
+++ b/ydb/library/yql/providers/solomon/provider/yql_solomon_provider.cpp
@@ -15,7 +15,8 @@ TDataProviderInitializer GetSolomonDataProviderInitializer(ISolomonGateway::TPtr
         TIntrusivePtr<IRandomProvider> randomProvider,
         TIntrusivePtr<TTypeAnnotationContext> typeCtx,
         const TOperationProgressWriter& progressWriter,
-        const TYqlOperationOptions& operationOptions)
+        const TYqlOperationOptions& operationOptions,
+        THiddenQueryAborter)
     {
         Y_UNUSED(sessionId);
         Y_UNUSED(userName);
diff --git a/ydb/library/yql/providers/ydb/provider/yql_ydb_provider.cpp b/ydb/library/yql/providers/ydb/provider/yql_ydb_provider.cpp
index 8bb64063b2..4abe4ca763 100644
--- a/ydb/library/yql/providers/ydb/provider/yql_ydb_provider.cpp
+++ b/ydb/library/yql/providers/ydb/provider/yql_ydb_provider.cpp
@@ -16,7 +16,8 @@ TDataProviderInitializer GetYdbDataProviderInitializer(
         TIntrusivePtr<IRandomProvider> randomProvider,
         TIntrusivePtr<TTypeAnnotationContext> typeCtx,
         const TOperationProgressWriter& progressWriter,
-        const TYqlOperationOptions& operationOptions)
+        const TYqlOperationOptions& operationOptions,
+        THiddenQueryAborter)
     {
         Y_UNUSED(sessionId);
         Y_UNUSED(userName);