Ticket #12540: t12540-r12262.2.diff

File t12540-r12262.2.diff, 53.4 KB (added by Russell Keith-Magee, 15 years ago)

Version 2 of a fix for database allocation problems. This time with a public DB Router API.

  • django/conf/global_settings.py 

    diff -r e165fea06e06 django/conf/global_settings.py
    a b  
    128128SEND_BROKEN_LINK_EMAILS = False
    129129
    130130# Database connection info.
     131# Legacy format
    131132DATABASE_ENGINE = ''           # 'postgresql_psycopg2', 'postgresql', 'mysql', 'sqlite3' or 'oracle'.
    132133DATABASE_NAME = ''             # Or path to database file if using sqlite3.
    133134DATABASE_USER = ''             # Not used with sqlite3.
    … …  
    136137DATABASE_PORT = ''             # Set to empty string for default. Not used with sqlite3.
    137138DATABASE_OPTIONS = {}          # Set to empty dictionary for default.
    138139
     140# New format
    139141DATABASES = {
    140142}
    141143
     144# Classes used to implement db routing behaviour
     145DATABASE_ROUTERS = []
     146
    142147# The email backend to use. For possible shortcuts see django.core.mail.
    143148# The default is to use the SMTP backend.
    144149# Third-party backends can be specified by providing a Python path

  • django/contrib/auth/models.py 

    diff -r e165fea06e06 django/contrib/auth/models.py
    a b  
    33
    44from django.contrib import auth
    55from django.core.exceptions import ImproperlyConfigured
    6 from django.db import models, DEFAULT_DB_ALIAS
     6from django.db import models
    77from django.db.models.manager import EmptyManager
    88from django.contrib.contenttypes.models import ContentType
    99from django.utils.encoding import smart_str

  • django/contrib/contenttypes/generic.py 

    diff -r e165fea06e06 django/contrib/contenttypes/generic.py
    a b  
    55from django.core.exceptions import ObjectDoesNotExist
    66from django.db import connection
    77from django.db.models import signals
    8 from django.db import models, DEFAULT_DB_ALIAS
     8from django.db import models
    99from django.db.models.fields.related import RelatedField, Field, ManyToManyRel
    1010from django.db.models.loading import get_model
    1111from django.forms import ModelForm
    … …  
    255255                    raise TypeError("'%s' instance expected" % self.model._meta.object_name)
    256256                setattr(obj, self.content_type_field_name, self.content_type)
    257257                setattr(obj, self.object_id_field_name, self.pk_val)
    258                 obj.save(using=self.instance._state.db)
     258                obj.save()
    259259        add.alters_data = True
    260260
    261261        def remove(self, *objs):

  • django/contrib/contenttypes/models.py 

    diff -r e165fea06e06 django/contrib/contenttypes/models.py
    a b  
    1 from django.db import models, DEFAULT_DB_ALIAS
     1from django.db import models
    22from django.utils.translation import ugettext_lazy as _
    33from django.utils.encoding import smart_unicode
    44

  • django/contrib/gis/db/models/sql/query.py 

    diff -r e165fea06e06 django/contrib/gis/db/models/sql/query.py
    a b  
    1 from django.db import connections, DEFAULT_DB_ALIAS
     1from django.db import connections
    22from django.db.models.query import sql
    33
    44from django.contrib.gis.db.models.fields import GeometryField

  • django/db/__init__.py 

    diff -r e165fea06e06 django/db/__init__.py
    a b  
    11from django.conf import settings
    22from django.core import signals
    33from django.core.exceptions import ImproperlyConfigured
    4 from django.db.utils import ConnectionHandler, load_backend
     4from django.db.utils import ConnectionHandler, ConnectionRouter, load_backend, DEFAULT_DB_ALIAS
    55from django.utils.functional import curry
    66
    77__all__ = ('backend', 'connection', 'connections', 'DatabaseError',
    88    'IntegrityError', 'DEFAULT_DB_ALIAS')
    99
    10 DEFAULT_DB_ALIAS = 'default'
     10
    1111
    1212# For backwards compatibility - Port any old database settings over to
    1313# the new values.
    … …  
    6060        database['ENGINE'] = full_engine
    6161
    6262connections = ConnectionHandler(settings.DATABASES)
    63 
     63router = ConnectionRouter(settings.DATABASE_ROUTERS)
    6464
    6565# `connection`, `DatabaseError` and `IntegrityError` are convenient aliases
    6666# for backend bits.

  • django/db/models/base.py 

    diff -r e165fea06e06 django/db/models/base.py
    a b  
    1010from django.db.models.query import delete_objects, Q
    1111from django.db.models.query_utils import CollectedObjects, DeferredAttribute
    1212from django.db.models.options import Options
    13 from django.db import connections, transaction, DatabaseError, DEFAULT_DB_ALIAS
     13from django.db import connections, router, transaction, DatabaseError, DEFAULT_DB_ALIAS
    1414from django.db.models import signals
    1515from django.db.models.loading import register_models, get_model
    1616from django.utils.translation import ugettext_lazy as _
    … …  
    439439        need for overrides of save() to pass around internal-only parameters
    440440        ('raw', 'cls', and 'origin').
    441441        """
    442         using = using or self._state.db or DEFAULT_DB_ALIAS
     442        using = using or router.db_for_write(self.__class__, self)
    443443        connection = connections[using]
    444444        assert not (force_insert and force_update)
    445445        if cls is None:
    … …  
    592592            parent_obj._collect_sub_objects(seen_objs)
    593593
    594594    def delete(self, using=None):
    595         using = using or self._state.db or DEFAULT_DB_ALIAS
     595        using = using or router.db_for_write(self.__class__, self)
    596596        connection = connections[using]
    597597        assert self._get_pk_val() is not None, "%s object can't be deleted because its %s attribute is set to None." % (self._meta.object_name, self._meta.pk.attname)
    598598
    … …  
    719719                    # no value, skip the lookup
    720720                    continue
    721721                if f.primary_key and not getattr(self, '_adding', False):
    722                     # no need to check for unique primary key when editting
     722                    # no need to check for unique primary key when editing
    723723                    continue
    724724                lookup_kwargs[str(field_name)] = lookup_value
    725725

  • django/db/models/fields/related.py 

    diff -r e165fea06e06 django/db/models/fields/related.py
    a b  
    1 from django.db import connection, transaction, DEFAULT_DB_ALIAS
     1from django.conf import settings
     2from django.db import connection, router, transaction
    23from django.db.backends import util
    34from django.db.models import signals, get_model
    45from django.db.models.fields import (AutoField, Field, IntegerField,
    … …  
    218219            raise ValueError('Cannot assign "%r": "%s.%s" must be a "%s" instance.' %
    219220                                (value, instance._meta.object_name,
    220221                                 self.related.get_accessor_name(), self.related.opts.object_name))
     222        elif value is not None:
     223            if instance._state.db is None:
     224                instance._state.db = router.db_for_write(instance.__class__, value)
     225            elif value._state.db is None:
     226                value._state.db = router.db_for_write(value.__class__, instance)
     227            elif value._state.db is not None and instance._state.db is not None:
     228                if not router.allow_relation(value, instance):
     229                    raise ValueError('Cannot assign "%r": instance is on database "%s", value is is on database "%s"' %
     230                                        (value, instance._state.db, value._state.db))
    221231
    222232        # Set the value of the related field to the value of the related object's related field
    223233        setattr(value, self.related.field.attname, getattr(instance, self.related.field.rel.get_related_field().attname))
    … …  
    260270            # If the related manager indicates that it should be used for
    261271            # related fields, respect that.
    262272            rel_mgr = self.field.rel.to._default_manager
    263             using = instance._state.db or DEFAULT_DB_ALIAS
     273            using = router.db_for_read(self.field.rel.to, instance)
    264274            if getattr(rel_mgr, 'use_for_related_fields', False):
    265275                rel_obj = rel_mgr.using(using).get(**params)
    266276            else:
    … …  
    281291            raise ValueError('Cannot assign "%r": "%s.%s" must be a "%s" instance.' %
    282292                                (value, instance._meta.object_name,
    283293                                 self.field.name, self.field.rel.to._meta.object_name))
    284         elif value is not None and value._state.db != instance._state.db:
     294        elif value is not None:
    285295            if instance._state.db is None:
    286                 instance._state.db = value._state.db
    287             else:#elif value._state.db is None:
    288                 value._state.db = instance._state.db
    289 #            elif value._state.db is not None and instance._state.db is not None:
    290 #                raise ValueError('Cannot assign "%r": instance is on database "%s", value is is on database "%s"' %
    291 #                                    (value, instance._state.db, value._state.db))
     296                instance._state.db = router.db_for_write(instance.__class__, value)
     297            elif value._state.db is None:
     298                value._state.db = router.db_for_write(value.__class__, instance)
     299            elif value._state.db is not None and instance._state.db is not None:
     300                if not router.allow_relation(value, instance):
     301                    raise ValueError('Cannot assign "%r": instance is on database "%s", value is is on database "%s"' %
     302                                        (value, instance._state.db, value._state.db))
    292303
    293304        # If we're setting the value of a OneToOneField to None, we need to clear
    294305        # out the cache on any old related object. Otherwise, deleting the
    … …  
    370381
    371382        class RelatedManager(superclass):
    372383            def get_query_set(self):
    373                 using = instance._state.db or DEFAULT_DB_ALIAS
     384                using = router.db_for_read(rel_model, instance)
    374385                return superclass.get_query_set(self).using(using).filter(**(self.core_filters))
    375386
    376387            def add(self, *objs):
    … …  
    378389                    if not isinstance(obj, self.model):
    379390                        raise TypeError("'%s' instance expected" % self.model._meta.object_name)
    380391                    setattr(obj, rel_field.name, instance)
    381                     obj.save(using=instance._state.db)
     392                    obj.save()
    382393            add.alters_data = True
    383394
    384395            def create(self, **kwargs):
    … …  
    390401                # Update kwargs with the related object that this
    391402                # ForeignRelatedObjectsDescriptor knows about.
    392403                kwargs.update({rel_field.name: instance})
    393                 using = instance._state.db or DEFAULT_DB_ALIAS
     404                using = router.db_for_write(rel_model, instance)
    394405                return super(RelatedManager, self).using(using).get_or_create(**kwargs)
    395406            get_or_create.alters_data = True
    396407
    … …  
    402413                        # Is obj actually part of this descriptor set?
    403414                        if getattr(obj, rel_field.attname) == val:
    404415                            setattr(obj, rel_field.name, None)
    405                             obj.save(using=instance._state.db)
     416                            obj.save()
    406417                        else:
    407418                            raise rel_field.rel.to.DoesNotExist("%r is not related to %r." % (obj, instance))
    408419                remove.alters_data = True
    … …  
    410421                def clear(self):
    411422                    for obj in self.all():
    412423                        setattr(obj, rel_field.name, None)
    413                         obj.save(using=instance._state.db)
     424                        obj.save()
    414425                clear.alters_data = True
    415426
    416427        manager = RelatedManager()
    … …  
    505516                new_ids = set()
    506517                for obj in objs:
    507518                    if isinstance(obj, self.model):
    508 #                        if obj._state.db != self.instance._state.db:
    509 #                            raise ValueError('Cannot add "%r": instance is on database "%s", value is is on database "%s"' %
    510 #                                                (obj, self.instance._state.db, obj._state.db))
     519                        if not router.allow_relation(obj, self.instance):
     520                           raise ValueError('Cannot add "%r": instance is on database "%s", value is is on database "%s"' %
     521                                               (obj, self.instance._state.db, obj._state.db))
    511522                        new_ids.add(obj.pk)
    512523                    elif isinstance(obj, Model):
    513524                        raise TypeError("'%s' instance expected" % self.model._meta.object_name)
    514525                    else:
    515526                        new_ids.add(obj)
    516                 vals = self.through._default_manager.using(self.instance._state.db).values_list(target_field_name, flat=True)
     527                db = router.db_for_read(self.through.__class__, self.instance)
     528                vals = self.through._default_manager.using(db).values_list(target_field_name, flat=True)
    517529                vals = vals.filter(**{
    518530                    source_field_name: self._pk_val,
    519531                    '%s__in' % target_field_name: new_ids,
    520532                })
    521533                new_ids = new_ids - set(vals)
    522534                # Add the ones that aren't there already
     535                db = router.db_for_write(self.through.__class__, self.instance)
    523536                for obj_id in new_ids:
    524                     self.through._default_manager.using(self.instance._state.db).create(**{
     537                    self.through._default_manager.using(db).create(**{
    525538                        '%s_id' % source_field_name: self._pk_val,
    526539                        '%s_id' % target_field_name: obj_id,
    527540                    })

  • django/db/models/manager.py 

    diff -r e165fea06e06 django/db/models/manager.py
    a b  
    11from django.utils import copycompat as copy
    2 
    3 from django.db import DEFAULT_DB_ALIAS
     2from django.conf import settings
    43from django.db.models.query import QuerySet, EmptyQuerySet, insert_query, RawQuerySet
    54from django.db.models import signals
    65from django.db.models.fields import FieldDoesNotExist
    76
     7
    88def ensure_default_manager(sender, **kwargs):
    99    """
    1010    Ensures that a Model subclass contains a default manager  and sets the
    … …  
    8787        mgr._inherited = True
    8888        return mgr
    8989
    90     def db_manager(self, alias):
     90    def db_manager(self, using):
    9191        obj = copy.copy(self)
    92         obj._db = alias
     92        obj._db = using
    9393        return obj
    9494
    9595    @property
    9696    def db(self):
    97         return self._db or DEFAULT_DB_ALIAS
     97        return self._db
    9898
    9999    #######################
    100100    # PROXIES TO QUERYSET #
    101101    #######################
    102102
    103103    def get_empty_query_set(self):
    104         return EmptyQuerySet(self.model)
     104        return EmptyQuerySet(self.model, using=self.db)
    105105
    106106    def get_query_set(self):
    107107        """Returns a new QuerySet object.  Subclasses can override this method
    108108        to easily customize the behavior of the Manager.
    109109        """
    110         qs = QuerySet(self.model)
    111         if self._db is not None:
    112             qs = qs.using(self._db)
    113         return qs
     110        return QuerySet(self.model, using=self.db)
    114111
    115112    def none(self):
    116113        return self.get_empty_query_set()

  • django/db/models/query.py 

    diff -r e165fea06e06 django/db/models/query.py
    a b  
    44
    55from copy import deepcopy
    66
    7 from django.db import connections, transaction, IntegrityError, DEFAULT_DB_ALIAS
     7from django.db import connections, router, transaction, IntegrityError
    88from django.db.models.aggregates import Aggregate
    99from django.db.models.fields import DateField
    1010from django.db.models.query_utils import Q, select_related_descend, CollectedObjects, CyclicDependency, deferred_class_factory, InvalidQuery
    … …  
    429429
    430430            if not seen_objs:
    431431                break
    432             delete_objects(seen_objs, del_query.db)
     432            delete_objects(seen_objs, del_query._db or router.db_for_write(self.model))
    433433
    434434        # Clear the result cache, in case this QuerySet gets reused.
    435435        self._result_cache = None
    … …  
    444444                "Cannot update a query once a slice has been taken."
    445445        query = self.query.clone(sql.UpdateQuery)
    446446        query.add_update_values(kwargs)
    447         if not transaction.is_managed(using=self.db):
    448             transaction.enter_transaction_management(using=self.db)
     447        db = self._db or router.db_for_write(self.model)
     448        if not transaction.is_managed(using=db):
     449            transaction.enter_transaction_management(using=db)
    449450            forced_managed = True
    450451        else:
    451452            forced_managed = False
    452453        try:
    453             rows = query.get_compiler(self.db).execute_sql(None)
     454            rows = query.get_compiler(db).execute_sql(None)
    454455            if forced_managed:
    455                 transaction.commit(using=self.db)
     456                transaction.commit(using=db)
    456457            else:
    457                 transaction.commit_unless_managed(using=self.db)
     458                transaction.commit_unless_managed(using=db)
    458459        finally:
    459460            if forced_managed:
    460                 transaction.leave_transaction_management(using=self.db)
     461                transaction.leave_transaction_management(using=db)
    461462        self._result_cache = None
    462463        return rows
    463464    update.alters_data = True
    … …  
    714715    @property
    715716    def db(self):
    716717        "Return the database that will be used if this query is executed now"
    717         return self._db or DEFAULT_DB_ALIAS
     718        return self._db or router.db_for_read(self.model)
    718719
    719720    ###################
    720721    # PRIVATE METHODS #
    … …  
    988989
    989990
    990991class EmptyQuerySet(QuerySet):
    991     def __init__(self, model=None, query=None):
    992         super(EmptyQuerySet, self).__init__(model, query)
     992    def __init__(self, model=None, query=None, using=None):
     993        super(EmptyQuerySet, self).__init__(model, query, using)
    993994        self._result_cache = []
    994995
    995996    def __and__(self, other):
    … …  
    12541255    @property
    12551256    def db(self):
    12561257        "Return the database that will be used if this query is executed now"
    1257         return self._db or DEFAULT_DB_ALIAS
     1258        return self._db or router.db_for_read(self.model)
    12581259
    12591260    def using(self, alias):
    12601261        """

  • django/db/utils.py 

    diff -r e165fea06e06 django/db/utils.py
    a b  
    55from django.core.exceptions import ImproperlyConfigured
    66from django.utils.importlib import import_module
    77
     8DEFAULT_DB_ALIAS = 'default'
     9
    810def load_backend(backend_name):
    911    try:
    1012        module = import_module('.base', 'django.db.backends.%s' % backend_name)
    … …  
    5557            conn = self.databases[alias]
    5658        except KeyError:
    5759            raise ConnectionDoesNotExist("The connection %s doesn't exist" % alias)
     60
    5861        conn.setdefault('ENGINE', 'django.db.backends.dummy')
    5962        if conn['ENGINE'] == 'django.db.backends.' or not conn['ENGINE']:
    6063            conn['ENGINE'] = 'django.db.backends.dummy'
    … …  
    8285
    8386    def all(self):
    8487        return [self[alias] for alias in self]
     88
     89class ConnectionRouter(object):
     90    def __init__(self, routers):
     91        self.routers = []
     92        for r in routers:
     93            if isinstance(r, basestring):
     94                module_name, klass_name = r.rsplit('.', 1)
     95                module = import_module(module_name)
     96                router = getattr(module, klass_name)()
     97            else:
     98                router = r
     99            self.routers.append(router)
     100
     101    def _router_func(action):
     102        def _route_db(self, model, instance=None):
     103            chosen_db = None
     104            for router in self.routers:
     105                chosen_db = getattr(router, action)(model, instance=instance)
     106                if chosen_db:
     107                    return chosen_db
     108            if instance:
     109                return instance._state.db or DEFAULT_DB_ALIAS
     110            return DEFAULT_DB_ALIAS
     111        return _route_db
     112
     113    db_for_read = _router_func('db_for_read')
     114    db_for_write = _router_func('db_for_write')
     115
     116    def allow_relation(self, obj1, obj2):
     117        for router in self.routers:
     118            allow = router.allow_relation(obj1, obj2)
     119            if allow is not None:
     120                return allow
     121        return obj1._state.db == obj2._state.db

  • tests/regressiontests/multiple_database/models.py 

    diff -r e165fea06e06 tests/regressiontests/multiple_database/models.py
    a b  
    22from django.contrib.auth.models import User
    33from django.contrib.contenttypes.models import ContentType
    44from django.contrib.contenttypes import generic
    5 from django.db import models, DEFAULT_DB_ALIAS
     5from django.db import models
    66
    77class Review(models.Model):
    88    source = models.CharField(max_length=100)

  • tests/regressiontests/multiple_database/tests.py 

    diff -r e165fea06e06 tests/regressiontests/multiple_database/tests.py
    a b  
    33
    44from django.conf import settings
    55from django.contrib.auth.models import User
    6 from django.db import connections
     6from django.db import connections, router, DEFAULT_DB_ALIAS
     7from django.db.utils import ConnectionRouter
    78from django.test import TestCase
    89
    910from models import Book, Person, Review, UserProfile
    … …  
    259260        self.assertEquals(list(Person.objects.using('other').filter(book__title='Dive into HTML5').values_list('name', flat=True)),
    260261                          [u'Mark Pilgrim'])
    261262
    262 #    def test_m2m_cross_database_protection(self):
    263 #        "Operations that involve sharing M2M objects across databases raise an error"
    264 #        # Create a book and author on the default database
    265 #        pro = Book.objects.create(title="Pro Django",
    266 #                                  published=datetime.date(2008, 12, 16))
     263    def test_m2m_cross_database_protection(self):
     264        "Operations that involve sharing M2M objects across databases raise an error"
     265        # Create a book and author on the default database
     266        pro = Book.objects.create(title="Pro Django",
     267                                  published=datetime.date(2008, 12, 16))
    267268
    268 #        marty = Person.objects.create(name="Marty Alchin")
     269        marty = Person.objects.create(name="Marty Alchin")
    269270
    270 #        # Create a book and author on the other database
    271 #        dive = Book.objects.using('other').create(title="Dive into Python",
    272 #                                                  published=datetime.date(2009, 5, 4))
     271        # Create a book and author on the other database
     272        dive = Book.objects.using('other').create(title="Dive into Python",
     273                                                  published=datetime.date(2009, 5, 4))
    273274
    274 #        mark = Person.objects.using('other').create(name="Mark Pilgrim")
    275 #        # Set a foreign key set with an object from a different database
    276 #        try:
    277 #            marty.book_set = [pro, dive]
    278 #            self.fail("Shouldn't be able to assign across databases")
    279 #        except ValueError:
    280 #            pass
     275        mark = Person.objects.using('other').create(name="Mark Pilgrim")
     276        # Set a foreign key set with an object from a different database
     277        try:
     278            marty.book_set = [pro, dive]
     279            self.fail("Shouldn't be able to assign across databases")
     280        except ValueError:
     281            pass
    281282
    282 #        # Add to an m2m with an object from a different database
    283 #        try:
    284 #            marty.book_set.add(dive)
    285 #            self.fail("Shouldn't be able to assign across databases")
    286 #        except ValueError:
    287 #            pass
     283        # Add to an m2m with an object from a different database
     284        try:
     285            marty.book_set.add(dive)
     286            self.fail("Shouldn't be able to assign across databases")
     287        except ValueError:
     288            pass
    288289
    289 #        # Set a m2m with an object from a different database
    290 #        try:
    291 #            marty.book_set = [pro, dive]
    292 #            self.fail("Shouldn't be able to assign across databases")
    293 #        except ValueError:
    294 #            pass
     290        # Set a m2m with an object from a different database
     291        try:
     292            marty.book_set = [pro, dive]
     293            self.fail("Shouldn't be able to assign across databases")
     294        except ValueError:
     295            pass
    295296
    296 #        # Add to a reverse m2m with an object from a different database
    297 #        try:
    298 #            dive.authors.add(marty)
    299 #            self.fail("Shouldn't be able to assign across databases")
    300 #        except ValueError:
    301 #            pass
     297        # Add to a reverse m2m with an object from a different database
     298        try:
     299            dive.authors.add(marty)
     300            self.fail("Shouldn't be able to assign across databases")
     301        except ValueError:
     302            pass
    302303
    303 #        # Set a reverse m2m with an object from a different database
    304 #        try:
    305 #            dive.authors = [mark, marty]
    306 #            self.fail("Shouldn't be able to assign across databases")
    307 #        except ValueError:
    308 #            pass
     304        # Set a reverse m2m with an object from a different database
     305        try:
     306            dive.authors = [mark, marty]
     307            self.fail("Shouldn't be able to assign across databases")
     308        except ValueError:
     309            pass
    309310
    310311    def test_foreign_key_separation(self):
    311312        "FK fields are constrained to a single database"
    … …  
    401402        self.assertEquals(list(Person.objects.using('other').filter(edited__title='Dive into Python').values_list('name', flat=True)),
    402403                          [])
    403404
    404 #    def test_foreign_key_cross_database_protection(self):
    405 #        "Operations that involve sharing FK objects across databases raise an error"
    406 #        # Create a book and author on the default database
    407 #        pro = Book.objects.create(title="Pro Django",
    408 #                                  published=datetime.date(2008, 12, 16))
     405    def test_foreign_key_cross_database_protection(self):
     406        "Operations that involve sharing FK objects across databases raise an error"
     407        # Create a book and author on the default database
     408        pro = Book.objects.create(title="Pro Django",
     409                                  published=datetime.date(2008, 12, 16))
    409410
    410 #        marty = Person.objects.create(name="Marty Alchin")
     411        marty = Person.objects.create(name="Marty Alchin")
    411412
    412 #        # Create a book and author on the other database
    413 #        dive = Book.objects.using('other').create(title="Dive into Python",
    414 #                                                  published=datetime.date(2009, 5, 4))
     413        # Create a book and author on the other database
     414        dive = Book.objects.using('other').create(title="Dive into Python",
     415                                                  published=datetime.date(2009, 5, 4))
    415416
    416 #        mark = Person.objects.using('other').create(name="Mark Pilgrim")
     417        mark = Person.objects.using('other').create(name="Mark Pilgrim")
    417418
    418 #        # Set a foreign key with an object from a different database
    419 #        try:
    420 #            dive.editor = marty
    421 #            self.fail("Shouldn't be able to assign across databases")
    422 #        except ValueError:
    423 #            pass
     419        # Set a foreign key with an object from a different database
     420        try:
     421            dive.editor = marty
     422            self.fail("Shouldn't be able to assign across databases")
     423        except ValueError:
     424            pass
    424425
    425 #        # Set a foreign key set with an object from a different database
    426 #        try:
    427 #            marty.edited = [pro, dive]
    428 #            self.fail("Shouldn't be able to assign across databases")
    429 #        except ValueError:
    430 #            pass
     426        # Set a foreign key set with an object from a different database
     427        try:
     428            marty.edited = [pro, dive]
     429            self.fail("Shouldn't be able to assign across databases")
     430        except ValueError:
     431            pass
    431432
    432 #        # Add to a foreign key set with an object from a different database
    433 #        try:
    434 #            marty.edited.add(dive)
    435 #            self.fail("Shouldn't be able to assign across databases")
    436 #        except ValueError:
    437 #            pass
     433        # Add to a foreign key set with an object from a different database
     434        try:
     435            marty.edited.add(dive)
     436            self.fail("Shouldn't be able to assign across databases")
     437        except ValueError:
     438            pass
    438439
    439 #        # BUT! if you assign a FK object when the base object hasn't
    440 #        # been saved yet, you implicitly assign the database for the
    441 #        # base object.
    442 #        chris = Person(name="Chris Mills")
    443 #        html5 = Book(title="Dive into HTML5", published=datetime.date(2010, 3, 15))
    444 #        # initially, no db assigned
    445 #        self.assertEquals(chris._state.db, None)
    446 #        self.assertEquals(html5._state.db, None)
     440        # BUT! if you assign a FK object when the base object hasn't
     441        # been saved yet, you implicitly assign the database for the
     442        # base object.
     443        chris = Person(name="Chris Mills")
     444        html5 = Book(title="Dive into HTML5", published=datetime.date(2010, 3, 15))
     445        # initially, no db assigned
     446        self.assertEquals(chris._state.db, None)
     447        self.assertEquals(html5._state.db, None)
    447448
    448 #        # old object comes from 'other', so the new object is set to use 'other'...
    449 #        dive.editor = chris
    450 #        html5.editor = mark
    451 #        # self.assertEquals(chris._state.db, 'other')
    452 #        self.assertEquals(html5._state.db, 'other')
    453 #        # ... but it isn't saved yet
    454 #        self.assertEquals(list(Person.objects.using('other').values_list('name',flat=True)),
    455 #                          [u'Mark Pilgrim'])
    456 #        self.assertEquals(list(Book.objects.using('other').values_list('title',flat=True)),
    457 #                           [u'Dive into Python'])
     449        # old object comes from 'other', so the new object is set to use 'other'...
     450        dive.editor = chris
     451        html5.editor = mark
     452        self.assertEquals(chris._state.db, 'other')
     453        self.assertEquals(html5._state.db, 'other')
     454        # ... but it isn't saved yet
     455        self.assertEquals(list(Person.objects.using('other').values_list('name',flat=True)),
     456                          [u'Mark Pilgrim'])
     457        self.assertEquals(list(Book.objects.using('other').values_list('title',flat=True)),
     458                           [u'Dive into Python'])
    458459
    459 #        # When saved (no using required), new objects goes to 'other'
    460 #        chris.save()
    461 #        html5.save()
    462 #        self.assertEquals(list(Person.objects.using('default').values_list('name',flat=True)),
    463 #                          [u'Marty Alchin'])
    464 #        self.assertEquals(list(Person.objects.using('other').values_list('name',flat=True)),
    465 #                          [u'Chris Mills', u'Mark Pilgrim'])
    466 #        self.assertEquals(list(Book.objects.using('default').values_list('title',flat=True)),
    467 #                          [u'Pro Django'])
    468 #        self.assertEquals(list(Book.objects.using('other').values_list('title',flat=True)),
    469 #                          [u'Dive into HTML5', u'Dive into Python'])
     460        # When saved (no using required), new objects goes to 'other'
     461        chris.save()
     462        html5.save()
     463        self.assertEquals(list(Person.objects.using('default').values_list('name',flat=True)),
     464                          [u'Marty Alchin'])
     465        self.assertEquals(list(Person.objects.using('other').values_list('name',flat=True)),
     466                          [u'Chris Mills', u'Mark Pilgrim'])
     467        self.assertEquals(list(Book.objects.using('default').values_list('title',flat=True)),
     468                          [u'Pro Django'])
     469        self.assertEquals(list(Book.objects.using('other').values_list('title',flat=True)),
     470                          [u'Dive into HTML5', u'Dive into Python'])
    470471
    471 #        # This also works if you assign the FK in the constructor
    472 #        water = Book(title="Dive into Water", published=datetime.date(2001, 1, 1), editor=mark)
    473 #        self.assertEquals(water._state.db, 'other')
    474 #        # ... but it isn't saved yet
    475 #        self.assertEquals(list(Book.objects.using('default').values_list('title',flat=True)),
    476 #                          [u'Pro Django'])
    477 #        self.assertEquals(list(Book.objects.using('other').values_list('title',flat=True)),
    478 #                          [u'Dive into HTML5', u'Dive into Python'])
     472        # This also works if you assign the FK in the constructor
     473        water = Book(title="Dive into Water", published=datetime.date(2001, 1, 1), editor=mark)
     474        self.assertEquals(water._state.db, 'other')
     475        # ... but it isn't saved yet
     476        self.assertEquals(list(Book.objects.using('default').values_list('title',flat=True)),
     477                          [u'Pro Django'])
     478        self.assertEquals(list(Book.objects.using('other').values_list('title',flat=True)),
     479                          [u'Dive into HTML5', u'Dive into Python'])
    479480
    480 #        # When saved, the new book goes to 'other'
    481 #        water.save()
    482 #        self.assertEquals(list(Book.objects.using('default').values_list('title',flat=True)),
    483 #                          [u'Pro Django'])
    484 #        self.assertEquals(list(Book.objects.using('other').values_list('title',flat=True)),
    485 #                          [u'Dive into HTML5', u'Dive into Python', u'Dive into Water'])
     481        # When saved, the new book goes to 'other'
     482        water.save()
     483        self.assertEquals(list(Book.objects.using('default').values_list('title',flat=True)),
     484                          [u'Pro Django'])
     485        self.assertEquals(list(Book.objects.using('other').values_list('title',flat=True)),
     486                          [u'Dive into HTML5', u'Dive into Python', u'Dive into Water'])
    486487
    487488    def test_generic_key_separation(self):
    488489        "Generic fields are constrained to a single database"
    … …  
    555556        self.assertEquals(list(Review.objects.using('other').filter(object_id=dive.pk).values_list('source', flat=True)),
    556557                          [u'Python Daily'])
    557558
    558 #    def test_generic_key_cross_database_protection(self):
    559 ##        "Operations that involve sharing FK objects across databases raise an error"
    560 ##        # Create a book and author on the default database
    561 ##        pro = Book.objects.create(title="Pro Django",
    562 ##                                  published=datetime.date(2008, 12, 16))
     559    def test_generic_key_cross_database_protection(self):
     560        "Operations that involve sharing generic key objects across databases raise an error"
     561        # Create a book and author on the default database
     562        pro = Book.objects.create(title="Pro Django",
     563                                  published=datetime.date(2008, 12, 16))
    563564
    564 ##        review1 = Review.objects.create(source="Python Monthly", content_object=pro)
     565        review1 = Review.objects.create(source="Python Monthly", content_object=pro)
    565566
    566 ##        # Create a book and author on the other database
    567 ##        dive = Book.objects.using('other').create(title="Dive into Python",
    568 ##                                                  published=datetime.date(2009, 5, 4))
     567        # Create a book and author on the other database
     568        dive = Book.objects.using('other').create(title="Dive into Python",
     569                                                  published=datetime.date(2009, 5, 4))
    569570
    570 ##        review2 = Review.objects.using('other').create(source="Python Weekly", content_object=dive)
     571        review2 = Review.objects.using('other').create(source="Python Weekly", content_object=dive)
    571572
    572 ##        # Set a foreign key with an object from a different database
    573 ##        try:
    574 ##            review1.content_object = dive
    575 ##            self.fail("Shouldn't be able to assign across databases")
    576 ##        except ValueError:
    577 ##            pass
     573        # Set a foreign key with an object from a different database
     574        try:
     575            review1.content_object = dive
     576            self.fail("Shouldn't be able to assign across databases")
     577        except ValueError:
     578            pass
    578579
    579 #        # Add to a foreign key set with an object from a different database
    580 #        try:
    581 #            dive.reviews.add(review1)
    582 #            self.fail("Shouldn't be able to assign across databases")
    583 #        except ValueError:
    584 #            pass
     580        # Add to a foreign key set with an object from a different database
     581        try:
     582            dive.reviews.add(review1)
     583            self.fail("Shouldn't be able to assign across databases")
     584        except ValueError:
     585            pass
    585586
    586 #        # BUT! if you assign a FK object when the base object hasn't
    587 #        # been saved yet, you implicitly assign the database for the
    588 #        # base object.
    589 #        review3 = Review(source="Python Daily")
    590 #        # initially, no db assigned
    591 #        self.assertEquals(review3._state.db, None)
     587        # BUT! if you assign a FK object when the base object hasn't
     588        # been saved yet, you implicitly assign the database for the
     589        # base object.
     590        review3 = Review(source="Python Daily")
     591        # initially, no db assigned
     592        self.assertEquals(review3._state.db, None)
    592593
    593 #        # Dive comes from 'other', so review3 is set to use 'other'...
    594 #        review3.content_object = dive
    595 #        self.assertEquals(review3._state.db, 'other')
    596 #        # ... but it isn't saved yet
    597 #        self.assertEquals(list(Review.objects.using('default').filter(object_id=pro.pk).values_list('source', flat=True)),
    598 #                          [u'Python Monthly'])
    599 #        self.assertEquals(list(Review.objects.using('other').filter(object_id=dive.pk).values_list('source',flat=True)),
    600 #                          [u'Python Weekly'])
     594        # Dive comes from 'other', so review3 is set to use 'other'...
     595        review3.content_object = dive
     596        self.assertEquals(review3._state.db, 'other')
     597        # ... but it isn't saved yet
     598        self.assertEquals(list(Review.objects.using('default').filter(object_id=pro.pk).values_list('source', flat=True)),
     599                          [u'Python Monthly'])
     600        self.assertEquals(list(Review.objects.using('other').filter(object_id=dive.pk).values_list('source',flat=True)),
     601                          [u'Python Weekly'])
    601602
    602 #        # When saved, John goes to 'other'
    603 #        review3.save()
    604 #        self.assertEquals(list(Review.objects.using('default').filter(object_id=pro.pk).values_list('source', flat=True)),
    605 #                          [u'Python Monthly'])
    606 #        self.assertEquals(list(Review.objects.using('other').filter(object_id=dive.pk).values_list('source',flat=True)),
    607 #                          [u'Python Daily', u'Python Weekly'])
     603        # When saved, John goes to 'other'
     604        review3.save()
     605        self.assertEquals(list(Review.objects.using('default').filter(object_id=pro.pk).values_list('source', flat=True)),
     606                          [u'Python Monthly'])
     607        self.assertEquals(list(Review.objects.using('other').filter(object_id=dive.pk).values_list('source',flat=True)),
     608                          [u'Python Daily', u'Python Weekly'])
    608609
    609610    def test_ordering(self):
    610611        "get_next_by_XXX commands stick to a single database"
    … …  
    630631        val = Book.objects.raw('SELECT id FROM "multiple_database_book"').using('other')
    631632        self.assertEqual(map(lambda o: o.pk, val), [dive.pk])
    632633
     634class CommonSourceRouter(object):
     635    def db_for_read(self, model, instance):
     636        if instance:
     637            return instance._state.db or DEFAULT_DB_ALIAS
     638        return DEFAULT_DB_ALIAS
     639
     640    def db_for_write(self, model, instance):
     641        return DEFAULT_DB_ALIAS
     642
     643    def allow_relation(self, obj1, obj2):
     644        return obj1._state.db in ('default', 'other') and obj2._state.db in ('default', 'other')
     645
     646class CommonSourceTestCase(TestCase):
     647    multi_db = True
     648
     649    def setUp(self):
     650        # Make the 'other' database appear to be a slave of the 'default'
     651        self.old_routers = router.routers
     652        router.routers = [CommonSourceRouter()]
     653
     654    def tearDown(self):
     655        # Restore the 'other' database as an independent database
     656        router.routers = self.old_routers
     657
     658    def test_foreign_key_cross_database_protection(self):
     659        "Foreign keys can cross databases if they two databases have a common source"
     660        # Create a book and author on the default database
     661        pro = Book.objects.create(title="Pro Django",
     662                                  published=datetime.date(2008, 12, 16))
     663
     664        marty = Person.objects.create(name="Marty Alchin")
     665
     666        # Create a book and author on the other database
     667        dive = Book.objects.using('other').create(title="Dive into Python",
     668                                                  published=datetime.date(2009, 5, 4))
     669
     670        mark = Person.objects.using('other').create(name="Mark Pilgrim")
     671
     672        # Set a foreign key with an object from a different database
     673        try:
     674            dive.editor = marty
     675        except ValueError:
     676            self.fail("Assignment across master/slave databases with a common source should be ok")
     677
     678        # Database assignments of original objects haven't changed...
     679        self.assertEquals(marty._state.db, 'default')
     680        self.assertEquals(pro._state.db, 'default')
     681        self.assertEquals(dive._state.db, 'other')
     682        self.assertEquals(mark._state.db, 'other')
     683
     684        # ... but they will when the affected object is saved.
     685        dive.save()
     686        self.assertEquals(dive._state.db, 'default')
     687
     688        # ...and the source database now has a copy of any object saved
     689        try:
     690            Book.objects.using('default').get(title='Dive into Python').delete()
     691        except Book.DoesNotExist:
     692            self.fail('Source database should have a copy of saved object')
     693
     694        # This isn't a real master-slave database, so restore the original from other
     695        dive = Book.objects.using('other').get(title='Dive into Python')
     696        self.assertEquals(dive._state.db, 'other')
     697
     698        # Set a foreign key set with an object from a different database
     699        try:
     700            marty.edited = [pro, dive]
     701        except ValueError:
     702            self.fail("Assignment across master/slave databases with a common source should be ok")
     703
     704        # Assignment implies a save, so database assignments of original objects have changed...
     705        self.assertEquals(marty._state.db, 'default')
     706        self.assertEquals(pro._state.db, 'default')
     707        self.assertEquals(dive._state.db, 'default')
     708        self.assertEquals(mark._state.db, 'other')
     709
     710        # ...and the source database now has a copy of any object saved
     711        try:
     712            Book.objects.using('default').get(title='Dive into Python').delete()
     713        except Book.DoesNotExist:
     714            self.fail('Source database should have a copy of saved object')
     715
     716        # This isn't a real master-slave database, so restore the original from other
     717        dive = Book.objects.using('other').get(title='Dive into Python')
     718        self.assertEquals(dive._state.db, 'other')
     719
     720        # Add to a foreign key set with an object from a different database
     721        try:
     722            marty.edited.add(dive)
     723        except ValueError:
     724            self.fail("Assignment across master/slave databases with a common source should be ok")
     725
     726        # Add implies a save, so database assignments of original objects have changed...
     727        self.assertEquals(marty._state.db, 'default')
     728        self.assertEquals(pro._state.db, 'default')
     729        self.assertEquals(dive._state.db, 'default')
     730        self.assertEquals(mark._state.db, 'other')
     731
     732        # ...and the source database now has a copy of any object saved
     733        try:
     734            Book.objects.using('default').get(title='Dive into Python').delete()
     735        except Book.DoesNotExist:
     736            self.fail('Source database should have a copy of saved object')
     737
     738        # This isn't a real master-slave database, so restore the original from other
     739        dive = Book.objects.using('other').get(title='Dive into Python')
     740
     741        # If you assign a FK object when the base object hasn't
     742        # been saved yet, you implicitly assign the database for the
     743        # base object.
     744        chris = Person(name="Chris Mills")
     745        html5 = Book(title="Dive into HTML5", published=datetime.date(2010, 3, 15))
     746        # initially, no db assigned
     747        self.assertEquals(chris._state.db, None)
     748        self.assertEquals(html5._state.db, None)
     749
     750        # old object comes from 'other', so the new object is set to use the
     751        # source of 'other'...
     752        self.assertEquals(dive._state.db, 'other')
     753        dive.editor = chris
     754        html5.editor = mark
     755
     756        self.assertEquals(dive._state.db, 'other')
     757        self.assertEquals(mark._state.db, 'other')
     758        self.assertEquals(chris._state.db, 'default')
     759        self.assertEquals(html5._state.db, 'default')
     760
     761        # This also works if you assign the FK in the constructor
     762        water = Book(title="Dive into Water", published=datetime.date(2001, 1, 1), editor=mark)
     763        self.assertEquals(water._state.db, 'default')
     764
     765    def test_m2m_cross_database_protection(self):
     766        "M2M relations can cross databases if the database share a source"
     767        # Create books and authors on the inverse to the usual database
     768        pro = Book.objects.using('other').create(pk=1, title="Pro Django",
     769                                                 published=datetime.date(2008, 12, 16))
     770
     771        marty = Person.objects.using('other').create(pk=1, name="Marty Alchin")
     772
     773        dive = Book.objects.using('default').create(pk=2, title="Dive into Python",
     774                                                    published=datetime.date(2009, 5, 4))
     775
     776        mark = Person.objects.using('default').create(pk=2, name="Mark Pilgrim")
     777
     778        # Now save back onto the usual databse.
     779        # This simulates master/slave - the objects exist on both database,
     780        # but the _state.db is as it is for all other tests.
     781        pro.save(using='default')
     782        marty.save(using='default')
     783        dive.save(using='other')
     784        mark.save(using='other')
     785
     786        # Check that we have 2 of both types of object on both databases
     787        self.assertEquals(Book.objects.using('default').count(), 2)
     788        self.assertEquals(Book.objects.using('other').count(), 2)
     789        self.assertEquals(Person.objects.using('default').count(), 2)
     790        self.assertEquals(Person.objects.using('other').count(), 2)
     791
     792        # Set a m2m set with an object from a different database
     793        try:
     794            marty.book_set = [pro, dive]
     795        except ValueError:
     796            self.fail("Assignment across master/slave databases with a common source should be ok")
     797
     798        # Database assignments don't change
     799        self.assertEquals(marty._state.db, 'default')
     800        self.assertEquals(pro._state.db, 'default')
     801        self.assertEquals(dive._state.db, 'other')
     802        self.assertEquals(mark._state.db, 'other')
     803
     804        # All m2m relations should be saved on the default database
     805        self.assertEquals(Book.authors.through.objects.using('default').count(), 2)
     806        self.assertEquals(Book.authors.through.objects.using('other').count(), 0)
     807
     808        # Reset relations
     809        Book.authors.through.objects.using('default').delete()
     810
     811        # Add to an m2m with an object from a different database
     812        try:
     813            marty.book_set.add(dive)
     814        except ValueError:
     815            self.fail("Assignment across master/slave databases with a common source should be ok")
     816
     817        # Database assignments don't change
     818        self.assertEquals(marty._state.db, 'default')
     819        self.assertEquals(pro._state.db, 'default')
     820        self.assertEquals(dive._state.db, 'other')
     821        self.assertEquals(mark._state.db, 'other')
     822
     823        # All m2m relations should be saved on the default database
     824        self.assertEquals(Book.authors.through.objects.using('default').count(), 1)
     825        self.assertEquals(Book.authors.through.objects.using('other').count(), 0)
     826
     827        # Reset relations
     828        Book.authors.through.objects.using('default').delete()
     829
     830        # Set a reverse m2m with an object from a different database
     831        try:
     832            dive.authors = [mark, marty]
     833        except ValueError:
     834            self.fail("Assignment across master/slave databases with a common source should be ok")
     835
     836        # Database assignments don't change
     837        self.assertEquals(marty._state.db, 'default')
     838        self.assertEquals(pro._state.db, 'default')
     839        self.assertEquals(dive._state.db, 'other')
     840        self.assertEquals(mark._state.db, 'other')
     841
     842        # All m2m relations should be saved on the default database
     843        self.assertEquals(Book.authors.through.objects.using('default').count(), 2)
     844        self.assertEquals(Book.authors.through.objects.using('other').count(), 0)
     845
     846        # Reset relations
     847        Book.authors.through.objects.using('default').delete()
     848
     849        self.assertEquals(Book.authors.through.objects.using('default').count(), 0)
     850        self.assertEquals(Book.authors.through.objects.using('other').count(), 0)
     851
     852        # Add to a reverse m2m with an object from a different database
     853        try:
     854            dive.authors.add(marty)
     855        except ValueError:
     856            self.fail("Assignment across master/slave databases with a common source should be ok")
     857
     858        # Database assignments don't change
     859        self.assertEquals(marty._state.db, 'default')
     860        self.assertEquals(pro._state.db, 'default')
     861        self.assertEquals(dive._state.db, 'other')
     862        self.assertEquals(mark._state.db, 'other')
     863
     864        # All m2m relations should be saved on the default database
     865        self.assertEquals(Book.authors.through.objects.using('default').count(), 1)
     866        self.assertEquals(Book.authors.through.objects.using('other').count(), 0)
     867
     868    def test_generic_key_cross_database_protection(self):
     869        "Generic Key operations can span databases if they share a source"
     870        # Create a book and author on the default database
     871        pro = Book.objects.create(title="Pro Django",
     872                                  published=datetime.date(2008, 12, 16))
     873
     874        review1 = Review.objects.create(source="Python Monthly", content_object=pro)
     875
     876        # Create a book and author on the other database
     877        dive = Book.objects.using('other').create(title="Dive into Python",
     878                                                  published=datetime.date(2009, 5, 4))
     879
     880        review2 = Review.objects.using('other').create(source="Python Weekly", content_object=dive)
     881
     882        # Set a generic foreign key with an object from a different database
     883        try:
     884            review1.content_object = dive
     885        except ValueError:
     886            self.fail("Assignment across master/slave databases with a common source should be ok")
     887
     888        # Database assignments of original objects haven't changed...
     889        self.assertEquals(pro._state.db, 'default')
     890        self.assertEquals(review1._state.db, 'default')
     891        self.assertEquals(dive._state.db, 'other')
     892        self.assertEquals(review2._state.db, 'other')
     893
     894        # ... but they will when the affected object is saved.
     895        dive.save()
     896        self.assertEquals(review1._state.db, 'default')
     897        self.assertEquals(dive._state.db, 'default')
     898
     899        # ...and the source database now has a copy of any object saved
     900        try:
     901            Book.objects.using('default').get(title='Dive into Python').delete()
     902        except Book.DoesNotExist:
     903            self.fail('Source database should have a copy of saved object')
     904
     905        # This isn't a real master-slave database, so restore the original from other
     906        dive = Book.objects.using('other').get(title='Dive into Python')
     907        self.assertEquals(dive._state.db, 'other')
     908
     909        # Add to a generic foreign key set with an object from a different database
     910        try:
     911            dive.reviews.add(review1)
     912        except ValueError:
     913            self.fail("Assignment across master/slave databases with a common source should be ok")
     914
     915        # Database assignments of original objects haven't changed...
     916        self.assertEquals(pro._state.db, 'default')
     917        self.assertEquals(review1._state.db, 'default')
     918        self.assertEquals(dive._state.db, 'other')
     919        self.assertEquals(review2._state.db, 'other')
     920
     921        # ... but they will when the affected object is saved.
     922        dive.save()
     923        self.assertEquals(dive._state.db, 'default')
     924
     925        # ...and the source database now has a copy of any object saved
     926        try:
     927            Book.objects.using('default').get(title='Dive into Python').delete()
     928        except Book.DoesNotExist:
     929            self.fail('Source database should have a copy of saved object')
     930
     931        # BUT! if you assign a FK object when the base object hasn't
     932        # been saved yet, you implicitly assign the database for the
     933        # base object.
     934        review3 = Review(source="Python Daily")
     935        # initially, no db assigned
     936        self.assertEquals(review3._state.db, None)
     937
     938        # Dive comes from 'other', so review3 is set to use the source of 'other'...
     939        review3.content_object = dive
     940        self.assertEquals(review3._state.db, 'default')
     941
    633942
    634943class UserProfileTestCase(TestCase):
    635944    def setUp(self):
Back to Top