SQLAlchemy: TINYINT types didn't reflect properly (databricks#315)

Signed-off-by: Jesse Whitehouse <[email protected]>

Author: TimTheinAtTabs

CHANGELOG.md

@@ -2,7 +2,8 @@
 
 # 3.1.0 (TBD)
 
-- SQLAlchemy: Added support for table and column comments (thanks @cbornet!)
+- SQLAlchemy dialect now supports table and column comments (thanks @cbornet!)
+- Fix: SQLAlchemy dialect now correctly reflects TINYINT types (thanks @TimTheinAtTabs!)
 - Fix: `server_hostname` URIs that included `https://` would raise an exception
 - Other: pinned to `pandas<=2.1` and `urllib3>=1.26` to avoid runtime errors in dbt-databricks (#330)
 

src/databricks/sqlalchemy/_parse.py

@@ -305,6 +305,7 @@ def get_comment_from_dte_output(dte_output: List[Dict[str, str]]) -> Optional[st
 GET_COLUMNS_TYPE_MAP = {
     "boolean": sqlalchemy.types.Boolean,
     "smallint": sqlalchemy.types.SmallInteger,
+    "tinyint": type_overrides.TINYINT,
     "int": sqlalchemy.types.Integer,
     "bigint": sqlalchemy.types.BigInteger,
     "float": sqlalchemy.types.Float,

src/databricks/sqlalchemy/_types.py

@@ -316,3 +316,7 @@ class TINYINT(sqlalchemy.types.TypeDecorator):
 
     impl = sqlalchemy.types.SmallInteger
     cache_ok = True
+
+@compiles(TINYINT, "databricks")
+def compile_tinyint(type_, compiler, **kw):
+    return "TINYINT"

src/databricks/sqlalchemy/test_local/test_types.py

@@ -111,7 +111,7 @@ def test_numeric_renders_as_decimal_with_precision(self):
         )
 
     def test_numeric_renders_as_decimal_with_precision_and_scale(self):
-        return self._assert_compiled_value_explicit(
+        self._assert_compiled_value_explicit(
             sqlalchemy.types.Numeric(10, 2), "DECIMAL(10, 2)"
         )
 
@@ -146,15 +146,16 @@ def test_bare_uppercase_types_compile(self, type_, expected):
         if isinstance(type_, type(sqlalchemy.types.ARRAY)):
             # ARRAY cannot be initialised without passing an item definition so we test separately
             # I preserve it in the uppercase_type_map for clarity
-            return True
-        return self._assert_compiled_value(type_, expected)
+            assert True
+        else:
+            self._assert_compiled_value(type_, expected)
 
     def test_array_string_renders_as_array_of_string(self):
         """SQLAlchemy's ARRAY type requires an item definition. And their docs indicate that they've only tested
         it with Postgres since that's the only first-class dialect with support for ARRAY.
 
         https://docs.sqlalchemy.org/en/20/core/type_basics.html#sqlalchemy.types.ARRAY
         """
-        return self._assert_compiled_value_explicit(
+        self._assert_compiled_value_explicit(
             sqlalchemy.types.ARRAY(sqlalchemy.types.String), "ARRAY<STRING>"
         )