Make sure the extra translation queries follow db routing (bug 905887)

apps/translations/tests/test_models.py

@@ -1,10 +1,17 @@
 # -*- coding: utf-8 -*-
-from django.conf import settings
+from contextlib import nested
+
 from django import test
+from django.conf import settings
+from django.db import connections
+from django.test.utils import override_settings
 from django.utils import translation
 from django.utils.functional import lazy
 
 import jinja2
+import mock
+from multidb.pinning import use_master
+from nose import SkipTest
 from nose.tools import eq_
 from test_utils import ExtraAppTestCase, trans_eq
 
@@ -346,6 +353,59 @@ def test_require_locale(self):
         eq_(unicode(obj.no_locale), 'blammo')
         eq_(obj.no_locale.locale, 'fr')
 
+    @override_settings(DEBUG=True)
+    def test_translations_reading_from_multiple_db(self):
+        # Pretend we are using a slave. Needs to be done inside the test because
+        # we need to copy default database.
+        settings.DATABASES['slave-1'] = settings.DATABASES['default'].copy()
+        settings.DATABASES['slave-2'] = settings.DATABASES['default'].copy()
+
+        # Make sure we are in a clean environnement.
+        for dbname in settings.DATABASES.keys():
+            connections[dbname].queries = []
+
+        with mock.patch('multidb.get_slave', lambda: 'slave-2'):
+            TranslatedModel.objects.get(pk=1)
+            eq_(len(connections['default'].queries), 0)
+            eq_(len(connections['slave-1'].queries), 0)
+            eq_(len(connections['slave-2'].queries), 3)
+
+    @override_settings(DEBUG=True)
+    def test_translations_reading_from_multiple_db_using(self):
+        raise SkipTest('Will need a django-queryset-transform patch to work')
+        # Pretend we are using a slave. Needs to be done inside the test because
+        # we need to copy default database.
+        settings.DATABASES['slave-1'] = settings.DATABASES['default'].copy()
+        settings.DATABASES['slave-2'] = settings.DATABASES['default'].copy()
+
+        # Make sure we are in a clean environnement.
+        for dbname in settings.DATABASES.keys():
+            connections[dbname].queries = []
+
+        with mock.patch('multidb.get_slave', lambda: 'slave-2'):
+            TranslatedModel.objects.no_cache().using('slave-1').get(pk=1)
+            eq_(len(connections['default'].queries), 0)
+            eq_(len(connections['slave-1'].queries), 3)
+            eq_(len(connections['slave-2'].queries), 0)
+
+    @override_settings(DEBUG=True)
+    def test_translations_reading_from_multiple_db_pinning(self):
+        # Pretend we are using a slave. Needs to be done inside the test because
+        # we need to copy default database.
+        settings.DATABASES['slave-1'] = settings.DATABASES['default'].copy()
+        settings.DATABASES['slave-2'] = settings.DATABASES['default'].copy()
+
+        # Make sure we are in a clean environnement.
+        for dbname in settings.DATABASES.keys():
+            connections[dbname].queries = []
+
+        with nested(mock.patch('multidb.get_slave', lambda: 'slave-2'),
+                    use_master):
+            TranslatedModel.objects.get(pk=1)
+            eq_(len(connections['default'].queries), 3)
+            eq_(len(connections['slave-1'].queries), 0)
+            eq_(len(connections['slave-2'].queries), 0)
+
 
 def test_translation_bool():
     t = lambda s: Translation(localized_string=s)

apps/translations/transformer.py

@@ -1,9 +1,7 @@
 from django.conf import settings
-from django.db import connections, models
+from django.db import connections, models, router
 from django.utils import translation
 
-import multidb
-
 from translations.models import Translation
 from translations.fields import TranslatedField
 
@@ -67,14 +65,16 @@ def get_trans(items):
     if not items:
         return
 
-    connection = connections[multidb.get_slave()]
-    cursor = connection.cursor()
-
     model = items[0].__class__
+    # FIXME: if we knew which db the queryset we are transforming used, we could
+    # make sure we are re-using the same one.
+    dbname = router.db_for_read(model)
+    connection = connections[dbname]
     sql, params = build_query(model, connection)
     item_dict = dict((item.pk, item) for item in items)
     ids = ','.join(map(str, item_dict.keys()))
 
+    cursor = connection.cursor()
     cursor.execute(sql.format(ids='(%s)' % ids), tuple(params))
     step = len(trans_fields)
     for row in cursor.fetchall():