SQLAlchemy 2: Add support for TINYINT (#265)

Closes #123

Signed-off-by: Jesse Whitehouse <jesse.whitehouse@databricks.com>
Signed-off-by: Sai Shree Pradhan <saishree.pradhan@databricks.com>

Author: Jesse

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
 
 ```

src/databricks/sqlalchemy/__init__.py

@@ -1 +1,4 @@
 from databricks.sqlalchemy.base import DatabricksDialect
+from databricks.sqlalchemy._types import TINYINT
+
+__all__ = ["TINYINT"]

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

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

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

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,
 }