feat: support custom build service account in remote_function

bigframes/functions/_function_client.py

@@ -77,6 +77,7 @@ def __init__(
  cloud_function_service_account=None,
  cloud_function_kms_key_name=None,
  cloud_function_docker_repository=None,
+ cloud_build_service_account=None,
  *,
  session: Session,
  ):
@@ -94,6 +95,7 @@ def __init__(
  self._cloud_function_service_account = cloud_function_service_account
  self._cloud_function_kms_key_name = cloud_function_kms_key_name
  self._cloud_function_docker_repository = cloud_function_docker_repository
+ self._cloud_build_service_account = cloud_build_service_account
 
  def _create_bq_connection(self) -> None:
  if self._bq_connection_manager:
@@ -452,6 +454,17 @@ def create_cloud_function(
  function.build_config.docker_repository = (
  self._cloud_function_docker_repository
  )
+
+ if self._cloud_build_service_account:
+ canonical_cloud_build_service_account = (
+ self._cloud_build_service_account
+ if "/" in self._cloud_build_service_account
+ else f"projects/{self._gcp_project_id}/serviceAccounts/{self._cloud_build_service_account}"
+ )
+ function.build_config.service_account = (
+ canonical_cloud_build_service_account
+ )
+
  function.service_config = functions_v2.ServiceConfig()
  if memory_mib is not None:
  function.service_config.available_memory = f"{memory_mib}Mi"

bigframes/functions/_function_session.py

@@ -263,6 +263,7 @@ def remote_function(
  cloud_function_ingress_settings: Literal[
  "all", "internal-only", "internal-and-gclb"
  ] = "internal-only",
+ cloud_build_service_account: Optional[str] = None,
  ):
  """Decorator to turn a user defined function into a BigQuery remote function.
 
@@ -453,6 +454,16 @@ def remote_function(
  If no setting is provided, `internal-only` will be used by default.
  See for more details
  https://cloud.google.com/functions/docs/networking/network-settings#ingress_settings.
+ cloud_build_service_account (str, Optional):
+ Service account in the fully qualified format
+ `projects/PROJECT_ID/serviceAccounts/SERVICE_ACCOUNT_EMAIL`, or
+ just the SERVICE_ACCOUNT_EMAIL. The latter would be interpreted
+ as belonging to the BigQuery DataFrames session project. This is
+ to be used by Cloud Build to build the function source code into
+ a deployable artifact. If not provided, the default Cloud Build
+ service account is used. See
+ https://cloud.google.com/build/docs/cloud-build-service-account
+ for more details.
  """
  # Some defaults may be used from the session if not provided otherwise.
  session = self._resolve_session(session)
@@ -599,6 +610,7 @@ def wrapper(func):
  else cloud_function_service_account,
  cloud_function_kms_key_name,
  cloud_function_docker_repository,
+ cloud_build_service_account=cloud_build_service_account,
  session=session, # type: ignore
  )
 

bigframes/pandas/__init__.py

@@ -89,6 +89,7 @@ def remote_function(
  cloud_function_ingress_settings: Literal[
  "all", "internal-only", "internal-and-gclb"
  ] = "internal-only",
+ cloud_build_service_account: Optional[str] = None,
 ):
  return global_session.with_default_session(
  bigframes.session.Session.remote_function,
@@ -108,6 +109,7 @@ def remote_function(
  cloud_function_vpc_connector=cloud_function_vpc_connector,
  cloud_function_memory_mib=cloud_function_memory_mib,
  cloud_function_ingress_settings=cloud_function_ingress_settings,
+ cloud_build_service_account=cloud_build_service_account,
  )
 
 

bigframes/session/__init__.py

@@ -1378,6 +1378,7 @@ def remote_function(
  cloud_function_ingress_settings: Literal[
  "all", "internal-only", "internal-and-gclb"
  ] = "internal-only",
+ cloud_build_service_account: Optional[str] = None,
  ):
  """Decorator to turn a user defined function into a BigQuery remote function. Check out
  the code samples at: https://cloud.google.com/bigquery/docs/remote-functions#bigquery-dataframes.
@@ -1553,6 +1554,16 @@ def remote_function(
  If no setting is provided, `internal-only` will be used by default.
  See for more details
  https://cloud.google.com/functions/docs/networking/network-settings#ingress_settings.
+ cloud_build_service_account (str, Optional):
+ Service account in the fully qualified format
+ `projects/PROJECT_ID/serviceAccounts/SERVICE_ACCOUNT_EMAIL`, or
+ just the SERVICE_ACCOUNT_EMAIL. The latter would be interpreted
+ as belonging to the BigQuery DataFrames session project. This is
+ to be used by Cloud Build to build the function source code into
+ a deployable artifact. If not provided, the default Cloud Build
+ service account is used. See
+ https://cloud.google.com/build/docs/cloud-build-service-account
+ for more details.
  Returns:
  collections.abc.Callable:
  A remote function object pointing to the cloud assets created
@@ -1581,6 +1592,7 @@ def remote_function(
  cloud_function_vpc_connector=cloud_function_vpc_connector,
  cloud_function_memory_mib=cloud_function_memory_mib,
  cloud_function_ingress_settings=cloud_function_ingress_settings,
+ cloud_build_service_account=cloud_build_service_account,
  )
 
  def udf(

tests/system/large/functions/test_remote_function.py

@@ -1342,7 +1342,7 @@ def test_remote_function_via_session_custom_sa(scalars_dfs):
  # For upfront convenience, the following set up has been statically created
  # in the project bigfrmames-dev-perf via cloud console:
  #
- # 1. Create a service account as per
+ # 1. Create a service account bigframes-dev-perf-1@bigframes-dev-perf.iam.gserviceaccount.com as per
  # https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console
  # 2. Give necessary roles as per
  # https://cloud.google.com/functions/docs/reference/iam/roles#additional-configuration
@@ -1395,6 +1395,80 @@ def square_num(x):
  )
 
 
+@pytest.mark.parametrize(
+ ("set_build_service_account"),
+ [
+ pytest.param(
+ "projects/bigframes-dev-perf/serviceAccounts/bigframes-dev-perf-1@bigframes-dev-perf.iam.gserviceaccount.com",
+ id="fully-qualified-sa",
+ ),
+ pytest.param(
+ "bigframes-dev-perf-1@bigframes-dev-perf.iam.gserviceaccount.com",
+ id="just-sa-email",
+ ),
+ ],
+)
+@pytest.mark.flaky(retries=2, delay=120)
+def test_remote_function_via_session_custom_build_sa(
+ scalars_dfs, set_build_service_account
+):
+ # TODO(shobs): Automate the following set-up during testing in the test project.
+ #
+ # For upfront convenience, the following set up has been statically created
+ # in the project bigfrmames-dev-perf via cloud console:
+ #
+ # 1. Create a service account bigframes-dev-perf-1@bigframes-dev-perf.iam.gserviceaccount.com as per
+ # https://cloud.google.com/iam/docs/service-accounts-create#iam-service-accounts-create-console
+ # 2. Give "Cloud Build Service Account (roles/cloudbuild.builds.builder)" role as per
+ # https://cloud.google.com/build/docs/cloud-build-service-account#default_permissions_of_the_legacy_service_account
+ #
+ project = "bigframes-dev-perf"
+ expected_build_service_account = "projects/bigframes-dev-perf/serviceAccounts/bigframes-dev-perf-1@bigframes-dev-perf.iam.gserviceaccount.com"
+
+ rf_session = bigframes.Session(context=bigframes.BigQueryOptions(project=project))
+
+ try:
+
+ # TODO(shobs): Figure out why the default ingress setting
+ # (internal-only) does not work here
+ @rf_session.remote_function(
+ input_types=[int],
+ output_type=int,
+ reuse=False,
+ cloud_function_service_account="default",
+ cloud_build_service_account=set_build_service_account,
+ cloud_function_ingress_settings="all",
+ )
+ def square_num(x):
+ if x is None:
+ return x
+ return x * x
+
+ # assert that the GCF is created with the intended SA
+ gcf = rf_session.cloudfunctionsclient.get_function(
+ name=square_num.bigframes_cloud_function
+ )
+ assert gcf.build_config.service_account == expected_build_service_account
+
+ # assert that the function works as expected on data
+ scalars_df, scalars_pandas_df = scalars_dfs
+
+ bf_int64_col = scalars_df["int64_col"]
+ bf_result_col = bf_int64_col.apply(square_num)
+ bf_result = bf_int64_col.to_frame().assign(result=bf_result_col).to_pandas()
+
+ pd_int64_col = scalars_pandas_df["int64_col"]
+ pd_result_col = pd_int64_col.apply(lambda x: x if x is None else x * x)
+ pd_result = pd_int64_col.to_frame().assign(result=pd_result_col)
+
+ assert_pandas_df_equal(bf_result, pd_result, check_dtype=False)
+ finally:
+ # clean up the gcp assets created for the remote function
+ cleanup_function_assets(
+ square_num, rf_session.bqclient, rf_session.cloudfunctionsclient
+ )
+
+
 def test_remote_function_throws_none_cloud_function_service_account(session):
  with pytest.raises(
  ValueError,