1 import sys

     2 import time

     3 

     4 from django.conf import settings

     5 from django.core.management import call_command

     6 

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

     8 # the test database.

     9 TEST_DATABASE_PREFIX = 'test_'

    10 

    11 class BaseDatabaseCreation(object):

    12     """

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

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

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

    16     destruction of test databases.

    17     """

    18     data_types = {}

    19 

    20     def __init__(self, connection):

    21         self.connection = connection

    22 

    23     def _digest(self, *args):

    24         """

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

    26         shorten identifying names.

    27         """

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

    29 

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

    31         """

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

    33             (list_of_sql, pending_references_dict)

    34         """

    35         from django.db import models

    36 

    37         opts = model._meta

    38         if not opts.managed or opts.proxy:

    39             return [], {}

    40         final_output = []

    41         table_output = []

    42         pending_references = {}

    43         qn = self.connection.ops.quote_name

    44         for f in opts.local_fields:

    45             col_type = f.db_type(connection=self.connection)

    46             tablespace = f.db_tablespace or opts.db_tablespace

    47             if col_type is None:

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

    49                 # database columns in this table.

    50                 continue

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

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

    53                 style.SQL_COLTYPE(col_type)]

    54             if not f.null:

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

    56             if f.primary_key:

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

    58             elif f.unique:

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

    60             if tablespace and f.unique:

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

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

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

    64             if f.rel:

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

    66                 if pending:

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

    68                 else:

    69                     field_output.extend(ref_output)

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

    71         for field_constraints in opts.unique_together:

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

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

    74 

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

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

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

    78         full_statement.append(')')

    79         if opts.db_tablespace:

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

    81         full_statement.append(';')

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

    83 

    84         if opts.has_auto_field:

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

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

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

    88             if autoinc_sql:

    89                 for stmt in autoinc_sql:

    90                     final_output.append(stmt)

    91 

    92         return final_output, pending_references

    93 

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

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

    96         qn = self.connection.ops.quote_name

    97         if field.rel.to in known_models:

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

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

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

   101                 self.connection.ops.deferrable_sql()

   102             ]

   103             pending = False

   104         else:

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

   106             # is related, so save it for later.

   107             output = []

   108             pending = True

   109 

   110         return output, pending

   111 

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

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

   114         from django.db.backends.util import truncate_name

   115 

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

   117             return []

   118         qn = self.connection.ops.quote_name

   119         final_output = []

   120         opts = model._meta

   121         if model in pending_references:

   122             for rel_class, f in pending_references[model]:

   123                 rel_opts = rel_class._meta

   124                 r_table = rel_opts.db_table

   125                 r_col = f.column

   126                 table = opts.db_table

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

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

   129                 # So we are careful with character usage here.

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

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

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

   133                     qn(r_col), qn(table), qn(col),

   134                     self.connection.ops.deferrable_sql()))

   135             del pending_references[model]

   136         return final_output

   137 

   138     def sql_for_many_to_many(self, model, style):

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

   140         import warnings

   141         warnings.warn(

   142             'Database creation API for m2m tables has been deprecated. M2M models are now automatically generated',

   143             PendingDeprecationWarning

   144         )

   145 

   146         output = []

   147         for f in model._meta.local_many_to_many:

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

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

   150         return output

   151 

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

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

   154         import warnings

   155         warnings.warn(

   156             'Database creation API for m2m tables has been deprecated. M2M models are now automatically generated',

   157             PendingDeprecationWarning

   158         )

   159 

   160         from django.db import models

   161         from django.db.backends.util import truncate_name

   162 

   163         output = []

   164         if f.auto_created:

   165             opts = model._meta

   166             qn = self.connection.ops.quote_name

   167             tablespace = f.db_tablespace or opts.db_tablespace

   168             if tablespace:

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

   170                 if sql:

   171                     tablespace_sql = ' ' + sql

   172                 else:

   173                     tablespace_sql = ''

   174             else:

   175                 tablespace_sql = ''

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

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

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

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

   180                 style.SQL_COLTYPE(models.AutoField(primary_key=True).db_type(connection=self.connection)),

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

   182                 tablespace_sql))

   183 

   184             deferred = []

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

   186             table_output.extend(inline_output)

   187 

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

   189                 (style.SQL_KEYWORD('UNIQUE'),

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

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

   192                 tablespace_sql))

   193             table_output.append(')')

   194             if opts.db_tablespace:

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

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

   197             table_output.append(';')

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

   199 

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

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

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

   203                 (qn(r_table),

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

   205                 qn(r_col), qn(table), qn(col),

   206                 self.connection.ops.deferrable_sql()))

   207 

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

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

   210             if autoinc_sql:

   211                 for stmt in autoinc_sql:

   212                     output.append(stmt)

   213         return output

   214 

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

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

   217         import warnings

   218         warnings.warn(

   219             'Database creation API for m2m tables has been deprecated. M2M models are now automatically generated',

   220             PendingDeprecationWarning

   221         )

   222 

   223         from django.db import models

   224         opts = model._meta

   225         qn = self.connection.ops.quote_name

   226 

   227         table_output = [

   228             '    %s %s %s %s (%s)%s,' %

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

   230                 style.SQL_COLTYPE(models.ForeignKey(model).db_type(connection=self.connection)),

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

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

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

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

   235             '    %s %s %s %s (%s)%s,' %

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

   237                 style.SQL_COLTYPE(models.ForeignKey(field.rel.to).db_type(connection=self.connection)),

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

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

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

   241                 self.connection.ops.deferrable_sql())

   242         ]

   243         deferred = []

   244 

   245         return table_output, deferred

   246 

   247     def sql_indexes_for_model(self, model, style):

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

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

   250             return []

   251         output = []

   252         for f in model._meta.local_fields:

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

   254         return output

   255 

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

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

   258         from django.db.backends.util import truncate_name

   259 

   260         if f.db_index and not f.unique:

   261             qn = self.connection.ops.quote_name

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

   263             if tablespace:

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

   265                 if sql:

   266                     tablespace_sql = ' ' + sql

   267                 else:

   268                     tablespace_sql = ''

   269             else:

   270                 tablespace_sql = ''

   271             i_name = '%s_%s' % (model._meta.db_table, self._digest(f.column))

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

   273                 style.SQL_TABLE(qn(truncate_name(i_name, self.connection.ops.max_name_length()))) + ' ' +

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

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

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

   277                 "%s;" % tablespace_sql]

   278         else:

   279             output = []

   280         return output

   281 

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

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

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

   285             return []

   286         # Drop the table now

   287         qn = self.connection.ops.quote_name

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

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

   290         if model in references_to_delete:

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

   292 

   293         if model._meta.has_auto_field:

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

   295             if ds:

   296                 output.append(ds)

   297         return output

   298 

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

   300         from django.db.backends.util import truncate_name

   301 

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

   303             return []

   304         output = []

   305         qn = self.connection.ops.quote_name

   306         for rel_class, f in references_to_delete[model]:

   307             table = rel_class._meta.db_table

   308             col = f.column

   309             r_table = model._meta.db_table

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

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

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

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

   314                 style.SQL_TABLE(qn(table)),

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

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

   317         del references_to_delete[model]

   318         return output

   319 

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

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

   322         import warnings

   323         warnings.warn(

   324             'Database creation API for m2m tables has been deprecated. M2M models are now automatically generated',

   325             PendingDeprecationWarning

   326         )

   327 

   328         qn = self.connection.ops.quote_name

   329         output = []

   330         if f.auto_created:

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

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

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

   334             if ds:

   335                 output.append(ds)

   336         return output

   337 

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

   339         """

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

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

   342         """

   343         if verbosity >= 1:

   344             print "Creating test database '%s'..." % self.connection.alias

   345 

   346         test_database_name = self._create_test_db(verbosity, autoclobber)

   347 

   348         self.connection.close()

   349         self.connection.settings_dict["NAME"] = test_database_name

   350         can_rollback = self._rollback_works()

   351         self.connection.settings_dict["SUPPORTS_TRANSACTIONS"] = can_rollback

   352 

   353         call_command('syncdb', verbosity=verbosity, interactive=False, database=self.connection.alias)

   354 

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

   356             from django.core.cache import parse_backend_uri

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

   358             call_command('createcachetable', cache_name)

   359 

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

   361         # the side effect of initializing the test database.

   362         cursor = self.connection.cursor()

   363 

   364         return test_database_name

   365 

   366     def _create_test_db(self, verbosity, autoclobber):

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

   368         suffix = self.sql_table_creation_suffix()

   369 

   370         if self.connection.settings_dict['TEST_NAME']:

   371             test_database_name = self.connection.settings_dict['TEST_NAME']

   372         else:

   373             test_database_name = TEST_DATABASE_PREFIX + self.connection.settings_dict['NAME']

   374 

   375         qn = self.connection.ops.quote_name

   376 

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

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

   379         # CREATE/DROP DATABASE statements within transactions.

   380         cursor = self.connection.cursor()

   381         self.set_autocommit()

   382         try:

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

   384         except Exception, e:

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

   386             if not autoclobber:

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

   388             if autoclobber or confirm == 'yes':

   389                 try:

   390                     if verbosity >= 1:

   391                         print "Destroying old test database..."

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

   393                     if verbosity >= 1:

   394                         print "Creating test database..."

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

   396                 except Exception, e:

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

   398                     sys.exit(2)

   399             else:

   400                 print "Tests cancelled."

   401                 sys.exit(1)

   402 

   403         return test_database_name

   404 

   405     def _rollback_works(self):

   406         cursor = self.connection.cursor()

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

   408         self.connection._commit()

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

   410         self.connection._rollback()

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

   412         count, = cursor.fetchone()

   413         cursor.execute('DROP TABLE ROLLBACK_TEST')

   414         self.connection._commit()

   415         return count == 0

   416 

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

   418         """

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

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

   421         """

   422         if verbosity >= 1:

   423             print "Destroying test database '%s'..." % self.connection.alias

   424         self.connection.close()

   425         test_database_name = self.connection.settings_dict['NAME']

   426         self.connection.settings_dict['NAME'] = old_database_name

   427 

   428         self._destroy_test_db(test_database_name, verbosity)

   429 

   430     def _destroy_test_db(self, test_database_name, verbosity):

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

   432         # Remove the test database to clean up after

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

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

   435         # connected to it.

   436         cursor = self.connection.cursor()

   437         self.set_autocommit()

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

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

   440         self.connection.close()

   441 

   442     def set_autocommit(self):

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

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

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

   446                 self.connection.connection.autocommit(True)

   447             else:

   448                 self.connection.connection.autocommit = True

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

   450             self.connection.connection.set_isolation_level(0)

   451 

   452     def sql_table_creation_suffix(self):

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

   454         return ''