Multi-Tenancy in Flask Application

30 June 2021 . 6 minute read

By Abhishek Aravindan

When I started developing a web application I had a requirement to make it multi-tenant.

So from the basics, I started researching how to implement multi-tenancy in Flask applications.

I went with the obvious one since I’m using PostgreSQL. Multiple Schema single database models like Apartment gem in RoR.

I failed to find libraries for my needs, Tutorials and documentation raised many confusions to do my requirement. So I started breaking up my requirements and solving each one.

The first thing was creating the schema and creating all tables in the new schema, Technically making a replica of the public schema as a new tenant which I have to do at runtime. So after a couple of research and coffees, I came up with this.

from flask import current_app as app
from app.models import UserModel
from app import db
import sqlalchemy
from flask_sqlalchemy import SQLAlchemyfrom sqlalchemy import create_engine, MetaData, Table
def method_name():
.....
engine = create_engine(app.config['SQLALCHEMY_DATABASE_URI'])
        if not engine.dialect.has_schema(engine, schema_name):
            engine.execute(sqlalchemy.schema.CreateSchema(schema_name))
        insp = sqlalchemy.inspect(engine)
        if engine.dialect.has_schema(engine, schema_name):
            for mapper in db.Model.registry.mappers:
                if not insp.has_table(mapper.class_, schema=schema_name):
                    meta = MetaData()
                    mapper.class_.__table__.schema = subdomain
                    mapper.class_.__table__.create(engine)
....

You can add the above schema in any method where you need to create the new tenant.

The main intention of the above code is to create a new schema > create the tables in the new schema(list of tables are fetching from our models)

Now, The second part we created schemas. What if we need to make changes to schemas, Like alteration of database structure like creating new columns, new tables, delete tables, columns etc…

For this, we need to modify a file that is generated by the flask.

Everyone must have seen the env.py in the migrations folder

So I’m going to modify some of the code,

I’m going for a loop to implement the migration for each schema in the database.

from sqlalchemy import engine_from_config, pool, MetaData, Table, ForeignKeyConstraint
from logging.config import fileConfig
from app import db
from app.models import UserModel
....
def run_migrations_online():
    """Run migrations in 'online' mode.
    In this scenario we need to create an Engine
    and associate a connection with the context.
    """
    engine = engine_from_config(
                config.get_section(config.config_ini_section),
                prefix='sqlalchemy.',
                poolclass=pool.NullPool)
connection = engine.connect()
    context.configure(
        connection=connection,
        target_metadata=target_metadata,
        include_schemas=True,
        include_object=include_schemas([None])
)
domains = ['public']
try:
connection.execute(get_schemas_query)]
        users = UserModel.query.all()
        for user in users:
            domains.append(user.schema)
        for domain in domains:
            print('migrating tenant -> '+ domain)
            connection.execute('set search_path to "{}"'.format(domain))
            with context.begin_transaction():
                context.run_migrations()
finally:
        connection.close()

I have saved my schema for each user in the “users” table. So I can fetch the list of schemas and migrate to each schema separately.

Note that before you do any migration make sure you run

> flask db stamp head

This will make your newly created schema up to date as the default schema (here public schema).

Now you can run the rest of the migration command

> flask db migrate
> flask db upgrade

Now the last requirement is to switch tenants while the application is running.

@app.before_request
def before_request():
    subdomain = request.host.split('.')[0]
    db.session.execute('set search_path to "{}"'.format(subdomain))

Before each request made to the server, the server fill fetch the schema name from the host URL and set search path to the desired schema

Conclusion

Here it is, there are plenty of ways to do this and I’m sure the above methods and code need to be improved. If anyone has suggestions please raise comments so I can improve the code and the article.

Thank you for reading.