Search code examples
pythonsqlalchemyalembicsqlmodel

Dynamically set sql-default value with table-name in SQLModel


I'm trying to create a base-class in SQLModel which looks like this:

class BaseModel(SQLModel):
    @declared_attr
    def __tablename__(cls) -> str:
        return cls.__name__

    guid: Optional[UUID] = Field(default=None, primary_key=True)

class SequencedBaseModel(BaseModel):
    sequence_id: str = Field(sa_column=Column(VARCHAR(50), server_default=text(f"SELECT '{TABLENAME}_' + convert(varchar(10), NEXT VALUE FOR dbo.sequence)")))

so I got a table like this:

class Project(SequencedBaseModel):
    ...

where alembic would generate a migration for a table Project with columns guid and sequence_id. The default-value for sequence-id is a sequence which is generated with

SELECT '{TABLENAME}_' + convert(varchar(10), NEXT VALUE FOR dbo.sequence)

and should insert into project-table the values Project_1, Project_2, ...

Any idea on how to set the tablename dynamically? I cannot use a constructor for setting the columns because alembic is ignoring them, I cannot access the __tablename__() function, or cls, because the columns are static...


Solution

  • Unfortunately, if you have an attribute that relies on a @declared_attr, it must also be a @declared_attr, since SqlAlchemy will wait until the whole mapping is completed and the classes get actual tables to be resolved (at least, that is my understanding). Now: declared_attr(s) are an SqlAlchemy concept, whereas the idea of Field is an SQLModel concept and they don't seem to talk to each other regarding this sort of """deferred""" attribute thingy ("""deferred""" in the sense that it won't be evaluated until the mapping is completed and the tables known, not deferred as in waiting until it's accessed)... at least, not that I know.

    You could (maybe? hopefully?) do something like what's recommended in this SQLModel GitHub issue: """Defer""" the SqlAlchemy Column and alias the SQLModel field:

    class SequencedBaseModel(BaseModel):
        sequence_id: str = Field(alias="sequence_id")
    
        @declared_attr
        def sequence_id(cls):
            return Column(
                'sequence_id',
                VARCHAR(50),
                server_default=text(f"SELECT '{cls.__tablename__}_'"
                                    f" + convert(varchar(10), NEXT VALUE FOR dbo.sequence)"))
    
    
    class Project(SequencedBaseModel, table=True):
        pass
    

    Running alembic revision --autogenerate -m "init" will produce a migration file with the proper __tablename__ + '_' (meaning: Product_) expanded in the server_default's SELECT...:

    def upgrade() -> None:
        op.create_table('Project',
        sa.Column('guid', sqlmodel.sql.sqltypes.GUID(), nullable=False),
        sa.Column('sequence_id', sa.VARCHAR(length=50), server_default=sa.text("SELECT 'Project_' + convert(varchar(10), NEXT VALUE FOR dbo.sequence)"), nullable=True),
        sa.PrimaryKeyConstraint('guid'),
        # ... 
        # ... 
    

    This assumes your alembic environment has been properly configured. I can't help but pointing out that alembic will generate the migration using sqlmodel.sql.sqltypes.GUID() as the column type for the guid attribute, so you'll need to make sure the package sqlmodel is imported on each migration file. Probably by editing the template script.py.mako as described in this link where it shows that you must add import sqlmodel # NEW .

    ⚠️ I tried to test this, but I don't exactly know where dbo.sequence comes from (a SQL server, maybe?). I used a PostgreSQL sequence (which I christened so75719072) to simulate it. This means that I can't confirm whether the syntax for the DEFAULT SELECT... will be valid in your situation. I'm quite suspicious about you being able to use the result of a SELECT as default value for a column but hopefully I'm wrong.

    import uuid
    from typing import Optional
    from uuid import UUID
    
    from sqlalchemy import Column, VARCHAR, text
    from sqlalchemy.orm import declared_attr
    from sqlmodel import SQLModel, Field, create_engine, Session, select
    
    
    class BaseModel(SQLModel):
        __table_args__ = {'schema': 'SO-75719072'}
    
        @declared_attr
        def __tablename__(cls) -> str:
            return cls.__name__
    
        guid: Optional[UUID] = Field(default=None, primary_key=True)
    
    
    class SequencedBaseModel(BaseModel):
        sequence_id: str = Field(alias="sequence_id")
    
        @declared_attr
        def sequence_id(cls):
            return Column(
                'sequence_id',
                VARCHAR(50),
                server_default=text(f"nextval('so75719072')"))
    
    
    class Project(SequencedBaseModel, table=True):
        pass
    
    
    if __name__ == "__main__":
        engine = create_engine(
            "postgresql+psycopg2://postgres:postgrespw@localhost:32768/stackoverflow")
    
        with Session(engine) as session:
            for i in range(3):
                proj1 = Project(guid=uuid.uuid4())
                session.add(proj1)
                session.commit()
    
        with Session(engine) as session:
            statement = select(Project).where(Project.sequence_id.in_(["1", "2", "3"]))
            for project in session.exec(statement):
                print(f"guid: {project.guid}")
    

    Produces the following output:

    guid: c5e5902d-e224-48f1-95f5-fa47a73f7b05
    guid: 1c25550b-258c-49c5-9acc-90ae7ad8460c
    guid: eb84e90c-9449-4974-8eb4-bad98728b0f9
    

    Which came from the following table in Postgres:

    # select * from "SO-75719072"."Project";
                     guid                 | sequence_id
    --------------------------------------+-------------
     c5e5902d-e224-48f1-95f5-fa47a73f7b05 | 1
     1c25550b-258c-49c5-9acc-90ae7ad8460c | 2
     eb84e90c-9449-4974-8eb4-bad98728b0f9 | 3
    (3 rows)