feat: custom query labels for compute options (#638)

* feat: Custom query labels for compute options

* Update docstring

* Update test

* Code example update.

* ignore type

* update format

Author: Genesis929

bigframes/_config/compute_options.py

@@ -15,7 +15,7 @@
 """Options for displaying objects."""
 
 import dataclasses
-from typing import Optional
+from typing import Any, Dict, Optional
 
 
 @dataclasses.dataclass
@@ -34,6 +34,26 @@ class ComputeOptions:
 
         >>> bpd.options.compute.maximum_bytes_billed = None  # reset option
 
+    To add multiple extra labels to a query configuration, use the `assign_extra_query_labels`
+    method with keyword arguments:
+
+        >>> bpd.options.compute.assign_extra_query_labels(test1=1, test2="abc")
+        >>> bpd.options.compute.extra_query_labels
+        {'test1': 1, 'test2': 'abc'}
+
+    Alternatively, you can add labels individually by directly accessing the `extra_query_labels`
+    dictionary:
+
+        >>> bpd.options.compute.extra_query_labels["test3"] = False
+        >>> bpd.options.compute.extra_query_labels
+        {'test1': 1, 'test2': 'abc', 'test3': False}
+
+    To remove a label from the configuration, use the `del` keyword on the desired label key:
+
+        >>> del bpd.options.compute.extra_query_labels["test1"]
+        >>> bpd.options.compute.extra_query_labels
+        {'test2': 'abc', 'test3': False}
+
     Attributes:
         maximum_bytes_billed (int, Options):
             Limits the bytes billed for query jobs. Queries that will have
@@ -44,7 +64,38 @@ class ComputeOptions:
             If enabled, large queries may be factored into multiple smaller queries
             in order to avoid generating queries that are too complex for the query
             engine to handle. However this comes at the cost of increase cost and latency.
+        extra_query_labels (Dict[str, Any], Options):
+            Stores additional custom labels for query configuration.
     """
 
     maximum_bytes_billed: Optional[int] = None
     enable_multi_query_execution: bool = False
+    extra_query_labels: Dict[str, Any] = dataclasses.field(
+        default_factory=dict, init=False
+    )
+
+    def assign_extra_query_labels(self, **kwargs: Any) -> None:
+        """
+        Assigns additional custom labels for query configuration. The method updates the
+        `extra_query_labels` dictionary with new labels provided through keyword arguments.
+
+        Args:
+            kwargs (Any):
+                Custom labels provided as keyword arguments. Each key-value pair
+                in `kwargs` represents a label name and its value.
+
+        Raises:
+            ValueError: If a key matches one of the reserved attribute names,
+                specifically 'maximum_bytes_billed' or 'enable_multi_query_execution',
+                to prevent conflicts with built-in settings.
+        """
+        reserved_keys = ["maximum_bytes_billed", "enable_multi_query_execution"]
+        for key in kwargs:
+            if key in reserved_keys:
+                raise ValueError(
+                    f"'{key}' is a reserved attribute name. Please use "
+                    "a different key for your custom labels to avoid "
+                    "conflicts with built-in settings."
+                )
+
+        self.extra_query_labels.update(kwargs)

bigframes/session/__init__.py

@@ -1699,6 +1699,12 @@ def _prepare_query_job_config(
                 bigframes.options.compute.maximum_bytes_billed
             )
 
+        current_labels = job_config.labels if job_config.labels else {}
+        for key, value in bigframes.options.compute.extra_query_labels.items():
+            if key not in current_labels:
+                current_labels[key] = value
+        job_config.labels = current_labels
+
         if self._bq_kms_key_name:
             job_config.destination_encryption_configuration = (
                 bigquery.EncryptionConfiguration(kms_key_name=self._bq_kms_key_name)

tests/system/small/test_session.py

@@ -405,6 +405,24 @@ def test_read_gbq_with_configuration(
     assert df.shape == (9, 3)
 
 
+def test_read_gbq_with_custom_global_labels(
+    session: bigframes.Session, scalars_table_id: str
+):
+    bigframes.options.compute.assign_extra_query_labels(test1=1, test2="abc")
+    bigframes.options.compute.extra_query_labels["test3"] = False
+
+    job_labels = session.read_gbq(scalars_table_id).query_job.labels  # type:ignore
+    expected_labels = {"test1": "1", "test2": "abc", "test3": "false"}
+
+    assert all(job_labels.get(key) == value for key, value in expected_labels.items())
+
+    del bigframes.options.compute.extra_query_labels["test1"]
+    del bigframes.options.compute.extra_query_labels["test2"]
+    del bigframes.options.compute.extra_query_labels["test3"]
+
+    assert len(bigframes.options.compute.extra_query_labels) == 0
+
+
 def test_read_gbq_model(session, penguins_linear_model_name):
     model = session.read_gbq_model(penguins_linear_model_name)
     assert isinstance(model, bigframes.ml.linear_model.LinearRegression)