Added unit tests

Author: jprakash-db

src/databricks/sql/parameters/__init__.py

@@ -12,4 +12,6 @@
     TimestampNTZParameter,
     TinyIntParameter,
     DecimalParameter,
+    MapParameter,
+    ArrayParameter,
 )

src/databricks/sql/parameters/native.py

@@ -467,7 +467,7 @@ def __init__(self, value: Sequence[Any], name: Optional[str] = None):
         self.name = name
         self.value = [dbsql_parameter_from_primitive(val) for val in value]
 
-    def as_tspark_param(self, named: bool) -> TSparkParameter:
+    def as_tspark_param(self, named: bool = None) -> TSparkParameter:
         """Returns a TSparkParameter object that can be passed to the DBR thrift server."""
 
         tsp = TSparkParameter(type=self._cast_expr())
@@ -516,7 +516,7 @@ def __init__(self, value: dict, name: Optional[str] = None):
             for item in (key, val)
         ]
 
-    def as_tspark_param(self, named: bool) -> TSparkParameter:
+    def as_tspark_param(self, named: bool = None) -> TSparkParameter:
         """Returns a TSparkParameter object that can be passed to the DBR thrift server."""
 
         tsp = TSparkParameter(type=self._cast_expr())

src/databricks/sql/utils.py

@@ -473,7 +473,7 @@ def escape_mapping(self, item):
     def escape_datetime(self, item, format, cutoff=0):
         dt_str = item.strftime(format)
         formatted = dt_str[:-cutoff] if cutoff and format.endswith(".%f") else dt_str
-        return "'{}'".format(formatted)
+        return "'{}'".format(formatted.strip())
 
     def escape_decimal(self, item):
         return str(item)

tests/unit/test_client.py

@@ -455,7 +455,7 @@ def test_execute_parameter_passthrough(self):
             ("SELECT %(x)s", "SELECT NULL", {"x": None}),
             ("SELECT %(int_value)d", "SELECT 48", {"int_value": 48}),
             ("SELECT %(float_value).2f", "SELECT 48.20", {"float_value": 48.2}),
-            ("SELECT %(iter)s", "SELECT (1,2,3,4,5)", {"iter": [1, 2, 3, 4, 5]}),
+            ("SELECT %(iter)s", "SELECT ARRAY(1,2,3,4,5)", {"iter": [1, 2, 3, 4, 5]}),
             (
                 "SELECT %(datetime)s",
                 "SELECT '2022-02-01 10:23:00.000000'",

tests/unit/test_param_escaper.py

@@ -120,23 +120,37 @@ def test_escape_date(self):
         assert pe.escape_datetime(INPUT, FORMAT) == OUTPUT
 
     def test_escape_sequence_integer(self):
-        assert pe.escape_sequence([1, 2, 3, 4]) == "(1,2,3,4)"
+        assert pe.escape_sequence([1, 2, 3, 4]) == "ARRAY(1,2,3,4)"
 
     def test_escape_sequence_float(self):
-        assert pe.escape_sequence([1.1, 2.2, 3.3, 4.4]) == "(1.1,2.2,3.3,4.4)"
+        assert pe.escape_sequence([1.1, 2.2, 3.3, 4.4]) == "ARRAY(1.1,2.2,3.3,4.4)"
 
     def test_escape_sequence_string(self):
         assert (
             pe.escape_sequence(["his", "name", "was", "robert", "palmer"])
-            == "('his','name','was','robert','palmer')"
+            == "ARRAY('his','name','was','robert','palmer')"
         )
 
     def test_escape_sequence_sequence_of_strings(self):
-        # This is not valid SQL.
         INPUT = [["his", "name"], ["was", "robert"], ["palmer"]]
-        OUTPUT = "(('his','name'),('was','robert'),('palmer'))"
+        OUTPUT = "ARRAY(ARRAY('his','name'),ARRAY('was','robert'),ARRAY('palmer'))"
 
         assert pe.escape_sequence(INPUT) == OUTPUT
+    
+    def test_escape_map_string_int(self):
+        INPUT = {"a": 1, "b": 2}
+        OUTPUT = "MAP('a',1,'b',2)"
+        assert pe.escape_mapping(INPUT) == OUTPUT
+
+    def test_escape_map_string_sequence_of_floats(self):
+        INPUT = {"a": [1.1, 2.2, 3.3], "b": [4.4, 5.5, 6.6]}
+        OUTPUT = "MAP('a',ARRAY(1.1,2.2,3.3),'b',ARRAY(4.4,5.5,6.6))"
+        assert pe.escape_mapping(INPUT) == OUTPUT
+    
+    def test_escape_sequence_of_map_int_string(self):
+        INPUT = [{1: "a", 2: "foo"}, {3: "b", 4: "bar"}]
+        OUTPUT = "ARRAY(MAP(1,'a',2,'foo'),MAP(3,'b',4,'bar'))"
+        assert pe.escape_sequence(INPUT) == OUTPUT
 
 
 class TestFullQueryEscaping(object):

tests/unit/test_parameters.py

@@ -21,10 +21,14 @@
     TimestampParameter,
     TinyIntParameter,
     VoidParameter,
+    MapParameter,
+    ArrayParameter,
 )
 from databricks.sql.parameters.native import (
     TDbsqlParameter,
+    TSparkParameter,
     TSparkParameterValue,
+    TSparkParameterValueArg,
     dbsql_parameter_from_primitive,
 )
 from databricks.sql.thrift_api.TCLIService import ttypes
@@ -112,6 +116,8 @@ class Primitive(Enum):
     DOUBLE = 3.14
     FLOAT = 3.15
     SMALLINT = 51
+    ARRAY = [1, 2, 3]
+    MAP = {"a": 1, "b": 2}
 
 
 class TestDbsqlParameter:
@@ -131,6 +137,8 @@ class TestDbsqlParameter:
             (TimestampParameter, Primitive.TIMESTAMP, "TIMESTAMP"),
             (TimestampNTZParameter, Primitive.TIMESTAMP, "TIMESTAMP_NTZ"),
             (TinyIntParameter, Primitive.INT, "TINYINT"),
+            (MapParameter, Primitive.MAP, "MAP"),
+            (ArrayParameter, Primitive.ARRAY, "ARRAY"),
         ),
     )
     def test_cast_expression(
@@ -165,6 +173,18 @@ def test_tspark_param_value(self, t: TDbsqlParameter, prim):
             assert output == None
         else:
             assert output == TSparkParameterValue(stringValue=str(prim.value))
+    
+    @pytest.mark.parametrize(
+            "base_type,input,expected_output",[
+                (ArrayParameter, [1,2,3], TSparkParameter(ordinal=True, name=None, type='ARRAY', value=None, arguments=[TSparkParameterValueArg(type='INT', value='1', arguments=None), TSparkParameterValueArg(type='INT', value='2', arguments=None), TSparkParameterValueArg(type='INT', value='3', arguments=None)])),
+                (MapParameter, {"a": 1, "b": 2}, TSparkParameter(ordinal=True, name=None, type='MAP', value=None, arguments=[TSparkParameterValueArg(type='STRING', value='a', arguments=None), TSparkParameterValueArg(type='INT', value='1', arguments=None), TSparkParameterValueArg(type='STRING', value='b', arguments=None), TSparkParameterValueArg(type='INT', value='2', arguments=None)])),
+                (ArrayParameter,[{"a":1,"b":2},{"c":3,"d":4}], TSparkParameter(ordinal=True, name=None, type='ARRAY', value=None, arguments=[TSparkParameterValueArg(type='MAP', value=None, arguments=[TSparkParameterValueArg(type='STRING', value='a', arguments=None), TSparkParameterValueArg(type='INT', value='1', arguments=None), TSparkParameterValueArg(type='STRING', value='b', arguments=None), TSparkParameterValueArg(type='INT', value='2', arguments=None)]), TSparkParameterValueArg(type='MAP', value=None, arguments=[TSparkParameterValueArg(type='STRING', value='c', arguments=None), TSparkParameterValueArg(type='INT', value='3', arguments=None), TSparkParameterValueArg(type='STRING', value='d', arguments=None), TSparkParameterValueArg(type='INT', value='4', arguments=None)])])),
+            ]
+    )
+    def test_complex_type_tspark_param(self,base_type,input,expected_output):
+        p = base_type(input)
+        tsp = p.as_tspark_param()
+        assert tsp == expected_output
 
     def test_tspark_param_named(self):
         p = dbsql_parameter_from_primitive(Primitive.INT.value, name="p")
@@ -192,6 +212,8 @@ def test_tspark_param_ordinal(self):
             (FloatParameter, Primitive.FLOAT),
             (VoidParameter, Primitive.NONE),
             (TimestampParameter, Primitive.TIMESTAMP),
+            (MapParameter, Primitive.MAP),
+            (ArrayParameter, Primitive.ARRAY),
         ),
     )
     def test_inference(self, _type: TDbsqlParameter, prim: Primitive):