From a8704872b09be739fccc1f7d9208f8e3d9af8395 Mon Sep 17 00:00:00 2001 From: jialuo Date: Wed, 10 Sep 2025 22:03:42 +0000 Subject: [PATCH] fix: Fix the potential invalid VPC egress configuration --- bigframes/functions/_function_client.py | 10 ++++- bigframes/functions/_function_session.py | 16 ++++++-- bigframes/pandas/__init__.py | 6 +-- bigframes/session/__init__.py | 6 +-- .../large/functions/test_remote_function.py | 40 +++++++++++++++++++ 5 files changed, 68 insertions(+), 10 deletions(-) diff --git a/bigframes/functions/_function_client.py b/bigframes/functions/_function_client.py index d994d6353a..641bf52dc9 100644 --- a/bigframes/functions/_function_client.py +++ b/bigframes/functions/_function_client.py @@ -25,9 +25,11 @@ import textwrap import types from typing import Any, cast, Optional, Sequence, Tuple, TYPE_CHECKING +import warnings import requests +import bigframes.exceptions as bfe import bigframes.formatting_helpers as bf_formatting import bigframes.functions.function_template as bff_template @@ -482,10 +484,16 @@ 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 is None: + msg = bfe.format_message( + "The 'vpc_connector_egress_settings' was not specified. Defaulting to 'private-ranges-only'.", + ) + warnings.warn(msg, category=UserWarning) + vpc_connector_egress_settings = "private-ranges-only" 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)}", + f"'{vpc_connector_egress_settings}' is 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, diff --git a/bigframes/functions/_function_session.py b/bigframes/functions/_function_session.py index 6b5c9bf071..9a38ef1957 100644 --- a/bigframes/functions/_function_session.py +++ b/bigframes/functions/_function_session.py @@ -245,9 +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_vpc_connector_egress_settings: Optional[ + Literal["all", "private-ranges-only", "unspecified"] + ] = None, cloud_function_memory_mib: Optional[int] = 1024, cloud_function_ingress_settings: Literal[ "all", "internal-only", "internal-and-gclb" @@ -514,6 +514,16 @@ def remote_function( " For more details see https://cloud.google.com/functions/docs/securing/cmek#before_you_begin.", ) + # A VPC connector is required to specify VPC egress settings. + if ( + cloud_function_vpc_connector_egress_settings is not None + and cloud_function_vpc_connector is None + ): + raise bf_formatting.create_exception_with_feedback_link( + ValueError, + "cloud_function_vpc_connector must be specified before cloud_function_vpc_connector_egress_settings.", + ) + if cloud_function_ingress_settings is None: cloud_function_ingress_settings = "internal-only" msg = bfe.format_message( diff --git a/bigframes/pandas/__init__.py b/bigframes/pandas/__init__.py index 9d4fc101f6..00c2c7bcd6 100644 --- a/bigframes/pandas/__init__.py +++ b/bigframes/pandas/__init__.py @@ -87,9 +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_vpc_connector_egress_settings: Optional[ + Literal["all", "private-ranges-only", "unspecified"] + ] = None, cloud_function_memory_mib: Optional[int] = 1024, cloud_function_ingress_settings: Literal[ "all", "internal-only", "internal-and-gclb" diff --git a/bigframes/session/__init__.py b/bigframes/session/__init__.py index 4c824256b1..f0cec864b4 100644 --- a/bigframes/session/__init__.py +++ b/bigframes/session/__init__.py @@ -1509,9 +1509,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_vpc_connector_egress_settings: Optional[ + Literal["all", "private-ranges-only", "unspecified"] + ] = None, cloud_function_memory_mib: Optional[int] = 1024, cloud_function_ingress_settings: Literal[ "all", "internal-only", "internal-and-gclb" diff --git a/tests/system/large/functions/test_remote_function.py b/tests/system/large/functions/test_remote_function.py index 22b623193d..068b9b5d30 100644 --- a/tests/system/large/functions/test_remote_function.py +++ b/tests/system/large/functions/test_remote_function.py @@ -1512,6 +1512,46 @@ def square_num(x): ) +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_function_no_vpc_connector(session): + def foo(x): + return x + + with pytest.raises( + ValueError, + match="^cloud_function_vpc_connector must be specified before cloud_function_vpc_connector_egress_settings", + ): + session.remote_function( + input_types=[int], + output_type=int, + reuse=False, + cloud_function_service_account="default", + cloud_function_vpc_connector=None, + cloud_function_vpc_connector_egress_settings="all", + cloud_function_ingress_settings="all", + )(foo) + + +@pytest.mark.flaky(retries=2, delay=120) +def test_remote_function_wrong_vpc_egress_value(session): + def foo(x): + return x + + with pytest.raises( + ValueError, + match="^'wrong-egress-value' is not one of the supported vpc egress settings values:", + ): + session.remote_function( + input_types=[int], + output_type=int, + reuse=False, + cloud_function_service_account="default", + cloud_function_vpc_connector="dummy-value", + cloud_function_vpc_connector_egress_settings="wrong-egress-value", + cloud_function_ingress_settings="all", + )(foo) + + @pytest.mark.parametrize( ("max_batching_rows"), [