Adding support for swapping databases at run-time using a context

manager/decorator. Fixes #531

playhouse/read_slave.py

@@ -21,6 +21,7 @@ class User(BaseModel):
 master_select = SelectQuery(User).where(...)
 """
 from peewee import *
+from peewee import ExecutionContext
 
 
 class ReadSlaveModel(Model):
@@ -44,3 +45,21 @@ def raw(cls, *args, **kwargs):
         if query._sql.lower().startswith('select'):
             query.database = cls._get_read_database()
         return query
+
+
+class Using(ExecutionContext):
+    def __init__(self, database, models, with_transaction=True):
+        super(Using, self).__init__(database, with_transaction)
+        self.models = models
+
+    def __enter__(self):
+        self._orig = []
+        for model in self.models:
+            self._orig.append(model._meta.database)
+            model._meta.database = self.database
+        return super(Using, self).__enter__()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        super(Using, self).__exit__(exc_type, exc_val, exc_tb)
+        for i, model in enumerate(self.models):
+            model._meta.database = self._orig[i]

playhouse/tests/test_read_slave.py

@@ -1,5 +1,6 @@
 from peewee import *
 from playhouse.read_slave import ReadSlaveModel
+from playhouse.read_slave import Using
 from playhouse.tests.base import database_initializer
 from playhouse.tests.base import ModelTestCase
 
@@ -31,6 +32,8 @@ class Slave2(QueryLogDatabase):
 slave1 = database_initializer.get_database('sqlite', db_class=Slave1)
 slave2 = database_initializer.get_database('sqlite', db_class=Slave2)
 
+# Models to use for testing read slaves.
+
 class BaseModel(ReadSlaveModel):
     class Meta:
         database = master
@@ -45,6 +48,87 @@ class Thing(BaseModel):
     class Meta:
         read_slaves = [slave2]
 
+# Regular models to use for testing `Using`.
+
+class BaseMasterOnly(Model):
+    class Meta:
+        database = master
+
+class A(BaseMasterOnly):
+    data = CharField()
+
+class B(BaseMasterOnly):
+    data = CharField()
+
+
+class TestUsing(ModelTestCase):
+    requires = [A, B]
+
+    def setUp(self):
+        super(TestUsing, self).setUp()
+        reset()
+
+    def assertDatabaseVerb(self, expected):
+        db_and_verb = [(db, sql.split()[0]) for db, sql in queries]
+        self.assertEqual(db_and_verb, expected)
+        reset()
+
+    def test_using_context(self):
+        models = [A, B]
+
+        with Using(slave1, models, False):
+            A.create(data='a1')
+            B.create(data='b1')
+
+        self.assertDatabaseVerb([
+            ('slave1', 'INSERT'),
+            ('slave1', 'INSERT')])
+
+        with Using(slave2, models, False):
+            A.create(data='a2')
+            B.create(data='b2')
+            a_obj = A.select().order_by(A.id).get()
+            self.assertEqual(a_obj.data, 'a1')
+
+        self.assertDatabaseVerb([
+            ('slave2', 'INSERT'),
+            ('slave2', 'INSERT'),
+            ('slave2', 'SELECT')])
+
+        with Using(master, models, False):
+            query = A.select().order_by(A.data.desc())
+            values = [a_obj.data for a_obj in query]
+
+        self.assertEqual(values, ['a2', 'a1'])
+        self.assertDatabaseVerb([('master', 'SELECT')])
+
+    def test_using_transactions(self):
+        with Using(slave1, [A]) as txn:
+            list(B.select())
+            A.create(data='a1')
+
+        B.create(data='b1')
+        self.assertDatabaseVerb([
+            ('slave1', 'BEGIN'),
+            ('master', 'SELECT'),
+            ('slave1', 'INSERT'),
+            ('master', 'INSERT')])
+
+        def fail_with_exc(data):
+            with Using(slave2, [A]):
+                A.create(data=data)
+                raise ValueError('xxx')
+
+        self.assertRaises(ValueError, fail_with_exc, 'a2')
+        self.assertDatabaseVerb([
+            ('slave2', 'BEGIN'),
+            ('slave2', 'INSERT')])
+
+        with Using(slave1, [A, B]):
+            a_objs = [a_obj.data for a_obj in A.select()]
+            self.assertEqual(a_objs, ['a1'])
+
+
 class TestMasterSlave(ModelTestCase):
     requires = [User, Thing]