From c71e081c8d0d43a8908d1fb5fbc65fd21e94c50c Mon Sep 17 00:00:00 2001
From: TimTheinAtTabs <100803283+TimTheinAtTabs@users.noreply.github.com>
Date: Thu, 25 Jan 2024 13:07:58 -0800
Subject: [PATCH] SQLAlchemy: TINYINT types didn't reflect properly (#315)

Signed-off-by: Jesse Whitehouse <jesse.whitehouse@databricks.com>
---
 CHANGELOG.md                                       | 3 ++-
 src/databricks/sqlalchemy/_parse.py                | 1 +
 src/databricks/sqlalchemy/_types.py                | 4 ++++
 src/databricks/sqlalchemy/test_local/test_types.py | 9 +++++----
 4 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f5356fdd..53f49fa2 100644
--- a/CHANGELOG.md
+++ b/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)
 
diff --git a/src/databricks/sqlalchemy/_parse.py b/src/databricks/sqlalchemy/_parse.py
index a80f37bb..6d38e1e6 100644
--- a/src/databricks/sqlalchemy/_parse.py
+++ b/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,
diff --git a/src/databricks/sqlalchemy/_types.py b/src/databricks/sqlalchemy/_types.py
index 1dc6d9e9..79f3d626 100644
--- a/src/databricks/sqlalchemy/_types.py
+++ b/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"
\ No newline at end of file
diff --git a/src/databricks/sqlalchemy/test_local/test_types.py b/src/databricks/sqlalchemy/test_local/test_types.py
index 73e28669..b91217ed 100644
--- a/src/databricks/sqlalchemy/test_local/test_types.py
+++ b/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,8 +146,9 @@ 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
@@ -155,6 +156,6 @@ def test_array_string_renders_as_array_of_string(self):
 
         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>"
         )