Skip to content

Commit

Permalink
BigQuery: Add support for array parameters to Cursor.execute() (#9189)
Browse files Browse the repository at this point in the history
* Add support for array params to Cursor.execute()

* Raise NotImplementedError for STRUCT-like values
  • Loading branch information
plamut authored Sep 10, 2019
1 parent 9e1b896 commit adfc6d3
Show file tree
Hide file tree
Showing 2 changed files with 214 additions and 25 deletions.
146 changes: 121 additions & 25 deletions bigquery/google/cloud/bigquery/dbapi/_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,27 +43,9 @@ def scalar_to_query_parameter(value, name=None):
:raises: :class:`~google.cloud.bigquery.dbapi.exceptions.ProgrammingError`
if the type cannot be determined.
"""
parameter_type = None
parameter_type = bigquery_scalar_type(value)

if isinstance(value, bool):
parameter_type = "BOOL"
elif isinstance(value, numbers.Integral):
parameter_type = "INT64"
elif isinstance(value, numbers.Real):
parameter_type = "FLOAT64"
elif isinstance(value, decimal.Decimal):
parameter_type = "NUMERIC"
elif isinstance(value, six.text_type):
parameter_type = "STRING"
elif isinstance(value, six.binary_type):
parameter_type = "BYTES"
elif isinstance(value, datetime.datetime):
parameter_type = "DATETIME" if value.tzinfo is None else "TIMESTAMP"
elif isinstance(value, datetime.date):
parameter_type = "DATE"
elif isinstance(value, datetime.time):
parameter_type = "TIME"
else:
if parameter_type is None:
raise exceptions.ProgrammingError(
"encountered parameter {} with value {} of unexpected type".format(
name, value
Expand All @@ -72,6 +54,46 @@ def scalar_to_query_parameter(value, name=None):
return bigquery.ScalarQueryParameter(name, parameter_type, value)


def array_to_query_parameter(value, name=None):
"""Convert an array-like value into a query parameter.
Args:
value (Sequence[Any]): The elements of the array (should not be a
string-like Sequence).
name (Optional[str]): Name of the query parameter.
Returns:
A query parameter corresponding with the type and value of the plain
Python object.
Raises:
:class:`~google.cloud.bigquery.dbapi.exceptions.ProgrammingError`
if the type of array elements cannot be determined.
"""
if not array_like(value):
raise exceptions.ProgrammingError(
"The value of parameter {} must be a sequence that is "
"not string-like.".format(name)
)

if not value:
raise exceptions.ProgrammingError(
"Encountered an empty array-like value of parameter {}, cannot "
"determine array elements type.".format(name)
)

# Assume that all elements are of the same type, and let the backend handle
# any type incompatibilities among the array elements
array_type = bigquery_scalar_type(value[0])
if array_type is None:
raise exceptions.ProgrammingError(
"Encountered unexpected first array element of parameter {}, "
"cannot determine array elements type.".format(name)
)

return bigquery.ArrayQueryParameter(name, array_type, value)


def to_query_parameters_list(parameters):
"""Converts a sequence of parameter values into query parameters.
Expand All @@ -81,7 +103,18 @@ def to_query_parameters_list(parameters):
:rtype: List[google.cloud.bigquery.query._AbstractQueryParameter]
:returns: A list of query parameters.
"""
return [scalar_to_query_parameter(value) for value in parameters]
result = []

for value in parameters:
if isinstance(value, collections_abc.Mapping):
raise NotImplementedError("STRUCT-like parameter values are not supported.")
elif array_like(value):
param = array_to_query_parameter(value)
else:
param = scalar_to_query_parameter(value)
result.append(param)

return result


def to_query_parameters_dict(parameters):
Expand All @@ -93,10 +126,21 @@ def to_query_parameters_dict(parameters):
:rtype: List[google.cloud.bigquery.query._AbstractQueryParameter]
:returns: A list of named query parameters.
"""
return [
scalar_to_query_parameter(value, name=name)
for name, value in six.iteritems(parameters)
]
result = []

for name, value in six.iteritems(parameters):
if isinstance(value, collections_abc.Mapping):
raise NotImplementedError(
"STRUCT-like parameter values are not supported "
"(parameter {}).".format(name)
)
elif array_like(value):
param = array_to_query_parameter(value, name=name)
else:
param = scalar_to_query_parameter(value, name=name)
result.append(param)

return result


def to_query_parameters(parameters):
Expand All @@ -115,3 +159,55 @@ def to_query_parameters(parameters):
return to_query_parameters_dict(parameters)

return to_query_parameters_list(parameters)


def bigquery_scalar_type(value):
"""Return a BigQuery name of the scalar type that matches the given value.
If the scalar type name could not be determined (e.g. for non-scalar
values), ``None`` is returned.
Args:
value (Any)
Returns:
Optional[str]: The BigQuery scalar type name.
"""
if isinstance(value, bool):
return "BOOL"
elif isinstance(value, numbers.Integral):
return "INT64"
elif isinstance(value, numbers.Real):
return "FLOAT64"
elif isinstance(value, decimal.Decimal):
return "NUMERIC"
elif isinstance(value, six.text_type):
return "STRING"
elif isinstance(value, six.binary_type):
return "BYTES"
elif isinstance(value, datetime.datetime):
return "DATETIME" if value.tzinfo is None else "TIMESTAMP"
elif isinstance(value, datetime.date):
return "DATE"
elif isinstance(value, datetime.time):
return "TIME"

return None


def array_like(value):
"""Determine if the given value is array-like.
Examples of array-like values (as interpreted by this function) are
sequences such as ``list`` and ``tuple``, but not strings and other
iterables such as sets.
Args:
value (Any)
Returns:
bool: ``True`` if the value is considered array-like, ``False`` otherwise.
"""
return isinstance(value, collections_abc.Sequence) and not isinstance(
value, (six.text_type, six.binary_type, bytearray)
)
93 changes: 93 additions & 0 deletions bigquery/tests/unit/test_dbapi__helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,61 @@ def test_scalar_to_query_parameter_w_special_floats(self):
self.assertTrue(math.isinf(inf_parameter.value))
self.assertEqual(inf_parameter.type_, "FLOAT64")

def test_array_to_query_parameter_valid_argument(self):
expected_types = [
([True, False], "BOOL"),
([123, -456, 0], "INT64"),
([1.25, 2.50], "FLOAT64"),
([decimal.Decimal("1.25")], "NUMERIC"),
([b"foo", b"bar"], "BYTES"),
([u"foo", u"bar"], "STRING"),
([datetime.date(2017, 4, 1), datetime.date(2018, 4, 1)], "DATE"),
([datetime.time(12, 34, 56), datetime.time(10, 20, 30)], "TIME"),
(
[
datetime.datetime(2012, 3, 4, 5, 6, 7),
datetime.datetime(2013, 1, 1, 10, 20, 30),
],
"DATETIME",
),
(
[
datetime.datetime(
2012, 3, 4, 5, 6, 7, tzinfo=google.cloud._helpers.UTC
),
datetime.datetime(
2013, 1, 1, 10, 20, 30, tzinfo=google.cloud._helpers.UTC
),
],
"TIMESTAMP",
),
]

for values, expected_type in expected_types:
msg = "value: {} expected_type: {}".format(values, expected_type)
parameter = _helpers.array_to_query_parameter(values)
self.assertIsNone(parameter.name, msg=msg)
self.assertEqual(parameter.array_type, expected_type, msg=msg)
self.assertEqual(parameter.values, values, msg=msg)
named_param = _helpers.array_to_query_parameter(values, name="my_param")
self.assertEqual(named_param.name, "my_param", msg=msg)
self.assertEqual(named_param.array_type, expected_type, msg=msg)
self.assertEqual(named_param.values, values, msg=msg)

def test_array_to_query_parameter_empty_argument(self):
with self.assertRaises(exceptions.ProgrammingError):
_helpers.array_to_query_parameter([])

def test_array_to_query_parameter_unsupported_sequence(self):
unsupported_iterables = [{10, 20, 30}, u"foo", b"bar", bytearray([65, 75, 85])]
for iterable in unsupported_iterables:
with self.assertRaises(exceptions.ProgrammingError):
_helpers.array_to_query_parameter(iterable)

def test_array_to_query_parameter_sequence_w_invalid_elements(self):
with self.assertRaises(exceptions.ProgrammingError):
_helpers.array_to_query_parameter([object(), 2, 7])

def test_to_query_parameters_w_dict(self):
parameters = {"somebool": True, "somestring": u"a-string-value"}
query_parameters = _helpers.to_query_parameters(parameters)
Expand All @@ -82,6 +137,23 @@ def test_to_query_parameters_w_dict(self):
),
)

def test_to_query_parameters_w_dict_array_param(self):
parameters = {"somelist": [10, 20]}
query_parameters = _helpers.to_query_parameters(parameters)

self.assertEqual(len(query_parameters), 1)
param = query_parameters[0]

self.assertEqual(param.name, "somelist")
self.assertEqual(param.array_type, "INT64")
self.assertEqual(param.values, [10, 20])

def test_to_query_parameters_w_dict_dict_param(self):
parameters = {"my_param": {"foo": "bar"}}

with self.assertRaises(NotImplementedError):
_helpers.to_query_parameters(parameters)

def test_to_query_parameters_w_list(self):
parameters = [True, u"a-string-value"]
query_parameters = _helpers.to_query_parameters(parameters)
Expand All @@ -92,3 +164,24 @@ def test_to_query_parameters_w_list(self):
sorted(query_parameter_tuples),
sorted([(None, "BOOL", True), (None, "STRING", u"a-string-value")]),
)

def test_to_query_parameters_w_list_array_param(self):
parameters = [[10, 20]]
query_parameters = _helpers.to_query_parameters(parameters)

self.assertEqual(len(query_parameters), 1)
param = query_parameters[0]

self.assertIsNone(param.name)
self.assertEqual(param.array_type, "INT64")
self.assertEqual(param.values, [10, 20])

def test_to_query_parameters_w_list_dict_param(self):
parameters = [{"foo": "bar"}]

with self.assertRaises(NotImplementedError):
_helpers.to_query_parameters(parameters)

def test_to_query_parameters_none_argument(self):
query_parameters = _helpers.to_query_parameters(None)
self.assertEqual(query_parameters, [])

0 comments on commit adfc6d3

Please sign in to comment.