Merge pull request pallets-eco#40 from EnTeQuAk/multidb-sessions

mitsuhiko · mitsuhiko · commit f23fdb9e24fb · 2011-08-01T03:38:59.000-07:00
Added support for multiple databases for the session.
diff --git a/flaskext/sqlalchemy.py b/flaskext/sqlalchemy.py
@@ -184,11 +184,12 @@ def after_rollback(self, session):
 class _SignallingSession(Session):
 
     def __init__(self, db, autocommit=False, autoflush=False, **options):
-        Session.__init__(self, autocommit=autocommit, autoflush=autoflush,
-                         extension=db.session_extensions,
-                         bind=db.engine, **options)
         self.app = db.get_app()
         self._model_changes = {}
+        Session.__init__(self, autocommit=autocommit, autoflush=autoflush,
+                         extension=db.session_extensions,
+                         bind=db.engine,
+                         binds=db.get_binds(self.app), **options)
 
     def get_bind(self, mapper, clause=None):
         # mapper is None if someone tries to just get a connection
@@ -756,6 +757,20 @@ def get_tables_for_bind(self, bind=None):
                 result.append(table)
         return result
 
+    def get_binds(self, app=None):
+        """Returns a dictionary with a table->engine mapping.
+
+        This is suitable for use of sessionmaker(binds=db.get_binds(app)).
+        """
+        app = self.get_app(app)
+        binds = [None] + list(app.config.get('SQLALCHEMY_BINDS') or ())
+        retval = {}
+        for bind in binds:
+            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 _execute_for_all_tables(self, app, bind, operation):
         app = self.get_app(app)
 
diff --git a/test_sqlalchemy.py b/test_sqlalchemy.py
@@ -204,6 +204,7 @@ def _remove_files():
         atexit.register(_remove_files)
 
         app = flask.Flask(__name__)
+        app.config['SQLALCHEMY_ENGINE'] = 'sqlite://'
         app.config['SQLALCHEMY_BINDS'] = {
             'foo':      'sqlite:///' + db1,
             'bar':      'sqlite:///' + db2
@@ -254,6 +255,13 @@ class Baz(db.Model):
         self.assertEqual(len(metadata.tables), 1)
         self.assert_('baz' in metadata.tables)
 
+        # do the session have the right binds set?
+        self.assertEqual(db.get_binds(app), {
+            Foo.__table__: db.get_engine(app, 'foo'),
+            Bar.__table__: db.get_engine(app, 'bar'),
+            Baz.__table__: db.get_engine(app, None)
+        })
+
 
 class DefaultQueryClassTestCase(unittest.TestCase):
 
