Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions bigframes/functions/_function_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,15 @@
}
)

# https://cloud.google.com/functions/docs/reference/rest/v2/projects.locations.functions#vpconnectoregresssettings
_VPC_EGRESS_SETTINGS_MAP = types.MappingProxyType(
{
"all": functions_v2.ServiceConfig.VpcConnectorEgressSettings.ALL_TRAFFIC,
"private-ranges-only": functions_v2.ServiceConfig.VpcConnectorEgressSettings.PRIVATE_RANGES_ONLY,
"unspecified": functions_v2.ServiceConfig.VpcConnectorEgressSettings.VPC_CONNECTOR_EGRESS_SETTINGS_UNSPECIFIED,
}
)

# BQ managed functions (@udf) currently only support Python 3.11.
_MANAGED_FUNC_PYTHON_VERSION = "python-3.11"

Expand Down Expand Up @@ -375,6 +384,7 @@ def create_cloud_function(
max_instance_count=None,
is_row_processor=False,
vpc_connector=None,
vpc_connector_egress_settings="private-ranges-only",
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

memory_mib=1024,
ingress_settings="internal-only",
):
Expand Down Expand Up @@ -472,6 +482,15 @@ def create_cloud_function(
function.service_config.max_instance_count = max_instance_count
if vpc_connector is not None:
function.service_config.vpc_connector = vpc_connector
if vpc_connector_egress_settings not in _VPC_EGRESS_SETTINGS_MAP:
raise bf_formatting.create_exception_with_feedback_link(
ValueError,
f"'{vpc_connector_egress_settings}' not one of the supported vpc egress settings values: {list(_VPC_EGRESS_SETTINGS_MAP)}",
)
function.service_config.vpc_connector_egress_settings = cast(
functions_v2.ServiceConfig.VpcConnectorEgressSettings,
_VPC_EGRESS_SETTINGS_MAP[vpc_connector_egress_settings],
)
function.service_config.service_account_email = (
self._cloud_function_service_account
)
Expand Down Expand Up @@ -532,6 +551,7 @@ def provision_bq_remote_function(
cloud_function_max_instance_count,
is_row_processor,
cloud_function_vpc_connector,
cloud_function_vpc_connector_egress_settings,
cloud_function_memory_mib,
cloud_function_ingress_settings,
bq_metadata,
Expand Down Expand Up @@ -580,6 +600,7 @@ def provision_bq_remote_function(
max_instance_count=cloud_function_max_instance_count,
is_row_processor=is_row_processor,
vpc_connector=cloud_function_vpc_connector,
vpc_connector_egress_settings=cloud_function_vpc_connector_egress_settings,
memory_mib=cloud_function_memory_mib,
ingress_settings=cloud_function_ingress_settings,
)
Expand Down
11 changes: 11 additions & 0 deletions bigframes/functions/_function_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -245,6 +245,9 @@ def remote_function(
cloud_function_timeout: Optional[int] = 600,
cloud_function_max_instances: Optional[int] = None,
cloud_function_vpc_connector: Optional[str] = None,
cloud_function_vpc_connector_egress_settings: Literal[
"all", "private-ranges-only", "unspecified"
] = "private-ranges-only",
cloud_function_memory_mib: Optional[int] = 1024,
cloud_function_ingress_settings: Literal[
"all", "internal-only", "internal-and-gclb"
Expand Down Expand Up @@ -425,6 +428,13 @@ def remote_function(
function. This is useful if your code needs access to data or
service(s) that are on a VPC network. See for more details
https://cloud.google.com/functions/docs/networking/connecting-vpc.
cloud_function_vpc_connector_egress_settings (str, Optional):
Egress settings for the VPC connector, controlling what outbound
traffic is routed through the VPC connector.
Options are: `all`, `private-ranges-only`, or `unspecified`.
If not specified, `private-ranges-only` is used by default.
See for more details
https://cloud.google.com/run/docs/configuring/vpc-connectors#egress-job.
cloud_function_memory_mib (int, Optional):
The amounts of memory (in mebibytes) to allocate for the cloud
function (2nd gen) created. This also dictates a corresponding
Expand Down Expand Up @@ -616,6 +626,7 @@ def wrapper(func):
cloud_function_max_instance_count=cloud_function_max_instances,
is_row_processor=is_row_processor,
cloud_function_vpc_connector=cloud_function_vpc_connector,
cloud_function_vpc_connector_egress_settings=cloud_function_vpc_connector_egress_settings,
cloud_function_memory_mib=cloud_function_memory_mib,
cloud_function_ingress_settings=cloud_function_ingress_settings,
bq_metadata=bqrf_metadata,
Expand Down
4 changes: 4 additions & 0 deletions bigframes/pandas/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,9 @@ def remote_function(
cloud_function_timeout: Optional[int] = 600,
cloud_function_max_instances: Optional[int] = None,
cloud_function_vpc_connector: Optional[str] = None,
cloud_function_vpc_connector_egress_settings: Literal[
"all", "private-ranges-only", "unspecified"
] = "private-ranges-only",
cloud_function_memory_mib: Optional[int] = 1024,
cloud_function_ingress_settings: Literal[
"all", "internal-only", "internal-and-gclb"
Expand All @@ -109,6 +112,7 @@ def remote_function(
cloud_function_timeout=cloud_function_timeout,
cloud_function_max_instances=cloud_function_max_instances,
cloud_function_vpc_connector=cloud_function_vpc_connector,
cloud_function_vpc_connector_egress_settings=cloud_function_vpc_connector_egress_settings,
cloud_function_memory_mib=cloud_function_memory_mib,
cloud_function_ingress_settings=cloud_function_ingress_settings,
cloud_build_service_account=cloud_build_service_account,
Expand Down
11 changes: 11 additions & 0 deletions bigframes/session/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -1510,6 +1510,9 @@ def remote_function(
cloud_function_timeout: Optional[int] = 600,
cloud_function_max_instances: Optional[int] = None,
cloud_function_vpc_connector: Optional[str] = None,
cloud_function_vpc_connector_egress_settings: Literal[
"all", "private-ranges-only", "unspecified"
] = "private-ranges-only",
cloud_function_memory_mib: Optional[int] = 1024,
cloud_function_ingress_settings: Literal[
"all", "internal-only", "internal-and-gclb"
Expand Down Expand Up @@ -1675,6 +1678,13 @@ def remote_function(
function. This is useful if your code needs access to data or
service(s) that are on a VPC network. See for more details
https://cloud.google.com/functions/docs/networking/connecting-vpc.
cloud_function_vpc_connector_egress_settings (str, Optional):
Egress settings for the VPC connector, controlling what outbound
traffic is routed through the VPC connector.
Options are: `all`, `private-ranges-only`, or `unspecified`.
If not specified, `private-ranges-only` is used by default.
See for more details
https://cloud.google.com/run/docs/configuring/vpc-connectors#egress-job.
cloud_function_memory_mib (int, Optional):
The amounts of memory (in mebibytes) to allocate for the cloud
function (2nd gen) created. This also dictates a corresponding
Expand Down Expand Up @@ -1732,6 +1742,7 @@ def remote_function(
cloud_function_timeout=cloud_function_timeout,
cloud_function_max_instances=cloud_function_max_instances,
cloud_function_vpc_connector=cloud_function_vpc_connector,
cloud_function_vpc_connector_egress_settings=cloud_function_vpc_connector_egress_settings,
cloud_function_memory_mib=cloud_function_memory_mib,
cloud_function_ingress_settings=cloud_function_ingress_settings,
cloud_build_service_account=cloud_build_service_account,
Expand Down
8 changes: 7 additions & 1 deletion tests/system/large/functions/test_remote_function.py
Original file line number Diff line number Diff line change
Expand Up @@ -1478,14 +1478,20 @@ def square_num(x):
reuse=False,
cloud_function_service_account="default",
cloud_function_vpc_connector=gcf_vpc_connector,
cloud_function_vpc_connector_egress_settings="all",
cloud_function_ingress_settings="all",
)(square_num)

# assert that the GCF is created with the intended vpc connector
gcf = rf_session.cloudfunctionsclient.get_function(
name=square_num_remote.bigframes_cloud_function
)

# assert that the GCF is created with the intended vpc connector and
# egress settings.
assert gcf.service_config.vpc_connector == gcf_vpc_connector
# The value is <VpcConnectorEgressSettings.ALL_TRAFFIC: 2> since we set
# cloud_function_vpc_connector_egress_settings="all" earlier.
assert gcf.service_config.vpc_connector_egress_settings == 2

# assert that the function works as expected on data
scalars_df, scalars_pandas_df = scalars_dfs
Expand Down