From b2e26d0d4ca5d1ae05d4163cf493558e2d313435 Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Fri, 20 Jun 2025 20:13:25 +0000 Subject: [PATCH 1/5] fix: update the iam binding update logic This is needed as some recent updates to a dependency (google-cloud-resource-manager or its dependencies) broke the existing logic, and we are seeing the error like this: --> policy.bindings.append(new_binding) TypeError: Expected a message object, but got {'role': 'roles/run.invoker', 'members': [...]} --- bigframes/clients.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/bigframes/clients.py b/bigframes/clients.py index f1f6d686fd..77489fda21 100644 --- a/bigframes/clients.py +++ b/bigframes/clients.py @@ -24,6 +24,7 @@ import google.api_core.exceptions import google.api_core.retry from google.cloud import bigquery_connection_v1, resourcemanager_v3 +import google.cloud.iam logger = logging.getLogger(__name__) @@ -172,10 +173,9 @@ def _ensure_iam_binding( return # Create a new binding - new_binding = { - "role": role, - "members": [service_account], - } # Use a dictionary to avoid problematic google.iam namespace package. + new_binding = google.iam.v1.policy_pb2.Binding( + role=role, members=[service_account] + ) # Use a dictionary to avoid problematic google.iam namespace package. policy.bindings.append(new_binding) request = { "resource": project, From 2b24e4e0f9d1eb0d31512cebe3155515276ca26b Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Fri, 20 Jun 2025 20:51:52 +0000 Subject: [PATCH 2/5] use right import --- bigframes/clients.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bigframes/clients.py b/bigframes/clients.py index 77489fda21..480d013189 100644 --- a/bigframes/clients.py +++ b/bigframes/clients.py @@ -24,7 +24,7 @@ import google.api_core.exceptions import google.api_core.retry from google.cloud import bigquery_connection_v1, resourcemanager_v3 -import google.cloud.iam +import google.iam.v1 logger = logging.getLogger(__name__) From 989992419f594c5fc7b80de120e2ea7edae3c95f Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Wed, 25 Jun 2025 23:13:44 +0000 Subject: [PATCH 3/5] add unit test mocking the iam update method --- tests/unit/test_clients.py | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/tests/unit/test_clients.py b/tests/unit/test_clients.py index 032512c26e..6f6344a042 100644 --- a/tests/unit/test_clients.py +++ b/tests/unit/test_clients.py @@ -12,6 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +from unittest import mock + +from google.cloud import bigquery_connection_v1, resourcemanager_v3 +from google.iam.v1 import policy_pb2 import pytest from bigframes import clients @@ -65,3 +69,25 @@ def test_get_canonical_bq_connection_id_invalid_path(): default_project="default-project", default_location="us", ) + + +def test_ensure_iam_binding(): + bq_connection_client = bigquery_connection_v1.ConnectionServiceClient() + resource_manager_client = mock.create_autospec( + resourcemanager_v3.ProjectsClient, instance=True + ) + resource_manager_client.get_iam_policy.return_value = policy_pb2.Policy( + bindings=[ + policy_pb2.Binding( + role="roles/test.role1", members=["serviceAccount:serviceAccount1"] + ) + ] + ) + bq_connection_manager = clients.BqConnectionManager( + bq_connection_client, resource_manager_client + ) + bq_connection_manager._IAM_WAIT_SECONDS = 0 # no need to wait in test + bq_connection_manager._ensure_iam_binding( + "test-project", "serviceAccount2", "roles/test.role2" + ) + resource_manager_client.set_iam_policy.assert_called_once() From 9d912c8b72b687ca55425b2f925fe10bfe995f4e Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Wed, 25 Jun 2025 23:19:32 +0000 Subject: [PATCH 4/5] use mock bq connection client --- tests/unit/test_clients.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_clients.py b/tests/unit/test_clients.py index 6f6344a042..9daa759838 100644 --- a/tests/unit/test_clients.py +++ b/tests/unit/test_clients.py @@ -72,7 +72,9 @@ def test_get_canonical_bq_connection_id_invalid_path(): def test_ensure_iam_binding(): - bq_connection_client = bigquery_connection_v1.ConnectionServiceClient() + bq_connection_client = mock.create_autospec( + bigquery_connection_v1.ConnectionServiceClient, instance=True + ) resource_manager_client = mock.create_autospec( resourcemanager_v3.ProjectsClient, instance=True ) From 5f6184319160aa0fb3a5ba0e2626f6b4b745d444 Mon Sep 17 00:00:00 2001 From: Shobhit Singh Date: Thu, 26 Jun 2025 00:26:31 +0000 Subject: [PATCH 5/5] import module rephrasing --- bigframes/clients.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/bigframes/clients.py b/bigframes/clients.py index 480d013189..e6ddd5c6cb 100644 --- a/bigframes/clients.py +++ b/bigframes/clients.py @@ -24,7 +24,7 @@ import google.api_core.exceptions import google.api_core.retry from google.cloud import bigquery_connection_v1, resourcemanager_v3 -import google.iam.v1 +from google.iam.v1 import policy_pb2 logger = logging.getLogger(__name__) @@ -173,9 +173,7 @@ def _ensure_iam_binding( return # Create a new binding - new_binding = google.iam.v1.policy_pb2.Binding( - role=role, members=[service_account] - ) # Use a dictionary to avoid problematic google.iam namespace package. + new_binding = policy_pb2.Binding(role=role, members=[service_account]) policy.bindings.append(new_binding) request = { "resource": project,