Check

Working e2e tests prototype

Author: jprakash-db

src/databricks/sql/client.py

@@ -789,7 +789,7 @@ def execute(
 
         :returns self
         """
-
+        
         param_approach = self._determine_parameter_approach(parameters)
         if param_approach == ParameterApproach.NONE:
             prepared_params = NO_NATIVE_PARAMS
@@ -808,7 +808,7 @@ def execute(
             prepared_operation, prepared_params = self._prepare_native_parameters(
                 transformed_operation, normalized_parameters, param_structure
             )
-
+        
         self._check_not_closed()
         self._close_and_clear_active_result_set()
         execute_response = self.thrift_backend.execute_command(

src/databricks/sql/utils.py

@@ -429,7 +429,7 @@ def user_friendly_error_message(self, no_retry_reason, attempt, elapsed):
 # Taken from PyHive
 class ParamEscaper:
     _DATE_FORMAT = "%Y-%m-%d"
-    _TIME_FORMAT = "%H:%M:%S.%f"
+    _TIME_FORMAT = "%H:%M:%S.%f %z"
     _DATETIME_FORMAT = "{} {}".format(_DATE_FORMAT, _TIME_FORMAT)
 
     def escape_args(self, parameters):
@@ -459,13 +459,15 @@ def escape_string(self, item):
 
     def escape_sequence(self, item):
         l = map(self.escape_item, item)
+        l = list(map(str, l))
         return "ARRAY(" + ",".join(l) + ")"
 
     def escape_mapping(self, item):
         l = map(
             self.escape_item,
             (element for key, value in item.items() for element in (key, value)),
         )
+        l = list(map(str, l))
         return "MAP(" + ",".join(l) + ")"
 
     def escape_datetime(self, item, format, cutoff=0):

tests/e2e/test_complex_types.py

@@ -1,5 +1,6 @@
 import pytest
 from numpy import ndarray
+from typing import Sequence
 
 from tests.e2e.test_driver import PySQLPytestTestCase
 
@@ -17,7 +18,10 @@ def table_fixture(self, connection_details):
                 CREATE TABLE IF NOT EXISTS pysql_test_complex_types_table (
                     array_col ARRAY<STRING>,
                     map_col MAP<STRING, INTEGER>,
-                    struct_col STRUCT<field1: STRING, field2: INTEGER>
+                    struct_col STRUCT<field1: STRING, field2: INTEGER>,
+                    array_array_col ARRAY<ARRAY<STRING>>,
+                    array_map_col ARRAY<MAP<STRING, INTEGER>>,
+                    map_array_col MAP<STRING, ARRAY<STRING>>
                 )
                 """
             )
@@ -28,7 +32,10 @@ def table_fixture(self, connection_details):
                 VALUES (
                     ARRAY('a', 'b', 'c'),
                     MAP('a', 1, 'b', 2, 'c', 3),
-                    NAMED_STRUCT('field1', 'a', 'field2', 1)
+                    NAMED_STRUCT('field1', 'a', 'field2', 1),
+                    ARRAY(ARRAY('a','b','c')),
+                    ARRAY(MAP('a', 1, 'b', 2, 'c', 3)),
+                    MAP('a', ARRAY('a', 'b', 'c'), 'b', ARRAY('d', 'e'))
                 )
                 """
             )
@@ -38,7 +45,7 @@ def table_fixture(self, connection_details):
 
     @pytest.mark.parametrize(
         "field,expected_type",
-        [("array_col", ndarray), ("map_col", list), ("struct_col", dict)],
+        [("array_col", ndarray), ("map_col", list), ("struct_col", dict), ("array_array_col", ndarray), ("array_map_col", ndarray), ("map_array_col", list)],
     )
     def test_read_complex_types_as_arrow(self, field, expected_type, table_fixture):
         """Confirms the return types of a complex type field when reading as arrow"""
@@ -47,10 +54,10 @@ def test_read_complex_types_as_arrow(self, field, expected_type, table_fixture):
             result = cursor.execute(
                 "SELECT * FROM pysql_test_complex_types_table LIMIT 1"
             ).fetchone()
-
+        
         assert isinstance(result[field], expected_type)
 
-    @pytest.mark.parametrize("field", [("array_col"), ("map_col"), ("struct_col")])
+    @pytest.mark.parametrize("field", [("array_col"), ("map_col"), ("struct_col"), ("array_array_col"), ("array_map_col"), ("map_array_col")])
     def test_read_complex_types_as_string(self, field, table_fixture):
         """Confirms the return type of a complex type that is returned as a string"""
         with self.cursor(

tests/e2e/test_parameterized_queries.py

@@ -26,6 +26,8 @@
     TimestampParameter,
     TinyIntParameter,
     VoidParameter,
+    ArrayParameter,
+    MapParameter,
 )
 from tests.e2e.test_driver import PySQLPytestTestCase
 
@@ -50,6 +52,8 @@ class Primitive(Enum):
     DOUBLE = 3.14
     FLOAT = 3.15
     SMALLINT = 51
+    ARRAYS = ["a", "b", "c"]
+    MAPS = {"a": 1, "b": 2, "c": 3}
 
 
 class PrimitiveExtra(Enum):
@@ -103,6 +107,8 @@ class TestParameterizedQueries(PySQLPytestTestCase):
         Primitive.BOOL: "boolean_col",
         Primitive.DATE: "date_col",
         Primitive.TIMESTAMP: "timestamp_col",
+        Primitive.ARRAYS: "array_col",
+        Primitive.MAPS: "map_col",
         Primitive.NONE: "null_col",
     }
 
@@ -134,7 +140,9 @@ def inline_table(self, connection_details):
             string_col STRING,
             boolean_col BOOLEAN,
             date_col DATE,
-            timestamp_col TIMESTAMP
+            timestamp_col TIMESTAMP,
+            array_col ARRAY<STRING>,
+            map_col MAP<STRING, INT>
             ) USING DELTA
         """
 
@@ -167,9 +175,9 @@ def _inline_roundtrip(self, params: dict, 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 pysql_e2e_inline_param_test_table (`{target_column}`) VALUES (%(p)s)"
-        SELECT_QUERY = f"SELECT {target_column} `col` FROM pysql_e2e_inline_param_test_table LIMIT 1"
-        DELETE_QUERY = "DELETE FROM pysql_e2e_inline_param_test_table"
+        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"
 
         with self.connection(extra_params={"use_inline_params": True}) as conn:
             with conn.cursor() as cursor:
@@ -229,10 +237,18 @@ def _eq(self, actual, expected: Primitive):
         If primitive is Primitive.DOUBLE than an extra quantize step is performed before
         making the assertion.
         """
-        if expected in (Primitive.DOUBLE, Primitive.FLOAT):
-            return self._quantize(actual) == self._quantize(expected.value)
+        actual_parsed = actual
+        expected_parsed = expected.value
 
-        return actual == expected.value
+        if expected in (Primitive.DOUBLE, Primitive.FLOAT):
+            actual_parsed = self._quantize(actual)
+            expected_parsed = self._quantize(expected.value)
+        elif expected == Primitive.ARRAYS:
+            actual_parsed = actual.tolist()
+        elif expected == Primitive.MAPS:
+            expected_parsed = list(expected.value.items())
+        
+        return actual_parsed == expected_parsed
 
     @pytest.mark.parametrize("primitive", Primitive)
     @pytest.mark.parametrize(
@@ -278,6 +294,8 @@ def test_primitive_single(
             (Primitive.SMALLINT, SmallIntParameter),
             (PrimitiveExtra.TIMESTAMP_NTZ, TimestampNTZParameter),
             (PrimitiveExtra.TINYINT, TinyIntParameter),
+            (Primitive.ARRAYS, ArrayParameter),
+            (Primitive.MAPS, MapParameter),
         ],
     )
     def test_dbsqlparameter_single(