fix type codes by using Thrift ttypes

Signed-off-by: varun-edachali-dbx <varun.edachali@databricks.com>

Author: varun-edachali-dbx

src/databricks/sql/backend/sea/utils/conversion.py

@@ -50,60 +50,65 @@ def _convert_decimal(
 
 class SqlType:
     """
-    SQL type constants
+    SQL type constants based on Thrift TTypeId values.
 
-    The list of types can be found in the SEA REST API Reference:
-    https://docs.databricks.com/api/workspace/statementexecution/executestatement
+    These correspond to the normalized type names that come from the SEA backend
+    after normalize_sea_type_to_thrift processing (lowercase, without _TYPE suffix).
     """
 
     # Numeric types
-    BYTE = "byte"
-    SHORT = "short"
-    INT = "int"
-    LONG = "long"
-    FLOAT = "float"
-    DOUBLE = "double"
-    DECIMAL = "decimal"
+    TINYINT = "tinyint"  # Maps to TTypeId.TINYINT_TYPE
+    SMALLINT = "smallint"  # Maps to TTypeId.SMALLINT_TYPE
+    INT = "int"  # Maps to TTypeId.INT_TYPE
+    BIGINT = "bigint"  # Maps to TTypeId.BIGINT_TYPE
+    FLOAT = "float"  # Maps to TTypeId.FLOAT_TYPE
+    DOUBLE = "double"  # Maps to TTypeId.DOUBLE_TYPE
+    DECIMAL = "decimal"  # Maps to TTypeId.DECIMAL_TYPE
 
     # Boolean type
-    BOOLEAN = "boolean"
+    BOOLEAN = "boolean"  # Maps to TTypeId.BOOLEAN_TYPE
 
     # Date/Time types
-    DATE = "date"
-    TIMESTAMP = "timestamp"
-    INTERVAL = "interval"
+    DATE = "date"  # Maps to TTypeId.DATE_TYPE
+    TIMESTAMP = "timestamp"  # Maps to TTypeId.TIMESTAMP_TYPE
+    INTERVAL_YEAR_MONTH = (
+        "interval_year_month"  # Maps to TTypeId.INTERVAL_YEAR_MONTH_TYPE
+    )
+    INTERVAL_DAY_TIME = "interval_day_time"  # Maps to TTypeId.INTERVAL_DAY_TIME_TYPE
 
     # String types
-    CHAR = "char"
-    STRING = "string"
+    CHAR = "char"  # Maps to TTypeId.CHAR_TYPE
+    VARCHAR = "varchar"  # Maps to TTypeId.VARCHAR_TYPE
+    STRING = "string"  # Maps to TTypeId.STRING_TYPE
 
     # Binary type
-    BINARY = "binary"
+    BINARY = "binary"  # Maps to TTypeId.BINARY_TYPE
 
     # Complex types
-    ARRAY = "array"
-    MAP = "map"
-    STRUCT = "struct"
+    ARRAY = "array"  # Maps to TTypeId.ARRAY_TYPE
+    MAP = "map"  # Maps to TTypeId.MAP_TYPE
+    STRUCT = "struct"  # Maps to TTypeId.STRUCT_TYPE
 
     # Other types
-    NULL = "null"
-    USER_DEFINED_TYPE = "user_defined_type"
+    NULL = "null"  # Maps to TTypeId.NULL_TYPE
+    UNION = "union"  # Maps to TTypeId.UNION_TYPE
+    USER_DEFINED = "user_defined"  # Maps to TTypeId.USER_DEFINED_TYPE
 
 
 class SqlTypeConverter:
     """
     Utility class for converting SQL types to Python types.
-    Based on the types supported by the Databricks SDK.
+    Based on the Thrift TTypeId types after normalization.
     """
 
     # SQL type to conversion function mapping
     # TODO: complex types
     TYPE_MAPPING: Dict[str, Callable] = {
         # Numeric types
-        SqlType.BYTE: lambda v: int(v),
-        SqlType.SHORT: lambda v: int(v),
+        SqlType.TINYINT: lambda v: int(v),
+        SqlType.SMALLINT: lambda v: int(v),
         SqlType.INT: lambda v: int(v),
-        SqlType.LONG: lambda v: int(v),
+        SqlType.BIGINT: lambda v: int(v),
         SqlType.FLOAT: lambda v: float(v),
         SqlType.DOUBLE: lambda v: float(v),
         SqlType.DECIMAL: _convert_decimal,
@@ -112,16 +117,18 @@ class SqlTypeConverter:
         # Date/Time types
         SqlType.DATE: lambda v: datetime.date.fromisoformat(v),
         SqlType.TIMESTAMP: lambda v: parser.parse(v),
-        SqlType.INTERVAL: lambda v: v,  # Keep as string for now
+        SqlType.INTERVAL_YEAR_MONTH: lambda v: v,  # Keep as string for now
+        SqlType.INTERVAL_DAY_TIME: lambda v: v,  # Keep as string for now
         # String types - no conversion needed
         SqlType.CHAR: lambda v: v,
+        SqlType.VARCHAR: lambda v: v,
         SqlType.STRING: lambda v: v,
         # Binary type
         SqlType.BINARY: lambda v: bytes.fromhex(v),
         # Other types
         SqlType.NULL: lambda v: None,
         # Complex types and user-defined types return as-is
-        SqlType.USER_DEFINED_TYPE: lambda v: v,
+        SqlType.USER_DEFINED: lambda v: v,
     }
 
     @staticmethod
@@ -135,7 +142,7 @@ def convert_value(
 
         Args:
             value: The string value to convert
-            sql_type: The SQL type (e.g., 'int', 'decimal')
+            sql_type: The SQL type (e.g., 'tinyint', 'decimal')
             **kwargs: Additional keyword arguments for the conversion function
 
         Returns:

tests/unit/test_client.py

@@ -265,9 +265,7 @@ def test_negative_fetch_throws_exception(self):
         mock_backend = Mock()
         mock_backend.fetch_results.return_value = (Mock(), False, 0)
 
-        result_set = ThriftResultSet(
-            Mock(), Mock(), mock_backend
-        )
+        result_set = ThriftResultSet(Mock(), Mock(), mock_backend)
 
         with self.assertRaises(ValueError) as e:
             result_set.fetchmany(-1)

tests/unit/test_downloader.py

@@ -26,12 +26,14 @@ class DownloaderTests(unittest.TestCase):
     def _setup_time_mock_for_download(self, mock_time, end_time):
         """Helper to setup time mock that handles logging system calls."""
         call_count = [0]
+
         def time_side_effect():
             call_count[0] += 1
             if call_count[0] <= 2:  # First two calls (validation, start_time)
                 return 1000
             else:  # All subsequent calls (logging, duration calculation)
                 return end_time
+
         mock_time.side_effect = time_side_effect
 
     @patch("time.time", return_value=1000)
@@ -104,7 +106,7 @@ def test_run_get_response_not_ok(self, mock_time):
     @patch("time.time")
     def test_run_uncompressed_successful(self, mock_time):
         self._setup_time_mock_for_download(mock_time, 1000.5)
-        
+
         http_client = DatabricksHttpClient.get_instance()
         file_bytes = b"1234567890" * 10
         settings = Mock(link_expiry_buffer_secs=0, download_timeout=0, use_proxy=False)
@@ -133,7 +135,7 @@ def test_run_uncompressed_successful(self, mock_time):
     @patch("time.time")
     def test_run_compressed_successful(self, mock_time):
         self._setup_time_mock_for_download(mock_time, 1000.2)
-        
+
         http_client = DatabricksHttpClient.get_instance()
         file_bytes = b"1234567890" * 10
         compressed_bytes = b'\x04"M\x18h@d\x00\x00\x00\x00\x00\x00\x00#\x14\x00\x00\x00\xaf1234567890\n\x00BP67890\x00\x00\x00\x00'

tests/unit/test_sea_conversion.py

@@ -18,10 +18,12 @@ class TestSqlTypeConverter:
     def test_convert_numeric_types(self):
         """Test converting numeric types."""
         # Test integer types
-        assert SqlTypeConverter.convert_value("123", SqlType.BYTE) == 123
-        assert SqlTypeConverter.convert_value("456", SqlType.SHORT) == 456
+        assert SqlTypeConverter.convert_value("123", SqlType.TINYINT) == 123
+        assert SqlTypeConverter.convert_value("456", SqlType.SMALLINT) == 456
         assert SqlTypeConverter.convert_value("789", SqlType.INT) == 789
-        assert SqlTypeConverter.convert_value("1234567890", SqlType.LONG) == 1234567890
+        assert (
+            SqlTypeConverter.convert_value("1234567890", SqlType.BIGINT) == 1234567890
+        )
 
         # Test floating point types
         assert SqlTypeConverter.convert_value("123.45", SqlType.FLOAT) == 123.45
@@ -80,11 +82,16 @@ def test_convert_datetime_types(self):
         assert timestamp_value.minute == 30
         assert timestamp_value.second == 45
 
-        # Test interval type (currently returns as string)
-        interval_value = SqlTypeConverter.convert_value(
-            "1 day 2 hours", SqlType.INTERVAL
+        # Test interval types (currently return as string)
+        interval_ym_value = SqlTypeConverter.convert_value(
+            "1-6", SqlType.INTERVAL_YEAR_MONTH
+        )
+        assert interval_ym_value == "1-6"
+
+        interval_dt_value = SqlTypeConverter.convert_value(
+            "1 day 2 hours", SqlType.INTERVAL_DAY_TIME
         )
-        assert interval_value == "1 day 2 hours"
+        assert interval_dt_value == "1 day 2 hours"
 
         # Test invalid date input
         result = SqlTypeConverter.convert_value("not_a_date", SqlType.DATE)
@@ -98,6 +105,10 @@ def test_convert_string_types(self):
             == "test string"
         )
         assert SqlTypeConverter.convert_value("test char", SqlType.CHAR) == "test char"
+        assert (
+            SqlTypeConverter.convert_value("test varchar", SqlType.VARCHAR)
+            == "test varchar"
+        )
 
     def test_convert_binary_type(self):
         """Test converting binary type."""
@@ -115,7 +126,7 @@ def test_convert_unsupported_type(self):
         # Should return the original value
         assert SqlTypeConverter.convert_value("test", "unsupported_type") == "test"
 
-        # Complex types should return as-is
+        # Complex types should return as-is (not yet implemented in TYPE_MAPPING)
         assert (
             SqlTypeConverter.convert_value("complex_value", SqlType.ARRAY)
             == "complex_value"

tests/unit/test_telemetry_retry.py

@@ -6,7 +6,8 @@
 from databricks.sql.telemetry.telemetry_client import TelemetryClientFactory
 from databricks.sql.auth.retry import DatabricksRetryPolicy
 
-PATCH_TARGET = 'urllib3.connectionpool.HTTPSConnectionPool._get_conn'
+PATCH_TARGET = "urllib3.connectionpool.HTTPSConnectionPool._get_conn"
+
 
 def create_mock_conn(responses):
     """Creates a mock connection object whose getresponse() method yields a series of responses."""
@@ -16,15 +17,18 @@ def create_mock_conn(responses):
         mock_http_response = MagicMock()
         mock_http_response.status = resp.get("status")
         mock_http_response.headers = resp.get("headers", {})
-        body = resp.get("body", b'{}')
+        body = resp.get("body", b"{}")
         mock_http_response.fp = io.BytesIO(body)
+
         def release():
             mock_http_response.fp.close()
+
         mock_http_response.release_conn = release
         mock_http_responses.append(mock_http_response)
     mock_conn.getresponse.side_effect = mock_http_responses
     return mock_conn
 
+
 class TestTelemetryClientRetries:
     @pytest.fixture(autouse=True)
     def setup_and_teardown(self):
@@ -49,28 +53,28 @@ def get_client(self, session_id, num_retries=3):
             host_url="test.databricks.com",
         )
         client = TelemetryClientFactory.get_telemetry_client(session_id)
-        
+
         retry_policy = DatabricksRetryPolicy(
             delay_min=0.01,
             delay_max=0.02,
             stop_after_attempts_duration=2.0,
-            stop_after_attempts_count=num_retries, 
+            stop_after_attempts_count=num_retries,
             delay_default=0.1,
             force_dangerous_codes=[],
-            urllib3_kwargs={'total': num_retries}
+            urllib3_kwargs={"total": num_retries},
         )
         adapter = client._http_client.session.adapters.get("https://")
         adapter.max_retries = retry_policy
         return client
 
     @pytest.mark.parametrize(
-    "status_code, description",
-    [
-        (401, "Unauthorized"),
-        (403, "Forbidden"),
-        (501, "Not Implemented"),
-        (200, "Success"),
-    ],
+        "status_code, description",
+        [
+            (401, "Unauthorized"),
+            (403, "Forbidden"),
+            (501, "Not Implemented"),
+            (200, "Success"),
+        ],
     )
     def test_non_retryable_status_codes_are_not_retried(self, status_code, description):
         """
@@ -80,7 +84,9 @@ def test_non_retryable_status_codes_are_not_retried(self, status_code, descripti
         client = self.get_client(f"session-{status_code}")
         mock_responses = [{"status": status_code}]
 
-        with patch(PATCH_TARGET, return_value=create_mock_conn(mock_responses)) as mock_get_conn:
+        with patch(
+            PATCH_TARGET, return_value=create_mock_conn(mock_responses)
+        ) as mock_get_conn:
             client.export_failure_log("TestError", "Test message")
             TelemetryClientFactory.close(client._session_id_hex)
 
@@ -92,16 +98,26 @@ def test_exceeds_retry_count_limit(self):
         Verifies that the client respects the Retry-After header and retries on 429, 502, 503.
         """
         num_retries = 3
-        expected_total_calls = num_retries + 1 
+        expected_total_calls = num_retries + 1
         retry_after = 1
         client = self.get_client("session-exceed-limit", num_retries=num_retries)
-        mock_responses = [{"status": 503, "headers": {"Retry-After": str(retry_after)}}, {"status": 429}, {"status": 502}, {"status": 503}]
-        
-        with patch(PATCH_TARGET, return_value=create_mock_conn(mock_responses)) as mock_get_conn:
+        mock_responses = [
+            {"status": 503, "headers": {"Retry-After": str(retry_after)}},
+            {"status": 429},
+            {"status": 502},
+            {"status": 503},
+        ]
+
+        with patch(
+            PATCH_TARGET, return_value=create_mock_conn(mock_responses)
+        ) as mock_get_conn:
             start_time = time.time()
             client.export_failure_log("TestError", "Test message")
             TelemetryClientFactory.close(client._session_id_hex)
             end_time = time.time()
-            
-            assert mock_get_conn.return_value.getresponse.call_count == expected_total_calls
-            assert end_time - start_time > retry_after
+
+            assert (
+                mock_get_conn.return_value.getresponse.call_count
+                == expected_total_calls
+            )
+            assert end_time - start_time > retry_after