Support YDB native UUID type

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Add opt-in support for YDB's native UUID type via `sa.UUID` or `types.YqlUUID` while keeping `sa.Uuid` stored as `Utf8`
+
 ## 0.1.24 ##
 * Keep `ydb_sqlalchemy.alembic` importable when the optional Alembic dependency is absent
 

docs/types.rst

@@ -93,6 +93,11 @@ The following table shows the complete mapping between YDB native types, YDB SQL
      - ``String`` / ``Text``
      - ``str``
      -
+   * - ``Uuid``
+     - :class:`~ydb_sqlalchemy.sqlalchemy.types.YqlUUID`
+     - ``UUID`` (SQLAlchemy 2.0+)
+     - ``uuid.UUID`` / ``str``
+     - Explicit opt-in; ``sa.Uuid`` keeps using ``Utf8``
    * - ``Date``
      - :class:`~ydb_sqlalchemy.sqlalchemy.datetime_types.YqlDate`
      - ``Date``
@@ -162,6 +167,27 @@ Most standard SQLAlchemy types work with YDB:
        is_active = Column(Boolean)
        price = Column(Float)
 
+UUID Types
+----------
+
+For compatibility with earlier releases, SQLAlchemy's generic ``Uuid`` type
+continues to use a YDB ``Utf8`` column. Use
+:class:`~ydb_sqlalchemy.sqlalchemy.types.YqlUUID` when native YDB ``Uuid``
+storage is required. On SQLAlchemy 2.0 and newer, the SQL-native ``UUID`` type
+is an equivalent explicit opt-in.
+
+.. code-block:: python
+
+   import sqlalchemy as sa
+   from ydb_sqlalchemy import types as ydb_types
+
+   native_uuid = sa.Column(ydb_types.YqlUUID())
+   native_uuid_as_text = sa.Column(ydb_types.YqlUUID(as_uuid=False))
+   native_uuid_sa20 = sa.Column(sa.UUID())
+
+   # Existing behavior: stored as YDB Utf8.
+   compatible_uuid = sa.Column(sa.Uuid())
+
 YDB-Specific Integer Types
 --------------------------
 

tests/integration/test_core.py

@@ -1,5 +1,6 @@
 import asyncio
 import datetime
+import uuid
 from decimal import Decimal
 from typing import NamedTuple
 
@@ -14,6 +15,8 @@
 from ydb_sqlalchemy import sqlalchemy as ydb_sa
 from ydb_sqlalchemy.sqlalchemy import types
 
+_UUID_TABLE_NAME = f"test_uuid_types_{uuid.uuid4().hex[:8]}"
+
 if sa.__version__ >= "2.":
     from sqlalchemy import NullPool
     from sqlalchemy import QueuePool
@@ -254,6 +257,13 @@ def define_tables(cls, metadata: sa.MetaData):
             Column("date", sa.Date),
             # Column("interval", sa.Interval),
         )
+        Table(
+            _UUID_TABLE_NAME,
+            metadata,
+            Column("id", Integer, primary_key=True),
+            Column("uuid_native", types.YqlUUID),
+            Column("uuid_str", sa.Uuid if not ydb_sa.OLD_SA else sa.String),
+        )
 
     def test_primitive_types(self, connection):
         table = self.tables.test_primitive_types
@@ -341,6 +351,44 @@ def test_datetime_types_timezone(self, connection):
             today,
         )
 
+    def test_native_uuid_types(self, connection):
+        table = self.tables[_UUID_TABLE_NAME]
+        uuid_value = uuid.uuid4()
+
+        statement = sa.insert(table).values(id=1, uuid_native=uuid_value)
+        connection.execute(statement)
+        row = connection.execute(sa.select(table.c.id, table.c.uuid_native).where(table.c.id == 1)).fetchone()
+        assert row == (1, uuid_value)
+
+        uuid_value_str = str(uuid_value)
+        statement = sa.insert(table).values(id=2, uuid_native=uuid_value_str)
+        connection.execute(statement)
+        row = connection.execute(sa.select(table.c.id, table.c.uuid_native).where(table.c.id == 2)).fetchone()
+        assert row == (2, uuid_value)
+
+    @pytest.mark.skipif(ydb_sa.OLD_SA, reason="sa.Uuid was added in SQLAlchemy 2.0")
+    def test_generic_uuid_keeps_utf8_storage(self, connection):
+        table = self.tables[_UUID_TABLE_NAME]
+        uuid_value = uuid.uuid4()
+
+        connection.execute(sa.insert(table).values(id=3, uuid_str=uuid_value))
+        row = connection.execute(sa.select(table.c.uuid_str).where(table.c.id == 3)).fetchone()
+        assert row == (uuid_value,)
+
+        table_description = connection.connection.driver_connection.describe(table.name)
+        column_types = {column.name: column.type for column in table_description.columns}
+        assert column_types["uuid_native"].item == ydb.PrimitiveType.UUID
+        assert column_types["uuid_str"].item == ydb.PrimitiveType.Utf8
+
+    def test_native_uuid_reflection(self, connection):
+        table = self.tables[_UUID_TABLE_NAME]
+        reflected_metadata = sa.MetaData()
+
+        reflected_metadata.reflect(connection, only=[table.name])
+
+        reflected_type = reflected_metadata.tables[table.name].c.uuid_native.type
+        assert isinstance(reflected_type, types.YqlUUID)
+
 
 class TestWithClause(TablesTest):
     __backend__ = True

tests/integration/test_suite.py

@@ -616,7 +616,6 @@ def test_concat_func(self, connection):
 if not OLD_SA:
     from sqlalchemy.testing.suite.test_types import NativeUUIDTest as _NativeUUIDTest
 
-    @pytest.mark.skip("uuid unsupported for columns")
     class NativeUUIDTest(_NativeUUIDTest):
         pass
 

ydb_sqlalchemy/sqlalchemy/__init__.py

@@ -77,6 +77,7 @@ def upsert(table):
     ydb.PrimitiveType.Interval: sa.INTEGER,
     ydb.PrimitiveType.Bool: sa.BOOLEAN,
     ydb.PrimitiveType.DyNumber: sa.TEXT,
+    ydb.PrimitiveType.UUID: types.YqlUUID,
 }
 
 DBAPI_COLUMN_TYPES = {
@@ -198,6 +199,7 @@ class YqlDialect(StrCompileDialect):
         sa.types.LargeBinary: types.Binary,
         sa.types.BLOB: types.Binary,
         sa.types.ARRAY: types.ListType,
+        **({sa.types.UUID: types.YqlUUID} if not OLD_SA else {}),
     }
 
     connection_characteristics = util.immutabledict(

ydb_sqlalchemy/sqlalchemy/compiler/base.py

@@ -68,6 +68,9 @@
 
 
 class BaseYqlTypeCompiler(StrSQLTypeCompiler):
+    def visit_UUID(self, type_: types.YqlUUID, **kw):
+        return "UUID"
+
     def visit_JSON(self, type_: Union[sa.JSON, types.YqlJSON], **kw):
         return "JSON"
 
@@ -177,6 +180,12 @@ def get_ydb_type(
         if isinstance(type_, sa.TypeDecorator):
             type_ = type_.impl
 
+        if isinstance(type_, types.YqlUUID):
+            ydb_type = ydb.PrimitiveType.UUID
+            if is_optional:
+                return ydb.OptionalType(ydb_type)
+            return ydb_type
+
         if isinstance(type_, (sa.Text, sa.String)):
             ydb_type = ydb.PrimitiveType.Utf8
 
@@ -297,6 +306,8 @@ def limit_clause(self, select, **kw):
 
     def render_literal_value(self, value, type_):
         if isinstance(value, str):
+            if isinstance(type_.dialect_impl(self.dialect), types.YqlUUID):
+                return super().render_literal_value(value, type_)
             for pattern, replacement in ESCAPE_RULES:
                 value = value.replace(pattern, replacement)
             return f"'{value}'"

ydb_sqlalchemy/sqlalchemy/compiler/sa20.py

@@ -24,6 +24,12 @@ def get_ydb_type(
         if isinstance(type_, sa.TypeDecorator):
             type_ = type_.impl
 
+        if isinstance(type_, sa.UUID):
+            ydb_type = ydb.PrimitiveType.UUID
+            if is_optional:
+                return ydb.OptionalType(ydb_type)
+            return ydb_type
+
         if isinstance(type_, sa.Uuid):
             ydb_type = ydb.PrimitiveType.Utf8
             if is_optional:

ydb_sqlalchemy/sqlalchemy/test_sqlalchemy.py

@@ -1,5 +1,9 @@
 from datetime import date
+import uuid
+
+import pytest
 import sqlalchemy as sa
+import ydb
 
 from . import YqlDialect, types
 
@@ -114,6 +118,54 @@ def compile_type(type_):
     assert compile_type(struct) == "Struct<a:Int32,b:List<Int32>>"
 
 
+def test_native_uuid_is_explicit_opt_in():
+    dialect = YqlDialect()
+    type_compiler = dialect.type_compiler
+
+    assert type_compiler.process(types.YqlUUID()) == "UUID"
+    assert type_compiler.get_ydb_type(types.YqlUUID(), is_optional=False) == ydb.PrimitiveType.UUID
+    assert type_compiler.get_ydb_type(types.YqlUUID(), is_optional=True).item == ydb.PrimitiveType.UUID
+
+    if not hasattr(sa, "Uuid"):
+        return
+
+    assert dialect.supports_native_uuid is False
+    assert type_compiler.process(sa.Uuid()) == "UTF8"
+    assert type_compiler.get_ydb_type(sa.Uuid(), is_optional=False) == ydb.PrimitiveType.Utf8
+    assert type_compiler.process(sa.UUID()) == "UUID"
+    assert type_compiler.get_ydb_type(sa.UUID(), is_optional=False) == ydb.PrimitiveType.UUID
+
+    dialect_impl = sa.UUID(as_uuid=False).dialect_impl(dialect)
+    assert isinstance(dialect_impl, types.YqlUUID)
+    assert dialect_impl.as_uuid is False
+
+
+def test_native_uuid_processors():
+    dialect = YqlDialect()
+    value = uuid.uuid4()
+    uuid_type = types.YqlUUID()
+
+    bind_processor = uuid_type.bind_processor(dialect)
+    assert bind_processor(None) is None
+    assert bind_processor(value) == value
+    assert bind_processor(str(value)) == value
+    with pytest.raises(ValueError):
+        bind_processor("not-a-uuid")
+
+    result_processor = uuid_type.result_processor(dialect, None)
+    assert result_processor(None) is None
+    assert result_processor(value) == value
+    assert result_processor(str(value)) == value
+    assert uuid_type.literal_processor(dialect)(value) == f'Uuid("{value}")'
+
+    text_uuid_type = types.YqlUUID(as_uuid=False)
+    assert text_uuid_type.bind_processor(dialect)(str(value)) == value
+    assert text_uuid_type.result_processor(dialect, None)(value) == str(value)
+
+    text_literal = sa.literal(str(value), text_uuid_type)
+    assert str(text_literal.compile(dialect=dialect, compile_kwargs={"literal_binds": True})) == f'Uuid("{value}")'
+
+
 def test_statement_prefixes_prepended_to_query():
     dialect = YqlDialect(_statement_prefixes_list=["PRAGMA DistinctOverKeys;"])
     result = dialect._apply_statement_prefixes_impl("SELECT 1")

ydb_sqlalchemy/sqlalchemy/types.py

@@ -1,4 +1,5 @@
 import decimal
+import uuid
 from typing import Any, Mapping, Type, Union
 
 from sqlalchemy import __version__ as sa_version
@@ -15,6 +16,62 @@
 from .json import YqlJSON  # noqa: F401
 
 
+_UUIDBase = getattr(types, "UUID", types.TypeEngine)
+
+
+class YqlUUID(_UUIDBase):
+    """Store UUID values using YDB's native ``Uuid`` type.
+
+    The dialect intentionally keeps SQLAlchemy's generic ``Uuid`` type mapped
+    to ``Utf8`` for backwards compatibility.  This type is the explicit opt-in
+    for native storage and is also used as the dialect implementation of
+    SQLAlchemy 2.x's SQL-native ``UUID`` type.
+    """
+
+    __visit_name__ = "UUID"
+
+    def __init__(self, as_uuid=True):
+        self.as_uuid = as_uuid
+        if hasattr(types, "UUID"):
+            super().__init__(as_uuid=as_uuid)
+        else:
+            super().__init__()
+
+    @property
+    def python_type(self):
+        return uuid.UUID if self.as_uuid else str
+
+    def bind_processor(self, dialect):
+        def process(value):
+            if value is None:
+                return None
+            if isinstance(value, str):
+                value = uuid.UUID(value)
+            return value
+
+        return process
+
+    def result_processor(self, dialect, coltype):
+        def process(value):
+            if value is None:
+                return None
+            if self.as_uuid:
+                return value if isinstance(value, uuid.UUID) else uuid.UUID(value)
+            return str(value)
+
+        return process
+
+    def literal_processor(self, dialect):
+        def process(value):
+            if value is None:
+                return None
+            if not isinstance(value, uuid.UUID):
+                value = uuid.UUID(value)
+            return f'Uuid("{value}")'
+
+        return process
+
+
 class UInt64(types.Integer):
     __visit_name__ = "uint64"