Cloned databases for running tests in parallel.

4 files changed, 146 insertions, 4 deletions

diff --git a/django/db/backends/base/creation.py b/django/db/backends/base/creation.py
index 82ccadcbe7..888dc792fe 100644
--- a/django/db/backends/base/creation.py
+++ b/django/db/backends/base/creation.py
@@ -190,13 +190,56 @@ class BaseDatabaseCreation(object):
 
         return test_database_name
 
-    def destroy_test_db(self, old_database_name, verbosity=1, keepdb=False):
+    def clone_test_db(self, number, verbosity=1, autoclobber=False, keepdb=False):
+        """
+        Clone a test database.
+        """
+        source_database_name = self.connection.settings_dict['NAME']
+
+        if verbosity >= 1:
+            test_db_repr = ''
+            action = 'Cloning test database'
+            if verbosity >= 2:
+                test_db_repr = " ('%s')" % source_database_name
+            if keepdb:
+                action = 'Using existing clone'
+            print("%s for alias '%s'%s..." % (action, self.connection.alias, test_db_repr))
+
+        # We could skip this call if keepdb is True, but we instead
+        # give it the keepdb param. See create_test_db for details.
+        self._clone_test_db(number, verbosity, keepdb)
+
+    def get_test_db_clone_settings(self, number):
+        """
+        Return a modified connection settings dict for the n-th clone of a DB.
+        """
+        # When this function is called, the test database has been created
+        # already and its name has been copied to settings_dict['NAME'] so
+        # we don't need to call _get_test_db_name.
+        orig_settings_dict = self.connection.settings_dict
+        new_settings_dict = orig_settings_dict.copy()
+        new_settings_dict['NAME'] = '{}_{}'.format(orig_settings_dict['NAME'], number)
+        return new_settings_dict
+
+    def _clone_test_db(self, number, verbosity, keepdb=False):
+        """
+        Internal implementation - duplicate the test db tables.
+        """
+        raise NotImplementedError(
+            "The database backend doesn't support cloning databases. "
+            "Disable the option to run tests in parallel processes.")
+
+    def destroy_test_db(self, old_database_name=None, verbosity=1, keepdb=False, number=None):
         """
         Destroy a test database, prompting the user for confirmation if the
         database already exists.
         """
         self.connection.close()
-        test_database_name = self.connection.settings_dict['NAME']
+        if number is None:
+            test_database_name = self.connection.settings_dict['NAME']
+        else:
+            test_database_name = self.get_test_db_clone_settings(number)['NAME']
+
         if verbosity >= 1:
             test_db_repr = ''
             action = 'Destroying'
@@ -213,8 +256,9 @@ class BaseDatabaseCreation(object):
             self._destroy_test_db(test_database_name, verbosity)
 
         # Restore the original database name
-        settings.DATABASES[self.connection.alias]["NAME"] = old_database_name
-        self.connection.settings_dict["NAME"] = old_database_name
+        if old_database_name is not None:
+            settings.DATABASES[self.connection.alias]["NAME"] = old_database_name
+            self.connection.settings_dict["NAME"] = old_database_name
 
     def _destroy_test_db(self, test_database_name, verbosity):
         """
diff --git a/django/db/backends/mysql/creation.py b/django/db/backends/mysql/creation.py
index aba9e8b841..a992e65952 100644
--- a/django/db/backends/mysql/creation.py
+++ b/django/db/backends/mysql/creation.py
@@ -1,5 +1,10 @@
+import subprocess
+import sys
+
 from django.db.backends.base.creation import BaseDatabaseCreation
 
+from .client import DatabaseClient
+
 
 class DatabaseCreation(BaseDatabaseCreation):
 
@@ -11,3 +16,34 @@ class DatabaseCreation(BaseDatabaseCreation):
         if test_settings['COLLATION']:
             suffix.append('COLLATE %s' % test_settings['COLLATION'])
         return ' '.join(suffix)
+
+    def _clone_test_db(self, number, verbosity, keepdb=False):
+        qn = self.connection.ops.quote_name
+        source_database_name = self.connection.settings_dict['NAME']
+        target_database_name = self.get_test_db_clone_settings(number)['NAME']
+
+        with self._nodb_connection.cursor() as cursor:
+            try:
+                cursor.execute("CREATE DATABASE %s" % qn(target_database_name))
+            except Exception as e:
+                if keepdb:
+                    return
+                try:
+                    if verbosity >= 1:
+                        print("Destroying old test database '%s'..." % self.connection.alias)
+                    cursor.execute("DROP DATABASE %s" % qn(target_database_name))
+                    cursor.execute("CREATE DATABASE %s" % qn(target_database_name))
+                except Exception as e:
+                    sys.stderr.write("Got an error recreating the test database: %s\n" % e)
+                    sys.exit(2)
+
+        dump_cmd = DatabaseClient.settings_to_cmd_args(self.connection.settings_dict)
+        dump_cmd[0] = 'mysqldump'
+        dump_cmd[-1] = source_database_name
+        load_cmd = DatabaseClient.settings_to_cmd_args(self.connection.settings_dict)
+        load_cmd[-1] = target_database_name
+
+        dump_proc = subprocess.Popen(dump_cmd, stdout=subprocess.PIPE)
+        load_proc = subprocess.Popen(load_cmd, stdin=dump_proc.stdout, stdout=subprocess.PIPE)
+        dump_proc.stdout.close()    # allow dump_proc to receive a SIGPIPE if load_proc exits.
+        load_proc.communicate()
diff --git a/django/db/backends/postgresql/creation.py b/django/db/backends/postgresql/creation.py
index 13eea0d3b7..dcb0430f1b 100644
--- a/django/db/backends/postgresql/creation.py
+++ b/django/db/backends/postgresql/creation.py
@@ -1,3 +1,5 @@
+import sys
+
 from django.db.backends.base.creation import BaseDatabaseCreation
 
 
@@ -11,3 +13,29 @@ class DatabaseCreation(BaseDatabaseCreation):
         if test_settings['CHARSET']:
             return "WITH ENCODING '%s'" % test_settings['CHARSET']
         return ''
+
+    def _clone_test_db(self, number, verbosity, keepdb=False):
+        # CREATE DATABASE ... WITH TEMPLATE ... requires closing connections
+        # to the template database.
+        self.connection.close()
+
+        qn = self.connection.ops.quote_name
+        source_database_name = self.connection.settings_dict['NAME']
+        target_database_name = self.get_test_db_clone_settings(number)['NAME']
+
+        with self._nodb_connection.cursor() as cursor:
+            try:
+                cursor.execute("CREATE DATABASE %s WITH TEMPLATE %s" % (
+                    qn(target_database_name), qn(source_database_name)))
+            except Exception as e:
+                if keepdb:
+                    return
+                try:
+                    if verbosity >= 1:
+                        print("Destroying old test database '%s'..." % self.connection.alias)
+                    cursor.execute("DROP DATABASE %s" % qn(target_database_name))
+                    cursor.execute("CREATE DATABASE %s WITH TEMPLATE %s" % (
+                        qn(target_database_name), qn(source_database_name)))
+                except Exception as e:
+                    sys.stderr.write("Got an error cloning the test database: %s\n" % e)
+                    sys.exit(2)
diff --git a/django/db/backends/sqlite3/creation.py b/django/db/backends/sqlite3/creation.py
index d8fde0c9b7..d52349e57f 100644
--- a/django/db/backends/sqlite3/creation.py
+++ b/django/db/backends/sqlite3/creation.py
@@ -1,4 +1,5 @@
 import os
+import shutil
 import sys
 
 from django.core.exceptions import ImproperlyConfigured
@@ -47,6 +48,39 @@ class DatabaseCreation(BaseDatabaseCreation):
                     sys.exit(1)
         return test_database_name
 
+    def get_test_db_clone_settings(self, number):
+        orig_settings_dict = self.connection.settings_dict
+        source_database_name = orig_settings_dict['NAME']
+        if self.connection.is_in_memory_db(source_database_name):
+            return orig_settings_dict
+        else:
+            new_settings_dict = orig_settings_dict.copy()
+            root, ext = os.path.splitext(orig_settings_dict['NAME'])
+            new_settings_dict['NAME'] = '{}_{}.{}'.format(root, number, ext)
+            return new_settings_dict
+
+    def _clone_test_db(self, number, verbosity, keepdb=False):
+        source_database_name = self.connection.settings_dict['NAME']
+        target_database_name = self.get_test_db_clone_settings(number)['NAME']
+        # Forking automatically makes a copy of an in-memory database.
+        if not self.connection.is_in_memory_db(source_database_name):
+            # Erase the old test database
+            if os.access(target_database_name, os.F_OK):
+                if keepdb:
+                    return
+                if verbosity >= 1:
+                    print("Destroying old test database '%s'..." % target_database_name)
+                try:
+                    os.remove(target_database_name)
+                except Exception as e:
+                    sys.stderr.write("Got an error deleting the old test database: %s\n" % e)
+                    sys.exit(2)
+            try:
+                shutil.copy(source_database_name, target_database_name)
+            except Exception as e:
+                sys.stderr.write("Got an error cloning the test database: %s\n" % e)
+                sys.exit(2)
+
     def _destroy_test_db(self, test_database_name, verbosity):
         if test_database_name and not self.connection.is_in_memory_db(test_database_name):
             # Remove the SQLite database file