Ticket #3070: django_management_sql.diff

django/core/management.py

@@ -3 +3 @@
 
 import django
 from django.core.exceptions import ImproperlyConfigured
-import os, re, shutil, sys, textwrap
+import os, re, shutil, sys, textwrap, glob
 from optparse import OptionParser
 from django.utils import termcolors
 
@@ -330 +330 @@
 get_sql_reset.help_doc = "Prints the DROP TABLE SQL, then the CREATE TABLE SQL, for the given app name(s)."
 get_sql_reset.args = APP_ARGS
 
-def get_sql_initial_data_for_model(model):
-    from django.db import models
+def flatten (listoflists):
+    return [item for sublist in listoflists for item in sublist]
+
+def get_sql_initial_files_for_model (app_dir, model):
     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'))
+    # Find custom SQL, if it's available.
+    # Part one: using model names as filename.
+    fnames = ["%s.%s.sql" % (opts.object_name.lower(), settings.DATABASE_ENGINE),
+              "%s.sql" % opts.object_name.lower()]
+    sql_files = [os.path.join(app_dir, fname) for fname in fnames]
+    # Part two: using model names as suffix.
+    fnames = flatten([glob.glob(os.path.join(app_dir, "*[._]%s" % fname))
+                        for fname in fnames])
+    sql_files.extend(fnames)
+    return sql_files
+
+
+def get_sql_initial_data_for_models(app_dir, models):
     output = []
+    sql_files = []
 
     # 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 model in models:
+        sql_files.extend(get_sql_initial_files_for_model(app_dir, model))
+    sql_files.sort()
     for sql_file in sql_files:
         if os.path.exists(sql_file):
             fp = open(sql_file, 'U')
@@ -360 +375 @@
 def get_sql_initial_data(app):
     "Returns a list of the initial INSERT 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(get_sql_initial_data_for_model(model))
-
-    return output
+    return get_sql_initial_data_for_models(app_dir, app_models)
 get_sql_initial_data.help_doc = "Prints the initial INSERT SQL statements for the given app name(s)."
 get_sql_initial_data.args = APP_ARGS
 
@@ -483 +493 @@
             if verbosity >= 1:
                 print "Creating table %s" % model._meta.db_table
             for statement in sql:
+                if verbosity >= 2:
+                    print "SQL:", statement
                 cursor.execute(statement)
             table_list.append(model._meta.db_table)
 
@@ -493 +505 @@
                     if verbosity >= 2:
                         print "Creating many-to-many tables for %s.%s model" % (app_name, model._meta.object_name)
                     for statement in sql:
+                        if verbosity >= 2:
+                            print "SQL:", statement
                         cursor.execute(statement)
 
         transaction.commit_unless_managed()
@@ -507 +521 @@
             app=app, created_models=created_models,
             verbosity=verbosity, interactive=interactive)
 
+        app_dir = os.path.normpath(os.path.join(os.path.dirname(app.__file__), 'sql'))
         # Install initial data for the app (but only if this is a model we've
         # just created)
-        for model in models.get_models(app):
-            if model in created_models:
-                initial_sql = get_sql_initial_data_for_model(model)
-                if initial_sql:
-                    if verbosity >= 1:
-                        print "Installing initial data for %s.%s model" % (app_name, model._meta.object_name)
-                    try:
-                        for sql in initial_sql:
-                            cursor.execute(sql)
-                    except Exception, e:
-                        sys.stderr.write("Failed to install initial SQL data for %s.%s model: %s" % \
-                                            (app_name, model._meta.object_name, e))
-                        transaction.rollback_unless_managed()
-                    else:
-                        transaction.commit_unless_managed()
+        
+        newmodels = [m for m in models.get_models(app) if m in created_models]
+        initial_sql = get_sql_initial_data_for_models(app_dir, newmodels)
+        if initial_sql:
+            modelnames = ["%s.%s" % (app_name, m.__name__) for m in newmodels]
+            strmodelnames = ", ".join(modelnames)
+            if verbosity >= 1:
+                print "Installing initial data for models", strmodelnames
+            try:
+                for sql in initial_sql:
+                    if verbosity >= 2:
+                        print "SQL:", sql
+                    cursor.execute(sql)
+            except Exception, e:
+                sys.stderr.write("Failed to install initial SQL data for models %s: %s" % \
+                                 (strmodelnames, e))
+                transaction.rollback_unless_managed()
+            else:
+                transaction.commit_unless_managed()
 
     # Install SQL indicies for all newly created models
     for app in models.get_apps():
@@ -536 +555 @@
                         print "Installing index for %s.%s model" % (app_name, model._meta.object_name)
                     try:
                         for sql in index_sql:
+                            if verbosity >= 2:
+                                print "SQL:", sql
                             cursor.execute(sql)
                     except Exception, e:
                         sys.stderr.write("Failed to install index for %s.%s model: %s" % \
@@ -597 +618 @@
     print '\n'.join(output)
 diffsettings.args = ""
 
-def install(app):
+def install(app, verbosity=1):
     "Executes the equivalent of 'get_sql_all' in the current database."
     from django.db import connection, transaction
 
@@ -613 +634 @@
     try:
         cursor = connection.cursor()
         for sql in sql_list:
+            if verbosity >= 2:
+                print "SQL:", sql
             cursor.execute(sql)
     except Exception, e:
         sys.stderr.write(style.ERROR("""Error: %s couldn't be installed. Possible reasons:
@@ -627 +650 @@
 install.help_doc = "Executes ``sqlall`` for the given app(s) in the current database."
 install.args = APP_ARGS
 
-def reset(app, interactive=True):
+def reset(app, interactive=True, verbosity=1):
     "Executes the equivalent of 'get_sql_reset' in the current database."
     from django.db import connection, transaction
     from django.conf import settings
@@ -654 +677 @@
         try:
             cursor = connection.cursor()
             for sql in sql_list:
+                if verbosity >= 2:
+                    print "SQL:", sql
                 cursor.execute(sql)
         except Exception, e:
             sys.stderr.write(style.ERROR("""Error: %s couldn't be reset. Possible reasons: