Support of replay of credentials (#10015)

2 files changed, 58 insertions, 1 deletions

diff --git a/ydb/library/yql/core/facade/yql_facade.cpp b/ydb/library/yql/core/facade/yql_facade.cpp
index 14ec1d4817..1b98b0a9ae 100644
--- a/ydb/library/yql/core/facade/yql_facade.cpp
+++ b/ydb/library/yql/core/facade/yql_facade.cpp
@@ -53,6 +53,8 @@ const TString ParametersLabel = "Parameters";
 const TString TranslationLabel = "Translation";
 const TString StaticUserFilesLabel = "UserFiles";
 const TString DynamicUserFilesLabel = "DynamicUserFiles";
+const TString StaticCredentialsLabel = "Credentials";
+const TString DynamicCredentialsLabel = "DynamicCredentials";
 
 class TUrlLoader : public IUrlLoader {
 public:
@@ -297,8 +299,33 @@ TProgram::TProgram(
         SessionId_ = CreateGuidAsString();
     }
 
+    if (QContext_.CanWrite()) {
+        NYT::TNode credListNode = NYT::TNode::CreateList();
+        Credentials_->ForEach([&](const TString name, const TCredential& cred) {
+            credListNode.Add(NYT::TNode()
+                ("Name", name)
+                ("Category", cred.Category)
+                ("Subcategory", cred.Subcategory));
+        });
+
+        auto credList = NYT::NodeToYsonString(credListNode, NYT::NYson::EYsonFormat::Binary);
+        QContext_.GetWriter()->Put({FacadeComponent, StaticCredentialsLabel}, credList).GetValueSync();
+    } else if (QContext_.CanRead()) {
+        Credentials_ = MakeIntrusive<TCredentials>();
+        for (const auto& label : {StaticCredentialsLabel, DynamicCredentialsLabel}) {
+            auto item = QContext_.GetReader()->Get({FacadeComponent, label}).GetValueSync();
+            if (item) {
+                auto node = NYT::NodeFromYsonString(item->Value);
+                for (const auto& c : node.AsList()) {
+                    Credentials_->AddCredential(c["Name"].AsString(), TCredential(
+                        c["Category"].AsString(),c["Subcategory"].AsString(),"REPLAY"));
+                }
+            }
+        }
+    }
+
     if (QContext_.CanWrite() && !SavedUserDataTable_.empty()) {
-        NYT::TNode userFilesNode;
+        NYT::TNode userFilesNode = NYT::TNode::CreateList();
         for (const auto& p : SavedUserDataTable_) {
             userFilesNode.Add(p.first.Alias());
         }
@@ -509,6 +536,19 @@ TString TProgram::GetSessionId() const {
 void TProgram::AddCredentials(const TVector<std::pair<TString, TCredential>>& credentials) {
     Y_ENSURE(!TypeCtx_, "TypeCtx_ already created");
 
+    if (QContext_.CanWrite()) {
+        NYT::TNode credListNode = NYT::TNode::CreateList();
+        for (const auto& c : credentials) {
+            credListNode.Add(NYT::TNode()
+                ("Name", c.first)
+                ("Category", c.second.Category)
+                ("Subcategory", c.second.Subcategory));
+        }
+
+        auto credList = NYT::NodeToYsonString(credListNode, NYT::NYson::EYsonFormat::Binary);
+        QContext_.GetWriter()->Put({FacadeComponent, DynamicCredentialsLabel}, credList).GetValueSync();
+    }
+
     for (const auto& credential : credentials) {
         Credentials_->AddCredential(credential.first, credential.second);
     }
diff --git a/ydb/library/yql/core/ut/yql_qplayer_ut.cpp b/ydb/library/yql/core/ut/yql_qplayer_ut.cpp
index 3783f33910..b0ba345c06 100644
--- a/ydb/library/yql/core/ut/yql_qplayer_ut.cpp
+++ b/ydb/library/yql/core/ut/yql_qplayer_ut.cpp
@@ -55,6 +55,7 @@ struct TRunSettings {
     TMaybe<TString> ParametersYson;
     THashMap<TString, TString> StaticFiles, DynamicFiles;
     bool ChangeTime = false;
+    TVector<std::pair<TString, TCredential>> Credentials;
 };
 
 TUserDataTable MakeUserTables(const THashMap<TString, TString>& map) {
@@ -107,6 +108,15 @@ bool RunProgram(bool replay, const TString& query, const TQContext& qContext, co
         factory.AddUserDataTable(MakeUserTables(runSettings.StaticFiles));
     }
 
+    if (!replay && !runSettings.Credentials.empty()) {
+        auto credentials = MakeIntrusive<TCredentials>();
+        for (const auto& x : runSettings.Credentials) {
+            credentials->AddCredential(x.first, x.second);
+        }
+        
+        factory.SetCredentials(credentials);
+    }
+
     TProgramPtr program = factory.Create("-stdin-", query, "", EHiddenMode::Disable, qContext);
     if (!replay && runSettings.ParametersYson) {
         program->SetParametersYson(*runSettings.ParametersYson);
@@ -289,4 +299,11 @@ SELECT State FROM plato.WalkFolders(``, $postHandler AS PostHandler);
         runSettings.ChangeTime = true;
         CheckProgram(s, runSettings);
     }
+
+    Y_UNIT_TEST(Credentials) {
+        auto s = "select SecureParam('token:foo');";
+        TRunSettings runSettings;
+        runSettings.Credentials.push_back(std::make_pair("foo",TCredential("C","S","V")));
+        CheckProgram(s, runSettings);
+    }
 }