diff --git a/CHANGELOG.md b/CHANGELOG.md index 29511a3..9b2ed5f 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/docs/types.rst b/docs/types.rst index 91b0a57..8b47176 100644 --- a/docs/types.rst +++ b/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 -------------------------- diff --git a/tests/integration/test_core.py b/tests/integration/test_core.py index 1940ec8..cb85e38 100644 --- a/tests/integration/test_core.py +++ b/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 diff --git a/tests/integration/test_suite.py b/tests/integration/test_suite.py index 5ef403a..78abda1 100644 --- a/tests/integration/test_suite.py +++ b/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 diff --git a/ydb_sqlalchemy/sqlalchemy/__init__.py b/ydb_sqlalchemy/sqlalchemy/__init__.py index 73198cd..afd7fac 100644 --- a/ydb_sqlalchemy/sqlalchemy/__init__.py +++ b/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( diff --git a/ydb_sqlalchemy/sqlalchemy/compiler/base.py b/ydb_sqlalchemy/sqlalchemy/compiler/base.py index e95154f..0ed1e7a 100644 --- a/ydb_sqlalchemy/sqlalchemy/compiler/base.py +++ b/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}'" diff --git a/ydb_sqlalchemy/sqlalchemy/compiler/sa20.py b/ydb_sqlalchemy/sqlalchemy/compiler/sa20.py index 702d7aa..7879dbc 100644 --- a/ydb_sqlalchemy/sqlalchemy/compiler/sa20.py +++ b/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: diff --git a/ydb_sqlalchemy/sqlalchemy/test_sqlalchemy.py b/ydb_sqlalchemy/sqlalchemy/test_sqlalchemy.py index 3f1de82..d2f56a1 100644 --- a/ydb_sqlalchemy/sqlalchemy/test_sqlalchemy.py +++ b/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>" +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") diff --git a/ydb_sqlalchemy/sqlalchemy/types.py b/ydb_sqlalchemy/sqlalchemy/types.py index 4e4a902..75ad5a2 100644 --- a/ydb_sqlalchemy/sqlalchemy/types.py +++ b/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"