app/django/core/management/sql.py
author Sverre Rabbelier <srabbelier@gmail.com>
Thu, 09 Oct 2008 23:48:20 +0000
changeset 294 1fdaab4a6ef2
parent 54 03e267d67478
child 323 ff1a9aa48cfd
permissions -rw-r--r--
Refactor existing code to use the new access module Instead of ending up with many different ways to do access control, we end up having only one centralized place wher access control is done. Patch by: Sverre Rabbelier Reviewed by: Pawel Solyga, Augie Fackler, Todd Larsen Reviewed at: http://codereviews.googleopensourceprograms.com/1601 Review id: 1601

from django.core.management.base import CommandError
import os
import re

try:
    set
except NameError:
    from sets import Set as set   # Python 2.3 fallback

def table_names():
    "Returns a list of all table names that exist in the database."
    from django.db import connection, get_introspection_module
    cursor = connection.cursor()
    return set(get_introspection_module().get_table_list(cursor))

def django_table_names(only_existing=False):
    """
    Returns a list of all table names that have associated Django models and
    are in INSTALLED_APPS.

    If only_existing is True, the resulting list will only include the tables
    that actually exist in the database.
    """
    from django.db import models
    tables = set()
    for app in models.get_apps():
        for model in models.get_models(app):
            tables.add(model._meta.db_table)
            tables.update([f.m2m_db_table() for f in model._meta.local_many_to_many])
    if only_existing:
        tables = [t for t in tables if t in table_names()]
    return tables

def installed_models(table_list):
    "Returns a set of all models that are installed, given a list of existing table names."
    from django.db import connection, models
    all_models = []
    for app in models.get_apps():
        for model in models.get_models(app):
            all_models.append(model)
    if connection.features.uses_case_insensitive_names:
        converter = lambda x: x.upper()
    else:
        converter = lambda x: x
    return set([m for m in all_models if converter(m._meta.db_table) in map(converter, table_list)])

def sequence_list():
    "Returns a list of information about all DB sequences for all models in all apps."
    from django.db import models

    apps = models.get_apps()
    sequence_list = []

    for app in apps:
        for model in models.get_models(app):
            for f in model._meta.local_fields:
                if isinstance(f, models.AutoField):
                    sequence_list.append({'table': model._meta.db_table, 'column': f.column})
                    break # Only one AutoField is allowed per model, so don't bother continuing.

            for f in model._meta.local_many_to_many:
                sequence_list.append({'table': f.m2m_db_table(), 'column': None})

    return sequence_list

def sql_create(app, style):
    "Returns a list of the CREATE TABLE SQL statements for the given app."
    from django.db import models
    from django.conf import settings

    if settings.DATABASE_ENGINE == 'dummy':
        # This must be the "dummy" database backend, which means the user
        # hasn't set DATABASE_ENGINE.
        raise CommandError("Django doesn't know which syntax to use for your SQL statements,\n" +
            "because you haven't specified the DATABASE_ENGINE setting.\n" +
            "Edit your settings file and change DATABASE_ENGINE to something like 'postgresql' or 'mysql'.")

    # Get installed models, so we generate REFERENCES right.
    # We trim models from the current app so that the sqlreset command does not
    # generate invalid SQL (leaving models out of known_models is harmless, so
    # we can be conservative).
    app_models = models.get_models(app)
    final_output = []
    known_models = set([model for model in installed_models(table_names()) if model not in app_models])
    pending_references = {}

    for model in app_models:
        output, references = sql_model_create(model, style, known_models)
        final_output.extend(output)
        for refto, refs in references.items():
            pending_references.setdefault(refto, []).extend(refs)
            if refto in known_models:
                final_output.extend(sql_for_pending_references(refto, style, pending_references))
        final_output.extend(sql_for_pending_references(model, style, pending_references))
        # Keep track of the fact that we've created the table for this model.
        known_models.add(model)

    # Create the many-to-many join tables.
    for model in app_models:
        final_output.extend(many_to_many_sql_for_model(model, style))

    # Handle references to tables that are from other apps
    # but don't exist physically.
    not_installed_models = set(pending_references.keys())
    if not_installed_models:
        alter_sql = []
        for model in not_installed_models:
            alter_sql.extend(['-- ' + sql for sql in
                sql_for_pending_references(model, style, pending_references)])
        if alter_sql:
            final_output.append('-- The following references should be added but depend on non-existent tables:')
            final_output.extend(alter_sql)

    return final_output

def sql_delete(app, style):
    "Returns a list of the DROP TABLE SQL statements for the given app."
    from django.db import connection, models, get_introspection_module
    from django.db.backends.util import truncate_name
    from django.contrib.contenttypes import generic
    introspection = get_introspection_module()

    # This should work even if a connection isn't available
    try:
        cursor = connection.cursor()
    except:
        cursor = None

    # Figure out which tables already exist
    if cursor:
        table_names = introspection.get_table_list(cursor)
    else:
        table_names = []
    if connection.features.uses_case_insensitive_names:
        table_name_converter = lambda x: x.upper()
    else:
        table_name_converter = lambda x: x

    output = []
    qn = connection.ops.quote_name

    # Output DROP TABLE statements for standard application tables.
    to_delete = set()

    references_to_delete = {}
    app_models = models.get_models(app)
    for model in app_models:
        if cursor and table_name_converter(model._meta.db_table) in table_names:
            # The table exists, so it needs to be dropped
            opts = model._meta
            for f in opts.local_fields:
                if f.rel and f.rel.to not in to_delete:
                    references_to_delete.setdefault(f.rel.to, []).append( (model, f) )

            to_delete.add(model)

    for model in app_models:
        if cursor and table_name_converter(model._meta.db_table) in table_names:
            # Drop the table now
            output.append('%s %s;' % (style.SQL_KEYWORD('DROP TABLE'),
                style.SQL_TABLE(qn(model._meta.db_table))))
            if connection.features.supports_constraints and model in references_to_delete:
                for rel_class, f in references_to_delete[model]:
                    table = rel_class._meta.db_table
                    col = f.column
                    r_table = model._meta.db_table
                    r_col = model._meta.get_field(f.rel.field_name).column
                    r_name = '%s_refs_%s_%x' % (col, r_col, abs(hash((table, r_table))))
                    output.append('%s %s %s %s;' % \
                        (style.SQL_KEYWORD('ALTER TABLE'),
                        style.SQL_TABLE(qn(table)),
                        style.SQL_KEYWORD(connection.ops.drop_foreignkey_sql()),
                        style.SQL_FIELD(truncate_name(r_name, connection.ops.max_name_length()))))
                del references_to_delete[model]
            if model._meta.has_auto_field:
                ds = connection.ops.drop_sequence_sql(model._meta.db_table)
                if ds:
                    output.append(ds)

    # Output DROP TABLE statements for many-to-many tables.
    for model in app_models:
        opts = model._meta
        for f in opts.local_many_to_many:
            if isinstance(f.rel, generic.GenericRel):
                continue
            if cursor and table_name_converter(f.m2m_db_table()) in table_names:
                output.append("%s %s;" % (style.SQL_KEYWORD('DROP TABLE'),
                    style.SQL_TABLE(qn(f.m2m_db_table()))))
                ds = connection.ops.drop_sequence_sql("%s_%s" % (model._meta.db_table, f.column))
                if ds:
                    output.append(ds)

    app_label = app_models[0]._meta.app_label

    # Close database connection explicitly, in case this output is being piped
    # directly into a database client, to avoid locking issues.
    if cursor:
        cursor.close()
        connection.close()

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

def sql_reset(app, style):
    "Returns a list of the DROP TABLE SQL, then the CREATE TABLE SQL, for the given module."
    return sql_delete(app, style) + sql_all(app, style)

def sql_flush(style, only_django=False):
    """
    Returns a list of the SQL statements used to flush the database.
    
    If only_django is True, then only table names that have associated Django
    models and are in INSTALLED_APPS will be included.
    """
    from django.db import connection
    if only_django:
        tables = django_table_names()
    else:
        tables = table_names()
    statements = connection.ops.sql_flush(style, tables, sequence_list())
    return statements

def sql_custom(app):
    "Returns a list of the custom table modifying SQL statements for the given app."
    from django.db.models import get_models
    output = []

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

    for model in app_models:
        output.extend(custom_sql_for_model(model))

    return output

def sql_indexes(app, style):
    "Returns a list of the CREATE INDEX SQL statements for all models in the given app."
    from django.db import models
    output = []
    for model in models.get_models(app):
        output.extend(sql_indexes_for_model(model, style))
    return output

def sql_all(app, style):
    "Returns a list of CREATE TABLE SQL, initial-data inserts, and CREATE INDEX SQL for the given module."
    return sql_create(app, style) + sql_custom(app) + sql_indexes(app, style)

def sql_model_create(model, style, known_models=set()):
    """
    Returns the SQL required to create a single model, as a tuple of:
        (list_of_sql, pending_references_dict)
    """
    from django.db import connection, models

    opts = model._meta
    final_output = []
    table_output = []
    pending_references = {}
    qn = connection.ops.quote_name
    inline_references = connection.features.inline_fk_references
    for f in opts.local_fields:
        col_type = f.db_type()
        tablespace = f.db_tablespace or opts.db_tablespace
        if col_type is None:
            # Skip ManyToManyFields, because they're not represented as
            # database columns in this table.
            continue
        # Make the definition (e.g. 'foo VARCHAR(30)') for this field.
        field_output = [style.SQL_FIELD(qn(f.column)),
            style.SQL_COLTYPE(col_type)]
        field_output.append(style.SQL_KEYWORD('%sNULL' % (not f.null and 'NOT ' or '')))
        if f.unique and (not f.primary_key or connection.features.allows_unique_and_pk):
            field_output.append(style.SQL_KEYWORD('UNIQUE'))
        if f.primary_key:
            field_output.append(style.SQL_KEYWORD('PRIMARY KEY'))
        if tablespace and connection.features.supports_tablespaces and (f.unique or f.primary_key) and connection.features.autoindexes_primary_keys:
            # We must specify the index tablespace inline, because we
            # won't be generating a CREATE INDEX statement for this field.
            field_output.append(connection.ops.tablespace_sql(tablespace, inline=True))
        if f.rel:
            if inline_references and f.rel.to in known_models:
                field_output.append(style.SQL_KEYWORD('REFERENCES') + ' ' + \
                    style.SQL_TABLE(qn(f.rel.to._meta.db_table)) + ' (' + \
                    style.SQL_FIELD(qn(f.rel.to._meta.get_field(f.rel.field_name).column)) + ')' +
                    connection.ops.deferrable_sql()
                )
            else:
                # We haven't yet created the table to which this field
                # is related, so save it for later.
                pr = pending_references.setdefault(f.rel.to, []).append((model, f))
        table_output.append(' '.join(field_output))
    if opts.order_with_respect_to:
        table_output.append(style.SQL_FIELD(qn('_order')) + ' ' + \
            style.SQL_COLTYPE(models.IntegerField().db_type()) + ' ' + \
            style.SQL_KEYWORD('NULL'))
    for field_constraints in opts.unique_together:
        table_output.append(style.SQL_KEYWORD('UNIQUE') + ' (%s)' % \
            ", ".join([style.SQL_FIELD(qn(opts.get_field(f).column)) for f in field_constraints]))

    full_statement = [style.SQL_KEYWORD('CREATE TABLE') + ' ' + style.SQL_TABLE(qn(opts.db_table)) + ' (']
    for i, line in enumerate(table_output): # Combine and add commas.
        full_statement.append('    %s%s' % (line, i < len(table_output)-1 and ',' or ''))
    full_statement.append(')')
    if opts.db_tablespace and connection.features.supports_tablespaces:
        full_statement.append(connection.ops.tablespace_sql(opts.db_tablespace))
    full_statement.append(';')
    final_output.append('\n'.join(full_statement))

    if opts.has_auto_field:
        # Add any extra SQL needed to support auto-incrementing primary keys.
        auto_column = opts.auto_field.db_column or opts.auto_field.name
        autoinc_sql = connection.ops.autoinc_sql(opts.db_table, auto_column)
        if autoinc_sql:
            for stmt in autoinc_sql:
                final_output.append(stmt)

    return final_output, pending_references

def sql_for_pending_references(model, style, pending_references):
    """
    Returns any ALTER TABLE statements to add constraints after the fact.
    """
    from django.db import connection
    from django.db.backends.util import truncate_name

    qn = connection.ops.quote_name
    final_output = []
    if connection.features.supports_constraints:
        opts = model._meta
        if model in pending_references:
            for rel_class, f in pending_references[model]:
                rel_opts = rel_class._meta
                r_table = rel_opts.db_table
                r_col = f.column
                table = opts.db_table
                col = opts.get_field(f.rel.field_name).column
                # For MySQL, r_name must be unique in the first 64 characters.
                # So we are careful with character usage here.
                r_name = '%s_refs_%s_%x' % (r_col, col, abs(hash((r_table, table))))
                final_output.append(style.SQL_KEYWORD('ALTER TABLE') + ' %s ADD CONSTRAINT %s FOREIGN KEY (%s) REFERENCES %s (%s)%s;' % \
                    (qn(r_table), truncate_name(r_name, connection.ops.max_name_length()),
                    qn(r_col), qn(table), qn(col),
                    connection.ops.deferrable_sql()))
            del pending_references[model]
    return final_output

def many_to_many_sql_for_model(model, style):
    from django.db import connection, models
    from django.contrib.contenttypes import generic
    from django.db.backends.util import truncate_name

    opts = model._meta
    final_output = []
    qn = connection.ops.quote_name
    inline_references = connection.features.inline_fk_references
    for f in opts.local_many_to_many:
        if not isinstance(f.rel, generic.GenericRel):
            tablespace = f.db_tablespace or opts.db_tablespace
            if tablespace and connection.features.supports_tablespaces and connection.features.autoindexes_primary_keys:
                tablespace_sql = ' ' + connection.ops.tablespace_sql(tablespace, inline=True)
            else:
                tablespace_sql = ''
            table_output = [style.SQL_KEYWORD('CREATE TABLE') + ' ' + \
                style.SQL_TABLE(qn(f.m2m_db_table())) + ' (']
            table_output.append('    %s %s %s%s,' %
                (style.SQL_FIELD(qn('id')),
                style.SQL_COLTYPE(models.AutoField(primary_key=True).db_type()),
                style.SQL_KEYWORD('NOT NULL PRIMARY KEY'),
                tablespace_sql))
            if inline_references:
                deferred = []
                table_output.append('    %s %s %s %s (%s)%s,' %
                    (style.SQL_FIELD(qn(f.m2m_column_name())),
                    style.SQL_COLTYPE(models.ForeignKey(model).db_type()),
                    style.SQL_KEYWORD('NOT NULL REFERENCES'),
                    style.SQL_TABLE(qn(opts.db_table)),
                    style.SQL_FIELD(qn(opts.pk.column)),
                    connection.ops.deferrable_sql()))
                table_output.append('    %s %s %s %s (%s)%s,' %
                    (style.SQL_FIELD(qn(f.m2m_reverse_name())),
                    style.SQL_COLTYPE(models.ForeignKey(f.rel.to).db_type()),
                    style.SQL_KEYWORD('NOT NULL REFERENCES'),
                    style.SQL_TABLE(qn(f.rel.to._meta.db_table)),
                    style.SQL_FIELD(qn(f.rel.to._meta.pk.column)),
                    connection.ops.deferrable_sql()))
            else:
                table_output.append('    %s %s %s,' %
                    (style.SQL_FIELD(qn(f.m2m_column_name())),
                    style.SQL_COLTYPE(models.ForeignKey(model).db_type()),
                    style.SQL_KEYWORD('NOT NULL')))
                table_output.append('    %s %s %s,' %
                    (style.SQL_FIELD(qn(f.m2m_reverse_name())),
                    style.SQL_COLTYPE(models.ForeignKey(f.rel.to).db_type()),
                    style.SQL_KEYWORD('NOT NULL')))
                deferred = [
                    (f.m2m_db_table(), f.m2m_column_name(), opts.db_table,
                        opts.pk.column),
                    ( f.m2m_db_table(), f.m2m_reverse_name(),
                        f.rel.to._meta.db_table, f.rel.to._meta.pk.column)
                    ]
            table_output.append('    %s (%s, %s)%s' %
                (style.SQL_KEYWORD('UNIQUE'),
                style.SQL_FIELD(qn(f.m2m_column_name())),
                style.SQL_FIELD(qn(f.m2m_reverse_name())),
                tablespace_sql))
            table_output.append(')')
            if opts.db_tablespace and connection.features.supports_tablespaces:
                # f.db_tablespace is only for indices, so ignore its value here.
                table_output.append(connection.ops.tablespace_sql(opts.db_tablespace))
            table_output.append(';')
            final_output.append('\n'.join(table_output))

            for r_table, r_col, table, col in deferred:
                r_name = '%s_refs_%s_%x' % (r_col, col,
                        abs(hash((r_table, table))))
                final_output.append(style.SQL_KEYWORD('ALTER TABLE') + ' %s ADD CONSTRAINT %s FOREIGN KEY (%s) REFERENCES %s (%s)%s;' % 
                (qn(r_table),
                truncate_name(r_name, connection.ops.max_name_length()),
                qn(r_col), qn(table), qn(col),
                connection.ops.deferrable_sql()))

            # Add any extra SQL needed to support auto-incrementing PKs
            autoinc_sql = connection.ops.autoinc_sql(f.m2m_db_table(), 'id')
            if autoinc_sql:
                for stmt in autoinc_sql:
                    final_output.append(stmt)

    return final_output

def custom_sql_for_model(model):
    from django.db import models
    from django.conf import settings

    opts = model._meta
    app_dir = os.path.normpath(os.path.join(os.path.dirname(models.get_app(model._meta.app_label).__file__), 'sql'))
    output = []

    # Some backends can't execute more than one SQL statement at a time,
    # so split into separate statements.
    statements = re.compile(r";[ \t]*$", re.M)

    # Find custom SQL, if it's available.
    sql_files = [os.path.join(app_dir, "%s.%s.sql" % (opts.object_name.lower(), settings.DATABASE_ENGINE)),
                 os.path.join(app_dir, "%s.sql" % opts.object_name.lower())]
    for sql_file in sql_files:
        if os.path.exists(sql_file):
            fp = open(sql_file, 'U')
            for statement in statements.split(fp.read().decode(settings.FILE_CHARSET)):
                # Remove any comments from the file
                statement = re.sub(ur"--.*[\n\Z]", "", statement)
                if statement.strip():
                    output.append(statement + u";")
            fp.close()

    return output

def sql_indexes_for_model(model, style):
    "Returns the CREATE INDEX SQL statements for a single model"
    from django.db import connection
    output = []

    qn = connection.ops.quote_name
    for f in model._meta.local_fields:
        if f.db_index and not ((f.primary_key or f.unique) and connection.features.autoindexes_primary_keys):
            unique = f.unique and 'UNIQUE ' or ''
            tablespace = f.db_tablespace or model._meta.db_tablespace
            if tablespace and connection.features.supports_tablespaces:
                tablespace_sql = ' ' + connection.ops.tablespace_sql(tablespace)
            else:
                tablespace_sql = ''
            output.append(
                style.SQL_KEYWORD('CREATE %sINDEX' % unique) + ' ' + \
                style.SQL_TABLE(qn('%s_%s' % (model._meta.db_table, f.column))) + ' ' + \
                style.SQL_KEYWORD('ON') + ' ' + \
                style.SQL_TABLE(qn(model._meta.db_table)) + ' ' + \
                "(%s)" % style.SQL_FIELD(qn(f.column)) + \
                "%s;" % tablespace_sql
            )
    return output

def emit_post_sync_signal(created_models, verbosity, interactive):
    from django.db import models
    from django.dispatch import dispatcher
    # Emit the post_sync signal for every application.
    for app in models.get_apps():
        app_name = app.__name__.split('.')[-2]
        if verbosity >= 2:
            print "Running post-sync handlers for application", app_name
        dispatcher.send(signal=models.signals.post_syncdb, sender=app,
            app=app, created_models=created_models,
            verbosity=verbosity, interactive=interactive)