How to run sql on a Django test database before syncdb

I've had a problem for some time with running tests on one of my older Django projects. In this project I had some models that wasn't being managed by Django. This mean that it was me who created the tables for the models instead of Django. This can be done pretty easily with the Meta class like this:

from django.db import models
class SomeModel(models.Model):
....

class Meta:
managed = False
db_table = 'table_name'

“Now why would you do such a thing?” you might ask. I needed to create the tsvector data type - a postgres data type that is designed to support full text search. Something that quickly would be a pain with the Django ORM, that I enjoy to use. But in this case I was prepared to sacrifice the ORM to get full control over the SQL queries when I needed, to take full advantage of my postgres database.

Anyways everything worked out as planed and I didn't have any problems with my Django model etc. Then came the time to write some unittests and this is where I ran into some serious problems. When you use the Django testing framework, it will create a test database for you, to run you tests on. I discovered that unmanaged models wont get their tables created along with the others. This wasn’t a big surprise and would not be an issue, most of the time. In my case I needed those tables to be created before some of my apps got their tables created, since I had foreign keys to my unmanaged models. I thought such a task would be pretty simple, since this was an option that Django gives so easily and freely. I soon found out that is was not going to quite that easy to get my tables created. I was pretty fortunate that the unmanaged tables didn't have FK relations to some of the managed ones which would have made this much more complex.

So the task at hand was pretty simple. How to run some custom sql before the syncdb command is executed in test database creation.

I have listed here the steps I did to achive this, using django 1.1. In django 1.2 the test runner is a class instead of a function, but I bet the same procedure would work.

### 1 - Defining your own test runner###
In the settings.py file, you can write a setting, if you want to overwrite the default test runner: django.test.simple.run_tests all you need to do is enter this in your settings.py:

TEST_RUNNER = 'path.to.run_tests'

What I did was to copy the default test runner and work with that.

### 2 - Digging into the code###
In the default test runner the interesting thing for us happens at line 190 in simple.py:

from django.db import connection
connection.creation.create_test_db(verbosity, autoclobber=not interactive)

create_test_db is a method on the BaseDatabaseCreation object defined in django.db.backends.creation. The code for the method looks like this:

def create_test_db(self, verbosity=1, autoclobber=False):
"""
Creates a test database, prompting the user for confirmation if the
database already exists. Returns the name of the test database created.
"""
if verbosity >= 1:
print "Creating test database..."

test_database_name = self._create_test_db(verbosity, autoclobber)

self.connection.close()
settings.DATABASE_NAME = test_database_name
self.connection.settings_dict["DATABASE_NAME"] = test_database_name
can_rollback = self._rollback_works()
settings.DATABASE_SUPPORTS_TRANSACTIONS = can_rollback
self.connection.settings_dict["DATABASE_SUPPORTS_TRANSACTIONS"] = can_rollback

call_command('syncdb', verbosity=verbosity, interactive=False)

if settings.CACHE_BACKEND.startswith('db://'):
from django.core.cache import parse_backend_uri
_, cache_name, _ = parse_backend_uri(settings.CACHE_BACKEND)
call_command('createcachetable', cache_name)

# Get a cursor (even though we don't need one yet). This has
# the side effect of initializing the test database.
cursor = self.connection.cursor()

return test_database_name

### 3 - Subclassing the DatabaseWrapper###
I wanted add my own code before call_command('syncdb', verbosity=verbosity, interactive=False) was executed, as this is here Django creates the tables by calling syncdb on the freshly created test database. Since I knew what database type I have, I didn't need some logic to fetch the correct one as Django does, so I went right ahead and made a quick subclass of the psycopg2 DatabaseWrapper:

from django.db.backends.postgresql_psycopg2.base import DatabaseWrapper
class MyDatabaseWrapper(DatabaseWrapper):
def __init__(self, *args, **kwargs):
super(FlexDatabaseWrapper, self).__init__(*args, **kwargs)
self.creation = MyDatabaseCreation(self)

All I do is to change the DatabaseCreation object that is used, with my own, which I again needed to subclass:

class FlexDatabaseCreation(DatabaseCreation):
def create_test_db(self, verbosity=1, autoclobber=False):
...

Then in my run_test I could now replace:

from django.db import connection
connection.creation.create_test_db(verbosity, autoclobber=not interactive)

with

connection = MyDatabaseWrapper({
'DATABASE_HOST': settings.DATABASE_HOST,
'DATABASE_NAME': settings.DATABASE_NAME,
'DATABASE_OPTIONS': settings.DATABASE_OPTIONS,
'DATABASE_PASSWORD': settings.DATABASE_PASSWORD,
'DATABASE_PORT': settings.DATABASE_PORT,
'DATABASE_USER': settings.DATABASE_USER,
'TIME_ZONE': settings.TIME_ZONE,
})
connection.creation.create_test_db(verbosity, autoclobber=not interactive)

### 4 - Making it all work ###
We now have control over the function that creates the test database with tables. But this also gave a little side problem. Because we don't use Django's connection, but our own, Django will keep on using it's regular db and not the test db. This is quite unfortunate at the result is running all the tests on a live / staging database. This is fixed with a few lines of code.

from django.db import connection
class FlexDatabaseCreation(DatabaseCreation):
def create_test_db(self, verbosity=1, autoclobber=False):
...
self.connection.settings_dict["DATABASE_NAME"] = test_database_name
connection.settings_dict["DATABASE_NAME"] = test_database_name # New code.
...
self.connection.settings_dict["DATABASE_SUPPORTS_TRANSACTIONS"] = can_rollback
connection.settings_dict["DATABASE_SUPPORTS_TRANSACTIONS"] = can_rollback # New code.
...

With this little addition we now have created a test database and make Django use it, we are really where we started, only now Django is using our code to do it instead. The last thing that is needed is to add the custom sql. I did it be reading a file where I have the SQL for the tables stored:

from django.db import connection
class FlexDatabaseCreation(DatabaseCreation):
def create_test_db(self, verbosity=1, autoclobber=False):
...
f = open('path/to/schema.sql')
sql = f.read()
f.close()
cursor.execute(sql)
cursor.connection.commit()

call_command('syncdb', verbosity=verbosity, interactive=False)

###5 - Success###
That's all there is too it. I would have loved if there was a signal pre syncdb, so you simply could setup the SQL to be run in a signal instead of having to grind through all this code. But in the end, it was a fun challenge to overcome and not a type of thing I’ve done much when it comes to Django.

Add new comment