feat: support bq connection path format (#1550)

* feat: support bq connection path format

For example, now a user specified bq connection like
"projects/project_id/locations/northamerica-northeast1/connections/conn-name"
would also be supported.

* include path format in tests with connection mismatch

* pass cloud_function_service_account="default" in more tests

Author: shobsi

bigframes/clients.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import logging
+import textwrap
 import time
 from typing import cast, Optional
 
@@ -28,21 +29,46 @@
 logger = logging.getLogger(__name__)
 
 
-def resolve_full_bq_connection_name(
-    connection_name: str, default_project: str, default_location: str
+def get_canonical_bq_connection_id(
+    connection_id: str, default_project: str, default_location: str
 ) -> str:
-    """Retrieve the full connection name of the form <PROJECT_NUMBER/PROJECT_ID>.<LOCATION>.<CONNECTION_ID>.
-    Use default project, location or connection_id when any of them are missing."""
-    if connection_name.count(".") == 2:
-        return connection_name
-
-    if connection_name.count(".") == 1:
-        return f"{default_project}.{connection_name}"
-
-    if connection_name.count(".") == 0:
-        return f"{default_project}.{default_location}.{connection_name}"
-
-    raise ValueError(f"Invalid connection name format: {connection_name}.")
+    """
+    Retrieve the full connection id of the form
+    <PROJECT_NUMBER/PROJECT_ID>.<LOCATION>.<CONNECTION_ID>.
+    Use default project, location or connection_id when any of them are missing.
+    """
+
+    if "/" in connection_id:
+        fields = connection_id.split("/")
+        if (
+            len(fields) == 6
+            and fields[0] == "projects"
+            and fields[2] == "locations"
+            and fields[4] == "connections"
+        ):
+            return ".".join((fields[1], fields[3], fields[5]))
+    else:
+        if connection_id.count(".") == 2:
+            return connection_id
+
+        if connection_id.count(".") == 1:
+            return f"{default_project}.{connection_id}"
+
+        if connection_id.count(".") == 0:
+            return f"{default_project}.{default_location}.{connection_id}"
+
+    raise ValueError(
+        textwrap.dedent(
+            f"""
+        Invalid connection id format: {connection_id}.
+        Only the following formats are supported:
+            <project-id>.<location>.<connection-id>,
+            <location>.<connection-id>,
+            <connection-id>,
+            projects/<project-id>/locations/<location>/connections/<connection-id>
+        """
+        ).strip()
+    )
 
 
 class BqConnectionManager:

bigframes/functions/_function_session.py

@@ -167,7 +167,7 @@ def _resolve_bigquery_connection_id(
         if not bigquery_connection:
             bigquery_connection = session._bq_connection  # type: ignore
 
-        bigquery_connection = clients.resolve_full_bq_connection_name(
+        bigquery_connection = clients.get_canonical_bq_connection_id(
             bigquery_connection,
             default_project=dataset_ref.project,
             default_location=bq_location,

bigframes/operations/blob.py

@@ -297,7 +297,7 @@ def _resolve_connection(self, connection: Optional[str] = None) -> str:
             ValueError: If the connection cannot be resolved to a valid string.
         """
         connection = connection or self._block.session._bq_connection
-        return clients.resolve_full_bq_connection_name(
+        return clients.get_canonical_bq_connection_id(
             connection,
             default_project=self._block.session._project,
             default_location=self._block.session._location,

bigframes/session/__init__.py

@@ -1775,8 +1775,8 @@ def _create_bq_connection(
         """Create the connection with the session settings and try to attach iam role to the connection SA.
         If any of project, location or connection isn't specified, use the session defaults. Returns fully-qualified connection name."""
         connection = self._bq_connection if not connection else connection
-        connection = bigframes.clients.resolve_full_bq_connection_name(
-            connection_name=connection,
+        connection = bigframes.clients.get_canonical_bq_connection_id(
+            connection_id=connection,
             default_project=self._project,
             default_location=self._location,
         )

tests/system/large/functions/test_remote_function.py

@@ -2819,3 +2819,33 @@ def featurize(x: int) -> list[float]:
         cleanup_function_assets(
             featurize, session.bqclient, session.cloudfunctionsclient
         )
+
+
+@pytest.mark.flaky(retries=2, delay=120)
+def test_remote_function_connection_path_format(
+    session, scalars_dfs, dataset_id, bq_cf_connection
+):
+    try:
+
+        @session.remote_function(
+            dataset=dataset_id,
+            bigquery_connection=f"projects/{session.bqclient.project}/locations/{session._location}/connections/{bq_cf_connection}",
+            reuse=False,
+            cloud_function_service_account="default",
+        )
+        def foo(x: int) -> int:
+            return x + 1
+
+        scalars_df, scalars_pandas_df = scalars_dfs
+
+        bf_int64_col = scalars_df["int64_too"]
+        bf_result = bf_int64_col.apply(foo).to_pandas()
+
+        pd_int64_col = scalars_pandas_df["int64_too"]
+        pd_result = pd_int64_col.apply(foo)
+
+        # ignore any dtype disparity
+        pandas.testing.assert_series_equal(pd_result, bf_result, check_dtype=False)
+    finally:
+        # clean up the gcp assets created for the remote function
+        cleanup_function_assets(foo, session.bqclient, session.cloudfunctionsclient)

tests/system/small/functions/test_remote_function.py

@@ -25,6 +25,7 @@
 import test_utils.prefixer
 
 import bigframes
+import bigframes.clients
 import bigframes.dtypes
 import bigframes.exceptions
 from bigframes.functions import _utils as bff_utils
@@ -93,6 +94,11 @@ def session_with_bq_connection(bq_cf_connection) -> bigframes.Session:
     return session
 
 
+def get_bq_connection_id_path_format(connection_id_dot_format):
+    fields = connection_id_dot_format.split(".")
+    return f"projects/{fields[0]}/locations/{fields[1]}/connections/{fields[2]}"
+
+
 @pytest.mark.flaky(retries=2, delay=120)
 def test_remote_function_direct_no_session_param(
     bigquery_client,
@@ -155,11 +161,8 @@ def square(x):
 
 
 @pytest.mark.flaky(retries=2, delay=120)
-def test_remote_function_direct_no_session_param_location_specified(
-    bigquery_client,
-    bigqueryconnection_client,
-    cloudfunctions_client,
-    resourcemanager_client,
+def test_remote_function_connection_w_location(
+    session,
     scalars_dfs,
     dataset_id_permanent,
     bq_cf_connection_location,
@@ -170,10 +173,7 @@ def square(x):
     square = bff.remote_function(
         input_types=int,
         output_type=int,
-        bigquery_client=bigquery_client,
-        bigquery_connection_client=bigqueryconnection_client,
-        cloud_functions_client=cloudfunctions_client,
-        resource_manager_client=resourcemanager_client,
+        session=session,
         dataset=dataset_id_permanent,
         bigquery_connection=bq_cf_connection_location,
         # See e2e tests for tests that actually deploy the Cloud Function.
@@ -210,11 +210,8 @@ def square(x):
 
 
 @pytest.mark.flaky(retries=2, delay=120)
-def test_remote_function_direct_no_session_param_location_mismatched(
-    bigquery_client,
-    bigqueryconnection_client,
-    cloudfunctions_client,
-    resourcemanager_client,
+def test_remote_function_connection_w_location_mismatched(
+    session,
     dataset_id_permanent,
     bq_cf_connection_location_mismatched,
 ):
@@ -223,32 +220,41 @@ def square(x):
         # connection doesn't match the location of the dataset.
         return x * x  # pragma: NO COVER
 
-    with pytest.raises(
-        ValueError,
-        match=re.escape("The location does not match BigQuery connection location:"),
-    ):
-        bff.remote_function(
-            input_types=int,
-            output_type=int,
-            bigquery_client=bigquery_client,
-            bigquery_connection_client=bigqueryconnection_client,
-            cloud_functions_client=cloudfunctions_client,
-            resource_manager_client=resourcemanager_client,
-            dataset=dataset_id_permanent,
-            bigquery_connection=bq_cf_connection_location_mismatched,
-            # See e2e tests for tests that actually deploy the Cloud Function.
-            reuse=True,
-            name=get_function_name(square),
-            cloud_function_service_account="default",
-        )(square)
+    bq_cf_connection_location_mismatched_path_fmt = get_bq_connection_id_path_format(
+        bigframes.clients.get_canonical_bq_connection_id(
+            bq_cf_connection_location_mismatched,
+            session.bqclient.project,
+            session._location,
+        )
+    )
+    connection_ids = [
+        bq_cf_connection_location_mismatched,
+        bq_cf_connection_location_mismatched_path_fmt,
+    ]
+
+    for connection_id in connection_ids:
+        with pytest.raises(
+            ValueError,
+            match=re.escape(
+                "The location does not match BigQuery connection location:"
+            ),
+        ):
+            bff.remote_function(
+                input_types=int,
+                output_type=int,
+                session=session,
+                dataset=dataset_id_permanent,
+                bigquery_connection=connection_id,
+                # See e2e tests for tests that actually deploy the Cloud Function.
+                reuse=True,
+                name=get_function_name(square),
+                cloud_function_service_account="default",
+            )(square)
 
 
 @pytest.mark.flaky(retries=2, delay=120)
-def test_remote_function_direct_no_session_param_location_project_specified(
-    bigquery_client,
-    bigqueryconnection_client,
-    cloudfunctions_client,
-    resourcemanager_client,
+def test_remote_function_connection_w_location_project(
+    session,
     scalars_dfs,
     dataset_id_permanent,
     bq_cf_connection_location_project,
@@ -259,10 +265,7 @@ def square(x):
     square = bff.remote_function(
         input_types=int,
         output_type=int,
-        bigquery_client=bigquery_client,
-        bigquery_connection_client=bigqueryconnection_client,
-        cloud_functions_client=cloudfunctions_client,
-        resource_manager_client=resourcemanager_client,
+        session=session,
         dataset=dataset_id_permanent,
         bigquery_connection=bq_cf_connection_location_project,
         # See e2e tests for tests that actually deploy the Cloud Function.
@@ -299,11 +302,8 @@ def square(x):
 
 
 @pytest.mark.flaky(retries=2, delay=120)
-def test_remote_function_direct_no_session_param_project_mismatched(
-    bigquery_client,
-    bigqueryconnection_client,
-    cloudfunctions_client,
-    resourcemanager_client,
+def test_remote_function_connection_w_project_mismatched(
+    session,
     dataset_id_permanent,
     bq_cf_connection_location_project_mismatched,
 ):
@@ -312,26 +312,38 @@ def square(x):
         # connection doesn't match the project of the dataset.
         return x * x  # pragma: NO COVER
 
-    with pytest.raises(
-        ValueError,
-        match=re.escape(
-            "The project_id does not match BigQuery connection gcp_project_id:"
-        ),
-    ):
-        bff.remote_function(
-            input_types=int,
-            output_type=int,
-            bigquery_client=bigquery_client,
-            bigquery_connection_client=bigqueryconnection_client,
-            cloud_functions_client=cloudfunctions_client,
-            resource_manager_client=resourcemanager_client,
-            dataset=dataset_id_permanent,
-            bigquery_connection=bq_cf_connection_location_project_mismatched,
-            # See e2e tests for tests that actually deploy the Cloud Function.
-            reuse=True,
-            name=get_function_name(square),
-            cloud_function_service_account="default",
-        )(square)
+    bq_cf_connection_location_project_mismatched_path_fmt = (
+        get_bq_connection_id_path_format(
+            bigframes.clients.get_canonical_bq_connection_id(
+                bq_cf_connection_location_project_mismatched,
+                session.bqclient.project,
+                session._location,
+            )
+        )
+    )
+    connection_ids = [
+        bq_cf_connection_location_project_mismatched,
+        bq_cf_connection_location_project_mismatched_path_fmt,
+    ]
+
+    for connection_id in connection_ids:
+        with pytest.raises(
+            ValueError,
+            match=re.escape(
+                "The project_id does not match BigQuery connection gcp_project_id:"
+            ),
+        ):
+            bff.remote_function(
+                input_types=int,
+                output_type=int,
+                session=session,
+                dataset=dataset_id_permanent,
+                bigquery_connection=connection_id,
+                # See e2e tests for tests that actually deploy the Cloud Function.
+                reuse=True,
+                name=get_function_name(square),
+                cloud_function_service_account="default",
+            )(square)
 
 
 @pytest.mark.flaky(retries=2, delay=120)