Skip to content

Commit 9fb3cb4

Browse files
authored
fix: revert dict back to protobuf in the iam binding update (#1838)
* 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': [...]} * use right import * add unit test mocking the iam update method * use mock bq connection client * import module rephrasing
1 parent 4da333e commit 9fb3cb4

File tree

2 files changed

+30
-4
lines changed

2 files changed

+30
-4
lines changed

bigframes/clients.py

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import google.api_core.exceptions
2525
import google.api_core.retry
2626
from google.cloud import bigquery_connection_v1, resourcemanager_v3
27+
from google.iam.v1 import policy_pb2
2728

2829
logger = logging.getLogger(__name__)
2930

@@ -172,10 +173,7 @@ def _ensure_iam_binding(
172173
return
173174

174175
# Create a new binding
175-
new_binding = {
176-
"role": role,
177-
"members": [service_account],
178-
} # Use a dictionary to avoid problematic google.iam namespace package.
176+
new_binding = policy_pb2.Binding(role=role, members=[service_account])
179177
policy.bindings.append(new_binding)
180178
request = {
181179
"resource": project,

tests/unit/test_clients.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,10 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
from unittest import mock
16+
17+
from google.cloud import bigquery_connection_v1, resourcemanager_v3
18+
from google.iam.v1 import policy_pb2
1519
import pytest
1620

1721
from bigframes import clients
@@ -65,3 +69,27 @@ def test_get_canonical_bq_connection_id_invalid_path():
6569
default_project="default-project",
6670
default_location="us",
6771
)
72+
73+
74+
def test_ensure_iam_binding():
75+
bq_connection_client = mock.create_autospec(
76+
bigquery_connection_v1.ConnectionServiceClient, instance=True
77+
)
78+
resource_manager_client = mock.create_autospec(
79+
resourcemanager_v3.ProjectsClient, instance=True
80+
)
81+
resource_manager_client.get_iam_policy.return_value = policy_pb2.Policy(
82+
bindings=[
83+
policy_pb2.Binding(
84+
role="roles/test.role1", members=["serviceAccount:serviceAccount1"]
85+
)
86+
]
87+
)
88+
bq_connection_manager = clients.BqConnectionManager(
89+
bq_connection_client, resource_manager_client
90+
)
91+
bq_connection_manager._IAM_WAIT_SECONDS = 0 # no need to wait in test
92+
bq_connection_manager._ensure_iam_binding(
93+
"test-project", "serviceAccount2", "roles/test.role2"
94+
)
95+
resource_manager_client.set_iam_policy.assert_called_once()

0 commit comments

Comments
 (0)