fix: improve robustness of managed udf code extraction (#1634)

* fix: improve robustness of managed udf code extraction

* resolve comments

* quick fix

Author: jialuoo

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 = (
             (

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))