Providing groups in feature activation

YQL-15975

6 files changed, 38 insertions, 12 deletions

diff --git a/ydb/library/yql/core/credentials/yql_credentials.h b/ydb/library/yql/core/credentials/yql_credentials.h
index b3aac32e3a..e076aff12c 100644
--- a/ydb/library/yql/core/credentials/yql_credentials.h
+++ b/ydb/library/yql/core/credentials/yql_credentials.h
@@ -6,6 +6,7 @@
 #include <util/generic/hash.h>
 
 #include <functional>
+#include <unordered_set>
 
 namespace NYql {
 
@@ -38,6 +39,14 @@ public:
     void SetUserCredentials(const TUserCredentials& userCredentials) {
         UserCredentials_ = userCredentials;
     }
+    
+    void SetGroups(std::unordered_set<TString>&& groups) {
+        Groups_ = std::move(groups);
+    }
+
+    const std::unordered_set<TString>& GetGroups() const {
+        return Groups_;
+    }
 
     const TCredential* FindCredential(const TStringBuf& name) const;
     TString FindCredentialContent(const TStringBuf& name1, const TStringBuf& name2, const TString& defaultContent) const;
@@ -50,7 +59,7 @@ public:
 private:
     THashMap<TString, TCredential> CredentialTable_;
     TUserCredentials UserCredentials_;
-
+    std::unordered_set<TString> Groups_;
 };
 
 } // namespace NYql
diff --git a/ydb/library/yql/providers/common/activation/yql_activation.cpp b/ydb/library/yql/providers/common/activation/yql_activation.cpp
index 3247e448fc..34f3c8667c 100644
--- a/ydb/library/yql/providers/common/activation/yql_activation.cpp
+++ b/ydb/library/yql/providers/common/activation/yql_activation.cpp
@@ -10,13 +10,19 @@
 namespace NYql::NConfig {
 
 template <class TActivation>
-ui32 GetPercentage(const TActivation& activation, const TString& userName) {
+ui32 GetPercentage(const TActivation& activation, const TString& userName, const std::unordered_set<std::string_view>& groups) {
     if (AnyOf(activation.GetIncludeUsers(), [&](const auto& user) { return user == userName; })) {
         return 100;
     }
+    if (AnyOf(activation.GetIncludeGroups(), [&](const auto& includeGroup) { return groups.contains(includeGroup); })) {
+        return 100;
+    }
     if (AnyOf(activation.GetExcludeUsers(), [&](const auto& user) { return user == userName; })) {
         return 0;
     }
+    if (AnyOf(activation.GetExcludeGroups(), [&](const auto& excludeGroup) { return groups.contains(excludeGroup); })) {
+        return 0;
+    }
     if ((userName.StartsWith("robot-") || userName.StartsWith("zomb-")) && activation.GetExcludeRobots()) {
         return 0;
     }
@@ -40,13 +46,13 @@ ui32 GetPercentage(const TActivation& activation, const TString& userName) {
 }
 
 template <class TActivation>
-bool Allow(const TActivation& activation, const TString& userName) {
-    ui32 percent = GetPercentage(activation, userName);
+bool Allow(const TActivation& activation, const TString& userName, const std::unordered_set<std::string_view>& groups) {
+    ui32 percent = GetPercentage(activation, userName, groups);
     const auto random = RandomNumber<ui8>(100);
     return random < percent;
 }
 
-template ui32 GetPercentage<NYql::TActivationPercentage>(const NYql::TActivationPercentage& activation, const TString& userName);
-template bool Allow<NYql::TActivationPercentage>(const NYql::TActivationPercentage& activation, const TString& userName);
+template ui32 GetPercentage<NYql::TActivationPercentage>(const NYql::TActivationPercentage& activation, const TString& userName, const std::unordered_set<std::string_view>& groups);
+template bool Allow<NYql::TActivationPercentage>(const NYql::TActivationPercentage& activation, const TString& userName, const std::unordered_set<std::string_view>& groups);
 
 } // namespace
diff --git a/ydb/library/yql/providers/common/activation/yql_activation.h b/ydb/library/yql/providers/common/activation/yql_activation.h
index 128a463f8b..d9bffbf9bd 100644
--- a/ydb/library/yql/providers/common/activation/yql_activation.h
+++ b/ydb/library/yql/providers/common/activation/yql_activation.h
@@ -1,13 +1,14 @@
 #pragma once
 
 #include <util/generic/string.h>
+#include <unordered_set>
 
 namespace NYql::NConfig {
 
 template <class TActivation>
-ui32 GetPercentage(const TActivation& activation, const TString& userName);
+ui32 GetPercentage(const TActivation& activation, const TString& userName, const std::unordered_set<std::string_view>& groups);
 
 template <class TActivation>
-bool Allow(const TActivation& activation, const TString& userName);
+bool Allow(const TActivation& activation, const TString& userName, const std::unordered_set<std::string_view>& groups);
 
 } // namespace
diff --git a/ydb/library/yql/providers/common/proto/gateways_config.proto b/ydb/library/yql/providers/common/proto/gateways_config.proto
index fcde76fe05..4fec8c450e 100644
--- a/ydb/library/yql/providers/common/proto/gateways_config.proto
+++ b/ydb/library/yql/providers/common/proto/gateways_config.proto
@@ -20,6 +20,8 @@ message TActivationPercentage {
     repeated string ExcludeUsers = 4;
     optional bool ExcludeRobots = 5 [default = true];
     repeated string Tags = 6;
+    repeated string IncludeGroups = 7;
+    repeated string ExcludeGroups = 8;
 }
 
 message TAttr {
diff --git a/ydb/library/yql/providers/config/yql_config_provider.cpp b/ydb/library/yql/providers/config/yql_config_provider.cpp
index 375feccc30..794fa85580 100644
--- a/ydb/library/yql/providers/config/yql_config_provider.cpp
+++ b/ydb/library/yql/providers/config/yql_config_provider.cpp
@@ -146,11 +146,15 @@ namespace {
         }
 
         bool Initialize(TExprContext& ctx) override {
-            auto filter = [this](const TCoreAttr& attr) {
+            std::unordered_set<std::string_view> groups;
+            if (Types.Credentials != nullptr) {
+                groups.insert(Types.Credentials->GetGroups().begin(), Types.Credentials->GetGroups().end());
+            }
+            auto filter = [this, groups = std::move(groups)](const TCoreAttr& attr) {
                 if (!attr.HasActivation() || !Username) {
                     return true;
                 }
-                if (NConfig::Allow(attr.GetActivation(), Username)) {
+                if (NConfig::Allow(attr.GetActivation(), Username, groups)) {
                     Statistics.Entries.emplace_back(TStringBuilder() << "Activation:" << attr.GetName(), 0, 0, 0, 0, 1);
                     return true;
                 }
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 7ae9849c2f..4877f4f458 100644
--- a/ydb/library/yql/providers/dq/provider/yql_dq_provider.cpp
+++ b/ydb/library/yql/providers/dq/provider/yql_dq_provider.cpp
@@ -78,11 +78,15 @@ TDataProviderInitializer GetDqDataProviderInitializer(
             }
 
             if (gatewaysConfig) {
-                auto filter = [username, state](const NYql::TAttr& attr) -> bool {
+                std::unordered_set<std::string_view> groups;
+                if (state->TypeCtx->Credentials != nullptr) {
+                    groups.insert(state->TypeCtx->Credentials->GetGroups().begin(), state->TypeCtx->Credentials->GetGroups().end());
+                }
+                auto filter = [username, state, groups = std::move(groups)](const NYql::TAttr& attr) -> bool {
                     if (!attr.HasActivation()) {
                         return true;
                     }
-                    if (NConfig::Allow(attr.GetActivation(), username)) {
+                    if (NConfig::Allow(attr.GetActivation(), username, groups)) {
                         with_lock(state->Mutex) {
                             state->Statistics[Max<ui32>()].Entries.emplace_back(TStringBuilder() << "Activation:" << attr.GetName(), 0, 0, 0, 0, 1);
                         }