Query parameters: implement support for binding NoneType parameters (#233)

Signed-off-by: Jesse Whitehouse <jesse.whitehouse@databricks.com>

Author: Jesse

src/databricks/sql/utils.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import copy
 import datetime
 import decimal
 from abc import ABC, abstractmethod
@@ -495,6 +494,7 @@ class DbSqlType(Enum):
     BOOLEAN = "BOOLEAN"
     INTERVAL_MONTH = "INTERVAL MONTH"
     INTERVAL_DAY = "INTERVAL DAY"
+    VOID = "VOID"
 
 
 class DbSqlParameter:
@@ -542,20 +542,41 @@ def infer_types(params: list[DbSqlParameter]):
         datetime.date: DbSqlType.DATE,
         bool: DbSqlType.BOOLEAN,
         Decimal: DbSqlType.DECIMAL,
+        type(None): DbSqlType.VOID,
     }
-    new_params = copy.deepcopy(params)
-    for param in new_params:
-        if not param.type:
-            if type(param.value) in type_lookup_table:
-                param.type = type_lookup_table[type(param.value)]
-            else:
-                raise ValueError("Parameter type cannot be inferred")
-
-        if param.type == DbSqlType.DECIMAL:
+
+    new_params = []
+
+    # cycle through each parameter we've been passed
+    for param in params:
+        _name: str = param.name
+        _value: Any = param.value
+        _type: Union[DbSqlType, DbsqlDynamicDecimalType, Enum, None]
+
+        if param.type:
+            _type = param.type
+        else:
+            # figure out what type to use
+            _type = type_lookup_table.get(type(_value), None)
+            if not _type:
+                raise ValueError(
+                    f"Could not infer parameter type from {type(param.value)} - {param.value}"
+                )
+
+        # Decimal require special handling because one column type in Databricks can have multiple precisions
+        if _type == DbSqlType.DECIMAL:
             cast_exp = calculate_decimal_cast_string(param.value)
-            param.type = DbsqlDynamicDecimalType(cast_exp)
+            _type = DbsqlDynamicDecimalType(cast_exp)
+
+        # VOID / NULL types must be passed in a unique way as TSparkParameters with no value
+        if _type == DbSqlType.VOID:
+            new_params.append(DbSqlParameter(name=_name, type=DbSqlType.VOID))
+            continue
+        else:
+            _value = str(param.value)
+
+        new_params.append(DbSqlParameter(name=_name, value=_value, type=_type))
 
-        param.value = str(param.value)
     return new_params
 
 
@@ -594,11 +615,15 @@ def named_parameters_to_tsparkparams(parameters: Union[List[Any], Dict[str, str]
         dbsql_params = named_parameters_to_dbsqlparams_v2(parameters)
     inferred_type_parameters = infer_types(dbsql_params)
     for param in inferred_type_parameters:
-        tspark_params.append(
-            TSparkParameter(
-                type=param.type.value,
-                name=param.name,
-                value=TSparkParameterValue(stringValue=param.value),
+        # The only way to pass a VOID/NULL to DBR is to declare TSparkParameter without declaring
+        # its value or type arguments. If we set these to NoneType, the request will fail with a
+        # thrift transport error
+        if param.type == DbSqlType.VOID:
+            this_tspark_param = TSparkParameter(name=param.name)
+        else:
+            this_tspark_param_value = TSparkParameterValue(stringValue=param.value)
+            this_tspark_param = TSparkParameter(
+                type=param.type.value, name=param.name, value=this_tspark_param_value
             )
-        )
+        tspark_params.append(this_tspark_param)
     return tspark_params

tests/e2e/common/parameterized_query_tests.py

@@ -34,6 +34,12 @@ def _quantize(self, input: Union[float, int], place_value=2) -> Decimal:
 
         return Decimal(str(input)).quantize(Decimal("0." + "0" * place_value))
 
+    def test_primitive_inferred_none(self):
+
+        params = {"p": None}
+        result = self._get_one_result(self.QUERY, params)
+        assert result.col == None
+
     def test_primitive_inferred_bool(self):
 
         params = {"p": True}
@@ -79,6 +85,12 @@ def test_primitive_inferred_decimal(self):
         result = self._get_one_result(self.QUERY, params)
         assert result.col == Decimal("1234.56")
 
+    def test_dbsqlparam_inferred_none(self):
+
+        params = [DbSqlParameter(name="p", value=None, type=None)]
+        result = self._get_one_result(self.QUERY, params)
+        assert result.col == None
+
     def test_dbsqlparam_inferred_bool(self):
 
         params = [DbSqlParameter(name="p", value=True, type=None)]
@@ -124,6 +136,12 @@ def test_dbsqlparam_inferred_decimal(self):
         result = self._get_one_result(self.QUERY, params)
         assert result.col == Decimal("1234.56")
 
+    def test_dbsqlparam_explicit_none(self):
+
+        params = [DbSqlParameter(name="p", value=None, type=DbSqlType.VOID)]
+        result = self._get_one_result(self.QUERY, params)
+        assert result.col == None
+    
     def test_dbsqlparam_explicit_bool(self):
 
         params = [DbSqlParameter(name="p", value=True, type=DbSqlType.BOOLEAN)]

tests/e2e/test_driver.py

@@ -341,7 +341,7 @@ def test_escape_single_quotes(self):
             assert rows[0]["col_1"] == "you're"
 
             # Test escape syntax in parameter
-            cursor.execute("SELECT * FROM {} WHERE {}.col_1 LIKE %(var)s".format(table_name, table_name), parameters={"var": "you're"})
+            cursor.execute("SELECT * FROM {} WHERE {}.col_1 LIKE :var".format(table_name, table_name), parameters={"var": "you're"})
             rows = cursor.fetchall()
             assert rows[0]["col_1"] == "you're"
 

tests/unit/test_parameters.py

@@ -4,7 +4,7 @@
     named_parameters_to_dbsqlparams_v1,
     named_parameters_to_dbsqlparams_v2,
     calculate_decimal_cast_string,
-    DbsqlDynamicDecimalType
+    DbsqlDynamicDecimalType,
 )
 from databricks.sql.thrift_api.TCLIService.ttypes import (
     TSparkParameter,
@@ -37,7 +37,9 @@ def test_conversion_e2e(self):
                 name="", type="FLOAT", value=TSparkParameterValue(stringValue="1.0")
             ),
             TSparkParameter(
-                name="", type="DECIMAL(2,1)", value=TSparkParameterValue(stringValue="1.0")
+                name="",
+                type="DECIMAL(2,1)",
+                value=TSparkParameterValue(stringValue="1.0"),
             ),
         ]
 
@@ -91,26 +93,45 @@ def test_infer_types_decimal(self):
         # The output decimal will have a dynamically calculated decimal type with a value of DECIMAL(2,1)
         input = DbSqlParameter("", Decimal("1.0"))
         output: List[DbSqlParameter] = infer_types([input])
-        
+
         x = output[0]
 
         assert x.value == "1.0"
         assert isinstance(x.type, DbsqlDynamicDecimalType)
         assert x.type.value == "DECIMAL(2,1)"
-       
 
-class TestCalculateDecimalCast(object):
+    def test_infer_types_none(self):
+
+        input = DbSqlParameter("", None)
+        output: List[DbSqlParameter] = infer_types([input])
+
+        x = output[0]
+
+        assert x.value == None
+        assert x.type == DbSqlType.VOID
+        assert x.type.value == "VOID"
 
+    def test_infer_types_unsupported(self):
+        class ArbitraryType:
+            pass
+
+        input = DbSqlParameter("", ArbitraryType())
+
+        with pytest.raises(ValueError, match="Could not infer parameter type from"):
+            infer_types([input])
+
+
+class TestCalculateDecimalCast(object):
     def test_38_38(self):
         input = Decimal(".12345678912345678912345678912345678912")
         output = calculate_decimal_cast_string(input)
         assert output == "DECIMAL(38,38)"
-    
+
     def test_18_9(self):
         input = Decimal("123456789.123456789")
         output = calculate_decimal_cast_string(input)
         assert output == "DECIMAL(18,9)"
-    
+
     def test_38_0(self):
         input = Decimal("12345678912345678912345678912345678912")
         output = calculate_decimal_cast_string(input)
@@ -119,4 +140,4 @@ def test_38_0(self):
     def test_6_2(self):
         input = Decimal("1234.56")
         output = calculate_decimal_cast_string(input)
-        assert output == "DECIMAL(6,2)"
+        assert output == "DECIMAL(6,2)"