Failing to render postgresql's partial index statement in case of comparing to StrEnum

AlexanderPodorov commented 1 year ago

Describe the bug

If the comparison with StrEnum is occurred inside the postgresql_where clause, SQLAlchemy fails to render the DDL for that index.

Optional link from https://docs.sqlalchemy.org which documents the behavior that is expected

No response

SQLAlchemy Version in Use

2.0.19

DBAPI (i.e. the database driver)

asyncpg

Database Vendor and Major Version

PostgreSQL 15

Python Version

3.11.4

Operating system

MacOS

To Reproduce

from enum import StrEnum

from sqlalchemy import Index, select
from sqlalchemy.dialects import postgresql
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
from sqlalchemy.schema import CreateIndex

class UserType(StrEnum):
    T1 = 't1'
    T2 = 't2'

class Base(DeclarativeBase):
    pass

class User(Base):
    __tablename__ = "a"

    id: Mapped[int] = mapped_column(primary_key=True)
    account_id: Mapped[int] = mapped_column()
    type: Mapped[str] = mapped_column()

    # works: str(UserType.T1)
    # does not work: UserType.T1
    __table_args__ = (
        Index(None, account_id, unique=True, postgresql_where=(type == UserType.T1)),
    )

# comparing to StrEnum member works fine with select.where(..)
print(select(User).where(User.type == UserType.T1))

# throws the error: if "(type == UserType.T1)" is used
# throws no error: if "(type == str(UserType.T1))" is used
print(CreateIndex(User.__table_args__[0]).compile(dialect=postgresql.dialect()))

Error

Traceback (most recent call last): File "/testtttt.py", line 37, in print(CreateIndex(User.table_args__[0]).compile(dialect=postgresql.dialect())) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/sqlalchemy/sql/elements.py", line 280, in compile return self._compiler(dialect, kw) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/sqlalchemy/sql/ddl.py", line 69, in _compiler return dialect.ddl_compiler(dialect, self, kw) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/sqlalchemy/sql/compiler.py", line 866, in init__ self.string = self.process(self.statement, compile_kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/sqlalchemy/sql/compiler.py", line 911, in process return obj._compiler_dispatch(self, kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/sqlalchemy/sql/visitors.py", line 143, in _compiler_dispatch return meth(self, kw) # type: ignore # noqa: E501 ^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/sqlalchemy/dialects/postgresql/base.py", line 2313, in visit_create_index where_compiled = self.sql_compiler.process( ^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/sqlalchemy/sql/compiler.py", line 911, in process return obj._compiler_dispatch(self, kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/sqlalchemy/sql/visitors.py", line 143, in _compiler_dispatch return meth(self, **kw) # type: ignore # noqa: E501 ^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/sqlalchemy/sql/compiler.py", line 3321, in visit_binary return self._generate_generic_binary( ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/sqlalchemy/sql/compiler.py", line 3376, in _generate_generic_binary

binary.right._compiler_dispatch( ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/sqlalchemy/sql/visitors.py", line 143, in _compiler_dispatch return meth(self, **kw) # type: ignore # noqa: E501 ^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/sqlalchemy/sql/compiler.py", line 3612, in visit_bindparam ret = self.render_literal_bindparam( ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/sqlalchemy/sql/compiler.py", line 3751, in render_literal_bindparam return self.render_literal_value(value, bindparam.type) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/sqlalchemy/dialects/postgresql/base.py", line 1865, in render_literal_value value = super().render_literalvalue(value, type) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/sqlalchemy/sql/compiler.py", line 3778, in render_literal_value raise exc.CompileError( sqlalchemy.exc.CompileError: No literal value renderer is available for literal value "<UserType.T1: 't1'>" with datatype NULL

Additional context

No response

zzzeek commented 1 year ago

you are using type before declarative has any chance to set up the actual column, so you have to give it a type:

class User(Base):
    __tablename__ = "a"

    id: Mapped[int] = mapped_column(primary_key=True)
    account_id: Mapped[int] = mapped_column()
    type: Mapped[str] = mapped_column(String)

    # works: str(UserType.T1)
    # does not work: UserType.T1
    __table_args__ = (
        Index(None, account_id, unique=True, postgresql_where=(type == UserType.T1)),
    )

zzzeek commented 1 year ago

there's nothing unexpected here, User.type can't possibly have its full typing information set up within the class body like that, so closing

CaselIT commented 1 year ago

Not sure if it could make sense to recheck the types when set parent executes?

zzzeek commented 1 year ago

the binary expression object is generated when that type == thing is executed. that's where it fails. there's not a mechanism to sweep through the structure a second time and correct all the datatypes and other configuration that was not there up front.

I can see two things that could happen here:

we document this kind of use
mapped_column() emits a warning if it has NULLTYPE and is used in an expression

CaselIT commented 1 year ago

Yes, my suggestion before was to either traverse the binary expression or something along these lines. But it's probably not great.

Documentation / warning could be useful though

zzzeek commented 1 year ago

there's a complexity / usefulness ratio here that would be way out of whack for "build a new system where the compiler automatically re-traverses a whole expression tree and (mutates it in place?) so that any column objects that came in without their final type can be corrected for" / "we are using declarative __table_args__ and for one particular column we need a binary expression pulled from the mapping and we didn't know to just add an explicit type to it for this use"

mapped_column() has some specific structure that allows it to be used in live expressions in this way before declarative has occurred, so having it warn for NULLTYPE should be straightforward

CaselIT commented 1 year ago

now that you write it I tend to agree :)

zzzeek commented 1 year ago

another way to do it woudl be that declarative does all the columns first, sets up their types, then when it does __table_args__, searches in all of them for expressions and does a rebuild of each expression into something new that would set up the types again. that would still be a big new complicated thing but not as bad.

but it's way less intrusive and complicated to just warn that this column doesnt have a type yet

CaselIT commented 1 year ago

searches in all of them for expressions and does a rebuild of each expression into something new that would set up the types again. that would still be a big new complicated thing but not as bad.

that was kind of my thinking, index set_parent would look into the expression and see if the types are the final ones or not

sqlalchemy / sqlalchemy