feat: support script options in query job config (#690)

google/cloud/bigquery/__init__.py

@@ -37,6 +37,7 @@
 from google.cloud.bigquery.dataset import Dataset
 from google.cloud.bigquery.dataset import DatasetReference
 from google.cloud.bigquery import enums
+from google.cloud.bigquery.enums import KeyResultStatementKind
 from google.cloud.bigquery.enums import SqlTypeNames
 from google.cloud.bigquery.enums import StandardSqlDataTypes
 from google.cloud.bigquery.exceptions import LegacyBigQueryStorageError
@@ -62,6 +63,7 @@
 from google.cloud.bigquery.job import QueryJobConfig
 from google.cloud.bigquery.job import QueryPriority
 from google.cloud.bigquery.job import SchemaUpdateOption
+from google.cloud.bigquery.job import ScriptOptions
 from google.cloud.bigquery.job import SourceFormat
 from google.cloud.bigquery.job import UnknownJob
 from google.cloud.bigquery.job import WriteDisposition
@@ -138,6 +140,7 @@
     "CSVOptions",
     "GoogleSheetsOptions",
     "ParquetOptions",
+    "ScriptOptions",
     "DEFAULT_RETRY",
     # Enum Constants
     "enums",
@@ -147,6 +150,7 @@
     "DeterminismLevel",
     "ExternalSourceFormat",
     "Encoding",
+    "KeyResultStatementKind",
     "QueryPriority",
     "SchemaUpdateOption",
     "SourceFormat",

google/cloud/bigquery/enums.py

@@ -142,6 +142,19 @@ class SourceFormat(object):
     """Specifies Orc format."""
 
 
+class KeyResultStatementKind:
+    """Determines which statement in the script represents the "key result".
+
+    The "key result" is used to populate the schema and query results of the script job.
+
+    https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#keyresultstatementkind
+    """
+
+    KEY_RESULT_STATEMENT_KIND_UNSPECIFIED = "KEY_RESULT_STATEMENT_KIND_UNSPECIFIED"
+    LAST = "LAST"
+    FIRST_SELECT = "FIRST_SELECT"
+
+
 _SQL_SCALAR_TYPES = frozenset(
     (
         "INT64",

google/cloud/bigquery/job/__init__.py

@@ -34,6 +34,7 @@
 from google.cloud.bigquery.job.query import QueryJobConfig
 from google.cloud.bigquery.job.query import QueryPlanEntry
 from google.cloud.bigquery.job.query import QueryPlanEntryStep
+from google.cloud.bigquery.job.query import ScriptOptions
 from google.cloud.bigquery.job.query import TimelineEntry
 from google.cloud.bigquery.enums import Compression
 from google.cloud.bigquery.enums import CreateDisposition
@@ -67,6 +68,7 @@
     "QueryJobConfig",
     "QueryPlanEntry",
     "QueryPlanEntryStep",
+    "ScriptOptions",
     "TimelineEntry",
     "Compression",
     "CreateDisposition",

google/cloud/bigquery/job/query.py

@@ -18,7 +18,7 @@
 import copy
 import re
 import typing
-from typing import Any, Dict, Union
+from typing import Any, Dict, Optional, Union
 
 from google.api_core import exceptions
 from google.api_core.future import polling as polling_future
@@ -28,6 +28,7 @@
 from google.cloud.bigquery.dataset import DatasetListItem
 from google.cloud.bigquery.dataset import DatasetReference
 from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration
+from google.cloud.bigquery.enums import KeyResultStatementKind
 from google.cloud.bigquery.external_config import ExternalConfig
 from google.cloud.bigquery import _helpers
 from google.cloud.bigquery.query import _query_param_from_api_repr
@@ -113,6 +114,82 @@ def _to_api_repr_table_defs(value):
     return {k: ExternalConfig.to_api_repr(v) for k, v in value.items()}
 
 
+class ScriptOptions:
+    """Options controlling the execution of scripts.
+
+    https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#ScriptOptions
+    """
+
+    def __init__(
+        self,
+        statement_timeout_ms: Optional[int] = None,
+        statement_byte_budget: Optional[int] = None,
+        key_result_statement: Optional[KeyResultStatementKind] = None,
+    ):
+        self._properties = {}
+        self.statement_timeout_ms = statement_timeout_ms
+        self.statement_byte_budget = statement_byte_budget
+        self.key_result_statement = key_result_statement
+
+    @classmethod
+    def from_api_repr(cls, resource: Dict[str, Any]) -> "ScriptOptions":
+        """Factory: construct instance from the JSON repr.
+
+        Args:
+            resource(Dict[str: Any]):
+                ScriptOptions representation returned from API.
+
+        Returns:
+            google.cloud.bigquery.ScriptOptions:
+                ScriptOptions sample parsed from ``resource``.
+        """
+        entry = cls()
+        entry._properties = copy.deepcopy(resource)
+        return entry
+
+    def to_api_repr(self) -> Dict[str, Any]:
+        """Construct the API resource representation."""
+        return copy.deepcopy(self._properties)
+
+    @property
+    def statement_timeout_ms(self) -> Union[int, None]:
+        """Timeout period for each statement in a script."""
+        return _helpers._int_or_none(self._properties.get("statementTimeoutMs"))
+
+    @statement_timeout_ms.setter
+    def statement_timeout_ms(self, value: Union[int, None]):
+        if value is not None:
+            value = str(value)
+        self._properties["statementTimeoutMs"] = value
+
+    @property
+    def statement_byte_budget(self) -> Union[int, None]:
+        """Limit on the number of bytes billed per statement.
+
+        Exceeding this budget results in an error.
+        """
+        return _helpers._int_or_none(self._properties.get("statementByteBudget"))
+
+    @statement_byte_budget.setter
+    def statement_byte_budget(self, value: Union[int, None]):
+        if value is not None:
+            value = str(value)
+        self._properties["statementByteBudget"] = value
+
+    @property
+    def key_result_statement(self) -> Union[KeyResultStatementKind, None]:
+        """Determines which statement in the script represents the "key result".
+
+        This is used to populate the schema and query results of the script job.
+        Default is ``KeyResultStatementKind.LAST``.
+        """
+        return self._properties.get("keyResultStatement")
+
+    @key_result_statement.setter
+    def key_result_statement(self, value: Union[KeyResultStatementKind, None]):
+        self._properties["keyResultStatement"] = value
+
+
 class QueryJobConfig(_JobConfig):
     """Configuration options for query jobs.
 
@@ -502,6 +579,23 @@ def schema_update_options(self):
     def schema_update_options(self, values):
         self._set_sub_prop("schemaUpdateOptions", values)
 
+    @property
+    def script_options(self) -> ScriptOptions:
+        """Connection properties which can modify the query behavior.
+
+        https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#scriptoptions
+        """
+        prop = self._get_sub_prop("scriptOptions")
+        if prop is not None:
+            prop = ScriptOptions.from_api_repr(prop)
+        return prop
+
+    @script_options.setter
+    def script_options(self, value: Union[ScriptOptions, None]):
+        if value is not None:
+            value = value.to_api_repr()
+        self._set_sub_prop("scriptOptions", value)
+
     def to_api_repr(self) -> dict:
         """Build an API representation of the query job config.
 

tests/unit/job/test_query_config.py

@@ -253,3 +253,59 @@ def test_from_api_repr_with_encryption(self):
         self.assertEqual(
             config.destination_encryption_configuration.kms_key_name, self.KMS_KEY_NAME
         )
+
+    def test_to_api_repr_with_script_options_none(self):
+        config = self._make_one()
+        config.script_options = None
+
+        resource = config.to_api_repr()
+
+        self.assertEqual(resource, {"query": {"scriptOptions": None}})
+        self.assertIsNone(config.script_options)
+
+    def test_to_api_repr_with_script_options(self):
+        from google.cloud.bigquery import KeyResultStatementKind
+        from google.cloud.bigquery import ScriptOptions
+
+        config = self._make_one()
+        config.script_options = ScriptOptions(
+            statement_timeout_ms=60,
+            statement_byte_budget=999,
+            key_result_statement=KeyResultStatementKind.FIRST_SELECT,
+        )
+
+        resource = config.to_api_repr()
+
+        expected_script_options_repr = {
+            "statementTimeoutMs": "60",
+            "statementByteBudget": "999",
+            "keyResultStatement": KeyResultStatementKind.FIRST_SELECT,
+        }
+        self.assertEqual(
+            resource, {"query": {"scriptOptions": expected_script_options_repr}}
+        )
+
+    def test_from_api_repr_with_script_options(self):
+        from google.cloud.bigquery import KeyResultStatementKind
+        from google.cloud.bigquery import ScriptOptions
+
+        resource = {
+            "query": {
+                "scriptOptions": {
+                    "statementTimeoutMs": "42",
+                    "statementByteBudget": "123",
+                    "keyResultStatement": KeyResultStatementKind.LAST,
+                },
+            },
+        }
+        klass = self._get_target_class()
+
+        config = klass.from_api_repr(resource)
+
+        script_options = config.script_options
+        self.assertIsInstance(script_options, ScriptOptions)
+        self.assertEqual(script_options.statement_timeout_ms, 42)
+        self.assertEqual(script_options.statement_byte_budget, 123)
+        self.assertEqual(
+            script_options.key_result_statement, KeyResultStatementKind.LAST
+        )