1 import sys

     2 import time

     3 try:

     4     set

     5 except NameError:

     6     # Python 2.3 compat

     7     from sets import Set as set

     8 

     9 from django.conf import settings

    10 from django.core.management import call_command

    11 

    12 # The prefix to put on the default database name when creating

    13 # the test database.

    14 TEST_DATABASE_PREFIX = 'test_'

    15 

    16 class BaseDatabaseCreation(object):

    17     """

    18     This class encapsulates all backend-specific differences that pertain to

    19     database *creation*, such as the column types to use for particular Django

    20     Fields, the SQL used to create and destroy tables, and the creation and

    21     destruction of test databases.

    22     """

    23     data_types = {}

    24 

    25     def __init__(self, connection):

    26         self.connection = connection

    27 

    28     def _digest(self, *args):

    29         """

    30         Generates a 32-bit digest of a set of arguments that can be used to

    31         shorten identifying names.

    32         """

    33         return '%x' % (abs(hash(args)) % 4294967296L)  # 2**32

    34 

    35     def sql_create_model(self, model, style, known_models=set()):

    36         """

    37         Returns the SQL required to create a single model, as a tuple of:

    38             (list_of_sql, pending_references_dict)

    39         """

    40         from django.db import models

    41 

    42         opts = model._meta

    43         if not opts.managed or opts.proxy:

    44             return [], {}

    45         final_output = []

    46         table_output = []

    47         pending_references = {}

    48         qn = self.connection.ops.quote_name

    49         for f in opts.local_fields:

    50             col_type = f.db_type()

    51             tablespace = f.db_tablespace or opts.db_tablespace

    52             if col_type is None:

    53                 # Skip ManyToManyFields, because they're not represented as

    54                 # database columns in this table.

    55                 continue

    56             # Make the definition (e.g. 'foo VARCHAR(30)') for this field.

    57             field_output = [style.SQL_FIELD(qn(f.column)),

    58                 style.SQL_COLTYPE(col_type)]

    59             if not f.null:

    60                 field_output.append(style.SQL_KEYWORD('NOT NULL'))

    61             if f.primary_key:

    62                 field_output.append(style.SQL_KEYWORD('PRIMARY KEY'))

    63             elif f.unique:

    64                 field_output.append(style.SQL_KEYWORD('UNIQUE'))

    65             if tablespace and f.unique:

    66                 # We must specify the index tablespace inline, because we

    67                 # won't be generating a CREATE INDEX statement for this field.

    68                 field_output.append(self.connection.ops.tablespace_sql(tablespace, inline=True))

    69             if f.rel:

    70                 ref_output, pending = self.sql_for_inline_foreign_key_references(f, known_models, style)

    71                 if pending:

    72                     pr = pending_references.setdefault(f.rel.to, []).append((model, f))

    73                 else:

    74                     field_output.extend(ref_output)

    75             table_output.append(' '.join(field_output))

    76         if opts.order_with_respect_to:

    77             table_output.append(style.SQL_FIELD(qn('_order')) + ' ' + \

    78                 style.SQL_COLTYPE(models.IntegerField().db_type()))

    79         for field_constraints in opts.unique_together:

    80             table_output.append(style.SQL_KEYWORD('UNIQUE') + ' (%s)' % \

    81                 ", ".join([style.SQL_FIELD(qn(opts.get_field(f).column)) for f in field_constraints]))

    82 

    83         full_statement = [style.SQL_KEYWORD('CREATE TABLE') + ' ' + style.SQL_TABLE(qn(opts.db_table)) + ' (']

    84         for i, line in enumerate(table_output): # Combine and add commas.

    85             full_statement.append('    %s%s' % (line, i < len(table_output)-1 and ',' or ''))

    86         full_statement.append(')')

    87         if opts.db_tablespace:

    88             full_statement.append(self.connection.ops.tablespace_sql(opts.db_tablespace))

    89         full_statement.append(';')

    90         final_output.append('\n'.join(full_statement))

    91 

    92         if opts.has_auto_field:

    93             # Add any extra SQL needed to support auto-incrementing primary keys.

    94             auto_column = opts.auto_field.db_column or opts.auto_field.name

    95             autoinc_sql = self.connection.ops.autoinc_sql(opts.db_table, auto_column)

    96             if autoinc_sql:

    97                 for stmt in autoinc_sql:

    98                     final_output.append(stmt)

    99 

   100         return final_output, pending_references

   101 

   102     def sql_for_inline_foreign_key_references(self, field, known_models, style):

   103         "Return the SQL snippet defining the foreign key reference for a field"

   104         qn = self.connection.ops.quote_name

   105         if field.rel.to in known_models:

   106             output = [style.SQL_KEYWORD('REFERENCES') + ' ' + \

   107                 style.SQL_TABLE(qn(field.rel.to._meta.db_table)) + ' (' + \

   108                 style.SQL_FIELD(qn(field.rel.to._meta.get_field(field.rel.field_name).column)) + ')' +

   109                 self.connection.ops.deferrable_sql()

   110             ]

   111             pending = False

   112         else:

   113             # We haven't yet created the table to which this field

   114             # is related, so save it for later.

   115             output = []

   116             pending = True

   117 

   118         return output, pending

   119 

   120     def sql_for_pending_references(self, model, style, pending_references):

   121         "Returns any ALTER TABLE statements to add constraints after the fact."

   122         from django.db.backends.util import truncate_name

   123 

   124         if not model._meta.managed or model._meta.proxy:

   125             return []

   126         qn = self.connection.ops.quote_name

   127         final_output = []

   128         opts = model._meta

   129         if model in pending_references:

   130             for rel_class, f in pending_references[model]:

   131                 rel_opts = rel_class._meta

   132                 r_table = rel_opts.db_table

   133                 r_col = f.column

   134                 table = opts.db_table

   135                 col = opts.get_field(f.rel.field_name).column

   136                 # For MySQL, r_name must be unique in the first 64 characters.

   137                 # So we are careful with character usage here.

   138                 r_name = '%s_refs_%s_%s' % (r_col, col, self._digest(r_table, table))

   139                 final_output.append(style.SQL_KEYWORD('ALTER TABLE') + ' %s ADD CONSTRAINT %s FOREIGN KEY (%s) REFERENCES %s (%s)%s;' % \

   140                     (qn(r_table), qn(truncate_name(r_name, self.connection.ops.max_name_length())),

   141                     qn(r_col), qn(table), qn(col),

   142                     self.connection.ops.deferrable_sql()))

   143             del pending_references[model]

   144         return final_output

   145 

   146     def sql_for_many_to_many(self, model, style):

   147         "Return the CREATE TABLE statments for all the many-to-many tables defined on a model"

   148         output = []

   149         for f in model._meta.local_many_to_many:

   150             if model._meta.managed or f.rel.to._meta.managed:

   151                 output.extend(self.sql_for_many_to_many_field(model, f, style))

   152         return output

   153 

   154     def sql_for_many_to_many_field(self, model, f, style):

   155         "Return the CREATE TABLE statements for a single m2m field"

   156         from django.db import models

   157         from django.db.backends.util import truncate_name

   158 

   159         output = []

   160         if f.creates_table:

   161             opts = model._meta

   162             qn = self.connection.ops.quote_name

   163             tablespace = f.db_tablespace or opts.db_tablespace

   164             if tablespace:

   165                 sql = self.connection.ops.tablespace_sql(tablespace, inline=True)

   166                 if sql:

   167                     tablespace_sql = ' ' + sql

   168                 else:

   169                     tablespace_sql = ''

   170             else:

   171                 tablespace_sql = ''

   172             table_output = [style.SQL_KEYWORD('CREATE TABLE') + ' ' + \

   173                 style.SQL_TABLE(qn(f.m2m_db_table())) + ' (']

   174             table_output.append('    %s %s %s%s,' %

   175                 (style.SQL_FIELD(qn('id')),

   176                 style.SQL_COLTYPE(models.AutoField(primary_key=True).db_type()),

   177                 style.SQL_KEYWORD('NOT NULL PRIMARY KEY'),

   178                 tablespace_sql))

   179 

   180             deferred = []

   181             inline_output, deferred = self.sql_for_inline_many_to_many_references(model, f, style)

   182             table_output.extend(inline_output)

   183 

   184             table_output.append('    %s (%s, %s)%s' %

   185                 (style.SQL_KEYWORD('UNIQUE'),

   186                 style.SQL_FIELD(qn(f.m2m_column_name())),

   187                 style.SQL_FIELD(qn(f.m2m_reverse_name())),

   188                 tablespace_sql))

   189             table_output.append(')')

   190             if opts.db_tablespace:

   191                 # f.db_tablespace is only for indices, so ignore its value here.

   192                 table_output.append(self.connection.ops.tablespace_sql(opts.db_tablespace))

   193             table_output.append(';')

   194             output.append('\n'.join(table_output))

   195 

   196             for r_table, r_col, table, col in deferred:

   197                 r_name = '%s_refs_%s_%s' % (r_col, col, self._digest(r_table, table))

   198                 output.append(style.SQL_KEYWORD('ALTER TABLE') + ' %s ADD CONSTRAINT %s FOREIGN KEY (%s) REFERENCES %s (%s)%s;' %

   199                 (qn(r_table),

   200                 qn(truncate_name(r_name, self.connection.ops.max_name_length())),

   201                 qn(r_col), qn(table), qn(col),

   202                 self.connection.ops.deferrable_sql()))

   203 

   204             # Add any extra SQL needed to support auto-incrementing PKs

   205             autoinc_sql = self.connection.ops.autoinc_sql(f.m2m_db_table(), 'id')

   206             if autoinc_sql:

   207                 for stmt in autoinc_sql:

   208                     output.append(stmt)

   209         return output

   210 

   211     def sql_for_inline_many_to_many_references(self, model, field, style):

   212         "Create the references to other tables required by a many-to-many table"

   213         from django.db import models

   214         opts = model._meta

   215         qn = self.connection.ops.quote_name

   216 

   217         table_output = [

   218             '    %s %s %s %s (%s)%s,' %

   219                 (style.SQL_FIELD(qn(field.m2m_column_name())),

   220                 style.SQL_COLTYPE(models.ForeignKey(model).db_type()),

   221                 style.SQL_KEYWORD('NOT NULL REFERENCES'),

   222                 style.SQL_TABLE(qn(opts.db_table)),

   223                 style.SQL_FIELD(qn(opts.pk.column)),

   224                 self.connection.ops.deferrable_sql()),

   225             '    %s %s %s %s (%s)%s,' %

   226                 (style.SQL_FIELD(qn(field.m2m_reverse_name())),

   227                 style.SQL_COLTYPE(models.ForeignKey(field.rel.to).db_type()),

   228                 style.SQL_KEYWORD('NOT NULL REFERENCES'),

   229                 style.SQL_TABLE(qn(field.rel.to._meta.db_table)),

   230                 style.SQL_FIELD(qn(field.rel.to._meta.pk.column)),

   231                 self.connection.ops.deferrable_sql())

   232         ]

   233         deferred = []

   234 

   235         return table_output, deferred

   236 

   237     def sql_indexes_for_model(self, model, style):

   238         "Returns the CREATE INDEX SQL statements for a single model"

   239         if not model._meta.managed or model._meta.proxy:

   240             return []

   241         output = []

   242         for f in model._meta.local_fields:

   243             output.extend(self.sql_indexes_for_field(model, f, style))

   244         return output

   245 

   246     def sql_indexes_for_field(self, model, f, style):

   247         "Return the CREATE INDEX SQL statements for a single model field"

   248         if f.db_index and not f.unique:

   249             qn = self.connection.ops.quote_name

   250             tablespace = f.db_tablespace or model._meta.db_tablespace

   251             if tablespace:

   252                 sql = self.connection.ops.tablespace_sql(tablespace)

   253                 if sql:

   254                     tablespace_sql = ' ' + sql

   255                 else:

   256                     tablespace_sql = ''

   257             else:

   258                 tablespace_sql = ''

   259             output = [style.SQL_KEYWORD('CREATE INDEX') + ' ' +

   260                 style.SQL_TABLE(qn('%s_%s' % (model._meta.db_table, f.column))) + ' ' +

   261                 style.SQL_KEYWORD('ON') + ' ' +

   262                 style.SQL_TABLE(qn(model._meta.db_table)) + ' ' +

   263                 "(%s)" % style.SQL_FIELD(qn(f.column)) +

   264                 "%s;" % tablespace_sql]

   265         else:

   266             output = []

   267         return output

   268 

   269     def sql_destroy_model(self, model, references_to_delete, style):

   270         "Return the DROP TABLE and restraint dropping statements for a single model"

   271         if not model._meta.managed or model._meta.proxy:

   272             return []

   273         # Drop the table now

   274         qn = self.connection.ops.quote_name

   275         output = ['%s %s;' % (style.SQL_KEYWORD('DROP TABLE'),

   276                               style.SQL_TABLE(qn(model._meta.db_table)))]

   277         if model in references_to_delete:

   278             output.extend(self.sql_remove_table_constraints(model, references_to_delete, style))

   279 

   280         if model._meta.has_auto_field:

   281             ds = self.connection.ops.drop_sequence_sql(model._meta.db_table)

   282             if ds:

   283                 output.append(ds)

   284         return output

   285 

   286     def sql_remove_table_constraints(self, model, references_to_delete, style):

   287         from django.db.backends.util import truncate_name

   288 

   289         if not model._meta.managed or model._meta.proxy:

   290             return []

   291         output = []

   292         qn = self.connection.ops.quote_name

   293         for rel_class, f in references_to_delete[model]:

   294             table = rel_class._meta.db_table

   295             col = f.column

   296             r_table = model._meta.db_table

   297             r_col = model._meta.get_field(f.rel.field_name).column

   298             r_name = '%s_refs_%s_%s' % (col, r_col, self._digest(table, r_table))

   299             output.append('%s %s %s %s;' % \

   300                 (style.SQL_KEYWORD('ALTER TABLE'),

   301                 style.SQL_TABLE(qn(table)),

   302                 style.SQL_KEYWORD(self.connection.ops.drop_foreignkey_sql()),

   303                 style.SQL_FIELD(truncate_name(r_name, self.connection.ops.max_name_length()))))

   304         del references_to_delete[model]

   305         return output

   306 

   307     def sql_destroy_many_to_many(self, model, f, style):

   308         "Returns the DROP TABLE statements for a single m2m field"

   309         qn = self.connection.ops.quote_name

   310         output = []

   311         if f.creates_table:

   312             output.append("%s %s;" % (style.SQL_KEYWORD('DROP TABLE'),

   313                 style.SQL_TABLE(qn(f.m2m_db_table()))))

   314             ds = self.connection.ops.drop_sequence_sql("%s_%s" % (model._meta.db_table, f.column))

   315             if ds:

   316                 output.append(ds)

   317         return output

   318 

   319     def create_test_db(self, verbosity=1, autoclobber=False):

   320         """

   321         Creates a test database, prompting the user for confirmation if the

   322         database already exists. Returns the name of the test database created.

   323         """

   324         if verbosity >= 1:

   325             print "Creating test database..."

   326 

   327         test_database_name = self._create_test_db(verbosity, autoclobber)

   328 

   329         self.connection.close()

   330         settings.DATABASE_NAME = test_database_name

   331         self.connection.settings_dict["DATABASE_NAME"] = test_database_name

   332         can_rollback = self._rollback_works()

   333         settings.DATABASE_SUPPORTS_TRANSACTIONS = can_rollback

   334         self.connection.settings_dict["DATABASE_SUPPORTS_TRANSACTIONS"] = can_rollback

   335 

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

   337 

   338         if settings.CACHE_BACKEND.startswith('db://'):

   339             from django.core.cache import parse_backend_uri

   340             _, cache_name, _ = parse_backend_uri(settings.CACHE_BACKEND)

   341             call_command('createcachetable', cache_name)

   342 

   343         # Get a cursor (even though we don't need one yet). This has

   344         # the side effect of initializing the test database.

   345         cursor = self.connection.cursor()

   346 

   347         return test_database_name

   348 

   349     def _create_test_db(self, verbosity, autoclobber):

   350         "Internal implementation - creates the test db tables."

   351         suffix = self.sql_table_creation_suffix()

   352 

   353         if settings.TEST_DATABASE_NAME:

   354             test_database_name = settings.TEST_DATABASE_NAME

   355         else:

   356             test_database_name = TEST_DATABASE_PREFIX + settings.DATABASE_NAME

   357 

   358         qn = self.connection.ops.quote_name

   359 

   360         # Create the test database and connect to it. We need to autocommit

   361         # if the database supports it because PostgreSQL doesn't allow

   362         # CREATE/DROP DATABASE statements within transactions.

   363         cursor = self.connection.cursor()

   364         self.set_autocommit()

   365         try:

   366             cursor.execute("CREATE DATABASE %s %s" % (qn(test_database_name), suffix))

   367         except Exception, e:

   368             sys.stderr.write("Got an error creating the test database: %s\n" % e)

   369             if not autoclobber:

   370                 confirm = raw_input("Type 'yes' if you would like to try deleting the test database '%s', or 'no' to cancel: " % test_database_name)

   371             if autoclobber or confirm == 'yes':

   372                 try:

   373                     if verbosity >= 1:

   374                         print "Destroying old test database..."

   375                     cursor.execute("DROP DATABASE %s" % qn(test_database_name))

   376                     if verbosity >= 1:

   377                         print "Creating test database..."

   378                     cursor.execute("CREATE DATABASE %s %s" % (qn(test_database_name), suffix))

   379                 except Exception, e:

   380                     sys.stderr.write("Got an error recreating the test database: %s\n" % e)

   381                     sys.exit(2)

   382             else:

   383                 print "Tests cancelled."

   384                 sys.exit(1)

   385 

   386         return test_database_name

   387 

   388     def _rollback_works(self):

   389         cursor = self.connection.cursor()

   390         cursor.execute('CREATE TABLE ROLLBACK_TEST (X INT)')

   391         self.connection._commit()

   392         cursor.execute('INSERT INTO ROLLBACK_TEST (X) VALUES (8)')

   393         self.connection._rollback()

   394         cursor.execute('SELECT COUNT(X) FROM ROLLBACK_TEST')

   395         count, = cursor.fetchone()

   396         cursor.execute('DROP TABLE ROLLBACK_TEST')

   397         self.connection._commit()

   398         return count == 0

   399 

   400     def destroy_test_db(self, old_database_name, verbosity=1):

   401         """

   402         Destroy a test database, prompting the user for confirmation if the

   403         database already exists. Returns the name of the test database created.

   404         """

   405         if verbosity >= 1:

   406             print "Destroying test database..."

   407         self.connection.close()

   408         test_database_name = settings.DATABASE_NAME

   409         settings.DATABASE_NAME = old_database_name

   410         self.connection.settings_dict["DATABASE_NAME"] = old_database_name

   411 

   412         self._destroy_test_db(test_database_name, verbosity)

   413 

   414     def _destroy_test_db(self, test_database_name, verbosity):

   415         "Internal implementation - remove the test db tables."

   416         # Remove the test database to clean up after

   417         # ourselves. Connect to the previous database (not the test database)

   418         # to do so, because it's not allowed to delete a database while being

   419         # connected to it.

   420         cursor = self.connection.cursor()

   421         self.set_autocommit()

   422         time.sleep(1) # To avoid "database is being accessed by other users" errors.

   423         cursor.execute("DROP DATABASE %s" % self.connection.ops.quote_name(test_database_name))

   424         self.connection.close()

   425 

   426     def set_autocommit(self):

   427         "Make sure a connection is in autocommit mode."

   428         if hasattr(self.connection.connection, "autocommit"):

   429             if callable(self.connection.connection.autocommit):

   430                 self.connection.connection.autocommit(True)

   431             else:

   432                 self.connection.connection.autocommit = True

   433         elif hasattr(self.connection.connection, "set_isolation_level"):

   434             self.connection.connection.set_isolation_level(0)

   435 

   436     def sql_table_creation_suffix(self):

   437         "SQL to append to the end of the test table creation statements"

   438         return ''

   439