diff --git a/google/cloud/bigquery/__init__.py b/google/cloud/bigquery/__init__.py index f031cd81d..94f87304a 100644 --- a/google/cloud/bigquery/__init__.py +++ b/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", diff --git a/google/cloud/bigquery/enums.py b/google/cloud/bigquery/enums.py index 787c2449d..edf991b6f 100644 --- a/google/cloud/bigquery/enums.py +++ b/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", diff --git a/google/cloud/bigquery/job/__init__.py b/google/cloud/bigquery/job/__init__.py index 4945841d9..cdab92e05 100644 --- a/google/cloud/bigquery/job/__init__.py +++ b/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", diff --git a/google/cloud/bigquery/job/query.py b/google/cloud/bigquery/job/query.py index f52f9c621..455ef4632 100644 --- a/google/cloud/bigquery/job/query.py +++ b/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. diff --git a/tests/unit/job/test_query_config.py b/tests/unit/job/test_query_config.py index db03d6a3b..109cf7e44 100644 --- a/tests/unit/job/test_query_config.py +++ b/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 + )