Merge pull request #63 from databricks/fix-uuid-canonical-binding

Bind UUIDs in canonical hyphenated form (fixes #50)

Author: sreekanth-db

src/databricks/sqlalchemy/_types.py

@@ -1,6 +1,7 @@
 from datetime import datetime, time, timezone
 from itertools import product
 from typing import Any, Union, Optional
+from uuid import UUID
 
 import sqlalchemy
 from sqlalchemy.engine.interfaces import Dialect
@@ -315,6 +316,60 @@ def process(value):
         return process
 
 
+class DatabricksUUID(sqlalchemy.types.Uuid):
+    """Bind UUIDs in their canonical 8-4-4-4-12 hyphenated form.
+
+    Databricks has no native UUID type, so SQLAlchemy's default ``Uuid``
+    bind/literal processors render the 32-character hex form without dashes
+    (e.g. ``1daa91d78d35468486d63fa89042c1f4``). That breaks equality against
+    UUIDs stored as canonical strings in Databricks. We coerce every input
+    through ``uuid.UUID`` so the wire value is always the canonical hyphenated
+    form regardless of whether the caller passed a ``UUID``, a hyphenated
+    string, or a dash-less hex string. The ``UUID(...)`` round-trip also
+    validates the input — any non-UUID string raises ``ValueError`` instead of
+    being silently injected into SQL, which is critical for ``literal_binds``
+    rendering safety.
+
+    With the default ``as_uuid=True``, the inherited ``result_processor``
+    parses both hyphenated and dash-less hex forms back into a ``UUID``
+    object, so reads of legacy hex-stored rows continue to work. With
+    ``as_uuid=False`` the result is returned as the raw column string —
+    callers who mix legacy hex-stored rows with the canonical form should
+    normalize on read themselves.
+    """
+
+    cache_ok = True
+
+    @staticmethod
+    def _canonical(value):
+        """Return the canonical hyphenated string for ``value``.
+
+        For UUID instances we rebuild a stdlib ``UUID`` from ``.int`` so a
+        subclass cannot smuggle an arbitrary string through an overridden
+        ``__str__`` — the canonical hyphenated form of ``value.int`` goes to
+        the wire, so no attacker-controlled string can escape the quotes.
+        """
+        if isinstance(value, UUID):
+            return str(UUID(int=value.int))
+        return str(UUID(str(value)))
+
+    def bind_processor(self, dialect):
+        def process(value):
+            if value is None:
+                return None
+            return self._canonical(value)
+
+        return process
+
+    def literal_processor(self, dialect):
+        def process(value):
+            if value is None:
+                return "NULL"
+            return "'%s'" % self._canonical(value)
+
+        return process
+
+
 class TINYINT(sqlalchemy.types.TypeDecorator):
     """Represents 1-byte signed integers
 

src/databricks/sqlalchemy/base.py

@@ -70,6 +70,7 @@ class DatabricksDialect(default.DefaultDialect):
         sqlalchemy.types.DateTime: dialect_type_impl.TIMESTAMP_NTZ,
         sqlalchemy.types.Time: dialect_type_impl.DatabricksTimeType,
         sqlalchemy.types.String: dialect_type_impl.DatabricksStringType,
+        sqlalchemy.types.Uuid: dialect_type_impl.DatabricksUUID,
     }
 
     # SQLAlchemy requires that a table with no primary key

tests/test_local/test_types.py

@@ -1,10 +1,18 @@
 import enum
+from uuid import UUID
 
 import pytest
 import sqlalchemy
+from sqlalchemy import Column, MetaData, Table, select
 
 from databricks.sqlalchemy.base import DatabricksDialect
-from databricks.sqlalchemy._types import TINYINT, TIMESTAMP, TIMESTAMP_NTZ, DatabricksVariant
+from databricks.sqlalchemy._types import (
+    DatabricksUUID,
+    DatabricksVariant,
+    TINYINT,
+    TIMESTAMP,
+    TIMESTAMP_NTZ,
+)
 
 
 class DatabricksDataType(enum.Enum):
@@ -161,3 +169,140 @@ def test_array_string_renders_as_array_of_string(self):
         self._assert_compiled_value_explicit(
             sqlalchemy.types.ARRAY(sqlalchemy.types.String), "ARRAY<STRING>"
         )
+
+
+class TestDatabricksUUID:
+    """Regression coverage for github.com/databricks/databricks-sqlalchemy/issues/50.
+
+    SQLAlchemy's default Uuid renders the 32-char hex form (no dashes) on backends
+    without a native UUID type, which breaks equality against UUIDs stored as
+    canonical 8-4-4-4-12 strings in Databricks.
+    """
+
+    dialect = DatabricksDialect()
+    HYPHENATED = "1daa91d7-8d35-4684-86d6-3fa89042c1f4"
+    HEX = "1daa91d78d35468486d63fa89042c1f4"
+    sample = UUID(HYPHENATED)
+
+    def test_bind_processor_preserves_hyphenated_form(self):
+        process = DatabricksUUID().bind_processor(self.dialect)
+        assert process(self.sample) == self.HYPHENATED
+
+    def test_bind_processor_handles_none(self):
+        process = DatabricksUUID().bind_processor(self.dialect)
+        assert process(None) is None
+
+    def test_literal_processor_renders_hyphenated_form(self):
+        process = DatabricksUUID().literal_processor(self.dialect)
+        assert process(self.sample) == "'%s'" % self.HYPHENATED
+
+    def test_literal_processor_handles_none(self):
+        process = DatabricksUUID().literal_processor(self.dialect)
+        assert process(None) == "NULL"
+
+    def test_result_processor_accepts_both_forms(self):
+        process = DatabricksUUID().result_processor(self.dialect, None)
+        assert process(self.HYPHENATED) == self.sample
+        assert process(self.HEX) == self.sample
+        assert process(None) is None
+
+    def test_dialect_routes_uuid_to_databricks_uuid(self):
+        """The colspecs entry is what makes a plain ``Uuid`` column use our impl."""
+        assert self.dialect.colspecs[sqlalchemy.types.Uuid] is DatabricksUUID
+
+    def test_uuid_where_clause_renders_with_dashes(self):
+        meta = MetaData()
+        users = Table(
+            "users", meta, Column("id", sqlalchemy.types.Uuid, primary_key=True)
+        )
+        stmt = select(users).where(users.c.id == self.sample)
+
+        literal_sql = str(
+            stmt.compile(dialect=self.dialect, compile_kwargs={"literal_binds": True})
+        )
+        assert "'%s'" % self.HYPHENATED in literal_sql
+        assert self.HEX not in literal_sql
+
+    def test_uuid_bound_param_wire_value_has_dashes(self):
+        meta = MetaData()
+        users = Table(
+            "users", meta, Column("id", sqlalchemy.types.Uuid, primary_key=True)
+        )
+        stmt = select(users).where(users.c.id == self.sample)
+        compiled = stmt.compile(dialect=self.dialect)
+
+        raw = compiled.construct_params()
+        processed = {
+            key: (
+                compiled._bind_processors[key](value)
+                if key in compiled._bind_processors
+                else value
+            )
+            for key, value in raw.items()
+        }
+        assert self.HYPHENATED in processed.values()
+
+    def test_bind_processor_normalizes_hex_string_to_canonical(self):
+        """A bare 32-char hex string (no dashes) must be coerced to canonical form."""
+        process = DatabricksUUID().bind_processor(self.dialect)
+        assert process(self.HEX) == self.HYPHENATED
+
+    def test_bind_processor_normalizes_hyphenated_string(self):
+        """A canonical hyphenated string passes through unchanged."""
+        process = DatabricksUUID().bind_processor(self.dialect)
+        assert process(self.HYPHENATED) == self.HYPHENATED
+
+    def test_bind_processor_rejects_non_uuid_string(self):
+        """Bad input must raise instead of being silently written to the column."""
+        process = DatabricksUUID().bind_processor(self.dialect)
+        with pytest.raises(ValueError):
+            process("not-a-uuid")
+
+    def test_literal_processor_rejects_injection_attempt(self):
+        """An attacker-controlled string must not be allowed to escape the quotes.
+
+        Before the input was normalized through ``UUID(...)``, a string like
+        ``abc' OR '1'='1`` would inject directly into ``WHERE id = 'abc' OR
+        '1'='1'`` whenever ``literal_binds=True`` was used.
+        """
+        process = DatabricksUUID().literal_processor(self.dialect)
+        with pytest.raises(ValueError):
+            process("abc' OR '1'='1")
+
+    def test_literal_processor_rejects_injection_via_uuid_subclass(self):
+        """A UUID subclass with a malicious ``__str__`` must not bypass escaping.
+
+        ``_canonical`` reconstructs every UUID through ``UUID(int=value.int)``
+        so a subclass cannot inject SQL via an overridden string conversion.
+        """
+
+        class EvilUUID(UUID):
+            def __str__(self):  # type: ignore[override]
+                return "abc' OR '1'='1"
+
+        process = DatabricksUUID().literal_processor(self.dialect)
+        rendered = process(EvilUUID(self.HYPHENATED))
+        assert rendered == "'%s'" % self.HYPHENATED
+        assert "OR" not in rendered
+
+    def test_literal_processor_rejects_injection_with_as_uuid_false(self):
+        """``Uuid(as_uuid=False)`` shares ``_canonical``; lock in coverage anyway."""
+        process = DatabricksUUID(as_uuid=False).literal_processor(self.dialect)
+        with pytest.raises(ValueError):
+            process("abc' OR '1'='1")
+
+    def test_as_uuid_false_round_trip_normalizes_hex_input(self):
+        """``Uuid(as_uuid=False)`` users sometimes pass hex-form strings.
+
+        Pre-fix, the dialect wrote those through unchanged, so a query against
+        canonically-stored data returned 0 rows. After the fix, both forms
+        normalize to canonical on the wire.
+        """
+        type_ = DatabricksUUID(as_uuid=False)
+        bind = type_.bind_processor(self.dialect)
+        result = type_.result_processor(self.dialect, None)
+
+        assert bind(self.HEX) == self.HYPHENATED
+        assert bind(self.HYPHENATED) == self.HYPHENATED
+        assert result(self.HYPHENATED) == self.HYPHENATED
+        assert result(self.HEX) == self.HYPHENATED