diff --git a/bigframes/functions/remote_function.py b/bigframes/functions/remote_function.py index 7be252406c..4223998c53 100644 --- a/bigframes/functions/remote_function.py +++ b/bigframes/functions/remote_function.py @@ -1146,21 +1146,14 @@ def try_delattr(attr): def read_gbq_function( function_name: str, - session: Optional[Session] = None, - bigquery_client: Optional[bigquery.Client] = None, + *, + session: Session, ): """ Read an existing BigQuery function and prepare it for use in future queries. """ - - # A BigQuery client is required to perform BQ operations - if not bigquery_client and session: - bigquery_client = session.bqclient - if not bigquery_client: - raise ValueError( - "A bigquery client must be provided, either directly or via session. " - f"{constants.FEEDBACK_LINK}" - ) + bigquery_client = session.bqclient + ibis_client = session.ibis_client try: routine_ref = get_routine_reference(function_name, bigquery_client, session) @@ -1192,8 +1185,10 @@ def read_gbq_function( # non-standard names for the arguments here. def func(*ignored_args, **ignored_kwargs): f"""Remote function {str(routine_ref)}.""" - # TODO(swast): Construct an ibis client from bigquery_client and - # execute node via a query. + nonlocal node # type: ignore + + expr = node(*ignored_args, **ignored_kwargs) # type: ignore + return ibis_client.execute(expr) # TODO: Move ibis logic to compiler step func.__name__ = routine_ref.routine_id diff --git a/tests/system/small/test_remote_function.py b/tests/system/small/test_remote_function.py index d2ee4411f4..bf3424d0f0 100644 --- a/tests/system/small/test_remote_function.py +++ b/tests/system/small/test_remote_function.py @@ -537,12 +537,12 @@ def add_one(x): @pytest.mark.flaky(retries=2, delay=120) -def test_read_gbq_function_detects_invalid_function(bigquery_client, dataset_id): +def test_read_gbq_function_detects_invalid_function(session, dataset_id): dataset_ref = bigquery.DatasetReference.from_string(dataset_id) with pytest.raises(ValueError) as e: rf.read_gbq_function( str(dataset_ref.routine("not_a_function")), - bigquery_client=bigquery_client, + session=session, ) assert "Unknown function" in str(e.value) @@ -550,6 +550,7 @@ def test_read_gbq_function_detects_invalid_function(bigquery_client, dataset_id) @pytest.mark.flaky(retries=2, delay=120) def test_read_gbq_function_like_original( + session, bigquery_client, bigqueryconnection_client, cloudfunctions_client, @@ -577,7 +578,7 @@ def square1(x): square2 = rf.read_gbq_function( function_name=square1.bigframes_remote_function, - bigquery_client=bigquery_client, + session=session, ) # The newly-created function (square1) should have a remote function AND a @@ -607,7 +608,14 @@ def square1(x): @pytest.mark.flaky(retries=2, delay=120) -def test_read_gbq_function_reads_udfs(bigquery_client, dataset_id): +def test_read_gbq_function_runs_existing_udf(session, bigquery_client, dataset_id): + func = session.read_gbq_function("bqutil.fn.cw_lower_case_ascii_only") + got = func("AURÉLIE") + assert got == "aurÉlie" + + +@pytest.mark.flaky(retries=2, delay=120) +def test_read_gbq_function_reads_udfs(session, bigquery_client, dataset_id): dataset_ref = bigquery.DatasetReference.from_string(dataset_id) arg = bigquery.RoutineArgument( name="x", @@ -633,7 +641,8 @@ def test_read_gbq_function_reads_udfs(bigquery_client, dataset_id): # Create the routine in BigQuery and read it back using read_gbq_function. bigquery_client.create_routine(routine, exists_ok=True) square = rf.read_gbq_function( - str(routine.reference), bigquery_client=bigquery_client + str(routine.reference), + session=session, ) # It should point to the named routine and yield the expected results. @@ -658,7 +667,9 @@ def test_read_gbq_function_reads_udfs(bigquery_client, dataset_id): @pytest.mark.flaky(retries=2, delay=120) -def test_read_gbq_function_enforces_explicit_types(bigquery_client, dataset_id): +def test_read_gbq_function_enforces_explicit_types( + session, bigquery_client, dataset_id +): dataset_ref = bigquery.DatasetReference.from_string(dataset_id) typed_arg = bigquery.RoutineArgument( name="x", @@ -702,18 +713,22 @@ def test_read_gbq_function_enforces_explicit_types(bigquery_client, dataset_id): bigquery_client.create_routine(neither_type_specified, exists_ok=True) rf.read_gbq_function( - str(both_types_specified.reference), bigquery_client=bigquery_client + str(both_types_specified.reference), + session=session, ) rf.read_gbq_function( - str(only_return_type_specified.reference), bigquery_client=bigquery_client + str(only_return_type_specified.reference), + session=session, ) with pytest.raises(ValueError): rf.read_gbq_function( - str(only_arg_type_specified.reference), bigquery_client=bigquery_client + str(only_arg_type_specified.reference), + session=session, ) with pytest.raises(ValueError): rf.read_gbq_function( - str(neither_type_specified.reference), bigquery_client=bigquery_client + str(neither_type_specified.reference), + session=session, )