databricks-sqlalchemy 1.0.2__py3-none-any.whl → 2.0.1__py3-none-any.whl

databricks/sqlalchemy/_parse.py

@@ -0,0 +1,385 @@
+from typing import List, Optional, Dict
+import re
+
+import sqlalchemy
+from sqlalchemy.engine import CursorResult
+from sqlalchemy.engine.interfaces import ReflectedColumn
+
+from databricks.sqlalchemy import _types as type_overrides
+
+"""
+This module contains helper functions that can parse the contents
+of metadata and exceptions received from DBR. These are mostly just
+wrappers around regexes.
+"""
+
+
+class DatabricksSqlAlchemyParseException(Exception):
+    pass
+
+
+def _match_table_not_found_string(message: str) -> bool:
+    """Return True if the message contains a substring indicating that a table was not found"""
+
+    DBR_LTE_12_NOT_FOUND_STRING = "Table or view not found"
+    DBR_GT_12_NOT_FOUND_STRING = "TABLE_OR_VIEW_NOT_FOUND"
+    return any(
+        [
+            DBR_LTE_12_NOT_FOUND_STRING in message,
+            DBR_GT_12_NOT_FOUND_STRING in message,
+        ]
+    )
+
+
+def _describe_table_extended_result_to_dict_list(
+    result: CursorResult,
+) -> List[Dict[str, str]]:
+    """Transform the CursorResult of DESCRIBE TABLE EXTENDED into a list of Dictionaries"""
+
+    rows_to_return = []
+    for row in result.all():
+        this_row = {"col_name": row.col_name, "data_type": row.data_type}
+        rows_to_return.append(this_row)
+
+    return rows_to_return
+
+
+def extract_identifiers_from_string(input_str: str) -> List[str]:
+    """For a string input resembling (`a`, `b`, `c`) return a list of identifiers ['a', 'b', 'c']"""
+
+    # This matches the valid character list contained in DatabricksIdentifierPreparer
+    pattern = re.compile(r"`([A-Za-z0-9_]+)`")
+    matches = pattern.findall(input_str)
+    return [i for i in matches]
+
+
+def extract_identifier_groups_from_string(input_str: str) -> List[str]:
+    """For a string input resembling :
+
+    FOREIGN KEY (`pname`, `pid`, `pattr`) REFERENCES `main`.`pysql_sqlalchemy`.`tb1` (`name`, `id`, `attr`)
+
+    Return ['(`pname`, `pid`, `pattr`)', '(`name`, `id`, `attr`)']
+    """
+    pattern = re.compile(r"\([`A-Za-z0-9_,\s]*\)")
+    matches = pattern.findall(input_str)
+    return [i for i in matches]
+
+
+def extract_three_level_identifier_from_constraint_string(input_str: str) -> dict:
+    """For a string input resembling :
+    FOREIGN KEY (`parent_user_id`) REFERENCES `main`.`pysql_dialect_compliance`.`users` (`user_id`)
+
+    Return a dict like
+        {
+            "catalog": "main",
+            "schema": "pysql_dialect_compliance",
+            "table": "users"
+        }
+
+    Raise a DatabricksSqlAlchemyParseException if a 3L namespace isn't found
+    """
+    pat = re.compile(r"REFERENCES\s+(.*?)\s*\(")
+    matches = pat.findall(input_str)
+
+    if not matches:
+        raise DatabricksSqlAlchemyParseException(
+            "3L namespace not found in constraint string"
+        )
+
+    first_match = matches[0]
+    parts = first_match.split(".")
+
+    def strip_backticks(input: str):
+        return input.replace("`", "")
+
+    try:
+        return {
+            "catalog": strip_backticks(parts[0]),
+            "schema": strip_backticks(parts[1]),
+            "table": strip_backticks(parts[2]),
+        }
+    except IndexError:
+        raise DatabricksSqlAlchemyParseException(
+            "Incomplete 3L namespace found in constraint string: " + ".".join(parts)
+        )
+
+
+def _parse_fk_from_constraint_string(constraint_str: str) -> dict:
+    """Build a dictionary of foreign key constraint information from a constraint string.
+
+    For example:
+
+    ```
+    FOREIGN KEY (`pname`, `pid`, `pattr`) REFERENCES `main`.`pysql_dialect_compliance`.`tb1` (`name`, `id`, `attr`)
+    ```
+
+    Return a dictionary like:
+
+    ```
+    {
+        "constrained_columns": ["pname", "pid", "pattr"],
+        "referred_table": "tb1",
+        "referred_schema": "pysql_dialect_compliance",
+        "referred_columns": ["name", "id", "attr"]
+    }
+    ```
+
+    Note that the constraint name doesn't appear in the constraint string so it will not
+    be present in the output of this function.
+    """
+
+    referred_table_dict = extract_three_level_identifier_from_constraint_string(
+        constraint_str
+    )
+    referred_table = referred_table_dict["table"]
+    referred_schema = referred_table_dict["schema"]
+
+    # _extracted is a tuple of two lists of identifiers
+    # we assume the first immediately follows "FOREIGN KEY" and the second
+    # immediately follows REFERENCES $tableName
+    _extracted = extract_identifier_groups_from_string(constraint_str)
+    constrained_columns_str, referred_columns_str = (
+        _extracted[0],
+        _extracted[1],
+    )
+
+    constrained_columns = extract_identifiers_from_string(constrained_columns_str)
+    referred_columns = extract_identifiers_from_string(referred_columns_str)
+
+    return {
+        "constrained_columns": constrained_columns,
+        "referred_table": referred_table,
+        "referred_columns": referred_columns,
+        "referred_schema": referred_schema,
+    }
+
+
+def build_fk_dict(
+    fk_name: str, fk_constraint_string: str, schema_name: Optional[str]
+) -> dict:
+    """
+    Given a foriegn key name and a foreign key constraint string, return a dictionary
+    with the following keys:
+
+    name
+        the name of the foreign key constraint
+    constrained_columns
+        a list of column names that make up the foreign key
+    referred_table
+        the name of the table that the foreign key references
+    referred_columns
+        a list of column names that are referenced by the foreign key
+    referred_schema
+        the name of the schema that the foreign key references.
+
+    referred schema will be None if the schema_name argument is None.
+    This is required by SQLAlchey's ComponentReflectionTest::test_get_foreign_keys
+    """
+
+    # The foreign key name is not contained in the constraint string so we
+    # need to add it manually
+    base_fk_dict = _parse_fk_from_constraint_string(fk_constraint_string)
+
+    if not schema_name:
+        schema_override_dict = dict(referred_schema=None)
+    else:
+        schema_override_dict = {}
+
+    # mypy doesn't like this method of conditionally adding a key to a dictionary
+    # while keeping everything immutable
+    complete_foreign_key_dict = {
+        "name": fk_name,
+        **base_fk_dict,
+        **schema_override_dict,  # type: ignore
+    }
+
+    return complete_foreign_key_dict
+
+
+def _parse_pk_columns_from_constraint_string(constraint_str: str) -> List[str]:
+    """Build a list of constrained columns from a constraint string returned by DESCRIBE TABLE EXTENDED
+
+    For example:
+
+    PRIMARY KEY (`id`, `name`, `email_address`)
+
+    Returns a list like
+
+    ["id", "name", "email_address"]
+    """
+
+    _extracted = extract_identifiers_from_string(constraint_str)
+
+    return _extracted
+
+
+def build_pk_dict(pk_name: str, pk_constraint_string: str) -> dict:
+    """Given a primary key name and a primary key constraint string, return a dictionary
+    with the following keys:
+
+    constrained_columns
+      A list of string column names that make up the primary key
+
+    name
+      The name of the primary key constraint
+    """
+
+    constrained_columns = _parse_pk_columns_from_constraint_string(pk_constraint_string)
+
+    return {"constrained_columns": constrained_columns, "name": pk_name}
+
+
+def match_dte_rows_by_value(dte_output: List[Dict[str, str]], match: str) -> List[dict]:
+    """Return a list of dictionaries containing only the col_name:data_type pairs where the `data_type`
+    value contains the match argument.
+
+    Today, DESCRIBE TABLE EXTENDED doesn't give a deterministic name to the fields
+    a constraint will be found in its output. So we cycle through its output looking
+    for a match. This is brittle. We could optionally make two roundtrips: the first
+    would query information_schema for the name of the constraint on this table, and
+    a second to DESCRIBE TABLE EXTENDED, at which point we would know the name of the
+    constraint. But for now we instead assume that Python list comprehension is faster
+    than a network roundtrip
+    """
+
+    output_rows = []
+
+    for row_dict in dte_output:
+        if match in row_dict["data_type"]:
+            output_rows.append(row_dict)
+
+    return output_rows
+
+
+def match_dte_rows_by_key(dte_output: List[Dict[str, str]], match: str) -> List[dict]:
+    """Return a list of dictionaries containing only the col_name:data_type pairs where the `col_name`
+    value contains the match argument.
+    """
+
+    output_rows = []
+
+    for row_dict in dte_output:
+        if match in row_dict["col_name"]:
+            output_rows.append(row_dict)
+
+    return output_rows
+
+
+def get_fk_strings_from_dte_output(dte_output: List[Dict[str, str]]) -> List[dict]:
+    """If the DESCRIBE TABLE EXTENDED output contains foreign key constraints, return a list of dictionaries,
+    one dictionary per defined constraint
+    """
+
+    output = match_dte_rows_by_value(dte_output, "FOREIGN KEY")
+
+    return output
+
+
+def get_pk_strings_from_dte_output(
+    dte_output: List[Dict[str, str]]
+) -> Optional[List[dict]]:
+    """If the DESCRIBE TABLE EXTENDED output contains primary key constraints, return a list of dictionaries,
+    one dictionary per defined constraint.
+
+    Returns None if no primary key constraints are found.
+    """
+
+    output = match_dte_rows_by_value(dte_output, "PRIMARY KEY")
+
+    return output
+
+
+def get_comment_from_dte_output(dte_output: List[Dict[str, str]]) -> Optional[str]:
+    """Returns the value of the first "Comment" col_name data in dte_output"""
+    output = match_dte_rows_by_key(dte_output, "Comment")
+    if not output:
+        return None
+    else:
+        return output[0]["data_type"]
+
+
+# The keys of this dictionary are the values we expect to see in a
+# TGetColumnsRequest's .TYPE_NAME attribute.
+# These are enumerated in ttypes.py as class TTypeId.
+# TODO: confirm that all types in TTypeId are included here.
+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,
+    "double": sqlalchemy.types.Float,
+    "string": sqlalchemy.types.String,
+    "varchar": sqlalchemy.types.String,
+    "char": sqlalchemy.types.String,
+    "binary": sqlalchemy.types.String,
+    "array": sqlalchemy.types.String,
+    "map": sqlalchemy.types.String,
+    "struct": sqlalchemy.types.String,
+    "uniontype": sqlalchemy.types.String,
+    "decimal": sqlalchemy.types.Numeric,
+    "timestamp": type_overrides.TIMESTAMP,
+    "timestamp_ntz": type_overrides.TIMESTAMP_NTZ,
+    "date": sqlalchemy.types.Date,
+}
+
+
+def parse_numeric_type_precision_and_scale(type_name_str):
+    """Return an intantiated sqlalchemy Numeric() type that preserves the precision and scale indicated
+    in the output from TGetColumnsRequest.
+
+    type_name_str
+      The value of TGetColumnsReq.TYPE_NAME.
+
+    If type_name_str is "DECIMAL(18,5) returns sqlalchemy.types.Numeric(18,5)
+    """
+
+    pattern = re.compile(r"DECIMAL\((\d+,\d+)\)")
+    match = re.search(pattern, type_name_str)
+    precision_and_scale = match.group(1)
+    precision, scale = tuple(precision_and_scale.split(","))
+
+    return sqlalchemy.types.Numeric(int(precision), int(scale))
+
+
+def parse_column_info_from_tgetcolumnsresponse(thrift_resp_row) -> ReflectedColumn:
+    """Returns a dictionary of the ReflectedColumn schema parsed from
+    a single of the result of a TGetColumnsRequest thrift RPC
+    """
+
+    pat = re.compile(r"^\w+")
+
+    # This method assumes a valid TYPE_NAME field in the response.
+    # TODO: add error handling in case TGetColumnsResponse format changes
+
+    _raw_col_type = re.search(pat, thrift_resp_row.TYPE_NAME).group(0).lower()  # type: ignore
+    _col_type = GET_COLUMNS_TYPE_MAP[_raw_col_type]
+
+    if _raw_col_type == "decimal":
+        final_col_type = parse_numeric_type_precision_and_scale(
+            thrift_resp_row.TYPE_NAME
+        )
+    else:
+        final_col_type = _col_type
+
+    # See comments about autoincrement in test_suite.py
+    # Since Databricks SQL doesn't currently support inline AUTOINCREMENT declarations
+    # the autoincrement must be manually declared with an Identity() construct in SQLAlchemy
+    # Other dialects can perform this extra Identity() step automatically. But that is not
+    # implemented in the Databricks dialect right now. So autoincrement is currently always False.
+    # It's not clear what IS_AUTO_INCREMENT in the thrift response actually reflects or whether
+    # it ever returns a `YES`.
+
+    # Per the guidance in SQLAlchemy's docstrings, we prefer to not even include an autoincrement
+    # key in this dictionary.
+    this_column = {
+        "name": thrift_resp_row.COLUMN_NAME,
+        "type": final_col_type,
+        "nullable": bool(thrift_resp_row.NULLABLE),
+        "default": thrift_resp_row.COLUMN_DEF,
+        "comment": thrift_resp_row.REMARKS or None,
+    }
+
+    # TODO: figure out how to return sqlalchemy.interfaces in a way that mypy respects
+    return this_column  # type: ignore

databricks/sqlalchemy/_types.py

@@ -0,0 +1,323 @@
+from datetime import datetime, time, timezone
+from itertools import product
+from typing import Any, Union, Optional
+
+import sqlalchemy
+from sqlalchemy.engine.interfaces import Dialect
+from sqlalchemy.ext.compiler import compiles
+
+from databricks.sql.utils import ParamEscaper
+
+
+def process_literal_param_hack(value: Any):
+    """This method is supposed to accept a Python type and return a string representation of that type.
+    But due to some weirdness in the way SQLAlchemy's literal rendering works, we have to return
+    the value itself because, by the time it reaches our custom type code, it's already been converted
+    into a string.
+
+    TimeTest
+    DateTimeTest
+    DateTimeTZTest
+
+    This dynamic only seems to affect the literal rendering of datetime and time objects.
+
+    All fail without this hack in-place. I'm not sure why. But it works.
+    """
+    return value
+
+
+@compiles(sqlalchemy.types.Enum, "databricks")
+@compiles(sqlalchemy.types.String, "databricks")
+@compiles(sqlalchemy.types.Text, "databricks")
+@compiles(sqlalchemy.types.Time, "databricks")
+@compiles(sqlalchemy.types.Unicode, "databricks")
+@compiles(sqlalchemy.types.UnicodeText, "databricks")
+@compiles(sqlalchemy.types.Uuid, "databricks")
+def compile_string_databricks(type_, compiler, **kw):
+    """
+    We override the default compilation for Enum(), String(), Text(), and Time() because SQLAlchemy
+    defaults to incompatible / abnormal compiled names
+
+      Enum -> VARCHAR
+      String -> VARCHAR[LENGTH]
+      Text -> VARCHAR[LENGTH]
+      Time -> TIME
+      Unicode -> VARCHAR[LENGTH]
+      UnicodeText -> TEXT
+      Uuid -> CHAR[32]
+
+    But all of these types will be compiled to STRING in Databricks SQL
+    """
+    return "STRING"
+
+
+@compiles(sqlalchemy.types.Integer, "databricks")
+def compile_integer_databricks(type_, compiler, **kw):
+    """
+    We need to override the default Integer compilation rendering because Databricks uses "INT" instead of "INTEGER"
+    """
+    return "INT"
+
+
+@compiles(sqlalchemy.types.LargeBinary, "databricks")
+def compile_binary_databricks(type_, compiler, **kw):
+    """
+    We need to override the default LargeBinary compilation rendering because Databricks uses "BINARY" instead of "BLOB"
+    """
+    return "BINARY"
+
+
+@compiles(sqlalchemy.types.Numeric, "databricks")
+def compile_numeric_databricks(type_, compiler, **kw):
+    """
+    We need to override the default Numeric compilation rendering because Databricks uses "DECIMAL" instead of "NUMERIC"
+
+    The built-in visit_DECIMAL behaviour captures the precision and scale. Here we're just mapping calls to compile Numeric
+    to the SQLAlchemy Decimal() implementation
+    """
+    return compiler.visit_DECIMAL(type_, **kw)
+
+
+@compiles(sqlalchemy.types.DateTime, "databricks")
+def compile_datetime_databricks(type_, compiler, **kw):
+    """
+    We need to override the default DateTime compilation rendering because Databricks uses "TIMESTAMP_NTZ" instead of "DATETIME"
+    """
+    return "TIMESTAMP_NTZ"
+
+
+@compiles(sqlalchemy.types.ARRAY, "databricks")
+def compile_array_databricks(type_, compiler, **kw):
+    """
+    SQLAlchemy's default ARRAY can't compile as it's only implemented for Postgresql.
+    The Postgres implementation works for Databricks SQL, so we duplicate that here.
+
+    :type_:
+        This is an instance of sqlalchemy.types.ARRAY which always includes an item_type attribute
+        which is itself an instance of TypeEngine
+
+    https://docs.sqlalchemy.org/en/20/core/type_basics.html#sqlalchemy.types.ARRAY
+    """
+
+    inner = compiler.process(type_.item_type, **kw)
+
+    return f"ARRAY<{inner}>"
+
+
+class TIMESTAMP_NTZ(sqlalchemy.types.TypeDecorator):
+    """Represents values comprising values of fields year, month, day, hour, minute, and second.
+    All operations are performed without taking any time zone into account.
+
+    Our dialect maps sqlalchemy.types.DateTime() to this type, which means that all DateTime()
+    objects are stored without tzinfo. To read and write timezone-aware datetimes use
+    databricks.sql.TIMESTAMP instead.
+
+    https://docs.databricks.com/en/sql/language-manual/data-types/timestamp-ntz-type.html
+    """
+
+    impl = sqlalchemy.types.DateTime
+
+    cache_ok = True
+
+    def process_result_value(self, value: Union[None, datetime], dialect):
+        if value is None:
+            return None
+        return value.replace(tzinfo=None)
+
+
+class TIMESTAMP(sqlalchemy.types.TypeDecorator):
+    """Represents values comprising values of fields year, month, day, hour, minute, and second,
+    with the session local time-zone.
+
+    Our dialect maps sqlalchemy.types.DateTime() to TIMESTAMP_NTZ, which means that all DateTime()
+    objects are stored without tzinfo. To read and write timezone-aware datetimes use
+    this type instead.
+
+    ```python
+    # This won't work
+    `Column(sqlalchemy.DateTime(timezone=True))`
+
+    # But this does
+    `Column(TIMESTAMP)`
+    ````
+
+    https://docs.databricks.com/en/sql/language-manual/data-types/timestamp-type.html
+    """
+
+    impl = sqlalchemy.types.DateTime
+
+    cache_ok = True
+
+    def process_result_value(self, value: Union[None, datetime], dialect):
+        if value is None:
+            return None
+
+        if not value.tzinfo:
+            return value.replace(tzinfo=timezone.utc)
+        return value
+
+    def process_bind_param(
+        self, value: Union[datetime, None], dialect
+    ) -> Optional[datetime]:
+        """pysql can pass datetime.datetime() objects directly to DBR"""
+        return value
+
+    def process_literal_param(
+        self, value: Union[datetime, None], dialect: Dialect
+    ) -> str:
+        """ """
+        return process_literal_param_hack(value)
+
+
+@compiles(TIMESTAMP, "databricks")
+def compile_timestamp_databricks(type_, compiler, **kw):
+    """
+    We need to override the default DateTime compilation rendering because Databricks uses "TIMESTAMP_NTZ" instead of "DATETIME"
+    """
+    return "TIMESTAMP"
+
+
+class DatabricksTimeType(sqlalchemy.types.TypeDecorator):
+    """Databricks has no native TIME type. So we store it as a string."""
+
+    impl = sqlalchemy.types.Time
+    cache_ok = True
+
+    BASE_FMT = "%H:%M:%S"
+    MICROSEC_PART = ".%f"
+    TIMEZONE_PART = "%z"
+
+    def _generate_fmt_string(self, ms: bool, tz: bool) -> str:
+        """Return a format string for datetime.strptime() that includes or excludes microseconds and timezone."""
+        _ = lambda x, y: x if y else ""
+        return f"{self.BASE_FMT}{_(self.MICROSEC_PART,ms)}{_(self.TIMEZONE_PART,tz)}"
+
+    @property
+    def allowed_fmt_strings(self):
+        """Time strings can be read with or without microseconds and with or without a timezone."""
+
+        if not hasattr(self, "_allowed_fmt_strings"):
+            ms_switch = tz_switch = [True, False]
+            self._allowed_fmt_strings = [
+                self._generate_fmt_string(x, y)
+                for x, y in product(ms_switch, tz_switch)
+            ]
+
+        return self._allowed_fmt_strings
+
+    def _parse_result_string(self, value: str) -> time:
+        """Parse a string into a time object. Try all allowed formats until one works."""
+        for fmt in self.allowed_fmt_strings:
+            try:
+                # We use timetz() here because we want to preserve the timezone information
+                # Calling .time() will strip the timezone information
+                return datetime.strptime(value, fmt).timetz()
+            except ValueError:
+                pass
+
+        raise ValueError(f"Could not parse time string {value}")
+
+    def _determine_fmt_string(self, value: time) -> str:
+        """Determine which format string to use to render a time object as a string."""
+        ms_bool = value.microsecond > 0
+        tz_bool = value.tzinfo is not None
+        return self._generate_fmt_string(ms_bool, tz_bool)
+
+    def process_bind_param(self, value: Union[time, None], dialect) -> Union[None, str]:
+        """Values sent to the database are converted to %:H:%M:%S strings."""
+        if value is None:
+            return None
+        fmt_string = self._determine_fmt_string(value)
+        return value.strftime(fmt_string)
+
+    # mypy doesn't like this workaround because TypeEngine wants process_literal_param to return a string
+    def process_literal_param(self, value, dialect) -> time:  # type: ignore
+        """ """
+        return process_literal_param_hack(value)
+
+    def process_result_value(
+        self, value: Union[None, str], dialect
+    ) -> Union[time, None]:
+        """Values received from the database are parsed into datetime.time() objects"""
+        if value is None:
+            return None
+
+        return self._parse_result_string(value)
+
+
+class DatabricksStringType(sqlalchemy.types.TypeDecorator):
+    """We have to implement our own String() type because SQLAlchemy's default implementation
+    wants to escape single-quotes with a doubled single-quote. Databricks uses a backslash for
+    escaping of literal strings. And SQLAlchemy's default escaping breaks Databricks SQL.
+    """
+
+    impl = sqlalchemy.types.String
+    cache_ok = True
+    pe = ParamEscaper()
+
+    def process_literal_param(self, value, dialect) -> str:
+        """SQLAlchemy's default string escaping for backslashes doesn't work for databricks. The logic here
+        implements the same logic as our legacy inline escaping logic.
+        """
+
+        return self.pe.escape_string(value)
+
+    def literal_processor(self, dialect):
+        """We manually override this method to prevent further processing of the string literal beyond
+        what happens in the process_literal_param() method.
+
+        The SQLAlchemy docs _specifically_ say to not override this method.
+
+        It appears that any processing that happens from TypeEngine.process_literal_param happens _before_
+        and _in addition to_ whatever the class's impl.literal_processor() method does. The String.literal_processor()
+        method performs a string replacement that doubles any single-quote in the contained string. This raises a syntax
+        error in Databricks. And it's not necessary because ParamEscaper() already implements all the escaping we need.
+
+        We should consider opening an issue on the SQLAlchemy project to see if I'm using it wrong.
+
+        See type_api.py::TypeEngine.literal_processor:
+
+        ```python
+            def process(value: Any) -> str:
+                return fixed_impl_processor(
+                    fixed_process_literal_param(value, dialect)
+                )
+        ```
+
+        That call to fixed_impl_processor wraps the result of fixed_process_literal_param (which is the
+        process_literal_param defined in our Databricks dialect)
+
+        https://docs.sqlalchemy.org/en/20/core/custom_types.html#sqlalchemy.types.TypeDecorator.literal_processor
+        """
+
+        def process(value):
+            """This is a copy of the default String.literal_processor() method but stripping away
+            its double-escaping behaviour for single-quotes.
+            """
+
+            _step1 = self.process_literal_param(value, dialect="databricks")
+            if dialect.identifier_preparer._double_percents:
+                _step2 = _step1.replace("%", "%%")
+            else:
+                _step2 = _step1
+
+            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
+
+
+@compiles(TINYINT, "databricks")
+def compile_tinyint(type_, compiler, **kw):
+    return "TINYINT"