Skip to content

Commit 05b4114

Browse files
authored
BigQuery: fix parsing for array parameter with struct type. (#4040)
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.
1 parent f9fc725 commit 05b4114

File tree

3 files changed

+438
-11
lines changed

3 files changed

+438
-11
lines changed

bigquery/google/cloud/bigquery/_helpers.py

Lines changed: 129 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
import base64
1818
from collections import OrderedDict
19+
import copy
1920
import datetime
2021

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

472+
def _key(self):
473+
"""A tuple key that uniquely describes this field.
474+
475+
Used to compute this instance's hashcode and evaluate equality.
476+
477+
Returns:
478+
tuple: The contents of this :class:`ScalarQueryParameter`.
479+
"""
480+
return (
481+
self.name,
482+
self.type_.upper(),
483+
self.value,
484+
)
485+
486+
def __eq__(self, other):
487+
if not isinstance(other, ScalarQueryParameter):
488+
return NotImplemented
489+
return self._key() == other._key()
490+
491+
def __ne__(self, other):
492+
return not self == other
493+
494+
def __repr__(self):
495+
return 'ScalarQueryParameter{}'.format(self._key())
496+
471497

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

509535
@classmethod
510-
def from_api_repr(cls, resource):
511-
"""Factory: construct parameter from JSON resource.
512-
513-
:type resource: dict
514-
:param resource: JSON mapping of parameter
536+
def _from_api_repr_struct(cls, resource):
537+
name = resource.get('name')
538+
converted = []
539+
# We need to flatten the array to use the StructQueryParameter
540+
# parse code.
541+
resource_template = {
542+
# The arrayType includes all the types of the fields of the STRUCT
543+
'parameterType': resource['parameterType']['arrayType']
544+
}
545+
for array_value in resource['parameterValue']['arrayValues']:
546+
struct_resource = copy.deepcopy(resource_template)
547+
struct_resource['parameterValue'] = array_value
548+
struct_value = StructQueryParameter.from_api_repr(struct_resource)
549+
converted.append(struct_value)
550+
return cls(name, 'STRUCT', converted)
515551

516-
:rtype: :class:`ArrayQueryParameter`
517-
:returns: instance
518-
"""
552+
@classmethod
553+
def _from_api_repr_scalar(cls, resource):
519554
name = resource.get('name')
520555
array_type = resource['parameterType']['arrayType']['type']
521556
values = [
@@ -526,14 +561,29 @@ def from_api_repr(cls, resource):
526561
_CELLDATA_FROM_JSON[array_type](value, None) for value in values]
527562
return cls(name, array_type, converted)
528563

564+
@classmethod
565+
def from_api_repr(cls, resource):
566+
"""Factory: construct parameter from JSON resource.
567+
568+
:type resource: dict
569+
:param resource: JSON mapping of parameter
570+
571+
:rtype: :class:`ArrayQueryParameter`
572+
:returns: instance
573+
"""
574+
array_type = resource['parameterType']['arrayType']['type']
575+
if array_type == 'STRUCT':
576+
return cls._from_api_repr_struct(resource)
577+
return cls._from_api_repr_scalar(resource)
578+
529579
def to_api_repr(self):
530580
"""Construct JSON API representation for the parameter.
531581
532582
:rtype: dict
533583
:returns: JSON mapping
534584
"""
535585
values = self.values
536-
if self.array_type == 'RECORD':
586+
if self.array_type == 'RECORD' or self.array_type == 'STRUCT':
537587
reprs = [value.to_api_repr() for value in values]
538588
a_type = reprs[0]['parameterType']
539589
a_values = [repr_['parameterValue'] for repr_ in reprs]
@@ -556,6 +606,31 @@ def to_api_repr(self):
556606
resource['name'] = self.name
557607
return resource
558608

609+
def _key(self):
610+
"""A tuple key that uniquely describes this field.
611+
612+
Used to compute this instance's hashcode and evaluate equality.
613+
614+
Returns:
615+
tuple: The contents of this :class:`ArrayQueryParameter`.
616+
"""
617+
return (
618+
self.name,
619+
self.array_type.upper(),
620+
self.values,
621+
)
622+
623+
def __eq__(self, other):
624+
if not isinstance(other, ArrayQueryParameter):
625+
return NotImplemented
626+
return self._key() == other._key()
627+
628+
def __ne__(self, other):
629+
return not self == other
630+
631+
def __repr__(self):
632+
return 'ArrayQueryParameter{}'.format(self._key())
633+
559634

560635
class StructQueryParameter(AbstractQueryParameter):
561636
"""Named / positional query parameters for struct values.
@@ -606,14 +681,32 @@ def from_api_repr(cls, resource):
606681
"""
607682
name = resource.get('name')
608683
instance = cls(name)
684+
type_resources = {}
609685
types = instance.struct_types
610686
for item in resource['parameterType']['structTypes']:
611687
types[item['name']] = item['type']['type']
688+
type_resources[item['name']] = item['type']
612689
struct_values = resource['parameterValue']['structValues']
613690
for key, value in struct_values.items():
614691
type_ = types[key]
615-
value = value['value']
616-
converted = _CELLDATA_FROM_JSON[type_](value, None)
692+
converted = None
693+
if type_ == 'STRUCT':
694+
struct_resource = {
695+
'name': key,
696+
'parameterType': type_resources[key],
697+
'parameterValue': value,
698+
}
699+
converted = StructQueryParameter.from_api_repr(struct_resource)
700+
elif type_ == 'ARRAY':
701+
struct_resource = {
702+
'name': key,
703+
'parameterType': type_resources[key],
704+
'parameterValue': value,
705+
}
706+
converted = ArrayQueryParameter.from_api_repr(struct_resource)
707+
else:
708+
value = value['value']
709+
converted = _CELLDATA_FROM_JSON[type_](value, None)
617710
instance.struct_values[key] = converted
618711
return instance
619712

@@ -651,6 +744,31 @@ def to_api_repr(self):
651744
resource['name'] = self.name
652745
return resource
653746

747+
def _key(self):
748+
"""A tuple key that uniquely describes this field.
749+
750+
Used to compute this instance's hashcode and evaluate equality.
751+
752+
Returns:
753+
tuple: The contents of this :class:`ArrayQueryParameter`.
754+
"""
755+
return (
756+
self.name,
757+
self.struct_types,
758+
self.struct_values,
759+
)
760+
761+
def __eq__(self, other):
762+
if not isinstance(other, StructQueryParameter):
763+
return NotImplemented
764+
return self._key() == other._key()
765+
766+
def __ne__(self, other):
767+
return not self == other
768+
769+
def __repr__(self):
770+
return 'StructQueryParameter{}'.format(self._key())
771+
654772

655773
class QueryParametersProperty(object):
656774
"""Custom property type, holding query parameter instances."""

bigquery/tests/system.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -887,6 +887,16 @@ def test_sync_query_w_query_params(self):
887887
name='friends', array_type='STRING',
888888
values=[phred_name, bharney_name])
889889
with_friends_param = StructQueryParameter(None, friends_param)
890+
top_left_param = StructQueryParameter(
891+
'top_left',
892+
ScalarQueryParameter('x', 'INT64', 12),
893+
ScalarQueryParameter('y', 'INT64', 102))
894+
bottom_right_param = StructQueryParameter(
895+
'bottom_right',
896+
ScalarQueryParameter('x', 'INT64', 22),
897+
ScalarQueryParameter('y', 'INT64', 92))
898+
rectangle_param = StructQueryParameter(
899+
'rectangle', top_left_param, bottom_right_param)
890900
examples = [
891901
{
892902
'sql': 'SELECT @question',
@@ -943,6 +953,14 @@ def test_sync_query_w_query_params(self):
943953
'expected': ({'_field_1': question, '_field_2': answer}),
944954
'query_parameters': [struct_param],
945955
},
956+
{
957+
'sql':
958+
'SELECT '
959+
'((@rectangle.bottom_right.x - @rectangle.top_left.x) '
960+
'* (@rectangle.top_left.y - @rectangle.bottom_right.y))',
961+
'expected': 100,
962+
'query_parameters': [rectangle_param],
963+
},
946964
{
947965
'sql': 'SELECT ?',
948966
'expected': [

0 commit comments

Comments
 (0)