Skip to content
This repository

HTTPS clone URL

Subversion checkout URL

You can clone with HTTPS or Subversion.

Download ZIP
Browse code

Turn SchemaEditor into a context manager

  • Loading branch information...
commit ce5bd42259bc95d372ab0d65dbae793e6251ea80 1 parent b31eea0
Andrew Godwin authored May 18, 2013
22  django/db/backends/schema.py
... ...
@@ -1,8 +1,10 @@
  1
+import sys
1 2
 import hashlib
2 3
 from django.db.backends.creation import BaseDatabaseCreation
3 4
 from django.db.backends.util import truncate_name
4 5
 from django.utils.log import getLogger
5 6
 from django.db.models.fields.related import ManyToManyField
  7
+from django.db.transaction import atomic
6 8
 
7 9
 logger = getLogger('django.db.backends.schema')
8 10
 
@@ -64,9 +66,7 @@ def start(self):
64 66
         Marks the start of a schema-altering run.
65 67
         """
66 68
         self.deferred_sql = []
67  
-        self.old_autocommit = self.connection.autocommit
68  
-        if self.connection.autocommit:
69  
-            self.connection.set_autocommit(False)
  69
+        atomic(self.connection.alias).__enter__()
70 70
 
71 71
     def commit(self):
72 72
         """
@@ -74,8 +74,7 @@ def commit(self):
74 74
         """
75 75
         for sql in self.deferred_sql:
76 76
             self.execute(sql)
77  
-        self.connection.commit()
78  
-        self.connection.set_autocommit(self.old_autocommit)
  77
+        atomic(self.connection.alias).__exit__(None, None, None)
79 78
 
80 79
     def rollback(self):
81 80
         """
@@ -83,8 +82,17 @@ def rollback(self):
83 82
         """
84 83
         if not self.connection.features.can_rollback_ddl:
85 84
             raise RuntimeError("Cannot rollback schema changes on this backend")
86  
-        self.connection.rollback()
87  
-        self.connection.set_autocommit(self.old_autocommit)
  85
+        atomic(self.connection.alias).__exit__(*sys.exc_info())
  86
+
  87
+    def __enter__(self):
  88
+        self.start()
  89
+        return self
  90
+
  91
+    def __exit__(self, exc_type, exc_value, traceback):
  92
+        if exc_type is None:
  93
+            self.commit()
  94
+        else:
  95
+            self.rollback()
88 96
 
89 97
     # Core utility functions
90 98
 
6  django/db/migrations/recorder.py
@@ -37,10 +37,8 @@ def ensure_schema(self):
37 37
         if self.Migration._meta.db_table in self.connection.introspection.get_table_list(self.connection.cursor()):
38 38
             return
39 39
         # Make the table
40  
-        editor = self.connection.schema_editor()
41  
-        editor.start()
42  
-        editor.create_model(self.Migration)
43  
-        editor.commit()
  40
+        with self.connection.schema_editor() as editor:
  41
+            editor.create_model(self.Migration)
44 42
 
45 43
     def applied_migrations(self):
46 44
         """

0 notes on commit ce5bd42

Please sign in to comment.
Something went wrong with that request. Please try again.