1 from django.core.management.base import CommandError

     2 import os

     3 import re

     4 

     5 try:

     6     set

     7 except NameError:

     8     from sets import Set as set   # Python 2.3 fallback

     9 

    10 def table_names():

    11     "Returns a list of all table names that exist in the database."

    12     from django.db import connection, get_introspection_module

    13     cursor = connection.cursor()

    14     return set(get_introspection_module().get_table_list(cursor))

    15 

    16 def django_table_names(only_existing=False):

    17     """

    18     Returns a list of all table names that have associated Django models and

    19     are in INSTALLED_APPS.

    20 

    21     If only_existing is True, the resulting list will only include the tables

    22     that actually exist in the database.

    23     """

    24     from django.db import models

    25     tables = set()

    26     for app in models.get_apps():

    27         for model in models.get_models(app):

    28             tables.add(model._meta.db_table)

    29             tables.update([f.m2m_db_table() for f in model._meta.local_many_to_many])

    30     if only_existing:

    31         tables = [t for t in tables if t in table_names()]

    32     return tables

    33 

    34 def installed_models(table_list):

    35     "Returns a set of all models that are installed, given a list of existing table names."

    36     from django.db import connection, models

    37     all_models = []

    38     for app in models.get_apps():

    39         for model in models.get_models(app):

    40             all_models.append(model)

    41     if connection.features.uses_case_insensitive_names:

    42         converter = lambda x: x.upper()

    43     else:

    44         converter = lambda x: x

    45     return set([m for m in all_models if converter(m._meta.db_table) in map(converter, table_list)])

    46 

    47 def sequence_list():

    48     "Returns a list of information about all DB sequences for all models in all apps."

    49     from django.db import models

    50 

    51     apps = models.get_apps()

    52     sequence_list = []

    53 

    54     for app in apps:

    55         for model in models.get_models(app):

    56             for f in model._meta.local_fields:

    57                 if isinstance(f, models.AutoField):

    58                     sequence_list.append({'table': model._meta.db_table, 'column': f.column})

    59                     break # Only one AutoField is allowed per model, so don't bother continuing.

    60 

    61             for f in model._meta.local_many_to_many:

    62                 sequence_list.append({'table': f.m2m_db_table(), 'column': None})

    63 

    64     return sequence_list

    65 

    66 def sql_create(app, style):

    67     "Returns a list of the CREATE TABLE SQL statements for the given app."

    68     from django.db import models

    69     from django.conf import settings

    70 

    71     if settings.DATABASE_ENGINE == 'dummy':

    72         # This must be the "dummy" database backend, which means the user

    73         # hasn't set DATABASE_ENGINE.

    74         raise CommandError("Django doesn't know which syntax to use for your SQL statements,\n" +

    75             "because you haven't specified the DATABASE_ENGINE setting.\n" +

    76             "Edit your settings file and change DATABASE_ENGINE to something like 'postgresql' or 'mysql'.")

    77 

    78     # Get installed models, so we generate REFERENCES right.

    79     # We trim models from the current app so that the sqlreset command does not

    80     # generate invalid SQL (leaving models out of known_models is harmless, so

    81     # we can be conservative).

    82     app_models = models.get_models(app)

    83     final_output = []

    84     known_models = set([model for model in installed_models(table_names()) if model not in app_models])

    85     pending_references = {}

    86 

    87     for model in app_models:

    88         output, references = sql_model_create(model, style, known_models)

    89         final_output.extend(output)

    90         for refto, refs in references.items():

    91             pending_references.setdefault(refto, []).extend(refs)

    92             if refto in known_models:

    93                 final_output.extend(sql_for_pending_references(refto, style, pending_references))

    94         final_output.extend(sql_for_pending_references(model, style, pending_references))

    95         # Keep track of the fact that we've created the table for this model.

    96         known_models.add(model)

    97 

    98     # Create the many-to-many join tables.

    99     for model in app_models:

   100         final_output.extend(many_to_many_sql_for_model(model, style))

   101 

   102     # Handle references to tables that are from other apps

   103     # but don't exist physically.

   104     not_installed_models = set(pending_references.keys())

   105     if not_installed_models:

   106         alter_sql = []

   107         for model in not_installed_models:

   108             alter_sql.extend(['-- ' + sql for sql in

   109                 sql_for_pending_references(model, style, pending_references)])

   110         if alter_sql:

   111             final_output.append('-- The following references should be added but depend on non-existent tables:')

   112             final_output.extend(alter_sql)

   113 

   114     return final_output

   115 

   116 def sql_delete(app, style):

   117     "Returns a list of the DROP TABLE SQL statements for the given app."

   118     from django.db import connection, models, get_introspection_module

   119     from django.db.backends.util import truncate_name

   120     from django.contrib.contenttypes import generic

   121     introspection = get_introspection_module()

   122 

   123     # This should work even if a connection isn't available

   124     try:

   125         cursor = connection.cursor()

   126     except:

   127         cursor = None

   128 

   129     # Figure out which tables already exist

   130     if cursor:

   131         table_names = introspection.get_table_list(cursor)

   132     else:

   133         table_names = []

   134     if connection.features.uses_case_insensitive_names:

   135         table_name_converter = lambda x: x.upper()

   136     else:

   137         table_name_converter = lambda x: x

   138 

   139     output = []

   140     qn = connection.ops.quote_name

   141 

   142     # Output DROP TABLE statements for standard application tables.

   143     to_delete = set()

   144 

   145     references_to_delete = {}

   146     app_models = models.get_models(app)

   147     for model in app_models:

   148         if cursor and table_name_converter(model._meta.db_table) in table_names:

   149             # The table exists, so it needs to be dropped

   150             opts = model._meta

   151             for f in opts.local_fields:

   152                 if f.rel and f.rel.to not in to_delete:

   153                     references_to_delete.setdefault(f.rel.to, []).append( (model, f) )

   154 

   155             to_delete.add(model)

   156 

   157     for model in app_models:

   158         if cursor and table_name_converter(model._meta.db_table) in table_names:

   159             # Drop the table now

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

   161                 style.SQL_TABLE(qn(model._meta.db_table))))

   162             if connection.features.supports_constraints and model in references_to_delete:

   163                 for rel_class, f in references_to_delete[model]:

   164                     table = rel_class._meta.db_table

   165                     col = f.column

   166                     r_table = model._meta.db_table

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

   168                     r_name = '%s_refs_%s_%x' % (col, r_col, abs(hash((table, r_table))))

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

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

   171                         style.SQL_TABLE(qn(table)),

   172                         style.SQL_KEYWORD(connection.ops.drop_foreignkey_sql()),

   173                         style.SQL_FIELD(truncate_name(r_name, connection.ops.max_name_length()))))

   174                 del references_to_delete[model]

   175             if model._meta.has_auto_field:

   176                 ds = connection.ops.drop_sequence_sql(model._meta.db_table)

   177                 if ds:

   178                     output.append(ds)

   179 

   180     # Output DROP TABLE statements for many-to-many tables.

   181     for model in app_models:

   182         opts = model._meta

   183         for f in opts.local_many_to_many:

   184             if isinstance(f.rel, generic.GenericRel):

   185                 continue

   186             if cursor and table_name_converter(f.m2m_db_table()) in table_names:

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

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

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

   190                 if ds:

   191                     output.append(ds)

   192 

   193     app_label = app_models[0]._meta.app_label

   194 

   195     # Close database connection explicitly, in case this output is being piped

   196     # directly into a database client, to avoid locking issues.

   197     if cursor:

   198         cursor.close()

   199         connection.close()

   200 

   201     return output[::-1] # Reverse it, to deal with table dependencies.

   202 

   203 def sql_reset(app, style):

   204     "Returns a list of the DROP TABLE SQL, then the CREATE TABLE SQL, for the given module."

   205     return sql_delete(app, style) + sql_all(app, style)

   206 

   207 def sql_flush(style, only_django=False):

   208     """

   209     Returns a list of the SQL statements used to flush the database.

   210     

   211     If only_django is True, then only table names that have associated Django

   212     models and are in INSTALLED_APPS will be included.

   213     """

   214     from django.db import connection

   215     if only_django:

   216         tables = django_table_names()

   217     else:

   218         tables = table_names()

   219     statements = connection.ops.sql_flush(style, tables, sequence_list())

   220     return statements

   221 

   222 def sql_custom(app):

   223     "Returns a list of the custom table modifying SQL statements for the given app."

   224     from django.db.models import get_models

   225     output = []

   226 

   227     app_models = get_models(app)

   228     app_dir = os.path.normpath(os.path.join(os.path.dirname(app.__file__), 'sql'))

   229 

   230     for model in app_models:

   231         output.extend(custom_sql_for_model(model))

   232 

   233     return output

   234 

   235 def sql_indexes(app, style):

   236     "Returns a list of the CREATE INDEX SQL statements for all models in the given app."

   237     from django.db import models

   238     output = []

   239     for model in models.get_models(app):

   240         output.extend(sql_indexes_for_model(model, style))

   241     return output

   242 

   243 def sql_all(app, style):

   244     "Returns a list of CREATE TABLE SQL, initial-data inserts, and CREATE INDEX SQL for the given module."

   245     return sql_create(app, style) + sql_custom(app) + sql_indexes(app, style)

   246 

   247 def sql_model_create(model, style, known_models=set()):

   248     """

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

   250         (list_of_sql, pending_references_dict)

   251     """

   252     from django.db import connection, models

   253 

   254     opts = model._meta

   255     final_output = []

   256     table_output = []

   257     pending_references = {}

   258     qn = connection.ops.quote_name

   259     inline_references = connection.features.inline_fk_references

   260     for f in opts.local_fields:

   261         col_type = f.db_type()

   262         tablespace = f.db_tablespace or opts.db_tablespace

   263         if col_type is None:

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

   265             # database columns in this table.

   266             continue

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

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

   269             style.SQL_COLTYPE(col_type)]

   270         field_output.append(style.SQL_KEYWORD('%sNULL' % (not f.null and 'NOT ' or '')))

   271         if f.unique and (not f.primary_key or connection.features.allows_unique_and_pk):

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

   273         if f.primary_key:

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

   275         if tablespace and connection.features.supports_tablespaces and (f.unique or f.primary_key) and connection.features.autoindexes_primary_keys:

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

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

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

   279         if f.rel:

   280             if inline_references and f.rel.to in known_models:

   281                 field_output.append(style.SQL_KEYWORD('REFERENCES') + ' ' + \

   282                     style.SQL_TABLE(qn(f.rel.to._meta.db_table)) + ' (' + \

   283                     style.SQL_FIELD(qn(f.rel.to._meta.get_field(f.rel.field_name).column)) + ')' +

   284                     connection.ops.deferrable_sql()

   285                 )

   286             else:

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

   288                 # is related, so save it for later.

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

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

   291     if opts.order_with_respect_to:

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

   293             style.SQL_COLTYPE(models.IntegerField().db_type()) + ' ' + \

   294             style.SQL_KEYWORD('NULL'))

   295     for field_constraints in opts.unique_together:

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

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

   298 

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

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

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

   302     full_statement.append(')')

   303     if opts.db_tablespace and connection.features.supports_tablespaces:

   304         full_statement.append(connection.ops.tablespace_sql(opts.db_tablespace))

   305     full_statement.append(';')

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

   307 

   308     if opts.has_auto_field:

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

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

   311         autoinc_sql = connection.ops.autoinc_sql(opts.db_table, auto_column)

   312         if autoinc_sql:

   313             for stmt in autoinc_sql:

   314                 final_output.append(stmt)

   315 

   316     return final_output, pending_references

   317 

   318 def sql_for_pending_references(model, style, pending_references):

   319     """

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

   321     """

   322     from django.db import connection

   323     from django.db.backends.util import truncate_name

   324 

   325     qn = connection.ops.quote_name

   326     final_output = []

   327     if connection.features.supports_constraints:

   328         opts = model._meta

   329         if model in pending_references:

   330             for rel_class, f in pending_references[model]:

   331                 rel_opts = rel_class._meta

   332                 r_table = rel_opts.db_table

   333                 r_col = f.column

   334                 table = opts.db_table

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

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

   337                 # So we are careful with character usage here.

   338                 r_name = '%s_refs_%s_%x' % (r_col, col, abs(hash((r_table, table))))

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

   340                     (qn(r_table), truncate_name(r_name, connection.ops.max_name_length()),

   341                     qn(r_col), qn(table), qn(col),

   342                     connection.ops.deferrable_sql()))

   343             del pending_references[model]

   344     return final_output

   345 

   346 def many_to_many_sql_for_model(model, style):

   347     from django.db import connection, models

   348     from django.contrib.contenttypes import generic

   349     from django.db.backends.util import truncate_name

   350 

   351     opts = model._meta

   352     final_output = []

   353     qn = connection.ops.quote_name

   354     inline_references = connection.features.inline_fk_references

   355     for f in opts.local_many_to_many:

   356         if not isinstance(f.rel, generic.GenericRel):

   357             tablespace = f.db_tablespace or opts.db_tablespace

   358             if tablespace and connection.features.supports_tablespaces and connection.features.autoindexes_primary_keys:

   359                 tablespace_sql = ' ' + connection.ops.tablespace_sql(tablespace, inline=True)

   360             else:

   361                 tablespace_sql = ''

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

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

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

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

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

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

   368                 tablespace_sql))

   369             if inline_references:

   370                 deferred = []

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

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

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

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

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

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

   377                     connection.ops.deferrable_sql()))

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

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

   380                     style.SQL_COLTYPE(models.ForeignKey(f.rel.to).db_type()),

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

   382                     style.SQL_TABLE(qn(f.rel.to._meta.db_table)),

   383                     style.SQL_FIELD(qn(f.rel.to._meta.pk.column)),

   384                     connection.ops.deferrable_sql()))

   385             else:

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

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

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

   389                     style.SQL_KEYWORD('NOT NULL')))

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

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

   392                     style.SQL_COLTYPE(models.ForeignKey(f.rel.to).db_type()),

   393                     style.SQL_KEYWORD('NOT NULL')))

   394                 deferred = [

   395                     (f.m2m_db_table(), f.m2m_column_name(), opts.db_table,

   396                         opts.pk.column),

   397                     ( f.m2m_db_table(), f.m2m_reverse_name(),

   398                         f.rel.to._meta.db_table, f.rel.to._meta.pk.column)

   399                     ]

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

   401                 (style.SQL_KEYWORD('UNIQUE'),

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

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

   404                 tablespace_sql))

   405             table_output.append(')')

   406             if opts.db_tablespace and connection.features.supports_tablespaces:

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

   408                 table_output.append(connection.ops.tablespace_sql(opts.db_tablespace))

   409             table_output.append(';')

   410             final_output.append('\n'.join(table_output))

   411 

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

   413                 r_name = '%s_refs_%s_%x' % (r_col, col,

   414                         abs(hash((r_table, table))))

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

   416                 (qn(r_table),

   417                 truncate_name(r_name, connection.ops.max_name_length()),

   418                 qn(r_col), qn(table), qn(col),

   419                 connection.ops.deferrable_sql()))

   420 

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

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

   423             if autoinc_sql:

   424                 for stmt in autoinc_sql:

   425                     final_output.append(stmt)

   426 

   427     return final_output

   428 

   429 def custom_sql_for_model(model):

   430     from django.db import models

   431     from django.conf import settings

   432 

   433     opts = model._meta

   434     app_dir = os.path.normpath(os.path.join(os.path.dirname(models.get_app(model._meta.app_label).__file__), 'sql'))

   435     output = []

   436 

   437     # Some backends can't execute more than one SQL statement at a time,

   438     # so split into separate statements.

   439     statements = re.compile(r";[ \t]*$", re.M)

   440 

   441     # Find custom SQL, if it's available.

   442     sql_files = [os.path.join(app_dir, "%s.%s.sql" % (opts.object_name.lower(), settings.DATABASE_ENGINE)),

   443                  os.path.join(app_dir, "%s.sql" % opts.object_name.lower())]

   444     for sql_file in sql_files:

   445         if os.path.exists(sql_file):

   446             fp = open(sql_file, 'U')

   447             for statement in statements.split(fp.read().decode(settings.FILE_CHARSET)):

   448                 # Remove any comments from the file

   449                 statement = re.sub(ur"--.*[\n\Z]", "", statement)

   450                 if statement.strip():

   451                     output.append(statement + u";")

   452             fp.close()

   453 

   454     return output

   455 

   456 def sql_indexes_for_model(model, style):

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

   458     from django.db import connection

   459     output = []

   460 

   461     qn = connection.ops.quote_name

   462     for f in model._meta.local_fields:

   463         if f.db_index and not ((f.primary_key or f.unique) and connection.features.autoindexes_primary_keys):

   464             unique = f.unique and 'UNIQUE ' or ''

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

   466             if tablespace and connection.features.supports_tablespaces:

   467                 tablespace_sql = ' ' + connection.ops.tablespace_sql(tablespace)

   468             else:

   469                 tablespace_sql = ''

   470             output.append(

   471                 style.SQL_KEYWORD('CREATE %sINDEX' % unique) + ' ' + \

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

   473                 style.SQL_KEYWORD('ON') + ' ' + \

   474                 style.SQL_TABLE(qn(model._meta.db_table)) + ' ' + \

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

   476                 "%s;" % tablespace_sql

   477             )

   478     return output

   479 

   480 def emit_post_sync_signal(created_models, verbosity, interactive):

   481     from django.db import models

   482     from django.dispatch import dispatcher

   483     # Emit the post_sync signal for every application.

   484     for app in models.get_apps():

   485         app_name = app.__name__.split('.')[-2]

   486         if verbosity >= 2:

   487             print "Running post-sync handlers for application", app_name

   488         dispatcher.send(signal=models.signals.post_syncdb, sender=app,

   489             app=app, created_models=created_models,

   490             verbosity=verbosity, interactive=interactive)