How can I extend every unique key constraint of all models with a common key across using Flask-SQLAlchemy and Flask-Migrate in Python?

Question

Want to extend every unique key constraint of all models with a common key across. Have tried multiple things but doesn't seem to working when I do

flask db init; flask db migrate -m "init"; flask db upgrade;

But it works and adds

"users_email_organization_id_key" UNIQUE CONSTRAINT, btree (email, organization_id)

when I run it as python app.py since I have db.create_all() as part of app.py

My code looks something like this. Hope someone can help, stuck here for quite some time now.

from flask import Flask
from flask_migrate import Migrate
from sqlalchemy import event
from sqlalchemy import Column, Integer, String, UniqueConstraint
from sqlalchemy.orm import declarative_base, relationship
from flask_sqlalchemy import SQLAlchemy
from sqlalchemy.orm import declared_attr, relationship, declarative_base
import os

app = Flask(__name__)
basedir = os.path.abspath(os.path.dirname(__file__))
app.config['SQLALCHEMY_DATABASE_URI'] = 'postgresql://localhost:5432/postgres'
db = SQLAlchemy(app)
migrate = Migrate(app, db)


class Organization(db.Model):
    __tablename__ = 'organizations'
    id = Column(Integer, primary_key=True)
    name = Column(String(50), unique=True)


class OrganizationMixin:
    organization_id = Column(Integer, db.ForeignKey('organizations.id'))

    @classmethod
    def extend_unique_constraints(cls):
        table = cls.__table__
        constraints_to_modify = []
        for constraint in table.constraints:
            if isinstance(constraint, UniqueConstraint):
                constraints_to_modify.append(constraint)

        for constraint in constraints_to_modify:
            table.constraints.remove(constraint)
            columns = list(constraint.columns)
            if 'organization_id' not in columns:
                columns.append(table.c.organization_id)
            uc = UniqueConstraint(*columns, name=constraint.name)
            table.append_constraint(uc)


class User(OrganizationMixin, db.Model):
    __tablename__ = 'users'
    id = Column(Integer, primary_key=True)
    username = Column(String(50), unique=True)
    email = Column(String(100), unique=True)


class Resource(OrganizationMixin, db.Model):
    __tablename__ = 'resources'
    id = Column(Integer, primary_key=True)
    name = Column(String(50), unique=True)


@event.listens_for(db.metadata, 'before_create')
def extend_unique_constraints(target, connection, **kwargs):
    models = [User, Resource]
    for model in models:
        if issubclass(model, OrganizationMixin):
            model.extend_unique_constraints()

with app.app_context():
    db.create_all()

if __name__ == '__main__':
    app.run(port=8000)

All things tried are in the question above.

score 0 · Answer 1 · answered May 27 '23 at 05:09

Alembic does not check or support changes like the ones you are trying to make programmatically with the extend_unique_constraints method, I guess thats why the changes are not reflected when using flask db migrate .

Lets try to constraint on individual columns and then trying to modify them later, directly define the multi-column UniqueConstraint

from sqlalchemy import UniqueConstraint

class User(OrganizationMixin, db.Model):
    __tablename__ = 'users'
    id = Column(Integer, primary_key=True)
    username = Column(String(50))  # removed unique=True
    email = Column(String(100))  # removed unique=True
    __table_args__ = (UniqueConstraint('username', 'organization_id', name='users_username_organization_id_key'),
                      UniqueConstraint('email', 'organization_id', name='users_email_organization_id_key'),)

class Resource(OrganizationMixin, db.Model):
    __tablename__ = 'resources'
    id = Column(Integer, primary_key=True)
    name = Column(String(50))  # removed unique=True
    __table_args__ = (UniqueConstraint('name', 'organization_id', name='resources_name_organization_id_key'),)

How can I extend every unique key constraint of all models with a common key across using Flask-SQLAlchemy and Flask-Migrate in Python?

1 Answers1