Skip to content

fix: revert dict back to protobuf in the iam binding update #1838

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
Jun 26, 2025
6 changes: 2 additions & 4 deletions bigframes/clients.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
import google.api_core.exceptions
import google.api_core.retry
from google.cloud import bigquery_connection_v1, resourcemanager_v3
from google.iam.v1 import policy_pb2

logger = logging.getLogger(__name__)

Expand Down Expand Up @@ -172,10 +173,7 @@ 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 = policy_pb2.Binding(role=role, members=[service_account])
policy.bindings.append(new_binding)
request = {
"resource": project,
Expand Down
28 changes: 28 additions & 0 deletions tests/unit/test_clients.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -65,3 +69,27 @@ def test_get_canonical_bq_connection_id_invalid_path():
default_project="default-project",
default_location="us",
)


def test_ensure_iam_binding():
bq_connection_client = mock.create_autospec(
bigquery_connection_v1.ConnectionServiceClient, instance=True
)
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()