BigQuery: Add support for array parameters to Cursor.execute() (googleapis#9189)

plamut · emar-kar · commit ca57b99d8236 · 2019-09-18T13:48:36.000+03:00
* Add support for array params to Cursor.execute()

* Raise NotImplementedError for STRUCT-like values
diff --git a/bigquery/google/cloud/bigquery/dbapi/_helpers.py b/bigquery/google/cloud/bigquery/dbapi/_helpers.py
@@ -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
@@ -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.
 
@@ -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):
@@ -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):
@@ -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)
+    )
diff --git a/bigquery/tests/unit/test_dbapi__helpers.py b/bigquery/tests/unit/test_dbapi__helpers.py
@@ -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)
@@ -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)
@@ -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, [])
