The code below is what I ended up with to have this functionality.
A few notes:
- I changed
get_table_for_bind
to bind all tables without an explicit __bind_key__
to all the binds. This is done in order to be able to call db.create_all()
or db.drop_all()
and create/drop the tables in all the DBs. In order for this to work and not break the default DB selection, when not specifying a specific bind, get_binds
was changed to map the None
bind again after the original implementation, to override the Table->Bind mapping.
- If you don't specify a
using_bind
everything should work with the default DB.
- SQLAlchemy mapped objects keep a reference to the session and state so you can't really add the same object to two DBs. I made a copy of the object before adding it in order to persist it in two DBs. Not sure if there is some better way to do this.
- I haven't fully tested this and this might break some other functionality I'm not using or not aware of.
flask-sqlalchemy overrides:
from flask_sqlalchemy import SQLAlchemy, SignallingSession, get_state
from flask_sqlalchemy._compat import itervalues
class UsingBindSignallingSession(SignallingSession):
def get_bind(self, mapper=None, clause=None):
if self._name:
_eng = get_state(self.app).db.get_engine(self.app,bind=self._name)
return _eng
else:
return super(UsingBindSignallingSession, self).get_bind(mapper, clause)
_name = None
def using_bind(self, name):
self._name = name
return self
class UsingBindSQLAlchemy(SQLAlchemy):
def create_session(self, options):
return UsingBindSignallingSession(self, **options)
def get_binds(self, app=None):
retval = super(UsingBindSQLAlchemy, self).get_binds(app)
# get the binds for None again in order to make sure that it is the default bind for tables
# without an explicit bind
bind = None
engine = self.get_engine(app, bind)
tables = self.get_tables_for_bind(bind)
retval.update(dict((table, engine) for table in tables))
return retval
def get_tables_for_bind(self, bind=None):
"""Returns a list of all tables relevant for a bind.
Tables without an explicit __bind_key__ will be bound to all binds.
"""
result = []
for table in itervalues(self.Model.metadata.tables):
# if we don't have an explicit __bind_key__ bind this table to all databases
if table.info.get('bind_key') == bind or table.info.get('bind_key') == None:
result.append(table)
return result
db = UsingBindSQLAlchemy()
Now you can do this:
# This is the default DB
SQLALCHEMY_DATABASE_URI=YOUR_MAIN_DB_URI_CONNECT_STRING
# Master1 and Master2
SQLALCHEMY_BINDS = { 'master1':YOUR_MASTER1_DB_URI_CONNECT_STRING, 'master2':YOUR_MASTER2_DB_URI_CONNECT_STRING }
# Tables without __bind_key__ will be dropped/created on all DBs (default, master1, master2)
db.drop_all()
db.create_all()
s = db.session().using_bind('master1')
s.add(SOME_OBJECT)
s.commit()
s = db.session().using_bind('master2')
s.add(SOME_OBJECT_CLONE) # a clone of the original object, before the first add()
s.commit()
# and the default DB, as always
db.session.add(SOME_OTHER_OBJECT)
db.session.commit()