Skip to content

Commit

Permalink
BigQuery: fix parsing for array parameter with struct type. (#4040)
Browse files Browse the repository at this point in the history
Adds special cases for loading an array query parameter resource that
contains structs.

Similarly, adds special cases for loading a struct query parameter when
it contains nested structs or arrays.
  • Loading branch information
tswast authored Sep 22, 2017
1 parent f9fc725 commit 05b4114
Show file tree
Hide file tree
Showing 3 changed files with 438 additions and 11 deletions.
140 changes: 129 additions & 11 deletions bigquery/google/cloud/bigquery/_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

import base64
from collections import OrderedDict
import copy
import datetime

from google.cloud._helpers import UTC
Expand Down Expand Up @@ -468,6 +469,31 @@ def to_api_repr(self):
resource['name'] = self.name
return resource

def _key(self):
"""A tuple key that uniquely describes this field.
Used to compute this instance's hashcode and evaluate equality.
Returns:
tuple: The contents of this :class:`ScalarQueryParameter`.
"""
return (
self.name,
self.type_.upper(),
self.value,
)

def __eq__(self, other):
if not isinstance(other, ScalarQueryParameter):
return NotImplemented
return self._key() == other._key()

def __ne__(self, other):
return not self == other

def __repr__(self):
return 'ScalarQueryParameter{}'.format(self._key())


class ArrayQueryParameter(AbstractQueryParameter):
"""Named / positional query parameters for array values.
Expand Down Expand Up @@ -507,15 +533,24 @@ def positional(cls, array_type, values):
return cls(None, array_type, values)

@classmethod
def from_api_repr(cls, resource):
"""Factory: construct parameter from JSON resource.
:type resource: dict
:param resource: JSON mapping of parameter
def _from_api_repr_struct(cls, resource):
name = resource.get('name')
converted = []
# We need to flatten the array to use the StructQueryParameter
# parse code.
resource_template = {
# The arrayType includes all the types of the fields of the STRUCT
'parameterType': resource['parameterType']['arrayType']
}
for array_value in resource['parameterValue']['arrayValues']:
struct_resource = copy.deepcopy(resource_template)
struct_resource['parameterValue'] = array_value
struct_value = StructQueryParameter.from_api_repr(struct_resource)
converted.append(struct_value)
return cls(name, 'STRUCT', converted)

:rtype: :class:`ArrayQueryParameter`
:returns: instance
"""
@classmethod
def _from_api_repr_scalar(cls, resource):
name = resource.get('name')
array_type = resource['parameterType']['arrayType']['type']
values = [
Expand All @@ -526,14 +561,29 @@ def from_api_repr(cls, resource):
_CELLDATA_FROM_JSON[array_type](value, None) for value in values]
return cls(name, array_type, converted)

@classmethod
def from_api_repr(cls, resource):
"""Factory: construct parameter from JSON resource.
:type resource: dict
:param resource: JSON mapping of parameter
:rtype: :class:`ArrayQueryParameter`
:returns: instance
"""
array_type = resource['parameterType']['arrayType']['type']
if array_type == 'STRUCT':
return cls._from_api_repr_struct(resource)
return cls._from_api_repr_scalar(resource)

def to_api_repr(self):
"""Construct JSON API representation for the parameter.
:rtype: dict
:returns: JSON mapping
"""
values = self.values
if self.array_type == 'RECORD':
if self.array_type == 'RECORD' or self.array_type == 'STRUCT':
reprs = [value.to_api_repr() for value in values]
a_type = reprs[0]['parameterType']
a_values = [repr_['parameterValue'] for repr_ in reprs]
Expand All @@ -556,6 +606,31 @@ def to_api_repr(self):
resource['name'] = self.name
return resource

def _key(self):
"""A tuple key that uniquely describes this field.
Used to compute this instance's hashcode and evaluate equality.
Returns:
tuple: The contents of this :class:`ArrayQueryParameter`.
"""
return (
self.name,
self.array_type.upper(),
self.values,
)

def __eq__(self, other):
if not isinstance(other, ArrayQueryParameter):
return NotImplemented
return self._key() == other._key()

def __ne__(self, other):
return not self == other

def __repr__(self):
return 'ArrayQueryParameter{}'.format(self._key())


class StructQueryParameter(AbstractQueryParameter):
"""Named / positional query parameters for struct values.
Expand Down Expand Up @@ -606,14 +681,32 @@ def from_api_repr(cls, resource):
"""
name = resource.get('name')
instance = cls(name)
type_resources = {}
types = instance.struct_types
for item in resource['parameterType']['structTypes']:
types[item['name']] = item['type']['type']
type_resources[item['name']] = item['type']
struct_values = resource['parameterValue']['structValues']
for key, value in struct_values.items():
type_ = types[key]
value = value['value']
converted = _CELLDATA_FROM_JSON[type_](value, None)
converted = None
if type_ == 'STRUCT':
struct_resource = {
'name': key,
'parameterType': type_resources[key],
'parameterValue': value,
}
converted = StructQueryParameter.from_api_repr(struct_resource)
elif type_ == 'ARRAY':
struct_resource = {
'name': key,
'parameterType': type_resources[key],
'parameterValue': value,
}
converted = ArrayQueryParameter.from_api_repr(struct_resource)
else:
value = value['value']
converted = _CELLDATA_FROM_JSON[type_](value, None)
instance.struct_values[key] = converted
return instance

Expand Down Expand Up @@ -651,6 +744,31 @@ def to_api_repr(self):
resource['name'] = self.name
return resource

def _key(self):
"""A tuple key that uniquely describes this field.
Used to compute this instance's hashcode and evaluate equality.
Returns:
tuple: The contents of this :class:`ArrayQueryParameter`.
"""
return (
self.name,
self.struct_types,
self.struct_values,
)

def __eq__(self, other):
if not isinstance(other, StructQueryParameter):
return NotImplemented
return self._key() == other._key()

def __ne__(self, other):
return not self == other

def __repr__(self):
return 'StructQueryParameter{}'.format(self._key())


class QueryParametersProperty(object):
"""Custom property type, holding query parameter instances."""
Expand Down
18 changes: 18 additions & 0 deletions bigquery/tests/system.py
Original file line number Diff line number Diff line change
Expand Up @@ -887,6 +887,16 @@ def test_sync_query_w_query_params(self):
name='friends', array_type='STRING',
values=[phred_name, bharney_name])
with_friends_param = StructQueryParameter(None, friends_param)
top_left_param = StructQueryParameter(
'top_left',
ScalarQueryParameter('x', 'INT64', 12),
ScalarQueryParameter('y', 'INT64', 102))
bottom_right_param = StructQueryParameter(
'bottom_right',
ScalarQueryParameter('x', 'INT64', 22),
ScalarQueryParameter('y', 'INT64', 92))
rectangle_param = StructQueryParameter(
'rectangle', top_left_param, bottom_right_param)
examples = [
{
'sql': 'SELECT @question',
Expand Down Expand Up @@ -943,6 +953,14 @@ def test_sync_query_w_query_params(self):
'expected': ({'_field_1': question, '_field_2': answer}),
'query_parameters': [struct_param],
},
{
'sql':
'SELECT '
'((@rectangle.bottom_right.x - @rectangle.top_left.x) '
'* (@rectangle.top_left.y - @rectangle.bottom_right.y))',
'expected': 100,
'query_parameters': [rectangle_param],
},
{
'sql': 'SELECT ?',
'expected': [
Expand Down
Loading

0 comments on commit 05b4114

Please sign in to comment.