Change the model for multi-categories - update of the upgrade script

1 files changed, 43 insertions, 4 deletions

diff --git a/chimere/scripts/upgrade.py b/chimere/scripts/upgrade.py
index a17007b..0649a52 100755
--- a/chimere/scripts/upgrade.py
+++ b/chimere/scripts/upgrade.py
@@ -13,7 +13,7 @@ from django.db import connection, transaction
 
 cursor = connection.cursor()
 
-from main.models import Area, Marker, Route, Icon
+from main.models import Area, Marker, Route, Icon, SubCategory
 from django.contrib.gis.geos import LineString
 
 # early versions before 0.1: urn field doesn't exist for area
@@ -35,6 +35,10 @@ def slugfy(text, separator):
 QUERY_CHECK_FIELD = """SELECT a.attname AS field FROM pg_class c, pg_attribute a
     WHERE c.relname = '%s' AND a.attnum > 0 AND a.attrelid = c.oid
           AND a.attname='%s';"""
+QUERY_CHECK_TABLE = """SELECT c.relname FROM pg_class c
+WHERE c.relname = '%s';"""
+
+
 query = QUERY_CHECK_FIELD % ('main_area', 'urn')
 cursor.execute(query)
 transaction.commit_unless_managed()
@@ -82,8 +86,7 @@ CREATE TABLE "subcategory_areas" (
 
 # early versions before 0.1: main_tinyurl table doesn't exist
 
-query = """SELECT c.relname FROM pg_class c
-WHERE c.relname = 'main_tinyurl';"""
+query = QUERY_CHECK_TABLE % 'main_tinyurl'
 cursor.execute(query)
 transaction.commit_unless_managed()
 
@@ -154,7 +157,43 @@ for cls, attr in ((Icon, "image"), (Marker, "picture"),
             obj.width = image.width
             obj.height = image.height
             obj.save()
-    print " * height and width of " + table + " corrected"
+        print " * height and width of " + table + " corrected"
+
+# changement from version 1.0 to 1.1: multiple selection of categories
+
+for cls in (Marker, Route):
+    table = cls._meta.db_table[len("main_"):]
+    query = QUERY_CHECK_TABLE % ('main_' + table + '_categories')
+    cursor.execute(query)
+    transaction.commit_unless_managed()
+
+    row = cursor.fetchone()
+    if row:
+        continue
+    query_create = """
+CREATE TABLE "main_%s_categories" (
+  "id" serial NOT NULL PRIMARY KEY,
+  "%s_id" integer NOT NULL REFERENCES "main_%s" ("id") DEFERRABLE INITIALLY DEFERRED,
+  "subcategory_id" integer NOT NULL REFERENCES "main_subcategory" ("id") DEFERRABLE INITIALLY DEFERRED,
+  UNIQUE ("%s_id", "subcategory_id"));
+""" % (table, table, table, table)
+    cursor.execute(query_create)
+    transaction.commit_unless_managed()
+    for obj in cls.objects.all():
+        query = "select subcategory_id from main_%s where id=%d" % (table,
+                                                                    obj.id)
+        cursor.execute(query)
+        transaction.commit_unless_managed()
+
+        row = cursor.fetchone()
+        if row:
+            obj.categories.add(SubCategory.objects.get(id=row[0]))
+            obj.save()
+    query = "ALTER TABLE main_%s DROP COLUMN subcategory_id;" % table
+    cursor.execute(query)
+    transaction.commit_unless_managed()
+    print " * main_%s_categories created" % table
+
 
 # early versions before 0.1: save route with wrong SRID
 # only errors with default SRID is managed adapt the script for your SRID