From 8198a292d9d2983ace8bba7dfd9090447b649858 Mon Sep 17 00:00:00 2001
From: Jesse Whitehouse <jesse.whitehouse@databricks.com>
Date: Sat, 28 Oct 2023 11:45:02 -0400
Subject: [PATCH] SQLAlchemy 2: Add support for TINYINT

Closes #123

Signed-off-by: Jesse Whitehouse <jesse.whitehouse@databricks.com>
---
 src/databricks/sqlalchemy/README.tests.md     |  2 +
 src/databricks/sqlalchemy/__init__.py         |  3 ++
 src/databricks/sqlalchemy/_types.py           | 12 +++++
 src/databricks/sqlalchemy/test/_extra.py      | 48 +++++++++++++++++++
 src/databricks/sqlalchemy/test/test_suite.py  |  1 +
 .../sqlalchemy/test_local/test_types.py       |  2 +
 6 files changed, 68 insertions(+)
 create mode 100644 src/databricks/sqlalchemy/test/_extra.py

diff --git a/src/databricks/sqlalchemy/README.tests.md b/src/databricks/sqlalchemy/README.tests.md
index 52ad17b9b..d005daa3c 100644
--- a/src/databricks/sqlalchemy/README.tests.md
+++ b/src/databricks/sqlalchemy/README.tests.md
@@ -17,6 +17,8 @@ We maintain three files of test cases that we import from the SQLAlchemy source
 
 In some cases, only certain tests in class should be skipped with a `SkipReason` or `FutureFeature` justification. In those cases, we import the class into `_regression.py`, then import it from there into one or both of `_future.py` and `_unsupported.py`. If a class needs to be "touched" by regression, unsupported, and future, the class will be imported in that order. If an entire class should be skipped, then we do not import it into `_regression.py` at all.
 
+We maintain `_extra.py` with test cases that depend on SQLAlchemy's reusable dialect test fixtures but which are specific to Databricks (e.g TinyIntegerTest).
+
 ## Running the reusable dialect tests
 
 ```
diff --git a/src/databricks/sqlalchemy/__init__.py b/src/databricks/sqlalchemy/__init__.py
index e1c0d7554..25584506f 100644
--- a/src/databricks/sqlalchemy/__init__.py
+++ b/src/databricks/sqlalchemy/__init__.py
@@ -1 +1,4 @@
 from databricks.sqlalchemy.base import DatabricksDialect
+from databricks.sqlalchemy._types import TINYINT
+
+__all__ = ["TINYINT"]
diff --git a/src/databricks/sqlalchemy/_types.py b/src/databricks/sqlalchemy/_types.py
index d2ea8c083..133abb299 100644
--- a/src/databricks/sqlalchemy/_types.py
+++ b/src/databricks/sqlalchemy/_types.py
@@ -212,3 +212,15 @@ def process(value):
             return "%s" % _step2
 
         return process
+
+
+class TINYINT(sqlalchemy.types.TypeDecorator):
+    """Represents 1-byte signed integers
+
+    Acts like a sqlalchemy SmallInteger() in Python but writes to a TINYINT field in Databricks
+
+    https://docs.databricks.com/en/sql/language-manual/data-types/tinyint-type.html
+    """
+
+    impl = sqlalchemy.types.SmallInteger
+    cache_ok = True
diff --git a/src/databricks/sqlalchemy/test/_extra.py b/src/databricks/sqlalchemy/test/_extra.py
new file mode 100644
index 000000000..f8e11bde6
--- /dev/null
+++ b/src/databricks/sqlalchemy/test/_extra.py
@@ -0,0 +1,48 @@
+"""Additional tests authored by Databricks that use SQLAlchemy's test fixtures
+"""
+
+from sqlalchemy.testing.suite.test_types import (
+    _LiteralRoundTripFixture,
+    fixtures,
+    testing,
+    eq_,
+    select,
+    Table,
+    Column,
+    config,
+)
+from databricks.sqlalchemy import TINYINT
+
+
+class TinyIntegerTest(_LiteralRoundTripFixture, fixtures.TestBase):
+    __backend__ = True
+
+    def test_literal(self, literal_round_trip):
+        literal_round_trip(TINYINT, [5], [5])
+
+    @testing.fixture
+    def integer_round_trip(self, metadata, connection):
+        def run(datatype, data):
+            int_table = Table(
+                "tiny_integer_table",
+                metadata,
+                Column(
+                    "id",
+                    TINYINT,
+                    primary_key=True,
+                    test_needs_autoincrement=False,
+                ),
+                Column("tiny_integer_data", datatype),
+            )
+
+            metadata.create_all(config.db)
+
+            connection.execute(int_table.insert(), {"id": 1, "integer_data": data})
+
+            row = connection.execute(select(int_table.c.integer_data)).first()
+
+            eq_(row, (data,))
+
+            assert isinstance(row[0], int)
+
+        return run
diff --git a/src/databricks/sqlalchemy/test/test_suite.py b/src/databricks/sqlalchemy/test/test_suite.py
index e96c6a063..a38be2e48 100644
--- a/src/databricks/sqlalchemy/test/test_suite.py
+++ b/src/databricks/sqlalchemy/test/test_suite.py
@@ -9,3 +9,4 @@
 from databricks.sqlalchemy.test._regression import *
 from databricks.sqlalchemy.test._unsupported import *
 from databricks.sqlalchemy.test._future import *
+from databricks.sqlalchemy.test._extra import TinyIntegerTest
diff --git a/src/databricks/sqlalchemy/test_local/test_types.py b/src/databricks/sqlalchemy/test_local/test_types.py
index fb66562a7..c29edfcec 100644
--- a/src/databricks/sqlalchemy/test_local/test_types.py
+++ b/src/databricks/sqlalchemy/test_local/test_types.py
@@ -4,6 +4,7 @@
 import sqlalchemy
 
 from databricks.sqlalchemy.base import DatabricksDialect
+from databricks.sqlalchemy._types import TINYINT
 
 
 class DatabricksDataType(enum.Enum):
@@ -127,6 +128,7 @@ def test_numeric_renders_as_decimal_with_precision_and_scale(self):
     sqlalchemy.types.INT: DatabricksDataType.INT,
     sqlalchemy.types.SMALLINT: DatabricksDataType.SMALLINT,
     sqlalchemy.types.TIMESTAMP: DatabricksDataType.TIMESTAMP,
+    TINYINT: DatabricksDataType.TINYINT,
 }