How to use Enum with SQLAlchemy and Alembic? - python

Here's my Post model:
class Post(Base):
__tablename__ = 'posts'
title = db.Column(db.String(120), nullable=False)
description = db.Column(db.String(2048), nullable=False)
I'd like to add Enum status to it. So, I've created a new Enum:
import enum
class PostStatus(enum.Enum):
DRAFT='draft'
APPROVE='approve'
PUBLISHED='published'
And added a new field to model:
class Post(Base):
...
status = db.Column(db.Enum(PostStatus), nullable=False, default=PostStatus.DRAFT.value, server_default=PostStatus.DRAFT.value)
After doing FLASK_APP=server.py flask db migrate, a such migration was generated:
def upgrade():
op.add_column('posts', sa.Column('status', sa.Enum('DRAFT', 'APPROVE', 'PUBLISHED', name='poststatus'), server_default='draft', nullable=False))
After trying to upgrade DB, I'm getting:
sqlalchemy.exc.ProgrammingError: (psycopg2.ProgrammingError) type "poststatus" does not exist
LINE 1: ALTER TABLE posts ADD COLUMN status poststatus DEFAULT 'draf...
^
[SQL: "ALTER TABLE posts ADD COLUMN status poststatus DEFAULT 'draft' NOT NULL"]
Why type poststatus was not created on DB-level automatically? In the similar migration it was.
How to specify server_default option properly? I need both ORM-level defaults and DB-level ones, because I'm altering existing rows, so ORM defaults are not applied.
Why real values in DB are 'DRAFT', 'APPROVE', 'PUBLISHED', but not draft, etc? I supposed there should be ENUM values, not names.
Thank you in advance.

Why real values in DB are 'DRAFT', 'APPROVE', 'PUBLISHED', but not draft, etc? I supposed there should be ENUM values, not names.
As Peter Bašista's already mentioned SQLAlchemy uses the enum names (DRAFT, APPROVE, PUBLISHED) in the database. I assume that was done because the enum values ("draft", "approve", ...) can be arbitrary types in Python and they are not guaranteed to be unique (unless #unique is used).
However since SQLAlchemy 1.2.3 the Enum class accepts a parameter values_callable which can be used to store enum values in the database:
status = db.Column(
db.Enum(PostStatus, values_callable=lambda obj: [e.value for e in obj]),
nullable=False,
default=PostStatus.DRAFT.value,
server_default=PostStatus.DRAFT.value
)
Why type poststatus was not created on DB-level automatically? In the similar migration it was.
I think basically you are hitting a limitation of alembic: It won't handle enums on PostgreSQL correctly in some cases. I suspect the main issue in your case is Autogenerate doesn't correctly handle postgresql enums #278.
I noticed that the type is created correctly if I use alembic.op.create_table so my workaround is basically:
enum_type = SQLEnum(PostStatus, values_callable=lambda enum: [e.value for e in enum])
op.create_table(
'_dummy',
sa.Column('id', Integer, primary_key=True),
sa.Column('status', enum_type)
)
op.drop_table('_dummy')
c_status = Column('status', enum_type, nullable=False)
add_column('posts', c_status)

Use the following function example in case you are using PostgreSQL:
from sqlalchemy.dialects import postgresql
from ... import PostStatus
from alembic import op
import sqlalchemy as sa
def upgrade():
post_status = postgresql.ENUM(PostStatus, name="status")
post_status.create(op.get_bind(), checkfirst=True)
op.add_column('posts', sa.Column('status', post_status))
def downgrade():
post_status = postgresql.ENUM(PostStatus, name="status")
post_status.drop(op.get_bind())

This and related StackOverflow threads resort to PostgreSQL dialect-specific typing. However, generic support may be easily achieved in an Alembic migration as follows.
First, import the Python enum, the SQLAlchemy Enum, and your SQLAlchemy declarative base wherever you're going to declare your custom SQLAlchemy Enum column type.
import enum
from sqlalchemy import Enum
from sqlalchemy.ext.declarative import declarative_base
Base = declarative_base()
Let's take OP's original Python enumerated class:
class PostStatus(enum.Enum):
DRAFT='draft'
APPROVE='approve'
PUBLISHED='published'
Now we create a SQLAlchemy Enum instantiation:
PostStatusType: Enum = Enum(
PostStatus,
name="post_status_type",
create_constraint=True,
metadata=Base.metadata,
validate_strings=True,
)
When you run your Alembic alembic revision --autogenerate -m "Revision Notes" and try to apply the revision with alembic upgrade head, you'll likely get an error about the type not existing. For example:
...
sqlalchemy.exc.ProgrammingError: (psycopg2.errors.UndefinedObject) type "post_status_type" does not exist
LINE 10: post_status post_status_type NOT NULL,
...
To fix this, import your SQLAlchemy Enum class and add the following to your upgrade() and downgrade() functions in the Alembic autogenerated revision script.
from myproject.database import PostStatusType
...
def upgrade() -> None:
PostStatusType.create(op.get_bind(), checkfirst=True)
... the remainder of the autogen code...
def downgrade() -> None:
...the autogen code...
PostStatusType.drop(op.get_bind(), checkfirst=True)
Finally, be sure to update the auto-generated sa.Column() declaration in the table(s) using the enumerated type to simply reference the SQLAlchemy Enum type instead of using Alembic's attempt to re-declare it. For example in def upgrade() -> None:
op.create_table(
"my_table",
sa.Column(
"post_status",
PostStatusType,
nullable=False,
),
)

I can only answer the third part of your question.
The documentation for the Enum type in SQLAlchemy states that:
Above, the string names of each element, e.g. “one”, “two”, “three”, are persisted to the database; the values of the Python Enum, here indicated as integers, are not used; the value of each enum can therefore be any kind of Python object whether or not it is persistable.
So, it is by SQLAlchemy design that Enum names, not values are persisted into the database.

Related

SQLAlchemy: ForeignKey across schemas

In my postgres server we have a database database with 2 schemas: public and api.
public has several tables, and I need to create a table in api with a foreign key to a table in public called model.
So it's:
-Schemas
--public
---tables
----models
--api
---tables
Using SQLAlchemy I have the following class:
from sqlalchemy import create_engine, MetaData, Table, Column
class __PostgresService:
def __init__(self):
self.__client = create_engine("postgresql://postgres#localhost:5432/database")
metadata = MetaData(self.__client, schema="public")
self.__table = Table("training", metadata,
Column("id", String, primary_key=True, nullable=False),
Column("model_id", ForeignKey("model.id"), nullable=False),
schema="api")
metadata.create_all()
postgres_service = __PostgresService()
However upon launch I receive the following error:
sqlalchemy.exc.NoReferencedTableError: Foreign key associated with column 'training.model_id' could not find table 'public.model' with which to generate a foreign key to target column 'id'
It seems it does look for the correct thing but can't find it? I'm very confused as to why this is happening, especially because the error refers to not finding "public", which is created by default by postgres, rather than "api" which I created myself in pgAdmin.
Am I missing some cruicial config?
The error you are getting means that you are trying to create a foreign key referencing a table that SQLAlchemy does not know about. You can tell sqlalchemy about it by creating a Table associated with the same MetaData describing the referenced table. You can also do this using sqlalchemy's reflection capabilities. For example:
from sqlalchemy import create_engine, MetaData, Table, Column
class __PostgresService:
def __init__(self):
self.__client = create_engine("postgresql://postgres#localhost:5432/database")
metadata = MetaData(self.__client, schema="public")
metadata.reflect(schema="public", only=["model"])
self.__table = Table("training", metadata,
Column("id", String, primary_key=True, nullable=False),
Column("model_id", ForeignKey("model.id"), nullable=False),
schema="api")
metadata.create_all()
postgres_service = __PostgresService()
By default, MetaData.create_all() will check for the existence of tables first, before creating them, but you can also specify the exact tables to create: metadata.create_all(tables=[self.__table])

do not know how to remove sqlalchemy 1.4 warning about cache_ok

I am using sqlalchemy 1.4.17 with postgres and have a pytest-asyncio test which calls a function that creates a record that contains a uuid.
async def create_user(session: AsyncSession, input_user_data):
new_user = model.User(**dict(input_user_data))
session.add(new_user)
await session.commit()
class User(Base):
__tablename__ = "user"
id = Column(GUID, primary_key=True,
server_default=DefaultClause(text("gen_random_uuid()")))
and it runs ok, but creates a warning
sys:1: SAWarning: TypeDecorator GUID() will not produce a cache key because the ``cache_ok`` flag is not set to True. Set this flag to True if this type object's state is safe to use in a cache key, or False to disable this warning.
and don't understand how to silence it. any help is appreciated!
thanks #snakecharmerb. that pointed me to what I was doing wrong. in case this helps anyone else, I was importing GUID from fastapi_utils and instead imported right from sqlalchemy
# from fastapi_utils.guid_type import GUID, GUID_SERVER_DEFAULT_POSTGRESQL
from sqlalchemy.dialects.postgresql import UUID
class User(Base):
__tablename__ = "user"
id = Column(UUID, primary_key=True,
server_default=DefaultClause(text("gen_random_uuid()")))
rather than change the fastapi utils lib, using the SQL alchemy typedecorator was much easier.
If for whatever reason, you use a custom TypeDecorator, for example:
class MyUUIDType(TypeDecorator):
impl = sqlalchemy.dialects.postgresql.UUID
...
Then you'll need to add cache_ok = True as a class member
class MyUUIDType(TypeDecorator):
impl = sqlalchemy.dialects.postgresql.UUID
cache_ok = True
...
Surely do this only if your column's TypeDecorator is cacheable.

Alembic Migrations on Multiple Models

I am attempting to create a revision with --autogenerate using Alembic for two Models, but am receiving a duplicate table keys error. Does, a schema need to be specified? If so, how can it be set? The documentation I've read says to use __table_args__ = {'schema': 'somename'}, but that hasn't helped. Any tips or suggestions are greatly appreciated.
My current setup is:
base.py
from sqlalchemy.ext.declarative import declarative_base
Base = declarative_base()
workspace.py
from sqlalchemy import Column, Integer, String
from base import Base
class WorkspaceModel(Base):
__tablename__ = 'workspaces'
id = Column(Integer, primary_key=True)
name = Column(String)
host.py
from sqlalchemy import Column, Integer, String
from base import Base
class HostModel(Base):
__tablename__ = 'hosts'
id = Column(Integer, primary_key=true)
ip = Column(String)
alembic/env.py
from host import HostModel
from workspace import WorkspaceModel
target_metadata = [HostModel.metadata, WorkspaceModel.metadata]
Error
ValueError: Duplicate table keys across multiple MetaData objects: "hosts", "workspaces"
To make it clear from what #esdotzed and #univerio said, you have to use a single Base.metadata - but still import the individual models.
In the original question, this is how the alembic/env.py should look like:
from base import Base
# This two won't be referenced, but *have* to be imported to populate `Base.metadata`
from host import HostModel
from workspace import WorkspaceModel
target_metadata = Base.metadata
If you didn't import both models, the autogenerated migration would end up deleting your whole database - because Base.metadata doesn't know about any model by itself.
quoting univerio's answer from the comment section:
target_metadata should just be target_metadata = Base.metadata
Using Base.metadata doesn't mean you can remove the imports from host import HostModel and from workspace import WorkspaceModel
It worked for me.
I just want to add to #mgarciaisaia answer, it will work but the thing is when I tried changing, for example, the max length of username field of User model and running alembic revision --autogenerate -m "test migration", alembic output a migration file with empty upgrade() and downgrade() functions!
Note: the following operations will erase your data from the database so please back them up beforehand!
In order to update the changes made to the original User model, I had to
Delete the first migration file
Rerun alembic revision --autogenerate -m "update user model" and alembic upgrade head again for the changes to appear inside upgrade() and downgrade() functions of the migration file.

sqlalchemy dynamic schema on entity at runtime

I'm using SQL Alchemy and have some schema's that are account specific. The name of the schema is derived using the account ID, so I don't have the name of the schema until I hit my application service or repository layer. I'm wondering if it's possible to run a query against an entity that has it's schema dynamically set at runtime?
I know I need to set the __table_args__['schema'] and have tried doing that using the type() built-in, but I always get the following error:
could not assemble any primary key columns for mapped table
I'm ready to give up and just write straight sql, but I really hate to do that. Any idea how this can be done? I'm using SA 0.99 and I do have a PK mapped.
Thanks
from sqlalchemy 1.1,
this can be done easily using using schema_translation_map.
https://docs.sqlalchemy.org/en/11/changelog/migration_11.html#multi-tenancy-schema-translation-for-table-objects
One option would be to reflect the particular account-dependent tables. Here is the SqlAlchemy Documentation on the matter.
Alternatively, You can create the table with a static schema attribute and update it as needed at runtime and run the queries you need to. I can't think of a non-messy way to do this. So here's the messy option
Use a loop to update the schema property in each table definition whenever the account is switched.
add all the tables that are account-specific to a list.
if the tables are expressed in the declarative syntax, then you have to modify the DeclarativeName.__table__.schema attribute. I'm not sure if you need to also modify DeclarativeName.__table_args__['schema'], but I guess it won't hurt.
If the tables are expressed in the old style Table syntax, then you have to modify the Table.schema attribute.
If you're using text for any relationships or foreign keys, then that will break, and you have to inspect each table for such hard coded usage and change them
example
user_id = Column(ForeignKey('my_schema.user.id')) needs to be written as user_id = Column(ForeignKey(User.id)). Then you can change the schema of User to my_new_schema. Otherwise, at query time sqlalchemy will be confused because the foreign key will point to my_schema.user.id while the query would point to my_new_schema.user.
I'm not sure if more complicated relationships can be expressed without the use of plain text, so I guess that's the limit to my proposed solution.
Here's an example I wrote up in the terminal:
>>> from sqlalchemy import Column, Table, Integer, String, select, ForeignKey
>>> from sqlalchemy.orm import relationship, backref
>>> from sqlalchemy.ext.declarative import declarative_base
>>> B = declarative_base()
>>>
>>> class User(B):
... __tablename__ = 'user'
... __table_args__ = {'schema': 'first_schema'}
... id = Column(Integer, primary_key=True)
... name = Column(String)
... email = Column(String)
...
>>> class Posts(B):
... __tablename__ = 'posts'
... __table_args__ = {'schema':'first_schema'}
... id = Column(Integer, primary_key=True)
... user_id = Column(ForeignKey(User.id))
... text = Column(String)
...
>>> str(select([User.id, Posts.text]).select_from(User.__table__.join(Posts)))
'SELECT first_schema."user".id, first_schema.posts.text \nFROM first_schema."user" JOIN first_schema.posts ON first_schema."user".id = first_schema.posts.user_id'
>>> account_specific = [User, Posts]
>>> for Tbl in account_specific:
... Tbl.__table__.schema = 'second_schema'
...
>>> str(select([User.id, Posts.text]).select_from(User.__table__.join(Posts)))
'SELECT second_schema."user".id, second_schema.posts.text \nFROM second_schema."user" JOIN second_schema.posts ON second_schema."user".id = second_schema.posts.user_id'
As you see the same query refers to the second_schema after I update the table's schema attribute.
edit: Although you can do what I did here, using the schema translation map as shown in the the answer below is the proper way to do it.
They are set statically. Foreign keys needs the same treatment, and I have an additional issue, in that I have multiple schemas that contain multiple tables so I did this:
from sqlalchemy.ext.declarative import declarative_base
staging_dbase = declarative_base()
model_dbase = declarative_base()
def adjust_schemas(staging, model):
for vv in staging_dbase.metadata.tables.values():
vv.schema = staging
for vv in model_dbase.metadata.tables.values():
vv.schema = model
def all_tables():
return staging_dbase.metadata.tables.union(model_dbase.metadata.tables)
Then in my startup code:
adjust_schemas(staging=staging_name, model=model_name)
You can mod this for a single declarative base.
I'm working on a project in which I have to create postgres schemas and tables dynamically and then insert data in proper schema. Here is something I have done maybe it will help someone:
import sqlalchemy
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from app.models.user import User
engine_uri = "postgres://someusername:somepassword#localhost:5432/users"
engine = create_engine(engine_uri, pool_pre_ping=True)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
def create_schema(schema_name: str):
"""
Creates a new postgres schema
- **schema_name**: name of the new schema to create
"""
if not engine.dialect.has_schema(engine, schema_name):
engine.execute(sqlalchemy.schema.CreateSchema(schema_name))
def create_tables(schema_name: str):
"""
Create new tables for postgres schema
- **schema_name**: schema in which tables are to be created
"""
if (
engine.dialect.has_schema(engine, schema_name) and
not engine.dialect.has_table(engine, str(User.__table__.name))
):
User.__table__.schema = schema_name
User.__table__.create(engine)
def add_data(schema_name: str):
"""
Add data to a particular postgres schema
- **schema_name**: schema in which data is to be added
"""
if engine.dialect.has_table(engine, str(User.__table__.name)):
db = SessionLocal()
db.connection(execution_options={
"schema_translate_map": {None: schema_name}},
)
user = User()
user.name = "Moin"
user.salary = 10000
db.add(user)
db.commit()

Getting SQLAlchemy to issue CREATE SCHEMA on create_all

I have a SqlAlchemy model with a schema argument like so:
Base = declarative_base()
class Road(Base):
__tablename__ = "roads"
__table_args__ = {'schema': 'my_schema'}
id = Column(Integer, primary_key=True)
When I use Base.metadata.create_all(engine) it correctly issues a CREATE TABLE with the schema name on the front like so CREATE TABLE my_schema.roads ( but Postgresql rightly complains that the schema doesn't exist.
Am I missing a step to get SqlAlchemy to issue the CREATE SCHEMA my_schema or do I have to call this manually?
I have done it manually on my db init script like so:
from sqlalchemy.schema import CreateSchema
engine.execute(CreateSchema('my_schema'))
But this seems less magical than I was expecting.
I ran into the same issue and believe the "cleanest" way of issuing the DDL is something like this:
from sqlalchemy import event
from sqlalchemy.schema import CreateSchema
event.listen(Base.metadata, 'before_create', CreateSchema('my_schema'))
This will ensure that before anything contained in the metadata of your base is created, you have the schema for it. This does, however, not check if the schema already exists.
You can do CreateSchema('my_schema').execute_if(callback_=check_schema) if you can be bothered to write the check_schema callback ("Controlling DDL Sequences" on should_create in docs). Or, as an easy way out, just use DDL("CREATE SCHEMA IF NOT EXISTS my_schema") instead (for Postgres):
from sqlalchemy import DDL
event.listen(Base.metadata, 'before_create', DDL("CREATE SCHEMA IF NOT EXISTS my_schema"))
I wrote a function that creates the declared schemas based on the accepted answer. It uses the schema value from the __table_args__ dict from each mapped class.
from sqlalchemy import event, DDL
# Import or write your mapped classes and configuration here
def init_db():
for mapper in Base.registry.mappers:
cls = mapper.class_
if issubclass(cls, Base):
table_args = getattr(cls, '__table_args__', None)
if table_args:
schema = table_args.get('schema')
if schema:
stmt = f"CREATE SCHEMA IF NOT EXISTS {schema}"
event.listen(Base.metadata, 'before_create', DDL(stmt))
Base.metadata.create_all(bind=engine)

Categories