import datetime

import pickle

import sys

from StringIO import StringIO

from django.conf import settings

from django.contrib.auth.models import User

from django.core import management

from django.db import connections, router, DEFAULT_DB_ALIAS

from django.db.utils import ConnectionRouter

from django.test import TestCase

from models import Book, Person, Pet, Review, UserProfile

try:

    # we only have these models if the user is using multi-db, it's safe the

    # run the tests without them though.

    from models import Article, article_using

except ImportError:

    pass

class QueryTestCase(TestCase):

    multi_db = True

    def test_db_selection(self):

        "Check that querysets will use the default databse by default"

        self.assertEquals(Book.objects.db, DEFAULT_DB_ALIAS)

        self.assertEquals(Book.objects.all().db, DEFAULT_DB_ALIAS)

        self.assertEquals(Book.objects.using('other').db, 'other')

        self.assertEquals(Book.objects.db_manager('other').db, 'other')

        self.assertEquals(Book.objects.db_manager('other').all().db, 'other')

    def test_default_creation(self):

        "Objects created on the default database don't leak onto other databases"

        # Create a book on the default database using create()

        Book.objects.create(title="Pro Django",

                            published=datetime.date(2008, 12, 16))

        # Create a book on the default database using a save

        dive = Book()

        dive.title="Dive into Python"

        dive.published = datetime.date(2009, 5, 4)

        dive.save()

        # Check that book exists on the default database, but not on other database

        try:

            Book.objects.get(title="Pro Django")

            Book.objects.using('default').get(title="Pro Django")

        except Book.DoesNotExist:

            self.fail('"Dive Into Python" should exist on default database')

        self.assertRaises(Book.DoesNotExist,

            Book.objects.using('other').get,

            title="Pro Django"

        )

        try:

            Book.objects.get(title="Dive into Python")

            Book.objects.using('default').get(title="Dive into Python")

        except Book.DoesNotExist:

            self.fail('"Dive into Python" should exist on default database')

        self.assertRaises(Book.DoesNotExist,

            Book.objects.using('other').get,

            title="Dive into Python"

        )

    def test_other_creation(self):

        "Objects created on another database don't leak onto the default database"

        # Create a book on the second database

        Book.objects.using('other').create(title="Pro Django",

                                           published=datetime.date(2008, 12, 16))

        # Create a book on the default database using a save

        dive = Book()

        dive.title="Dive into Python"

        dive.published = datetime.date(2009, 5, 4)

        dive.save(using='other')

        # Check that book exists on the default database, but not on other database

        try:

            Book.objects.using('other').get(title="Pro Django")

        except Book.DoesNotExist:

            self.fail('"Dive Into Python" should exist on other database')

        self.assertRaises(Book.DoesNotExist,

            Book.objects.get,

            title="Pro Django"

        )

        self.assertRaises(Book.DoesNotExist,

            Book.objects.using('default').get,

            title="Pro Django"

        )

        try:

            Book.objects.using('other').get(title="Dive into Python")

        except Book.DoesNotExist:

            self.fail('"Dive into Python" should exist on other database')

        self.assertRaises(Book.DoesNotExist,

            Book.objects.get,

            title="Dive into Python"

        )

        self.assertRaises(Book.DoesNotExist,

            Book.objects.using('default').get,

            title="Dive into Python"

        )

    def test_basic_queries(self):

        "Queries are constrained to a single database"

        dive = Book.objects.using('other').create(title="Dive into Python",

                                                  published=datetime.date(2009, 5, 4))

        dive =  Book.objects.using('other').get(published=datetime.date(2009, 5, 4))

        self.assertEqual(dive.title, "Dive into Python")

        self.assertRaises(Book.DoesNotExist, Book.objects.using('default').get, published=datetime.date(2009, 5, 4))

        dive = Book.objects.using('other').get(title__icontains="dive")

        self.assertEqual(dive.title, "Dive into Python")

        self.assertRaises(Book.DoesNotExist, Book.objects.using('default').get, title__icontains="dive")

        dive = Book.objects.using('other').get(title__iexact="dive INTO python")

        self.assertEqual(dive.title, "Dive into Python")

        self.assertRaises(Book.DoesNotExist, Book.objects.using('default').get, title__iexact="dive INTO python")

        dive =  Book.objects.using('other').get(published__year=2009)

        self.assertEqual(dive.title, "Dive into Python")

        self.assertEqual(dive.published, datetime.date(2009, 5, 4))

        self.assertRaises(Book.DoesNotExist, Book.objects.using('default').get, published__year=2009)

        years = Book.objects.using('other').dates('published', 'year')

        self.assertEqual([o.year for o in years], [2009])

        years = Book.objects.using('default').dates('published', 'year')

        self.assertEqual([o.year for o in years], [])

        months = Book.objects.using('other').dates('published', 'month')

        self.assertEqual([o.month for o in months], [5])

        months = Book.objects.using('default').dates('published', 'month')

        self.assertEqual([o.month for o in months], [])

    def test_m2m_separation(self):

        "M2M fields are constrained to a single database"

        # Create a book and author on the default database

        pro = Book.objects.create(title="Pro Django",

                                  published=datetime.date(2008, 12, 16))

        marty = Person.objects.create(name="Marty Alchin")

        # Create a book and author on the other database

        dive = Book.objects.using('other').create(title="Dive into Python",

                                                  published=datetime.date(2009, 5, 4))

        mark = Person.objects.using('other').create(name="Mark Pilgrim")

        # Save the author relations

        pro.authors = [marty]

        dive.authors = [mark]

        # Inspect the m2m tables directly.

        # There should be 1 entry in each database

        self.assertEquals(Book.authors.through.objects.using('default').count(), 1)

        self.assertEquals(Book.authors.through.objects.using('other').count(), 1)

        # Check that queries work across m2m joins

        self.assertEquals(list(Book.objects.using('default').filter(authors__name='Marty Alchin').values_list('title', flat=True)),

                          [u'Pro Django'])

        self.assertEquals(list(Book.objects.using('other').filter(authors__name='Marty Alchin').values_list('title', flat=True)),

                          [])

        self.assertEquals(list(Book.objects.using('default').filter(authors__name='Mark Pilgrim').values_list('title', flat=True)),

                          [])

        self.assertEquals(list(Book.objects.using('other').filter(authors__name='Mark Pilgrim').values_list('title', flat=True)),

                          [u'Dive into Python'])

        # Reget the objects to clear caches

        dive = Book.objects.using('other').get(title="Dive into Python")

        mark = Person.objects.using('other').get(name="Mark Pilgrim")

        # Retrive related object by descriptor. Related objects should be database-baound

        self.assertEquals(list(dive.authors.all().values_list('name', flat=True)),

                          [u'Mark Pilgrim'])

        self.assertEquals(list(mark.book_set.all().values_list('title', flat=True)),

                          [u'Dive into Python'])

    def test_m2m_forward_operations(self):

        "M2M forward manipulations are all constrained to a single DB"

        # Create a book and author on the other database

        dive = Book.objects.using('other').create(title="Dive into Python",

                                                  published=datetime.date(2009, 5, 4))

        mark = Person.objects.using('other').create(name="Mark Pilgrim")

        # Save the author relations

        dive.authors = [mark]

        # Add a second author

        john = Person.objects.using('other').create(name="John Smith")

        self.assertEquals(list(Book.objects.using('other').filter(authors__name='John Smith').values_list('title', flat=True)),

                          [])

        dive.authors.add(john)

        self.assertEquals(list(Book.objects.using('other').filter(authors__name='Mark Pilgrim').values_list('title', flat=True)),

                          [u'Dive into Python'])

        self.assertEquals(list(Book.objects.using('other').filter(authors__name='John Smith').values_list('title', flat=True)),

                          [u'Dive into Python'])

        # Remove the second author

        dive.authors.remove(john)

        self.assertEquals(list(Book.objects.using('other').filter(authors__name='Mark Pilgrim').values_list('title', flat=True)),

                          [u'Dive into Python'])

        self.assertEquals(list(Book.objects.using('other').filter(authors__name='John Smith').values_list('title', flat=True)),

                          [])

        # Clear all authors

        dive.authors.clear()

        self.assertEquals(list(Book.objects.using('other').filter(authors__name='Mark Pilgrim').values_list('title', flat=True)),

                          [])

        self.assertEquals(list(Book.objects.using('other').filter(authors__name='John Smith').values_list('title', flat=True)),

                          [])

        # Create an author through the m2m interface

        dive.authors.create(name='Jane Brown')

        self.assertEquals(list(Book.objects.using('other').filter(authors__name='Mark Pilgrim').values_list('title', flat=True)),

                          [])

        self.assertEquals(list(Book.objects.using('other').filter(authors__name='Jane Brown').values_list('title', flat=True)),

                          [u'Dive into Python'])

    def test_m2m_reverse_operations(self):

        "M2M reverse manipulations are all constrained to a single DB"

        # Create a book and author on the other database

        dive = Book.objects.using('other').create(title="Dive into Python",

                                                  published=datetime.date(2009, 5, 4))

        mark = Person.objects.using('other').create(name="Mark Pilgrim")

        # Save the author relations

        dive.authors = [mark]

        # Create a second book on the other database

        grease = Book.objects.using('other').create(title="Greasemonkey Hacks",

                                                    published=datetime.date(2005, 11, 1))

        # Add a books to the m2m

        mark.book_set.add(grease)

        self.assertEquals(list(Person.objects.using('other').filter(book__title='Dive into Python').values_list('name', flat=True)),

                          [u'Mark Pilgrim'])

        self.assertEquals(list(Person.objects.using('other').filter(book__title='Greasemonkey Hacks').values_list('name', flat=True)),

                          [u'Mark Pilgrim'])

        # Remove a book from the m2m

        mark.book_set.remove(grease)

        self.assertEquals(list(Person.objects.using('other').filter(book__title='Dive into Python').values_list('name', flat=True)),

                          [u'Mark Pilgrim'])

        self.assertEquals(list(Person.objects.using('other').filter(book__title='Greasemonkey Hacks').values_list('name', flat=True)),

                          [])

        # Clear the books associated with mark

        mark.book_set.clear()

        self.assertEquals(list(Person.objects.using('other').filter(book__title='Dive into Python').values_list('name', flat=True)),

                          [])

        self.assertEquals(list(Person.objects.using('other').filter(book__title='Greasemonkey Hacks').values_list('name', flat=True)),

                          [])

        # Create a book through the m2m interface

        mark.book_set.create(title="Dive into HTML5", published=datetime.date(2020, 1, 1))

        self.assertEquals(list(Person.objects.using('other').filter(book__title='Dive into Python').values_list('name', flat=True)),

                          [])

        self.assertEquals(list(Person.objects.using('other').filter(book__title='Dive into HTML5').values_list('name', flat=True)),

                          [u'Mark Pilgrim'])

    def test_m2m_cross_database_protection(self):

        "Operations that involve sharing M2M objects across databases raise an error"

        # Create a book and author on the default database

        pro = Book.objects.create(title="Pro Django",

                                  published=datetime.date(2008, 12, 16))

        marty = Person.objects.create(name="Marty Alchin")

        # Create a book and author on the other database

        dive = Book.objects.using('other').create(title="Dive into Python",

                                                  published=datetime.date(2009, 5, 4))

        mark = Person.objects.using('other').create(name="Mark Pilgrim")

        # Set a foreign key set with an object from a different database

        try:

            marty.book_set = [pro, dive]

            self.fail("Shouldn't be able to assign across databases")

        except ValueError:

            pass

        # Add to an m2m with an object from a different database

        try:

            marty.book_set.add(dive)

            self.fail("Shouldn't be able to assign across databases")

        except ValueError:

            pass

        # Set a m2m with an object from a different database

        try:

            marty.book_set = [pro, dive]

            self.fail("Shouldn't be able to assign across databases")

        except ValueError:

            pass

        # Add to a reverse m2m with an object from a different database

        try:

            dive.authors.add(marty)

            self.fail("Shouldn't be able to assign across databases")

        except ValueError:

            pass

        # Set a reverse m2m with an object from a different database

        try:

            dive.authors = [mark, marty]

            self.fail("Shouldn't be able to assign across databases")

        except ValueError:

            pass

    def test_m2m_deletion(self):

        "Cascaded deletions of m2m relations issue queries on the right database"

        # Create a book and author on the other database

        dive = Book.objects.using('other').create(title="Dive into Python",

                                                  published=datetime.date(2009, 5, 4))

        mark = Person.objects.using('other').create(name="Mark Pilgrim")

        dive.authors = [mark]

        # Check the initial state

        self.assertEquals(Person.objects.using('default').count(), 0)

        self.assertEquals(Book.objects.using('default').count(), 0)

        self.assertEquals(Book.authors.through.objects.using('default').count(), 0)

        self.assertEquals(Person.objects.using('other').count(), 1)

        self.assertEquals(Book.objects.using('other').count(), 1)

        self.assertEquals(Book.authors.through.objects.using('other').count(), 1)

        # Delete the object on the other database

        dive.delete(using='other')

        self.assertEquals(Person.objects.using('default').count(), 0)

        self.assertEquals(Book.objects.using('default').count(), 0)

        self.assertEquals(Book.authors.through.objects.using('default').count(), 0)

        # The person still exists ...

        self.assertEquals(Person.objects.using('other').count(), 1)

        # ... but the book has been deleted

        self.assertEquals(Book.objects.using('other').count(), 0)

        # ... and the relationship object has also been deleted.

        self.assertEquals(Book.authors.through.objects.using('other').count(), 0)

        # Now try deletion in the reverse direction. Set up the relation again

        dive = Book.objects.using('other').create(title="Dive into Python",

                                                  published=datetime.date(2009, 5, 4))

        dive.authors = [mark]

        # Check the initial state

        self.assertEquals(Person.objects.using('default').count(), 0)

        self.assertEquals(Book.objects.using('default').count(), 0)

        self.assertEquals(Book.authors.through.objects.using('default').count(), 0)

        self.assertEquals(Person.objects.using('other').count(), 1)

        self.assertEquals(Book.objects.using('other').count(), 1)

        self.assertEquals(Book.authors.through.objects.using('other').count(), 1)

        # Delete the object on the other database

        mark.delete(using='other')

        self.assertEquals(Person.objects.using('default').count(), 0)

        self.assertEquals(Book.objects.using('default').count(), 0)

        self.assertEquals(Book.authors.through.objects.using('default').count(), 0)

        # The person has been deleted ...

        self.assertEquals(Person.objects.using('other').count(), 0)

        # ... but the book still exists

        self.assertEquals(Book.objects.using('other').count(), 1)

        # ... and the relationship object has been deleted.

        self.assertEquals(Book.authors.through.objects.using('other').count(), 0)

    def test_foreign_key_separation(self):

        "FK fields are constrained to a single database"

        # Create a book and author on the default database

        pro = Book.objects.create(title="Pro Django",

                                  published=datetime.date(2008, 12, 16))

        marty = Person.objects.create(name="Marty Alchin")

        george = Person.objects.create(name="George Vilches")

        # Create a book and author on the other database

        dive = Book.objects.using('other').create(title="Dive into Python",

                                                  published=datetime.date(2009, 5, 4))

        mark = Person.objects.using('other').create(name="Mark Pilgrim")

        chris = Person.objects.using('other').create(name="Chris Mills")

        # Save the author's favourite books

        pro.editor = george

        pro.save()

        dive.editor = chris

        dive.save()

        pro = Book.objects.using('default').get(title="Pro Django")

        self.assertEquals(pro.editor.name, "George Vilches")

        dive = Book.objects.using('other').get(title="Dive into Python")

        self.assertEquals(dive.editor.name, "Chris Mills")

        # Check that queries work across foreign key joins

        self.assertEquals(list(Person.objects.using('default').filter(edited__title='Pro Django').values_list('name', flat=True)),

                          [u'George Vilches'])

        self.assertEquals(list(Person.objects.using('other').filter(edited__title='Pro Django').values_list('name', flat=True)),

                          [])

        self.assertEquals(list(Person.objects.using('default').filter(edited__title='Dive into Python').values_list('name', flat=True)),

                          [])

        self.assertEquals(list(Person.objects.using('other').filter(edited__title='Dive into Python').values_list('name', flat=True)),

                          [u'Chris Mills'])

        # Reget the objects to clear caches

        chris = Person.objects.using('other').get(name="Chris Mills")

        dive = Book.objects.using('other').get(title="Dive into Python")

        # Retrive related object by descriptor. Related objects should be database-baound

        self.assertEquals(list(chris.edited.values_list('title', flat=True)),

                          [u'Dive into Python'])

    def test_foreign_key_reverse_operations(self):

        "FK reverse manipulations are all constrained to a single DB"

        dive = Book.objects.using('other').create(title="Dive into Python",

                                                       published=datetime.date(2009, 5, 4))

        mark = Person.objects.using('other').create(name="Mark Pilgrim")

        chris = Person.objects.using('other').create(name="Chris Mills")

        # Save the author relations

        dive.editor = chris

        dive.save()

        # Add a second book edited by chris

        html5 = Book.objects.using('other').create(title="Dive into HTML5", published=datetime.date(2010, 3, 15))

        self.assertEquals(list(Person.objects.using('other').filter(edited__title='Dive into HTML5').values_list('name', flat=True)),

                          [])

        chris.edited.add(html5)

        self.assertEquals(list(Person.objects.using('other').filter(edited__title='Dive into HTML5').values_list('name', flat=True)),

                          [u'Chris Mills'])

        self.assertEquals(list(Person.objects.using('other').filter(edited__title='Dive into Python').values_list('name', flat=True)),

                          [u'Chris Mills'])

        # Remove the second editor

        chris.edited.remove(html5)

        self.assertEquals(list(Person.objects.using('other').filter(edited__title='Dive into HTML5').values_list('name', flat=True)),

                          [])

        self.assertEquals(list(Person.objects.using('other').filter(edited__title='Dive into Python').values_list('name', flat=True)),

                          [u'Chris Mills'])

        # Clear all edited books

        chris.edited.clear()

        self.assertEquals(list(Person.objects.using('other').filter(edited__title='Dive into HTML5').values_list('name', flat=True)),

                          [])

        self.assertEquals(list(Person.objects.using('other').filter(edited__title='Dive into Python').values_list('name', flat=True)),

                          [])

        # Create an author through the m2m interface

        chris.edited.create(title='Dive into Water', published=datetime.date(2010, 3, 15))

        self.assertEquals(list(Person.objects.using('other').filter(edited__title='Dive into HTML5').values_list('name', flat=True)),

                          [])

        self.assertEquals(list(Person.objects.using('other').filter(edited__title='Dive into Water').values_list('name', flat=True)),

                          [u'Chris Mills'])

        self.assertEquals(list(Person.objects.using('other').filter(edited__title='Dive into Python').values_list('name', flat=True)),

                          [])

    def test_foreign_key_cross_database_protection(self):

        "Operations that involve sharing FK objects across databases raise an error"

        # Create a book and author on the default database

        pro = Book.objects.create(title="Pro Django",

                                  published=datetime.date(2008, 12, 16))

        marty = Person.objects.create(name="Marty Alchin")

        # Create a book and author on the other database

        dive = Book.objects.using('other').create(title="Dive into Python",

                                                  published=datetime.date(2009, 5, 4))

        mark = Person.objects.using('other').create(name="Mark Pilgrim")

        # Set a foreign key with an object from a different database

        try:

            dive.editor = marty

            self.fail("Shouldn't be able to assign across databases")

        except ValueError:

            pass

        # Set a foreign key set with an object from a different database

        try:

            marty.edited = [pro, dive]

            self.fail("Shouldn't be able to assign across databases")

        except ValueError:

            pass

        # Add to a foreign key set with an object from a different database

        try:

            marty.edited.add(dive)

            self.fail("Shouldn't be able to assign across databases")

        except ValueError:

            pass

        # BUT! if you assign a FK object when the base object hasn't

        # been saved yet, you implicitly assign the database for the

        # base object.

        chris = Person(name="Chris Mills")

        html5 = Book(title="Dive into HTML5", published=datetime.date(2010, 3, 15))

        # initially, no db assigned

        self.assertEquals(chris._state.db, None)

        self.assertEquals(html5._state.db, None)

        # old object comes from 'other', so the new object is set to use 'other'...

        dive.editor = chris

        html5.editor = mark

        self.assertEquals(chris._state.db, 'other')

        self.assertEquals(html5._state.db, 'other')

        # ... but it isn't saved yet

        self.assertEquals(list(Person.objects.using('other').values_list('name',flat=True)),

                          [u'Mark Pilgrim'])