Show ast on case of compilation error

ydb-platform · Jan 26, 2024 · 4c8d2a9 · 4c8d2a9
1 parent 65c46e5
commit 4c8d2a9
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 15 deletions.
diff --git a/ydb/core/kqp/compile_service/kqp_compile_actor.cpp b/ydb/core/kqp/compile_service/kqp_compile_actor.cpp
@@ -372,6 +372,14 @@ class TKqpCompileActor : public TActorBootstrapped<TKqpCompileActor> {
         PassAway();
     }
 
+    void FillPreparedQuery(std::unique_ptr<NKikimrKqp::TPreparedQuery> preparingQuery, NKikimrKqp::EQueryType queryType) {
+        auto preparedQueryHolder = std::make_shared<TPreparedQueryHolder>(
+            preparingQuery.release(), AppData()->FunctionRegistry);
+        preparedQueryHolder->MutableLlvmSettings().Fill(Config, queryType);
+        KqpCompileResult->PreparedQuery = preparedQueryHolder;
+        KqpCompileResult->AllowCache = CanCacheQuery(KqpCompileResult->PreparedQuery->GetPhysicalQuery());
+    }
+
     void Handle(TEvKqp::TEvContinueProcess::TPtr &ev, const TActorContext &ctx) {
         Y_ENSURE(!ev->Get()->QueryId);
 
@@ -403,16 +411,10 @@ class TKqpCompileActor : public TActorBootstrapped<TKqpCompileActor> {
 
         if (status == Ydb::StatusIds::SUCCESS) {
             YQL_ENSURE(kqpResult.PreparingQuery);
-            {
-                auto preparedQueryHolder = std::make_shared<TPreparedQueryHolder>(
-                    kqpResult.PreparingQuery.release(), AppData()->FunctionRegistry);
-                preparedQueryHolder->MutableLlvmSettings().Fill(Config, queryType);
-                KqpCompileResult->PreparedQuery = preparedQueryHolder;
-                KqpCompileResult->AllowCache = CanCacheQuery(KqpCompileResult->PreparedQuery->GetPhysicalQuery());
-
-                if (AstResult) {
-                    KqpCompileResult->Ast = AstResult->Ast;
-                }
+            FillPreparedQuery(std::move(kqpResult.PreparingQuery), queryType);
+
+            if (AstResult) {
+                KqpCompileResult->Ast = AstResult->Ast;
             }
 
             auto now = TInstant::Now();
@@ -423,6 +425,10 @@ class TKqpCompileActor : public TActorBootstrapped<TKqpCompileActor> {
                 << ", self: " << ctx.SelfID
                 << ", duration: " << duration);
         } else {
+            if (kqpResult.PreparingQuery) {
+                FillPreparedQuery(std::move(kqpResult.PreparingQuery), queryType);
+            }
+
             LOG_ERROR_S(ctx, NKikimrServices::KQP_COMPILE_ACTOR, "Compilation failed"
                 << ", self: " << ctx.SelfID
                 << ", status: " << Ydb::StatusIds_StatusCode_Name(status)

diff --git a/ydb/core/kqp/host/kqp_host.cpp b/ydb/core/kqp/host/kqp_host.cpp
@@ -323,6 +323,7 @@ class TAsyncPrepareYqlResult : public TKqpAsyncResultBase<IKqpHost::TQueryResult
         TIntrusivePtr<TKikimrQueryContext> queryCtx, const TKqpQueryRef& query, TMaybe<TSqlVersion> sqlVersion)
         : TKqpAsyncResultBase(queryRoot, exprCtx, transformer)
         , QueryCtx(queryCtx)
+        , ExprCtx(exprCtx)
         , QueryText(query.Text)
         , SqlVersion(sqlVersion) {}
 
@@ -342,8 +343,18 @@ class TAsyncPrepareYqlResult : public TKqpAsyncResultBase<IKqpHost::TQueryResult
         prepareResult.QueryAst = prepareResult.PreparingQuery->GetPhysicalQuery().GetQueryAst();
     }
 
+    void FillPartialResult(TResult& prepareResult) const override {
+        YQL_ENSURE(QueryCtx->PrepareOnly);
+
+        if (auto exprRoot = GetExprRoot()) {
+            prepareResult.PreparingQuery = std::move(QueryCtx->PreparingQuery);
+            prepareResult.PreparingQuery->MutablePhysicalQuery()->SetQueryAst(KqpExprToPrettyString(*GetExprRoot(), ExprCtx));
+        }
+    }
+
 private:
     TIntrusivePtr<TKikimrQueryContext> QueryCtx;
+    NYql::TExprContext& ExprCtx;
     TString QueryText;
     TMaybe<TSqlVersion> SqlVersion;
 };

diff --git a/ydb/core/kqp/host/kqp_host_impl.h b/ydb/core/kqp/host/kqp_host_impl.h
@@ -34,7 +34,9 @@ class TKqpAsyncResultBase : public NYql::IKikimrAsyncResult<TResult> {
         YQL_ENSURE(HasResult());
 
         if (Status.GetValue() == NYql::IGraphTransformer::TStatus::Error) {
-            return NYql::NCommon::ResultFromErrors<TResult>(ExprCtx.IssueManager.GetIssues());
+            TResult result = NYql::NCommon::ResultFromErrors<TResult>(ExprCtx.IssueManager.GetIssues());
+            FillPartialResult(result);
+            return result;
         }
 
         YQL_ENSURE(Status.GetValue() == NYql::IGraphTransformer::TStatus::Ok);
@@ -83,6 +85,8 @@ class TKqpAsyncResultBase : public NYql::IKikimrAsyncResult<TResult> {
 protected:
     virtual void FillResult(TResult& result) const = 0;
 
+    virtual void FillPartialResult(TResult&) const {}
+
     NYql::TExprNode::TPtr GetExprRoot() const { return ExprRoot; }
     NYql::TExprContext& GetExprContext() const { return ExprCtx; }
     NYql::IGraphTransformer& GetTransformer() const { return Transformer; }

diff --git a/ydb/core/kqp/session_actor/kqp_session_actor.cpp b/ydb/core/kqp/session_actor/kqp_session_actor.cpp
@@ -1688,10 +1688,17 @@ class TKqpSessionActor : public TActorBootstrapped<TKqpSessionActor> {
 
             const auto& phyQuery = QueryState->PreparedQuery->GetPhysicalQuery();
             FillColumnsMeta(phyQuery, response);
-        } else if (compileResult->Status == Ydb::StatusIds::TIMEOUT && QueryState->QueryDeadlines.CancelAt) {
-            // The compile timeout cause cancelation execution of request.
-            // So in case of cancel after we can reply with canceled status
-            ev.SetYdbStatus(Ydb::StatusIds::CANCELLED);
+        } else {
+            if (compileResult->Status == Ydb::StatusIds::TIMEOUT && QueryState->QueryDeadlines.CancelAt) {
+                // The compile timeout cause cancelation execution of request.
+                // So in case of cancel after we can reply with canceled status
+                ev.SetYdbStatus(Ydb::StatusIds::CANCELLED);
+            }
+
+            auto& preparedQuery = compileResult->PreparedQuery;
+            if (preparedQuery && QueryState->ReportStats() && QueryState->GetStatsMode() >= Ydb::Table::QueryStatsCollection::STATS_COLLECTION_FULL) {
+                response.SetQueryAst(preparedQuery->GetPhysicalQuery().GetQueryAst());
+            }
         }
     }