pydantic-sqlalchemy icon indicating copy to clipboard operation
pydantic-sqlalchemy copied to clipboard

sqlalchemy_to_pydantic is not compatible with some TypeDecorator

Open lootr opened this issue 3 years ago • 0 comments

How to reproduce

I used a TypeDecorator SQLAlchemy class (see uuid_type.py below) to create a handler of python uuid.UUID for SQLAlchemy, inspired by sqlalchemy_utils.types.uuid.UUIDType and providing a dialect agnostic using load_dialect_impl() associated with impl = sqlalchemy.types.TypeEngine.

Extract of documentation of TypeDecorator:

The class-level impl attribute is required, and can reference any TypeEngine class. Alternatively, the load_dialect_impl() method can be used to provide different type classes based on the dialect given; in this case, the impl variable can reference TypeEngine as a placeholder.

Then, I called a SQLAlchemy base model using the type decorator.

uuid_type.py

import uuid

from sqlalchemy import types, util
from sqlalchemy.dialects import mssql, postgresql


class UUIDType(types.TypeDecorator):
    """
    Stores a UUID in the database natively when it can and falls back to
    a BINARY(16) or a CHAR(32) when it can't.
    """
    impl = types.TypeEngine

    python_type = uuid.UUID

    cache_ok = True

    def __init__(self, binary=True, native=True):
        """
        :param binary: Whether to use a BINARY(16) or CHAR(32) fallback.
        """
        self.binary = binary
        self.native = native

    def __repr__(self):
        return util.generic_repr(self)

    def load_dialect_impl(self, dialect):
        if self.native and dialect.name in ('postgresql', 'cockroachdb'):
            # Use the native UUID type.
            return dialect.type_descriptor(postgresql.UUID())

        if dialect.name == 'mssql' and self.native:
            # Use the native UNIQUEIDENTIFIER type.
            return dialect.type_descriptor(mssql.UNIQUEIDENTIFIER())

        else:
            # Fallback to either a BINARY or a CHAR.
            kind = types.BINARY(16) if self.binary else types.CHAR(32)
            return dialect.type_descriptor(kind)

    @staticmethod
    def _coerce(value):
        if value and not isinstance(value, uuid.UUID):
            try:
                value = uuid.UUID(value)

            except (TypeError, ValueError):
                value = uuid.UUID(bytes=value)

        return value

    def process_literal_param(self, value, dialect):
        return "'{}'".format(value) if value else value

    def process_bind_param(self, value, dialect):
        if value is None:
            return value

        if not isinstance(value, uuid.UUID):
            value = self._coerce(value)

        if self.native and dialect.name in (
            'postgresql',
            'mssql',
            'cockroachdb'
        ):
            return str(value)

        return value.bytes if self.binary else value.hex

    def process_result_value(self, value, dialect):
        if value is None:
            return value

        if self.native and dialect.name in (
            'postgresql',
            'mssql',
            'cockroachdb'
        ):
            if isinstance(value, uuid.UUID):
                # Some drivers convert PostgreSQL's uuid values to
                # Python's uuid.UUID objects by themselves
                return value
            return uuid.UUID(value)

        return uuid.UUID(bytes=value) if self.binary else uuid.UUID(value)

model.py

from sqlalchemy import Column
from sqlalchemy.orm import declarative_base
from pydantic_sqlalchemy import sqlalchemy_to_pydantic

from .uuid_type import UUIDType

Base = declarative_base()


class ExampleMapper(Base):
    uuid = Column(UUIDType)


ExampleModel = sqlalchemy_to_pydantic(ExampleMapper)

Running model.py raises a RuntimeError as sqlalchemy_to_pydantic attempt to infer type of decorator from impl (whose value is TypeGeneric and associated impl is a property) instead of python_type (which is correctly set).

Fix proposal

The following patch proposal handle this specific case and fallback into python_type condition when TypeEngine is used as impl.

5a6
> from sqlalchemy import types
25c26,28
<                 if hasattr(column.type, "impl"):
---
>                 # TypeEngine is a placeholder when impl is abstract
>                 if (hasattr(column.type, "impl")
>                     and column.type.impl is not types.TypeEngine):

lootr avatar Mar 18 '22 08:03 lootr