Merge bec2570 into b49b9cc

ydb-platform · Aug 30, 2024 · c8580e3 · c8580e3
2 parents b49b9cc + bec2570
commit c8580e3
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 3 deletions.
diff --git a/ydb/library/yql/providers/yt/gateway/native/yql_yt_spec.cpp b/ydb/library/yql/providers/yt/gateway/native/yql_yt_spec.cpp
@@ -66,7 +66,8 @@ void FillSpec(NYT::TNode& spec,
     const TTransactionCache::TEntry::TPtr& entry,
     double extraCpu,
     const TMaybe<double>& secondExtraCpu,
-    EYtOpProps opProps)
+    EYtOpProps opProps,
+    const TSet<TString>& addSecTags)
 {
     auto& cluster = execCtx.Cluster_;
 
@@ -512,6 +513,13 @@ void FillSpec(NYT::TNode& spec,
     if (opProps.HasFlags(EYtOpProp::WithReducer)) {
         spec["reducer"]["environment"]["TMPDIR"] = ".";
     }
+
+    if (!addSecTags.empty()) {
+        spec["additional_security_tags"] = NYT::TNode::CreateList();
+        for (const auto& tag : addSecTags) {
+            spec["additional_security_tags"].Add(NYT::TNode(tag));
+        }
+    }
 }
 
 void FillSecureVault(NYT::TNode& spec, const IYtGateway::TSecureParams& secureParams) {

diff --git a/ydb/library/yql/providers/yt/gateway/native/yql_yt_spec.h b/ydb/library/yql/providers/yt/gateway/native/yql_yt_spec.h
@@ -46,7 +46,8 @@ void FillSpec(NYT::TNode& spec,
     const TTransactionCache::TEntry::TPtr& entry,
     double extraCpu,
     const TMaybe<double>& secondExtraCpu,
-    EYtOpProps opProps = 0);
+    EYtOpProps opProps = 0,
+    const TSet<TString>& addSecTags = {});
 
 void FillSecureVault(NYT::TNode& spec, const IYtGateway::TSecureParams& secureParams);
 
@@ -67,6 +68,7 @@ void FillOperationOptionsImpl(NYT::TOperationOptions& opOpts,
 
 namespace NPrivate {
     Y_HAS_MEMBER(SecureParams);
+    Y_HAS_MEMBER(AdditionalSecurityTags);
 }
 
 template <class TOptions>
@@ -77,7 +79,11 @@ inline void FillSpec(NYT::TNode& spec,
     const TMaybe<double>& secondExtraCpu,
     EYtOpProps opProps = 0)
 {
-    FillSpec(spec, execCtx, execCtx.Options_.Config(), entry, extraCpu, secondExtraCpu, opProps);
+    TSet<TString> addSecTags = {};
+    if constexpr (NPrivate::THasAdditionalSecurityTags<TOptions>::value) {
+        addSecTags = execCtx.Options_.AdditionalSecurityTags();
+    }
+    FillSpec(spec, execCtx, execCtx.Options_.Config(), entry, extraCpu, secondExtraCpu, opProps, addSecTags);
     if constexpr (NPrivate::THasSecureParams<TOptions>::value) {
         FillSecureVault(spec, execCtx.Options_.SecureParams());
     }
@@ -91,6 +97,7 @@ inline void FillOperationSpec(NYT::TUserOperationSpecBase<TDerived>& spec, const
     if (auto val = execCtx->Options_.Config()->CoreDumpPath.Get()) {
         spec.CoreTablePath(*val);
     }
+
 }
 
 template <class TExecParamsPtr>

diff --git a/ydb/library/yql/providers/yt/provider/yql_yt_datasink_exec.cpp b/ydb/library/yql/providers/yt/provider/yql_yt_datasink_exec.cpp
@@ -252,6 +252,30 @@ class TYtDataSinkExecTransformer : public TExecTransformerBase {
                 << ", cache mode: " << queryCacheMode;
         }
 
+        TSet<TString> addSecTags;
+        for (size_t pos = 0; pos < optimizedNode->ChildrenSize(); pos++) {
+            auto childPtr = optimizedNode->ChildPtr(pos);
+            if (childPtr->Type() == TExprNode::Lambda) {
+                VisitExpr(childPtr, [&addSecTags](const TExprNode::TPtr& node) -> bool {
+                    if (TYtTableContent::Match(node.Get())) {
+                        VisitExpr(node, [&addSecTags](const TExprNode::TPtr& node) -> bool {
+                            if (TYtTableBase::Match(node.Get())) {
+                                if (auto stat = TYtTableBaseInfo::GetStat(TExprBase(node))) {
+                                    for (const auto& tag : stat->SecurityTags) {
+                                        addSecTags.insert(tag);
+                                    }
+                                }
+                                return false;
+                            }
+                            return true;
+                        });
+                        return false;
+                    }
+                    return true;
+                });
+            }
+        }
+
         YQL_CLOG(DEBUG, ProviderYt) << "Executing " << input->Content() << " (UniqueId=" << input->UniqueId() << ")";
 
         return State_->Gateway->Run(optimizedNode, ctx,
@@ -265,6 +289,7 @@ class TYtDataSinkExecTransformer : public TExecTransformerBase {
                 .OptLLVM(State_->Types->OptLLVM.GetOrElse(TString()))
                 .OperationHash(operationHash)
                 .SecureParams(secureParams)
+                .AdditionalSecurityTags(addSecTags)
             );
     }
 

diff --git a/ydb/library/yql/providers/yt/provider/yql_yt_gateway.h b/ydb/library/yql/providers/yt/provider/yql_yt_gateway.h
@@ -383,6 +383,7 @@ class IYtGateway : public TThrRefBase {
         OPTION_FIELD(TString, OptLLVM)
         OPTION_FIELD(TString, OperationHash)
         OPTION_FIELD(TSecureParams, SecureParams)
+        OPTION_FIELD_DEFAULT(TSet<TString>, AdditionalSecurityTags, {})
     };
 
     struct TRunResult : public NCommon::TOperationResult {