address comments

Author: shivam2680

tests/e2e/test_variant_types.py

@@ -1,6 +1,7 @@
 import pytest
 from datetime import datetime
 import json
+
 try:
     import pyarrow
 except ImportError:
@@ -9,6 +10,8 @@
 from tests.e2e.test_driver import PySQLPytestTestCase
 from tests.e2e.common.predicates import pysql_supports_arrow
 
+
+@pytest.mark.skipif(not pysql_supports_arrow(), reason="Requires arrow support")
 class TestVariantTypes(PySQLPytestTestCase):
     """Tests for the proper detection and handling of VARIANT type columns"""
 
@@ -17,7 +20,7 @@ def variant_table(self, connection_details):
         """A pytest fixture that creates a test table and cleans up after tests"""
         self.arguments = connection_details.copy()
         table_name = "pysql_test_variant_types_table"
-        
+
         with self.cursor() as cursor:
             try:
                 # Create the table with variant columns
@@ -30,7 +33,7 @@ def variant_table(self, connection_details):
                     )
                     """
                 )
-                
+
                 # Insert test records with different variant values
                 cursor.execute(
                     """
@@ -44,37 +47,45 @@ def variant_table(self, connection_details):
             finally:
                 cursor.execute(f"DROP TABLE IF EXISTS {table_name}")
 
-    @pytest.mark.skipif(not pysql_supports_arrow(), reason="Requires arrow support")
     def test_variant_type_detection(self, variant_table):
         """Test that VARIANT type columns are properly detected in schema"""
         with self.cursor() as cursor:
             cursor.execute(f"SELECT * FROM {variant_table} LIMIT 0")
-            
+
             # Verify column types in description
-            assert cursor.description[0][1] == 'int', "Integer column type not correctly identified"
-            assert cursor.description[1][1] == 'variant', "VARIANT column type not correctly identified"
-            assert cursor.description[2][1] == 'string', "String column type not correctly identified"
+            assert (
+                cursor.description[0][1] == "int"
+            ), "Integer column type not correctly identified"
+            assert (
+                cursor.description[1][1] == "variant"
+            ), "VARIANT column type not correctly identified"
+            assert (
+                cursor.description[2][1] == "string"
+            ), "String column type not correctly identified"
 
-    @pytest.mark.skipif(not pysql_supports_arrow(), reason="Requires arrow support")
     def test_variant_data_retrieval(self, variant_table):
         """Test that VARIANT data is properly retrieved and can be accessed as JSON"""
         with self.cursor() as cursor:
             cursor.execute(f"SELECT * FROM {variant_table} ORDER BY id")
             rows = cursor.fetchall()
-            
+
             # First row should have a JSON object
             json_obj = rows[0][1]
-            assert isinstance(json_obj, str), "VARIANT column should be returned as string"
+            assert isinstance(
+                json_obj, str
+            ), "VARIANT column should be returned as string"
 
             parsed = json.loads(json_obj)
-            assert parsed.get('name') == 'John'
-            assert parsed.get('age') == 30
+            assert parsed.get("name") == "John"
+            assert parsed.get("age") == 30
 
             # Second row should have a JSON array
             json_array = rows[1][1]
-            assert isinstance(json_array, str), "VARIANT array should be returned as string"
-            
+            assert isinstance(
+                json_array, str
+            ), "VARIANT array should be returned as string"
+
             # Parsing to verify it's valid JSON array
             parsed_array = json.loads(json_array)
             assert isinstance(parsed_array, list)
-            assert parsed_array == [1, 2, 3, 4]
+            assert parsed_array == [1, 2, 3, 4]

tests/unit/test_thrift_backend.py

@@ -2330,7 +2330,7 @@ def test_execute_command_sets_complex_type_fields_correctly(
                 [],
                 auth_provider=AuthProvider(),
                 ssl_options=SSLOptions(),
-            http_client=MagicMock(),
+                http_client=MagicMock(),
                 **complex_arg_types,
             )
             thrift_backend.execute_command(
@@ -2356,148 +2356,85 @@ def test_execute_command_sets_complex_type_fields_correctly(
                 t_execute_statement_req.useArrowNativeTypes.intervalTypesAsArrow
             )
 
-    def test_col_to_description_with_variant_type(self):
-        # Test variant type detection from Arrow field metadata
-        col = ttypes.TColumnDesc(
-            columnName="variant_col",
-            typeDesc=self._make_type_desc(ttypes.TTypeId.STRING_TYPE),
-        )
-        
-        # Create a field with variant type in metadata
-        field = pyarrow.field(
-            "variant_col", 
-            pyarrow.string(),
-            metadata={b'Spark:DataType:SqlName': b'VARIANT'}
-        )
-        
-        result = ThriftDatabricksClient._col_to_description(col, field)
-        
-        # Verify the result has variant as the type
-        self.assertEqual(result[0], "variant_col")  # Column name
-        self.assertEqual(result[1], "variant")      # Type name (should be variant instead of string)
-        self.assertIsNone(result[2])                # No display size
-        self.assertIsNone(result[3])                # No internal size
-        self.assertIsNone(result[4])                # No precision
-        self.assertIsNone(result[5])                # No scale
-        self.assertIsNone(result[6])                # No null ok
-
-    def test_col_to_description_without_variant_type(self):
-        # Test normal column without variant type
-        col = ttypes.TColumnDesc(
-            columnName="normal_col",
-            typeDesc=self._make_type_desc(ttypes.TTypeId.STRING_TYPE),
-        )
-        
-        # Create a normal field without variant metadata
-        field = pyarrow.field(
-            "normal_col", 
-            pyarrow.string(),
-            metadata={}
-        )
-        
-        result = ThriftDatabricksClient._col_to_description(col, field)
-        
-        # Verify the result has string as the type (unchanged)
-        self.assertEqual(result[0], "normal_col")  # Column name
-        self.assertEqual(result[1], "string")      # Type name (should be string)
-        self.assertIsNone(result[2])               # No display size
-        self.assertIsNone(result[3])               # No internal size
-        self.assertIsNone(result[4])               # No precision
-        self.assertIsNone(result[5])               # No scale
-        self.assertIsNone(result[6])               # No null ok
-
-    def test_col_to_description_with_null_field(self):
-        # Test handling of null field
-        col = ttypes.TColumnDesc(
-            columnName="missing_field",
-            typeDesc=self._make_type_desc(ttypes.TTypeId.STRING_TYPE),
-        )
-        
-        # Pass None as the field
-        result = ThriftDatabricksClient._col_to_description(col, None)
-        
-        # Verify the result has string as the type (unchanged)
-        self.assertEqual(result[0], "missing_field")  # Column name
-        self.assertEqual(result[1], "string")        # Type name (should be string)
-        self.assertIsNone(result[2])                # No display size
-        self.assertIsNone(result[3])                # No internal size
-        self.assertIsNone(result[4])                # No precision
-        self.assertIsNone(result[5])                # No scale
-        self.assertIsNone(result[6])                # No null ok
-
-    def test_hive_schema_to_description_with_arrow_schema(self):
-        # Create a table schema with regular and variant columns
-        columns = [
-            ttypes.TColumnDesc(
-                columnName="regular_col",
-                typeDesc=self._make_type_desc(ttypes.TTypeId.STRING_TYPE),
-            ),
-            ttypes.TColumnDesc(
-                columnName="variant_col",
-                typeDesc=self._make_type_desc(ttypes.TTypeId.STRING_TYPE),
-            ),
+    @unittest.skipIf(pyarrow is None, "Requires pyarrow")
+    def test_col_to_description(self):
+        test_cases = [
+            ("variant_col", {b"Spark:DataType:SqlName": b"VARIANT"}, "variant"),
+            ("normal_col", {}, "string"),
+            ("weird_field", {b"Spark:DataType:SqlName": b"Some unexpected value"}, "string"),
+            ("missing_field", None, "string"),  # None field case
         ]
-        t_table_schema = ttypes.TTableSchema(columns=columns)
 
-        # Create an Arrow schema with one variant column
-        fields = [
-            pyarrow.field("regular_col", pyarrow.string()),
-            pyarrow.field(
-                "variant_col", 
-                pyarrow.string(),
-                metadata={b'Spark:DataType:SqlName': b'VARIANT'}
-            )
-        ]
-        arrow_schema = pyarrow.schema(fields)
-        schema_bytes = arrow_schema.serialize().to_pybytes()
-        
-        # Get the description
-        description = ThriftDatabricksClient._hive_schema_to_description(t_table_schema, schema_bytes)
-        
-        # Verify regular column type
-        self.assertEqual(description[0][0], "regular_col")
-        self.assertEqual(description[0][1], "string")
-        
-        # Verify variant column type
-        self.assertEqual(description[1][0], "variant_col")
-        self.assertEqual(description[1][1], "variant")
+        for column_name, field_metadata, expected_type in test_cases:
+            with self.subTest(column_name=column_name, expected_type=expected_type):
+                col = ttypes.TColumnDesc(
+                    columnName=column_name,
+                    typeDesc=self._make_type_desc(ttypes.TTypeId.STRING_TYPE),
+                )
 
-    def test_hive_schema_to_description_with_null_schema_bytes(self):
-        # Create a simple table schema
-        columns = [
-            ttypes.TColumnDesc(
-                columnName="regular_col",
-                typeDesc=self._make_type_desc(ttypes.TTypeId.STRING_TYPE),
+                field = (
+                    None
+                    if field_metadata is None
+                    else pyarrow.field(column_name, pyarrow.string(), metadata=field_metadata)
+                )
+
+                result = ThriftDatabricksClient._col_to_description(col, field)
+
+                self.assertEqual(result[0], column_name)
+                self.assertEqual(result[1], expected_type)
+                self.assertIsNone(result[2])
+                self.assertIsNone(result[3])
+                self.assertIsNone(result[4])
+                self.assertIsNone(result[5])
+                self.assertIsNone(result[6])
+
+    @unittest.skipIf(pyarrow is None, "Requires pyarrow")
+    def test_hive_schema_to_description(self):
+        test_cases = [
+            (
+                [
+                    ("regular_col", ttypes.TTypeId.STRING_TYPE),
+                    ("variant_col", ttypes.TTypeId.STRING_TYPE),
+                ],
+                [
+                    ("regular_col", {}),
+                    ("variant_col", {b"Spark:DataType:SqlName": b"VARIANT"}),
+                ],
+                [("regular_col", "string"), ("variant_col", "variant")],
+            ),
+            (
+                [("regular_col", ttypes.TTypeId.STRING_TYPE)],
+                None,  # No arrow schema
+                [("regular_col", "string")],
             ),
         ]
-        t_table_schema = ttypes.TTableSchema(columns=columns)
-        
-        # Get the description with null schema_bytes
-        description = ThriftDatabricksClient._hive_schema_to_description(t_table_schema, None)
 
-        # Verify column type remains unchanged
-        self.assertEqual(description[0][0], "regular_col")
-        self.assertEqual(description[0][1], "string")
+        for columns, arrow_fields, expected_types in test_cases:
+            with self.subTest(arrow_fields=arrow_fields is not None):
+                t_table_schema = ttypes.TTableSchema(
+                    columns=[
+                        ttypes.TColumnDesc(
+                            columnName=name, typeDesc=self._make_type_desc(col_type)
+                        )
+                        for name, col_type in columns
+                    ]
+                )
 
-    def test_col_to_description_with_malformed_metadata(self):
-        # Test handling of malformed metadata
-        col = ttypes.TColumnDesc(
-            columnName="weird_field",
-            typeDesc=self._make_type_desc(ttypes.TTypeId.STRING_TYPE),
-        )
-        
-        # Create a field with malformed metadata
-        field = pyarrow.field(
-            "weird_field", 
-            pyarrow.string(),
-            metadata={b'Spark:DataType:SqlName': b'Some unexpected value'}
-        )
-        
-        result = ThriftDatabricksClient._col_to_description(col, field)
-        
-        # Verify the type remains unchanged
-        self.assertEqual(result[0], "weird_field")  # Column name
-        self.assertEqual(result[1], "string")       # Type name (should remain string)
+                schema_bytes = None
+                if arrow_fields:
+                    fields = [
+                        pyarrow.field(name, pyarrow.string(), metadata=metadata)
+                        for name, metadata in arrow_fields
+                    ]
+                    schema_bytes = pyarrow.schema(fields).serialize().to_pybytes()
+
+                description = ThriftDatabricksClient._hive_schema_to_description(
+                    t_table_schema, schema_bytes
+                )
+
+                for i, (expected_name, expected_type) in enumerate(expected_types):
+                    self.assertEqual(description[i][0], expected_name)
+                    self.assertEqual(description[i][1], expected_type)
 
 
 if __name__ == "__main__":