diff --git a/bigframes/functions/_function_client.py b/bigframes/functions/_function_client.py index 8a591f6916..0cc3d52c38 100644 --- a/bigframes/functions/_function_client.py +++ b/bigframes/functions/_function_client.py @@ -19,6 +19,7 @@ import logging import os import random +import re import shutil import string import tempfile @@ -272,7 +273,10 @@ def provision_bq_managed_function( # i.e. there are no references to variables or imports outside the # body. udf_code = textwrap.dedent(inspect.getsource(func)) - udf_code = udf_code[udf_code.index("def") :] + match = re.search(r"^def ", udf_code, flags=re.MULTILINE) + if match is None: + raise ValueError("The UDF is not defined correctly.") + udf_code = udf_code[match.start() :] with_connection_clause = ( ( diff --git a/tests/system/large/functions/test_managed_function.py b/tests/system/large/functions/test_managed_function.py index a15bce83ad..9eba1907e6 100644 --- a/tests/system/large/functions/test_managed_function.py +++ b/tests/system/large/functions/test_managed_function.py @@ -171,7 +171,12 @@ def featurize(x: int) -> list[float]: def test_managed_function_series_apply(session, dataset_id, scalars_dfs): try: - @session.udf(dataset=dataset_id, name=prefixer.create_prefix()) + # An explicit name with "def" in it is used to test the robustness of + # the user code extraction logic, which depends on that term. + bq_name = f"{prefixer.create_prefix()}_def_to_test_code_extraction" + assert "def" in bq_name, "The substring 'def' was not found in 'bq_name'" + + @session.udf(dataset=dataset_id, name=bq_name) def foo(x: int) -> bytes: return bytes(abs(x))