disable current iam has been fixed

ref:7adbe57fb5f39bf28444b11cf4040f75d1c8ab4a

2 files changed, 60 insertions, 21 deletions

diff --git a/ydb/core/yq/libs/control_plane_storage/ydb_control_plane_storage_impl.h b/ydb/core/yq/libs/control_plane_storage/ydb_control_plane_storage_impl.h
index c5fa4d318b..b9029cd567 100644
--- a/ydb/core/yq/libs/control_plane_storage/ydb_control_plane_storage_impl.h
+++ b/ydb/core/yq/libs/control_plane_storage/ydb_control_plane_storage_impl.h
@@ -577,7 +577,7 @@ private:
         TResultSetParser parser(resultSet);
         while (parser.TryNextRow()) {
             T entity;
-            Y_VERIFY(entity.ParseFromString(*parser.ColumnParser(columnName).GetOptionalString())); // TODO: move to run actor
+            Y_VERIFY(entity.ParseFromString(*parser.ColumnParser(columnName).GetOptionalString()));
             const TString name = entity.content().name();
             if (auto it = entities.find(name); it != entities.end()) {
                 const auto visibility = entity.content().acl().visibility();
@@ -592,6 +592,17 @@ private:
         return entities;
     }
 
+    template<typename T>
+    TVector<T> GetEntities(const TResultSet& resultSet, const TString& columnName)
+    {
+        TVector<T> entities;
+        TResultSetParser parser(resultSet);
+        while (parser.TryNextRow()) {
+            Y_VERIFY(entities.emplace_back().ParseFromString(*parser.ColumnParser(columnName).GetOptionalString()));
+        }
+        return entities;
+    }
+
     template<class ResponseEvent, class Result, class RequestEventPtr>
     TFuture<bool> SendResponse(const TString& name,
         NActors::TActorSystem* actorSystem,
diff --git a/ydb/core/yq/libs/control_plane_storage/ydb_control_plane_storage_queries.cpp b/ydb/core/yq/libs/control_plane_storage/ydb_control_plane_storage_queries.cpp
index cec2f1b7fb..04826e1a48 100644
--- a/ydb/core/yq/libs/control_plane_storage/ydb_control_plane_storage_queries.cpp
+++ b/ydb/core/yq/libs/control_plane_storage/ydb_control_plane_storage_queries.cpp
@@ -196,28 +196,43 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvCreateQuery
 
         if (request.execute_mode() != YandexQuery::SAVE) {
             // TODO: move to run actor priority selection
-            if (permissions.Check(TPermissions::BINDINGS_USE)) {
-                auto bindings = GetEntitiesWithVisibilityPriority<YandexQuery::Binding>(resultSets[resultSets.size() - 1], BINDING_COLUMN_NAME);
-                for (const auto& [_, binding]: bindings) {
-                    if (!Config.AvailableBindings.contains(binding.content().setting().binding_case())) {
-                        continue;
-                    }
 
-                    *queryInternal.add_binding() = binding;
+            TSet<TString> disabledConnections;
+            for (const auto& connection: GetEntities<YandexQuery::Connection>(resultSets[resultSets.size() - 2], CONNECTION_COLUMN_NAME)) {
+                if (!Config.AvailableConnections.contains(connection.content().setting().connection_case())) {
+                    disabledConnections.insert(connection.meta().id());
+                    continue;
+                }
+
+                if (GetIamAuth(connection) == YandexQuery::IamAuth::kCurrentIam && Config.Proto.GetDisableCurrentIam()) {
+                    disabledConnections.insert(connection.meta().id());
+                    continue;
                 }
             }
 
             if (permissions.Check(TPermissions::CONNECTIONS_USE)) {
                 auto connections = GetEntitiesWithVisibilityPriority<YandexQuery::Connection>(resultSets[resultSets.size() - 2], CONNECTION_COLUMN_NAME);
                 for (const auto& [_, connection]: connections) {
-                    if (!Config.AvailableConnections.contains(connection.content().setting().connection_case())) {
+                    if (disabledConnections.contains(connection.meta().id())) {
+                        continue;
+                    }
+                    *queryInternal.add_connection() = connection;
+                }
+            }
+
+            if (permissions.Check(TPermissions::BINDINGS_USE)) {
+                auto bindings = GetEntitiesWithVisibilityPriority<YandexQuery::Binding>(resultSets[resultSets.size() - 1], BINDING_COLUMN_NAME);
+                for (const auto& [_, binding]: bindings) {
+                    if (!Config.AvailableBindings.contains(binding.content().setting().binding_case())) {
                         continue;
                     }
 
-                    if (GetIamAuth(connection) == YandexQuery::IamAuth::kCurrentIam && Config.Proto.GetDisableCurrentIam()) {
+                    if (disabledConnections.contains(binding.content().connection_id())) {
                         continue;
                     }
-                    *queryInternal.add_connection() = connection;
+
+                    Cerr << "PAM: " << binding.content().connection_id() << Endl;
+                    *queryInternal.add_binding() = binding;
                 }
             }
         }
@@ -830,30 +845,43 @@ void TYdbControlPlaneStorageActor::Handle(TEvControlPlaneStorage::TEvModifyQuery
 
             internal.clear_binding();
             internal.clear_connection();
+
             // TODO: move to run actor priority selection
-            if (permissions.Check(TPermissions::BINDINGS_USE)) {
-                auto bindings = GetEntitiesWithVisibilityPriority<YandexQuery::Binding>(resultSets[resultSets.size() - 2], BINDING_COLUMN_NAME);
-                for (const auto& [_, binding]: bindings) {
-                    if (!Config.AvailableBindings.contains(binding.content().setting().binding_case())) {
-                        continue;
-                    }
+            TSet<TString> disabledConnections;
+            for (const auto& connection: GetEntities<YandexQuery::Connection>(resultSets[resultSets.size() - 3], CONNECTION_COLUMN_NAME)) {
+                if (!Config.AvailableConnections.contains(connection.content().setting().connection_case())) {
+                    disabledConnections.insert(connection.meta().id());
+                    continue;
+                }
 
-                    *internal.add_binding() = binding;
+                if (GetIamAuth(connection) == YandexQuery::IamAuth::kCurrentIam && Config.Proto.GetDisableCurrentIam()) {
+                    disabledConnections.insert(connection.meta().id());
+                    continue;
                 }
             }
 
             if (permissions.Check(TPermissions::CONNECTIONS_USE)) {
                 auto connections = GetEntitiesWithVisibilityPriority<YandexQuery::Connection>(resultSets[resultSets.size() - 3], CONNECTION_COLUMN_NAME);
                 for (const auto& [_, connection]: connections) {
-                    if (!Config.AvailableConnections.contains(connection.content().setting().connection_case())) {
+                    if (disabledConnections.contains(connection.meta().id())) {
+                        continue;
+                    }
+                    *internal.add_connection() = connection;
+                }
+            }
+
+            if (permissions.Check(TPermissions::BINDINGS_USE)) {
+                auto bindings = GetEntitiesWithVisibilityPriority<YandexQuery::Binding>(resultSets[resultSets.size() - 2], BINDING_COLUMN_NAME);
+                for (const auto& [_, binding]: bindings) {
+                    if (!Config.AvailableBindings.contains(binding.content().setting().binding_case())) {
                         continue;
                     }
 
-                    if (GetIamAuth(connection) == YandexQuery::IamAuth::kCurrentIam && Config.Proto.GetDisableCurrentIam()) {
+                    if (disabledConnections.contains(binding.content().connection_id())) {
                         continue;
                     }
 
-                    *internal.add_connection() = connection;
+                    *internal.add_binding() = binding;
                 }
             }
         }