align expected types with databricks sdk

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

Author: varun-edachali-dbx

src/databricks/sql/conversion.py …databricks/sql/backend/sea/conversion.pysrc/databricks/sql/conversion.py renamed to src/databricks/sql/backend/sea/conversion.py src/databricks/sql/conversion.py renamed to src/databricks/sql/backend/sea/conversion.py

@@ -15,89 +15,75 @@
 
 
 class SqlType:
-    """SQL type constants for improved maintainability."""
+    """
+    SQL type constants
+    """
 
     # Numeric types
-    TINYINT = "tinyint"
-    SMALLINT = "smallint"
+    BYTE = "byte"
+    SHORT = "short"
     INT = "int"
-    INTEGER = "integer"
-    BIGINT = "bigint"
+    LONG = "long"
     FLOAT = "float"
-    REAL = "real"
     DOUBLE = "double"
     DECIMAL = "decimal"
-    NUMERIC = "numeric"
 
-    # Boolean types
+    # Boolean type
     BOOLEAN = "boolean"
-    BIT = "bit"
 
     # Date/Time types
     DATE = "date"
-    TIME = "time"
     TIMESTAMP = "timestamp"
-    TIMESTAMP_NTZ = "timestamp_ntz"
-    TIMESTAMP_LTZ = "timestamp_ltz"
-    TIMESTAMP_TZ = "timestamp_tz"
+    INTERVAL = "interval"
 
     # String types
     CHAR = "char"
-    VARCHAR = "varchar"
     STRING = "string"
-    TEXT = "text"
 
-    # Binary types
+    # Binary type
     BINARY = "binary"
-    VARBINARY = "varbinary"
 
     # Complex types
     ARRAY = "array"
     MAP = "map"
     STRUCT = "struct"
 
+    # Other types
+    NULL = "null"
+    USER_DEFINED_TYPE = "user_defined_type"
+
     @classmethod
     def is_numeric(cls, sql_type: str) -> bool:
         """Check if the SQL type is a numeric type."""
         return sql_type.lower() in (
-            cls.TINYINT,
-            cls.SMALLINT,
+            cls.BYTE,
+            cls.SHORT,
             cls.INT,
-            cls.INTEGER,
-            cls.BIGINT,
+            cls.LONG,
             cls.FLOAT,
-            cls.REAL,
             cls.DOUBLE,
             cls.DECIMAL,
-            cls.NUMERIC,
         )
 
     @classmethod
     def is_boolean(cls, sql_type: str) -> bool:
         """Check if the SQL type is a boolean type."""
-        return sql_type.lower() in (cls.BOOLEAN, cls.BIT)
+        return sql_type.lower() == cls.BOOLEAN
 
     @classmethod
     def is_datetime(cls, sql_type: str) -> bool:
         """Check if the SQL type is a date/time type."""
-        return sql_type.lower() in (
-            cls.DATE,
-            cls.TIME,
-            cls.TIMESTAMP,
-            cls.TIMESTAMP_NTZ,
-            cls.TIMESTAMP_LTZ,
-            cls.TIMESTAMP_TZ,
-        )
+        return sql_type.lower() in (cls.DATE, cls.TIMESTAMP, cls.INTERVAL)
 
     @classmethod
     def is_string(cls, sql_type: str) -> bool:
         """Check if the SQL type is a string type."""
-        return sql_type.lower() in (cls.CHAR, cls.VARCHAR, cls.STRING, cls.TEXT)
+        return sql_type.lower() in (cls.CHAR, cls.STRING)
 
     @classmethod
     def is_binary(cls, sql_type: str) -> bool:
         """Check if the SQL type is a binary type."""
-        return sql_type.lower() in (cls.BINARY, cls.VARBINARY)
+        return sql_type.lower() == cls.BINARY
 
     @classmethod
     def is_complex(cls, sql_type: str) -> bool:
@@ -107,25 +93,25 @@ def is_complex(cls, sql_type: str) -> bool:
             sql_type.startswith(cls.ARRAY)
             or sql_type.startswith(cls.MAP)
             or sql_type.startswith(cls.STRUCT)
+            or sql_type == cls.USER_DEFINED_TYPE
         )
 
 
 class SqlTypeConverter:
     """
     Utility class for converting SQL types to Python types.
-    Based on the JDBC ConverterHelper implementation.
+    Based on the types supported by the Databricks SDK.
     """
 
     # SQL type to conversion function mapping
+    # TODO: complex types
     TYPE_MAPPING: Dict[str, Callable] = {
         # Numeric types
-        SqlType.TINYINT: lambda v: int(v),
-        SqlType.SMALLINT: lambda v: int(v),
+        SqlType.BYTE: lambda v: int(v),
+        SqlType.SHORT: lambda v: int(v),
         SqlType.INT: lambda v: int(v),
-        SqlType.INTEGER: lambda v: int(v),
-        SqlType.BIGINT: lambda v: int(v),
+        SqlType.LONG: lambda v: int(v),
         SqlType.FLOAT: lambda v: float(v),
-        SqlType.REAL: lambda v: float(v),
         SqlType.DOUBLE: lambda v: float(v),
         SqlType.DECIMAL: lambda v, p=None, s=None: (
             decimal.Decimal(v).quantize(
@@ -134,31 +120,21 @@ class SqlTypeConverter:
             if p is not None and s is not None
             else decimal.Decimal(v)
         ),
-        SqlType.NUMERIC: lambda v, p=None, s=None: (
-            decimal.Decimal(v).quantize(
-                decimal.Decimal(f'0.{"0" * s}'), context=decimal.Context(prec=p)
-            )
-            if p is not None and s is not None
-            else decimal.Decimal(v)
-        ),
-        # Boolean types
+        # Boolean type
         SqlType.BOOLEAN: lambda v: v.lower() in ("true", "t", "1", "yes", "y"),
-        SqlType.BIT: lambda v: v.lower() in ("true", "t", "1", "yes", "y"),
         # Date/Time types
         SqlType.DATE: lambda v: datetime.date.fromisoformat(v),
-        SqlType.TIME: lambda v: datetime.time.fromisoformat(v),
         SqlType.TIMESTAMP: lambda v: parser.parse(v),
-        SqlType.TIMESTAMP_NTZ: lambda v: parser.parse(v).replace(tzinfo=None),
-        SqlType.TIMESTAMP_LTZ: lambda v: parser.parse(v).astimezone(tz=None),
-        SqlType.TIMESTAMP_TZ: lambda v: parser.parse(v),
+        SqlType.INTERVAL: 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,
-        SqlType.TEXT: lambda v: v,
-        # Binary types
+        # Binary type
         SqlType.BINARY: lambda v: bytes.fromhex(v),
-        SqlType.VARBINARY: 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,
     }
 
     @staticmethod
@@ -180,6 +156,7 @@ def convert_value(
         Returns:
             The converted value in the appropriate Python type
         """
+
         if value is None:
             return None
 
@@ -190,7 +167,7 @@ def convert_value(
 
         converter_func = SqlTypeConverter.TYPE_MAPPING[sql_type]
         try:
-            if sql_type in (SqlType.DECIMAL, SqlType.NUMERIC):
+            if sql_type == SqlType.DECIMAL:
                 return converter_func(value, precision, scale)
             else:
                 return converter_func(value)

src/databricks/sql/result_set.py

@@ -6,7 +6,7 @@
 
 from databricks.sql.backend.sea.backend import SeaDatabricksClient
 from databricks.sql.backend.sea.models.base import ResultData, ResultManifest
-from databricks.sql.conversion import SqlTypeConverter
+from databricks.sql.backend.sea.conversion import SqlTypeConverter
 
 try:
     import pyarrow

tests/unit/test_type_conversion.py tests/unit/test_conversion.pytests/unit/test_type_conversion.py renamed to tests/unit/test_conversion.py

@@ -6,7 +6,7 @@
 from datetime import date, datetime, time
 from decimal import Decimal
 
-from databricks.sql.conversion import SqlType, SqlTypeConverter
+from databricks.sql.backend.sea.conversion import SqlType, SqlTypeConverter
 
 
 class TestSqlType(unittest.TestCase):
@@ -15,48 +15,40 @@ class TestSqlType(unittest.TestCase):
     def test_is_numeric(self):
         """Test the is_numeric method."""
         self.assertTrue(SqlType.is_numeric(SqlType.INT))
-        self.assertTrue(SqlType.is_numeric(SqlType.TINYINT))
-        self.assertTrue(SqlType.is_numeric(SqlType.SMALLINT))
-        self.assertTrue(SqlType.is_numeric(SqlType.BIGINT))
+        self.assertTrue(SqlType.is_numeric(SqlType.BYTE))
+        self.assertTrue(SqlType.is_numeric(SqlType.SHORT))
+        self.assertTrue(SqlType.is_numeric(SqlType.LONG))
         self.assertTrue(SqlType.is_numeric(SqlType.FLOAT))
         self.assertTrue(SqlType.is_numeric(SqlType.DOUBLE))
         self.assertTrue(SqlType.is_numeric(SqlType.DECIMAL))
-        self.assertTrue(SqlType.is_numeric(SqlType.NUMERIC))
         self.assertFalse(SqlType.is_numeric(SqlType.BOOLEAN))
         self.assertFalse(SqlType.is_numeric(SqlType.STRING))
         self.assertFalse(SqlType.is_numeric(SqlType.DATE))
 
     def test_is_boolean(self):
         """Test the is_boolean method."""
         self.assertTrue(SqlType.is_boolean(SqlType.BOOLEAN))
-        self.assertTrue(SqlType.is_boolean(SqlType.BIT))
         self.assertFalse(SqlType.is_boolean(SqlType.INT))
         self.assertFalse(SqlType.is_boolean(SqlType.STRING))
 
     def test_is_datetime(self):
         """Test the is_datetime method."""
         self.assertTrue(SqlType.is_datetime(SqlType.DATE))
-        self.assertTrue(SqlType.is_datetime(SqlType.TIME))
         self.assertTrue(SqlType.is_datetime(SqlType.TIMESTAMP))
-        self.assertTrue(SqlType.is_datetime(SqlType.TIMESTAMP_NTZ))
-        self.assertTrue(SqlType.is_datetime(SqlType.TIMESTAMP_LTZ))
-        self.assertTrue(SqlType.is_datetime(SqlType.TIMESTAMP_TZ))
+        self.assertTrue(SqlType.is_datetime(SqlType.INTERVAL))
         self.assertFalse(SqlType.is_datetime(SqlType.INT))
         self.assertFalse(SqlType.is_datetime(SqlType.STRING))
 
     def test_is_string(self):
         """Test the is_string method."""
         self.assertTrue(SqlType.is_string(SqlType.CHAR))
-        self.assertTrue(SqlType.is_string(SqlType.VARCHAR))
         self.assertTrue(SqlType.is_string(SqlType.STRING))
-        self.assertTrue(SqlType.is_string(SqlType.TEXT))
         self.assertFalse(SqlType.is_string(SqlType.INT))
         self.assertFalse(SqlType.is_string(SqlType.DATE))
 
     def test_is_binary(self):
         """Test the is_binary method."""
         self.assertTrue(SqlType.is_binary(SqlType.BINARY))
-        self.assertTrue(SqlType.is_binary(SqlType.VARBINARY))
         self.assertFalse(SqlType.is_binary(SqlType.INT))
         self.assertFalse(SqlType.is_binary(SqlType.STRING))
 
@@ -75,9 +67,9 @@ class TestSqlTypeConverter(unittest.TestCase):
     def test_numeric_conversions(self):
         """Test numeric type conversions."""
         self.assertEqual(SqlTypeConverter.convert_value("123", SqlType.INT), 123)
-        self.assertEqual(SqlTypeConverter.convert_value("123", SqlType.TINYINT), 123)
-        self.assertEqual(SqlTypeConverter.convert_value("123", SqlType.SMALLINT), 123)
-        self.assertEqual(SqlTypeConverter.convert_value("123", SqlType.BIGINT), 123)
+        self.assertEqual(SqlTypeConverter.convert_value("123", SqlType.BYTE), 123)
+        self.assertEqual(SqlTypeConverter.convert_value("123", SqlType.SHORT), 123)
+        self.assertEqual(SqlTypeConverter.convert_value("123", SqlType.LONG), 123)
         self.assertEqual(
             SqlTypeConverter.convert_value("123.45", SqlType.FLOAT), 123.45
         )
@@ -113,9 +105,6 @@ def test_datetime_conversions(self):
             SqlTypeConverter.convert_value("2023-01-15", SqlType.DATE),
             date(2023, 1, 15),
         )
-        self.assertEqual(
-            SqlTypeConverter.convert_value("14:30:45", SqlType.TIME), time(14, 30, 45)
-        )
         self.assertEqual(
             SqlTypeConverter.convert_value("2023-01-15 14:30:45", SqlType.TIMESTAMP),
             datetime(2023, 1, 15, 14, 30, 45),
@@ -124,15 +113,19 @@ def test_datetime_conversions(self):
     def test_string_conversions(self):
         """Test string type conversions."""
         self.assertEqual(SqlTypeConverter.convert_value("test", SqlType.STRING), "test")
+        self.assertEqual(SqlTypeConverter.convert_value("test", SqlType.CHAR), "test")
+
+    def test_binary_conversions(self):
+        """Test binary type conversions."""
+        hex_str = "68656c6c6f"  # "hello" in hex
+        expected_bytes = b"hello"
+
         self.assertEqual(
-            SqlTypeConverter.convert_value("test", SqlType.VARCHAR), "test"
+            SqlTypeConverter.convert_value(hex_str, SqlType.BINARY), expected_bytes
         )
-        self.assertEqual(SqlTypeConverter.convert_value("test", SqlType.CHAR), "test")
-        self.assertEqual(SqlTypeConverter.convert_value("test", SqlType.TEXT), "test")
 
     def test_error_handling(self):
         """Test error handling in conversions."""
-        # Test invalid conversions - should return original value
         self.assertEqual(SqlTypeConverter.convert_value("abc", SqlType.INT), "abc")
         self.assertEqual(SqlTypeConverter.convert_value("abc", SqlType.FLOAT), "abc")
         self.assertEqual(SqlTypeConverter.convert_value("abc", SqlType.DECIMAL), "abc")
@@ -155,6 +148,10 @@ def test_complex_type_handling(self):
         self.assertEqual(
             SqlTypeConverter.convert_value('{"a": 1}', "struct<a:int>"), '{"a": 1}'
         )
+        self.assertEqual(
+            SqlTypeConverter.convert_value('{"a": 1}', SqlType.USER_DEFINED_TYPE),
+            '{"a": 1}',
+        )
 
 
 if __name__ == "__main__":