Skip to content
Permalink
Browse files
feat: support script options in query job config (#690)
  • Loading branch information
plamut committed Jun 7, 2021
1 parent d034a4d commit 1259e16394784315368e8be959c1ac097782b62e
Showing with 170 additions and 1 deletion.
  1. +4 −0 google/cloud/bigquery/__init__.py
  2. +13 −0 google/cloud/bigquery/enums.py
  3. +2 −0 google/cloud/bigquery/job/__init__.py
  4. +95 −1 google/cloud/bigquery/job/query.py
  5. +56 −0 tests/unit/job/test_query_config.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",
@@ -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",
@@ -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",
@@ -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.
@@ -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.