YDB FQ: handle exception in YQL Generic Provider when IAM service is not available (#9092)

Author: vitalyisaev2

ydb/library/yql/providers/generic/actors/ut/yql_generic_lookup_actor_ut.cpp

@@ -78,7 +78,7 @@ Y_UNIT_TEST_SUITE(GenericProviderLookupActor) {
         dsi.Setdatabase("some_db");
         dsi.Setuse_tls(true);
         dsi.set_protocol(::NYql::NConnector::NApi::EProtocol::NATIVE);
-        auto token = dsi.mutable_credentials() -> mutable_token();
+        auto token = dsi.mutable_credentials()->mutable_token();
         token->Settype("IAM");
         token->Setvalue("TEST_TOKEN");
 

ydb/library/yql/providers/generic/actors/yql_generic_lookup_actor.cpp

@@ -152,7 +152,14 @@ namespace NYql::NDq {
             Y_ABORT_UNLESS(response.splits_size() == 1);
             auto& split = response.splits(0);
             NConnector::NApi::TReadSplitsRequest readRequest;
-            *readRequest.mutable_data_source_instance() = GetDataSourceInstanceWithToken();
+
+            *readRequest.mutable_data_source_instance() = LookupSource.data_source_instance();
+            auto error = TokenProvider->MaybeFillToken(*readRequest.mutable_data_source_instance());
+            if (error) {
+                SendError(TActivationContext::ActorSystem(), SelfId(), std::move(error));
+                return;
+            }
+
             *readRequest.add_splits() = split;
             readRequest.Setformat(NConnector::NApi::TReadSplitsRequest_EFormat::TReadSplitsRequest_EFormat_ARROW_IPC_STREAMING);
             Connector->ReadSplits(readRequest).Subscribe([actorSystem = TActivationContext::ActorSystem(), selfId = SelfId()](const NConnector::TReadSplitsStreamIteratorAsyncResult& asyncResult) {
@@ -194,9 +201,16 @@ namespace NYql::NDq {
             YQL_CLOG(DEBUG, ProviderGeneric) << "ActorId=" << SelfId() << " Got LookupRequest for " << request.size() << " keys";
             Y_ABORT_IF(InProgress);
             Y_ABORT_IF(request.size() == 0 || request.size() > MaxKeysInRequest);
+
             Request = std::move(request);
             NConnector::NApi::TListSplitsRequest splitRequest;
-            *splitRequest.add_selects() = CreateSelect();
+
+            auto error = FillSelect(*splitRequest.add_selects());
+            if (error) {
+                SendError(TActivationContext::ActorSystem(), SelfId(), std::move(error));
+                return;
+            };
+
             splitRequest.Setmax_split_count(1);
             Connector->ListSplits(splitRequest).Subscribe([actorSystem = TActivationContext::ActorSystem(), selfId = SelfId()](const NConnector::TListSplitsStreamIteratorAsyncResult& asyncResult) {
                 auto result = ExtractFromConstFuture(asyncResult);
@@ -285,6 +299,12 @@ namespace NYql::NDq {
             SendError(actorSystem, selfId, NConnector::ErrorFromGRPCStatus(status));
         }
 
+        static void SendError(NActors::TActorSystem* actorSystem, const NActors::TActorId& selfId, TString error) {
+            NConnector::NApi::TError dst;
+            *dst.mutable_message() = error;
+            SendError(actorSystem, selfId, std::move(dst));
+        }
+
     private:
         enum class EColumnDestination {
             Key,
@@ -314,17 +334,13 @@ namespace NYql::NDq {
             return result;
         }
 
-        NYql::NConnector::NApi::TDataSourceInstance GetDataSourceInstanceWithToken() const {
+        TString FillSelect(NConnector::NApi::TSelect& select) {
             auto dsi = LookupSource.data_source_instance();
-            // Note: returned token may be stale and we have no way to check or recover here
-            // Consider to redesign ICredentialsProvider
-            TokenProvider->MaybeFillToken(dsi);
-            return dsi;
-        }
-
-        NConnector::NApi::TSelect CreateSelect() {
-            NConnector::NApi::TSelect select;
-            *select.mutable_data_source_instance() = GetDataSourceInstanceWithToken();
+            auto error = TokenProvider->MaybeFillToken(dsi);
+            if (error) {
+                return error;
+            }
+            *select.mutable_data_source_instance() = dsi;
 
             for (ui32 i = 0; i != SelectResultType->GetMembersCount(); ++i) {
                 auto c = select.mutable_what()->add_items()->mutable_column();
@@ -349,7 +365,7 @@ namespace NYql::NDq {
                 *disjunction.mutable_operands()->Add()->mutable_conjunction() = conjunction;
             }
             *select.mutable_where()->mutable_filter_typed()->mutable_disjunction() = disjunction;
-            return select;
+            return {};
         }
 
     private:

ydb/library/yql/providers/generic/actors/yql_generic_read_actor.cpp

@@ -59,7 +59,14 @@ namespace NYql::NDq {
 
         void Bootstrap() {
             Become(&TGenericReadActor::StateFunc);
-            InitSplitsListing();
+            auto issue = InitSplitsListing();
+            if (issue) {
+                return NotifyComputeActorWithIssue(
+                    TActivationContext::ActorSystem(),
+                    ComputeActorId_,
+                    InputIndex_,
+                    std::move(*issue));
+            };
         }
 
         static constexpr char ActorName[] = "GENERIC_READ_ACTOR";
@@ -79,13 +86,18 @@ namespace NYql::NDq {
 
         // ListSplits
 
-        void InitSplitsListing() {
+        TMaybe<TIssue> InitSplitsListing() {
             YQL_CLOG(DEBUG, ProviderGeneric) << "Start splits listing";
 
             // Prepare request
             NConnector::NApi::TListSplitsRequest request;
             NConnector::NApi::TSelect select = Source_.select(); // copy TSelect from source
-            TokenProvider_->MaybeFillToken(*select.mutable_data_source_instance());
+
+            auto error = TokenProvider_->MaybeFillToken(*select.mutable_data_source_instance());
+            if (error) {
+                return TIssue(error);
+            }
+
             *request.mutable_selects()->Add() = std::move(select);
 
             // Initialize stream
@@ -100,6 +112,8 @@ namespace NYql::NDq {
                         TEvListSplitsIterator>(
                         actorSystem, selfId, computeActorId, inputIndex, future);
                 });
+
+            return Nothing();
         }
 
         void Handle(TEvListSplitsIterator::TPtr& ev) {
@@ -145,7 +159,16 @@ namespace NYql::NDq {
             // Server sent EOF, now we are ready to start splits reading
             if (NConnector::GrpcStatusEndOfStream(status)) {
                 YQL_CLOG(DEBUG, ProviderGeneric) << "Handle :: EvListSplitsFinished :: last message was reached, start data reading";
-                return InitSplitsReading();
+                auto issue = InitSplitsReading();
+                if (issue) {
+                    return NotifyComputeActorWithIssue(
+                        TActivationContext::ActorSystem(),
+                        ComputeActorId_,
+                        InputIndex_,
+                        std::move(*issue));
+                }
+
+                return;
             }
 
             // Server temporary failure
@@ -163,27 +186,31 @@ namespace NYql::NDq {
         }
 
         // ReadSplits
-        void InitSplitsReading() {
+        TMaybe<TIssue> InitSplitsReading() {
             YQL_CLOG(DEBUG, ProviderGeneric) << "Start splits reading";
 
             if (Splits_.empty()) {
                 YQL_CLOG(WARN, ProviderGeneric) << "Accumulated empty list of splits";
                 ReadSplitsFinished_ = true;
-                return NotifyComputeActorWithData();
+                NotifyComputeActorWithData();
+                return Nothing();
             }
 
             // Prepare request
             NConnector::NApi::TReadSplitsRequest request;
             request.set_format(NConnector::NApi::TReadSplitsRequest::ARROW_IPC_STREAMING);
             request.mutable_splits()->Reserve(Splits_.size());
 
-            std::for_each(
-                Splits_.cbegin(), Splits_.cend(),
-                [&](const NConnector::NApi::TSplit& split) {
-                    NConnector::NApi::TSplit splitCopy = split;
-                    TokenProvider_->MaybeFillToken(*splitCopy.mutable_select()->mutable_data_source_instance());
-                    *request.mutable_splits()->Add() = std::move(split);
-                });
+            for (const auto& split : Splits_) {
+                NConnector::NApi::TSplit splitCopy = split;
+
+                auto error = TokenProvider_->MaybeFillToken(*splitCopy.mutable_select()->mutable_data_source_instance());
+                if (error) {
+                    return TIssue(std::move(error));
+                }
+
+                *request.mutable_splits()->Add() = std::move(splitCopy);
+            }
 
             // Start streaming
             Client_->ReadSplits(request).Subscribe(
@@ -197,6 +224,8 @@ namespace NYql::NDq {
                         TEvReadSplitsIterator>(
                         actorSystem, selfId, computeActorId, inputIndex, future);
                 });
+
+            return Nothing();
         }
 
         void Handle(TEvReadSplitsIterator::TPtr& ev) {
@@ -308,8 +337,8 @@ namespace NYql::NDq {
 
         static void NotifyComputeActorWithError(
             TActorSystem* actorSystem,
-            const NActors::TActorId computeActorId,
-            const ui64 inputIndex,
+            NActors::TActorId computeActorId,
+            ui64 inputIndex,
             const NConnector::NApi::TError& error) {
             actorSystem->Send(computeActorId,
                               new TEvAsyncInputError(
@@ -319,6 +348,19 @@ namespace NYql::NDq {
             return;
         }
 
+        static void NotifyComputeActorWithIssue(
+            TActorSystem* actorSystem,
+            NActors::TActorId computeActorId,
+            ui64 inputIndex,
+            TIssue issue) {
+            actorSystem->Send(computeActorId,
+                              new TEvAsyncInputError(
+                                  inputIndex,
+                                  TIssues{std::move(issue)},
+                                  NDqProto::StatusIds::StatusCode::StatusIds_StatusCode_INTERNAL_ERROR));
+            return;
+        }
+
         i64 GetAsyncInputData(NKikimr::NMiniKQL::TUnboxedValueBatch& buffer,
                               TMaybe<TInstant>&,
                               bool& finished,

ydb/library/yql/providers/generic/actors/yql_generic_token_provider.cpp

@@ -1,6 +1,7 @@
 #include "yql_generic_token_provider.h"
 
 #include <ydb/library/yql/providers/common/structured_token/yql_token_builder.h>
+#include <ydb/library/yql/utils/log/log.h>
 
 namespace NYql::NDq {
     TGenericTokenProvider::TGenericTokenProvider(const TString& staticIamToken)
@@ -9,59 +10,63 @@ namespace NYql::NDq {
     }
 
     TGenericTokenProvider::TGenericTokenProvider(
-        const TString& serviceAccountId,
-        const TString& ServiceAccountIdSignature,
-        const ISecuredServiceAccountCredentialsFactory::TPtr& credentialsFactory)
-    {
+        const TString& serviceAccountId, const TString& ServiceAccountIdSignature,
+        const ISecuredServiceAccountCredentialsFactory::TPtr& credentialsFactory) {
         Y_ENSURE(!serviceAccountId.Empty(), "No service account provided");
         Y_ENSURE(!ServiceAccountIdSignature.Empty(), "No service account signature provided");
         Y_ENSURE(credentialsFactory, "CredentialsFactory is not initialized");
 
         auto structuredTokenJSON =
-            TStructuredTokenBuilder()
-                .SetServiceAccountIdAuth(serviceAccountId, ServiceAccountIdSignature)
-                .ToJson();
+            TStructuredTokenBuilder().SetServiceAccountIdAuth(serviceAccountId, ServiceAccountIdSignature).ToJson();
 
         Y_ENSURE(structuredTokenJSON, "empty structured token");
 
-        auto credentialsProviderFactory = CreateCredentialsProviderFactoryForStructuredToken(credentialsFactory, structuredTokenJSON, false);
+        auto credentialsProviderFactory =
+            CreateCredentialsProviderFactoryForStructuredToken(credentialsFactory, structuredTokenJSON, false);
         CredentialsProvider_ = credentialsProviderFactory->CreateProvider();
     }
 
-    void TGenericTokenProvider::MaybeFillToken(NConnector::NApi::TDataSourceInstance& dsi) const {
+    TString TGenericTokenProvider::MaybeFillToken(NConnector::NApi::TDataSourceInstance& dsi) const {
         // 1. Don't need tokens if basic auth is set
         if (dsi.credentials().has_basic()) {
-            return;
+            return {};
         }
 
         *dsi.mutable_credentials()->mutable_token()->mutable_type() = "IAM";
 
         // 2. If static IAM-token has been provided, use it
         if (!StaticIAMToken_.empty()) {
             *dsi.mutable_credentials()->mutable_token()->mutable_value() = StaticIAMToken_;
-            return;
+            return {};
         }
 
         // 3. Otherwise use credentials provider to get token
         Y_ENSURE(CredentialsProvider_, "CredentialsProvider is not initialized");
 
-        auto iamToken = CredentialsProvider_->GetAuthInfo();
+        TString iamToken;
+        try {
+            iamToken = CredentialsProvider_->GetAuthInfo();
+        } catch (const std::exception& e) {
+            YQL_CLOG(ERROR, ProviderGeneric) << "MaybeFillToken: " << e.what();
+            return TString(e.what());
+        }
+
         Y_ENSURE(iamToken, "CredentialsProvider returned empty IAM token");
 
         *dsi.mutable_credentials()->mutable_token()->mutable_value() = std::move(iamToken);
+        return {};
     }
 
     TGenericTokenProvider::TPtr
-    CreateGenericTokenProvider(
-        const TString& staticIamToken,
-        const TString& serviceAccountId, const TString& serviceAccountIdSignature,
-        const ISecuredServiceAccountCredentialsFactory::TPtr& credentialsFactory)
-    {
+    CreateGenericTokenProvider(const TString& staticIamToken, const TString& serviceAccountId,
+                               const TString& serviceAccountIdSignature,
+                               const ISecuredServiceAccountCredentialsFactory::TPtr& credentialsFactory) {
         if (!staticIamToken.Empty()) {
             return std::make_unique<TGenericTokenProvider>(staticIamToken);
         }
         if (!serviceAccountId.Empty()) {
-            return std::make_unique<TGenericTokenProvider>(serviceAccountId, serviceAccountIdSignature, credentialsFactory);
+            return std::make_unique<TGenericTokenProvider>(serviceAccountId, serviceAccountIdSignature,
+                                                           credentialsFactory);
         }
         return std::make_unique<TGenericTokenProvider>();
     }

ydb/library/yql/providers/generic/actors/yql_generic_token_provider.h

@@ -3,6 +3,7 @@
 #include <ydb/library/yql/providers/common/token_accessor/client/factory.h>
 #include <ydb/library/yql/providers/generic/connector/api/service/protos/connector.pb.h>
 #include <ydb/library/yql/providers/generic/proto/source.pb.h>
+#include <ydb/library/yql/public/issue/yql_issue.h>
 
 namespace NYql::NDq {
     // When accessing external data sources using authentication via tokens,
@@ -19,7 +20,9 @@ namespace NYql::NDq {
             const TString& ServiceAccountIdSignature,
             const ISecuredServiceAccountCredentialsFactory::TPtr& credentialsFactory);
 
-        void MaybeFillToken(NConnector::NApi::TDataSourceInstance& dsi) const;
+        // MaybeFillToken sets IAM-token within DataSourceInstance.
+        // Returns string containing error, if it happened.
+        TString MaybeFillToken(NConnector::NApi::TDataSourceInstance& dsi) const;
 
     private:
         TString StaticIAMToken_;