Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
feat: support script options in query job config (#690)
  • Loading branch information
plamut committed Jun 7, 2021
1 parent d034a4d commit 1259e16
Show file tree
Hide file tree
Showing 5 changed files with 170 additions and 1 deletion.
4 changes: 4 additions & 0 deletions google/cloud/bigquery/__init__.py
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -138,6 +140,7 @@
"CSVOptions",
"GoogleSheetsOptions",
"ParquetOptions",
"ScriptOptions",
"DEFAULT_RETRY",
# Enum Constants
"enums",
Expand All @@ -147,6 +150,7 @@
"DeterminismLevel",
"ExternalSourceFormat",
"Encoding",
"KeyResultStatementKind",
"QueryPriority",
"SchemaUpdateOption",
"SourceFormat",
Expand Down
13 changes: 13 additions & 0 deletions google/cloud/bigquery/enums.py
Expand Up @@ -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",
Expand Down
2 changes: 2 additions & 0 deletions google/cloud/bigquery/job/__init__.py
Expand Up @@ -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
Expand Down Expand Up @@ -67,6 +68,7 @@
"QueryJobConfig",
"QueryPlanEntry",
"QueryPlanEntryStep",
"ScriptOptions",
"TimelineEntry",
"Compression",
"CreateDisposition",
Expand Down
96 changes: 95 additions & 1 deletion google/cloud/bigquery/job/query.py
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down
56 changes: 56 additions & 0 deletions tests/unit/job/test_query_config.py
Expand Up @@ -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
)

0 comments on commit 1259e16

Please sign in to comment.