feat: option to use bq connection without check

bigframes/_config/bigquery_options.py

@@ -33,19 +33,22 @@ class BigQueryOptions:
 
     def __init__(
         self,
+        *,
         credentials: Optional[google.auth.credentials.Credentials] = None,
         project: Optional[str] = None,
         location: Optional[str] = None,
-        bq_connection: Optional[str] = None,
         use_regional_endpoints: bool = False,
+        bq_connection: Optional[str] = None,
+        skip_bq_connection_check: bool = False,
         application_name: Optional[str] = None,
         kms_key_name: Optional[str] = None,
     ):
         self._credentials = credentials
         self._project = project
         self._location = location
-        self._bq_connection = bq_connection
         self._use_regional_endpoints = use_regional_endpoints
+        self._bq_connection = bq_connection
+        self._skip_bq_connection_check = skip_bq_connection_check
         self._application_name = application_name
         self._kms_key_name = kms_key_name
         self._session_started = False
@@ -105,14 +108,16 @@ def project(self, value: Optional[str]):
 
     @property
     def bq_connection(self) -> Optional[str]:
-        """Name of the BigQuery connection to use. Should be of the form <PROJECT_NUMBER/PROJECT_ID>.<LOCATION>.<CONNECTION_ID>.
+        """Name of the BigQuery connection to use. Should be of the form
+        <PROJECT_NUMBER/PROJECT_ID>.<LOCATION>.<CONNECTION_ID>.
 
         You should either have the connection already created in the
         <code>location</code> you have chosen, or you should have the Project IAM
         Admin role to enable the service to create the connection for you if you
         need it.
 
-        If this option isn't provided, or project or location aren't provided, session will use its default project/location/connection_id as default connection.
+        If this option isn't provided, or project or location aren't provided,
+        session will use its default project/location/connection_id as default connection.
         """
         return self._bq_connection
 
@@ -122,6 +127,26 @@ def bq_connection(self, value: Optional[str]):
             raise ValueError(SESSION_STARTED_MESSAGE.format(attribute="bq_connection"))
         self._bq_connection = value
 
+    @property
+    def skip_bq_connection_check(self) -> bool:
+        """Forcibly use the BigQuery connection.
+
+        Setting this flag to True would avoid creating the BigQuery connection
+        and checking or setting IAM permissions on it. So if the BigQuery
+        connection (default or user-provided) does not exist, or it does not have
+        necessary permissions set up to support BigQuery DataFrames operations,
+        then a runtime error will be reported.
+        """
+        return self._skip_bq_connection_check
+
+    @skip_bq_connection_check.setter
+    def skip_bq_connection_check(self, value: bool):
+        if self._session_started and self._skip_bq_connection_check != value:
+            raise ValueError(
+                SESSION_STARTED_MESSAGE.format(attribute="skip_bq_connection_check")
+            )
+        self._skip_bq_connection_check = value
+
     @property
     def use_regional_endpoints(self) -> bool:
         """Flag to connect to regional API endpoints.

bigframes/clients.py

@@ -73,12 +73,6 @@ def create_bq_connection(
             iam_role:
                 str of the IAM role that the service account of the created connection needs to aquire. E.g. 'run.invoker', 'aiplatform.user'
         """
-        # TODO(shobs): The below command to enable BigQuery Connection API needs
-        # to be automated. Disabling for now since most target users would not
-        # have the privilege to enable API in a project.
-        # log("Making sure BigQuery Connection API is enabled")
-        # if os.system("gcloud services enable bigqueryconnection.googleapis.com"):
-        #    raise ValueError("Failed to enable BigQuery Connection API")
         # If the intended connection does not exist then create it
         service_account_id = self._get_service_account_if_connection_exists(
             project_id, location, connection_id

bigframes/functions/remote_function.py

@@ -126,9 +126,8 @@ def __init__(
         bq_location,
         bq_dataset,
         bq_client,
-        bq_connection_client,
         bq_connection_id,
-        cloud_resource_manager_client,
+        bq_connection_manager,
         cloud_function_service_account,
         cloud_function_kms_key_name,
         cloud_function_docker_repository,
@@ -140,9 +139,7 @@ def __init__(
         self._bq_dataset = bq_dataset
         self._bq_client = bq_client
         self._bq_connection_id = bq_connection_id
-        self._bq_connection_manager = clients.BqConnectionManager(
-            bq_connection_client, cloud_resource_manager_client
-        )
+        self._bq_connection_manager = bq_connection_manager
         self._cloud_function_service_account = cloud_function_service_account
         self._cloud_function_kms_key_name = cloud_function_kms_key_name
         self._cloud_function_docker_repository = cloud_function_docker_repository
@@ -152,12 +149,13 @@ def create_bq_remote_function(
     ):
         """Create a BigQuery remote function given the artifacts of a user defined
         function and the http endpoint of a corresponding cloud function."""
-        self._bq_connection_manager.create_bq_connection(
-            self._gcp_project_id,
-            self._bq_location,
-            self._bq_connection_id,
-            "run.invoker",
-        )
+        if self._bq_connection_manager:
+            self._bq_connection_manager.create_bq_connection(
+                self._gcp_project_id,
+                self._bq_location,
+                self._bq_connection_id,
+                "run.invoker",
+            )
 
         # Create BQ function
         # https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#create_a_remote_function_2
@@ -816,6 +814,17 @@ def remote_function(
             " For more details see https://cloud.google.com/functions/docs/securing/cmek#before_you_begin"
         )
 
+    skip_bq_connection_check = (
+        False if session is None else session._skip_bq_connection_check
+    )
+    bq_connection_manager = (
+        None
+        if skip_bq_connection_check
+        else clients.BqConnectionManager(
+            bigquery_connection_client, resource_manager_client
+        )
+    )
+
     def wrapper(f):
         if not callable(f):
             raise TypeError("f must be callable, got {}".format(f))
@@ -832,9 +841,8 @@ def wrapper(f):
             bq_location,
             dataset_ref.dataset_id,
             bigquery_client,
-            bigquery_connection_client,
             bq_connection_id,
-            resource_manager_client,
+            bq_connection_manager,
             cloud_function_service_account,
             cloud_function_kms_key_name,
             cloud_function_docker_repository,

bigframes/ml/llm.py

@@ -73,12 +73,16 @@ def __init__(
     ):
         self.model_name = model_name
         self.session = session or bpd.get_global_session()
-        self._bq_connection_manager = clients.BqConnectionManager(
-            self.session.bqconnectionclient, self.session.resourcemanagerclient
+        self._bq_connection_manager = (
+            None
+            if self.session._skip_bq_connection_check
+            else clients.BqConnectionManager(
+                self.session.bqconnectionclient, self.session.resourcemanagerclient
+            )
         )
 
         connection_name = connection_name or self.session._bq_connection
-        self.connection_name = self._bq_connection_manager.resolve_full_connection_name(
+        self.connection_name = clients.BqConnectionManager.resolve_full_connection_name(
             connection_name,
             default_project=self.session._project,
             default_location=self.session._location,
@@ -93,17 +97,19 @@ def _create_bqml_model(self):
             raise ValueError(
                 "Must provide connection_name, either in constructor or through session options."
             )
-        connection_name_parts = self.connection_name.split(".")
-        if len(connection_name_parts) != 3:
-            raise ValueError(
-                f"connection_name must be of the format <PROJECT_NUMBER/PROJECT_ID>.<LOCATION>.<CONNECTION_ID>, got {self.connection_name}."
+
+        if self._bq_connection_manager:
+            connection_name_parts = self.connection_name.split(".")
+            if len(connection_name_parts) != 3:
+                raise ValueError(
+                    f"connection_name must be of the format <PROJECT_NUMBER/PROJECT_ID>.<LOCATION>.<CONNECTION_ID>, got {self.connection_name}."
+                )
+            self._bq_connection_manager.create_bq_connection(
+                project_id=connection_name_parts[0],
+                location=connection_name_parts[1],
+                connection_id=connection_name_parts[2],
+                iam_role="aiplatform.user",
             )
-        self._bq_connection_manager.create_bq_connection(
-            project_id=connection_name_parts[0],
-            location=connection_name_parts[1],
-            connection_id=connection_name_parts[2],
-            iam_role="aiplatform.user",
-        )
 
         if self.model_name not in _TEXT_GENERATOR_ENDPOINTS:
             raise ValueError(
@@ -289,12 +295,16 @@ def __init__(
         self.model_name = model_name
         self.version = version
         self.session = session or bpd.get_global_session()
-        self._bq_connection_manager = clients.BqConnectionManager(
-            self.session.bqconnectionclient, self.session.resourcemanagerclient
+        self._bq_connection_manager = (
+            None
+            if self.session._skip_bq_connection_check
+            else clients.BqConnectionManager(
+                self.session.bqconnectionclient, self.session.resourcemanagerclient
+            )
         )
 
         connection_name = connection_name or self.session._bq_connection
-        self.connection_name = self._bq_connection_manager.resolve_full_connection_name(
+        self.connection_name = clients.BqConnectionManager.resolve_full_connection_name(
             connection_name,
             default_project=self.session._project,
             default_location=self.session._location,
@@ -309,17 +319,19 @@ def _create_bqml_model(self):
             raise ValueError(
                 "Must provide connection_name, either in constructor or through session options."
             )
-        connection_name_parts = self.connection_name.split(".")
-        if len(connection_name_parts) != 3:
-            raise ValueError(
-                f"connection_name must be of the format <PROJECT_NUMBER/PROJECT_ID>.<LOCATION>.<CONNECTION_ID>, got {self.connection_name}."
+
+        if self._bq_connection_manager:
+            connection_name_parts = self.connection_name.split(".")
+            if len(connection_name_parts) != 3:
+                raise ValueError(
+                    f"connection_name must be of the format <PROJECT_NUMBER/PROJECT_ID>.<LOCATION>.<CONNECTION_ID>, got {self.connection_name}."
+                )
+            self._bq_connection_manager.create_bq_connection(
+                project_id=connection_name_parts[0],
+                location=connection_name_parts[1],
+                connection_id=connection_name_parts[2],
+                iam_role="aiplatform.user",
             )
-        self._bq_connection_manager.create_bq_connection(
-            project_id=connection_name_parts[0],
-            location=connection_name_parts[1],
-            connection_id=connection_name_parts[2],
-            iam_role="aiplatform.user",
-        )
 
         if self.model_name not in _EMBEDDING_GENERATOR_ENDPOINTS:
             raise ValueError(
@@ -437,12 +449,16 @@ def __init__(
         connection_name: Optional[str] = None,
     ):
         self.session = session or bpd.get_global_session()
-        self._bq_connection_manager = clients.BqConnectionManager(
-            self.session.bqconnectionclient, self.session.resourcemanagerclient
+        self._bq_connection_manager = (
+            None
+            if self.session._skip_bq_connection_check
+            else clients.BqConnectionManager(
+                self.session.bqconnectionclient, self.session.resourcemanagerclient
+            )
         )
 
         connection_name = connection_name or self.session._bq_connection
-        self.connection_name = self._bq_connection_manager.resolve_full_connection_name(
+        self.connection_name = clients.BqConnectionManager.resolve_full_connection_name(
             connection_name,
             default_project=self.session._project,
             default_location=self.session._location,
@@ -457,17 +473,19 @@ def _create_bqml_model(self):
             raise ValueError(
                 "Must provide connection_name, either in constructor or through session options."
             )
-        connection_name_parts = self.connection_name.split(".")
-        if len(connection_name_parts) != 3:
-            raise ValueError(
-                f"connection_name must be of the format <PROJECT_NUMBER/PROJECT_ID>.<LOCATION>.<CONNECTION_ID>, got {self.connection_name}."
+
+        if self._bq_connection_manager:
+            connection_name_parts = self.connection_name.split(".")
+            if len(connection_name_parts) != 3:
+                raise ValueError(
+                    f"connection_name must be of the format <PROJECT_NUMBER/PROJECT_ID>.<LOCATION>.<CONNECTION_ID>, got {self.connection_name}."
+                )
+            self._bq_connection_manager.create_bq_connection(
+                project_id=connection_name_parts[0],
+                location=connection_name_parts[1],
+                connection_id=connection_name_parts[2],
+                iam_role="aiplatform.user",
             )
-        self._bq_connection_manager.create_bq_connection(
-            project_id=connection_name_parts[0],
-            location=connection_name_parts[1],
-            connection_id=connection_name_parts[2],
-            iam_role="aiplatform.user",
-        )
 
         options = {"endpoint": _GEMINI_PRO_ENDPOINT}
 

bigframes/ml/remote.py

@@ -62,11 +62,15 @@ def __init__(
         self.output = output
         self.session = session or bpd.get_global_session()
 
-        self._bq_connection_manager = clients.BqConnectionManager(
-            self.session.bqconnectionclient, self.session.resourcemanagerclient
+        self._bq_connection_manager = (
+            None
+            if self.session._skip_bq_connection_check
+            else clients.BqConnectionManager(
+                self.session.bqconnectionclient, self.session.resourcemanagerclient
+            )
         )
         connection_name = connection_name or self.session._bq_connection
-        self.connection_name = self._bq_connection_manager.resolve_full_connection_name(
+        self.connection_name = clients.BqConnectionManager.resolve_full_connection_name(
             connection_name,
             default_project=self.session._project,
             default_location=self.session._location,
@@ -81,17 +85,19 @@ def _create_bqml_model(self):
             raise ValueError(
                 "Must provide connection_name, either in constructor or through session options."
             )
-        connection_name_parts = self.connection_name.split(".")
-        if len(connection_name_parts) != 3:
-            raise ValueError(
-                f"connection_name must be of the format <PROJECT_NUMBER/PROJECT_ID>.<LOCATION>.<CONNECTION_ID>, got {self.connection_name}."
+
+        if self._bq_connection_manager:
+            connection_name_parts = self.connection_name.split(".")
+            if len(connection_name_parts) != 3:
+                raise ValueError(
+                    f"connection_name must be of the format <PROJECT_NUMBER/PROJECT_ID>.<LOCATION>.<CONNECTION_ID>, got {self.connection_name}."
+                )
+            self._bq_connection_manager.create_bq_connection(
+                project_id=connection_name_parts[0],
+                location=connection_name_parts[1],
+                connection_id=connection_name_parts[2],
+                iam_role="aiplatform.user",
             )
-        self._bq_connection_manager.create_bq_connection(
-            project_id=connection_name_parts[0],
-            location=connection_name_parts[1],
-            connection_id=connection_name_parts[2],
-            iam_role="aiplatform.user",
-        )
 
         options = {
             "endpoint": self.endpoint,

bigframes/session/__init__.py

@@ -196,6 +196,7 @@ def __init__(
 
         # Resolve the BQ connection for remote function and Vertex AI integration
         self._bq_connection = context.bq_connection or _BIGFRAMES_DEFAULT_CONNECTION_ID
+        self._skip_bq_connection_check = context._skip_bq_connection_check
 
         # Now that we're starting the session, don't allow the options to be
         # changed.