More tests

Author: jprakash-db

tests/e2e/test_parameterized_queries.py

@@ -5,8 +5,11 @@
 from typing import Dict, List, Type, Union
 from unittest.mock import patch
 
+import time
+import numpy as np
 import pytest
 import pytz
+from numpy.random.mtrand import Sequence
 
 from databricks.sql.parameters.native import (
     BigIntegerParameter,
@@ -112,6 +115,7 @@ class TestParameterizedQueries(PySQLPytestTestCase):
         Primitive.NONE: "null_col",
     }
 
+
     def _get_inline_table_column(self, value):
         return self.inline_type_map[Primitive(value)]
 
@@ -142,7 +146,9 @@ def inline_table(self, connection_details):
             date_col DATE,
             timestamp_col TIMESTAMP,
             array_col ARRAY<STRING>,
-            map_col MAP<STRING, INT>
+            map_col MAP<STRING, INT>,
+            array_map_col ARRAY<MAP<STRING,INT>>,
+            map_array_col MAP<INT,ARRAY<STRING>>
             ) USING DELTA
         """
 
@@ -163,7 +169,7 @@ def patch_server_supports_native_params(self, supports_native_params: bool = Tru
             finally:
                 pass
 
-    def _inline_roundtrip(self, params: dict, paramstyle: ParamStyle):
+    def _inline_roundtrip(self, params: dict, paramstyle: ParamStyle, target_column):
         """This INSERT, SELECT, DELETE dance is necessary because simply selecting
         ```
         "SELECT %(param)s"
@@ -174,7 +180,6 @@ def _inline_roundtrip(self, params: dict, paramstyle: ParamStyle):
         :paramstyle:
             This is a no-op but is included to make the test-code easier to read.
         """
-        target_column = self._get_inline_table_column(params.get("p"))
         INSERT_QUERY = f"INSERT INTO ___________________first.jprakash.pysql_e2e_inline_param_test_table (`{target_column}`) VALUES (%(p)s)"
         SELECT_QUERY = f"SELECT {target_column} `col` FROM ___________________first.jprakash.pysql_e2e_inline_param_test_table LIMIT 1"
         DELETE_QUERY = "DELETE FROM ___________________first.jprakash.pysql_e2e_inline_param_test_table"
@@ -220,7 +225,7 @@ def _get_one_result(
         if approach == ParameterApproach.INLINE:
             # inline mode always uses ParamStyle.PYFORMAT
             # inline mode doesn't support positional parameters
-            return self._inline_roundtrip(params, paramstyle=ParamStyle.PYFORMAT)
+            return self._inline_roundtrip(params, paramstyle=ParamStyle.PYFORMAT,target_column=self._get_inline_table_column(params.get("p")))
         elif approach == ParameterApproach.NATIVE:
             # native mode can use either ParamStyle.NAMED or ParamStyle.PYFORMAT
             # native mode can use either ParameterStructure.NAMED or ParameterStructure.POSITIONAL
@@ -250,6 +255,57 @@ def _eq(self, actual, expected: Primitive):
 
         return actual_parsed == expected_parsed
 
+    def _parse_to_common_type(self, value):
+        """
+        Function to convert the :value passed into a common python datatype for comparison
+
+        Convertion fyi
+        MAP Datatype on server is returned as a list of tuples
+            Ex:
+                {"a":1,"b":2} -> [("a",1),("b",2)]
+        
+        ARRAY Datatype on server is returned as a numpy array
+            Ex:
+                ["a","b","c"] -> np.array(["a","b","c"],dtype=object)
+        
+        Primitive datatype on server is returned as a numpy primitive
+            Ex:
+                1 -> np.int64(1)
+                2 -> np.int32(2)
+        """
+        if value is None:
+            return None
+        elif isinstance(value, (Sequence,np.ndarray)) and not isinstance(value, (str, bytes)):
+            return tuple(value)
+        elif isinstance(value,dict):
+            return tuple(value.items())
+        elif isinstance(value, np.generic):
+            return value.item()
+        else:
+            return value
+
+    def _recursive_compare(self,actual, expected):
+        """
+        Function to compare the :actual and :expected values, recursively checks and ensures that all the data matches till the leaf level
+
+        Note: Complex datatype like MAP is not returned as a dictionary but as a list of tuples
+        """
+        actual_parsed = self._parse_to_common_type(actual)
+        expected_parsed = self._parse_to_common_type(expected)
+
+        # Check if types are the same
+        if type(actual_parsed) != type(expected_parsed):
+            return False
+
+        # Handle lists or tuples
+        if isinstance(actual_parsed, (list, tuple)):
+            if len(actual_parsed) != len(expected_parsed):
+                return False
+            return all(self._recursive_compare(o1, o2) for o1, o2 in zip(actual_parsed, expected_parsed))
+
+        return actual_parsed==expected_parsed
+
+
     @pytest.mark.parametrize("primitive", Primitive)
     @pytest.mark.parametrize(
         "approach,paramstyle,parameter_structure", approach_paramstyle_combinations
@@ -379,6 +435,43 @@ def test_readme_example(self):
         assert len(result) == 10
         assert result[0].p == "foo"
 
+    @pytest.mark.parametrize(
+        "col_name,data",[
+        ("array_map_col", [{"a": 1, "b": 2}, {"c": 3, "d": 4}]),
+        ("map_array_col", {1: ["a", "b"], 2: ["c", "d"]}),
+    ])
+    def test_inline_recursive_complex_type(self,col_name,data):
+        params = {'p':data}
+        result = self._inline_roundtrip(params=params,paramstyle=ParamStyle.PYFORMAT,target_column=col_name)
+        assert self._recursive_compare(result.col, data)
+    
+    
+    @pytest.mark.parametrize(
+        "description,data",
+        [
+            ("ARRAY<MAP<STRING,INT>>",[{"a": 1, "b": 2}, {"c": 3, "d": 4}]),
+            ("MAP<INT,ARRAY<STRING>>",{1: ["a", "b"], 2: ["c", "d"]}),
+            ("ARRAY<ARRAY<INT>>",[[1,2,3],[1,2,3]]),
+            ("ARRAY<ARRAY<ARRAY<INT>>>",[[[1,2,3],[1,2,3]],[[1,2,3],[1,2,3]]]),
+            ("MAP<STRING,MAP<STRING,STRING>>",{"a":{"b":"c","d":"e"},"f":{"g":"h","i":"j"}})
+        ],
+    )
+    @pytest.mark.parametrize(
+        "paramstyle,parameter_structure",
+        [
+            (ParamStyle.NONE, ParameterStructure.POSITIONAL),
+            (ParamStyle.PYFORMAT, ParameterStructure.NAMED),
+            (ParamStyle.NAMED, ParameterStructure.NAMED),
+        ]
+    )
+    def test_native_recursive_complex_type(self,description,data,paramstyle,parameter_structure):
+        if(paramstyle==ParamStyle.NONE):
+            params=[data]
+        else:
+            params={'p':data}
+        result = self._native_roundtrip(parameters=params,paramstyle=paramstyle,parameter_structure=parameter_structure)
+        assert self._recursive_compare(result.col, data)
+
 
 class TestInlineParameterSyntax(PySQLPytestTestCase):
     """The inline parameter approach uses pyformat markers"""