- assurances that context.connection is safe to use by column default functions, helps proposal for [ticket:703]

2 files changed, 30 insertions, 14 deletions

diff --git a/lib/sqlalchemy/engine/base.py b/lib/sqlalchemy/engine/base.py
index 642eeac62..ff2da1165 100644
--- a/lib/sqlalchemy/engine/base.py
+++ b/lib/sqlalchemy/engine/base.py
@@ -275,8 +275,14 @@ class ExecutionContext(object):
     ExecutionContext should have these datamembers:
     
         connection
-            Connection object which initiated the call to the
-            dialect to create this ExecutionContext.
+            Connection object which can be freely used by default value generators
+            to execute SQL.  This Connection should reference the same underlying
+            connection/transactional resources of root_connection.
+            
+        root_connection
+            Connection object which is the source of this ExecutionContext.  This
+            Connection may have close_with_result=True set, in which case it can
+            only be used once.
 
         dialect
             dialect which created this ExecutionContext.
@@ -515,12 +521,13 @@ class Connection(Connectable):
     The Connection object is **not** threadsafe.
     """
 
-    def __init__(self, engine, connection=None, close_with_result=False):
+    def __init__(self, engine, connection=None, close_with_result=False, _branch=False):
         self.__engine = engine
         self.__connection = connection or engine.raw_connection()
         self.__transaction = None
         self.__close_with_result = close_with_result
         self.__savepoint_seq = 0
+        self.__branch = _branch
 
     def _get_connection(self):
         try:
@@ -530,9 +537,14 @@ class Connection(Connectable):
 
     def _branch(self):
         """return a new Connection which references this Connection's 
-        engine and connection; but does not have close_with_result enabled."""
+        engine and connection; but does not have close_with_result enabled,
+        and also whose close() method does nothing.
+        
+        This is used to execute "sub" statements within a single execution,
+        usually an INSERT statement.
+        """
         
-        return Connection(self.__engine, self.__connection)
+        return Connection(self.__engine, self.__connection, _branch=True)
         
     engine = property(lambda s:s.__engine, doc="The Engine with which this Connection is associated.")
     dialect = property(lambda s:s.__engine.dialect, doc="Dialect used by this Connection.")
@@ -686,7 +698,8 @@ class Connection(Connectable):
             c = self.__connection
         except AttributeError:
             return
-        self.__connection.close()
+        if not self.__branch:
+            self.__connection.close()
         self.__connection = None
         del self.__connection
 
@@ -757,7 +770,7 @@ class Connection(Connectable):
         else:
             self.__execute(context)
         self._autocommit(context.statement)
-
+        
     def __execute(self, context):
         if context.parameters is None:
             if context.dialect.positional:
@@ -1124,7 +1137,8 @@ class ResultProxy(object):
             self._rowcount = context.get_rowcount()
             self.close()
             
-    connection = property(lambda self:self.context.connection)
+    connection = property(lambda self:self.context.root_connection)
+    
     def _get_rowcount(self):
         if self._rowcount is not None:
             return self._rowcount
@@ -1510,9 +1524,7 @@ class DefaultRunner(schema.SchemaVisitor):
 
     def __init__(self, context):
         self.context = context
-        # branch the connection so it doesnt close after result
-        self.connection = context.connection._branch()
-        
+        self.connection = self.context._connection._branch()
     dialect = property(lambda self:self.context.dialect)
     
     def get_column_default(self, column):
diff --git a/lib/sqlalchemy/engine/default.py b/lib/sqlalchemy/engine/default.py
index a2e159639..185387177 100644
--- a/lib/sqlalchemy/engine/default.py
+++ b/lib/sqlalchemy/engine/default.py
@@ -145,7 +145,7 @@ class DefaultDialect(base.Dialect):
 class DefaultExecutionContext(base.ExecutionContext):
     def __init__(self, dialect, connection, compiled=None, statement=None, parameters=None):
         self.dialect = dialect
-        self.connection = connection
+        self._connection = connection
         self.compiled = compiled
         self._postfetch_cols = util.Set()
         
@@ -172,11 +172,15 @@ class DefaultExecutionContext(base.ExecutionContext):
             self.statement = self.statement.encode(self.dialect.encoding)
             
         self.cursor = self.create_cursor()
-        
+    
     engine = property(lambda s:s.connection.engine)
     isinsert = property(lambda s:s.compiled and s.compiled.isinsert)
     isupdate = property(lambda s:s.compiled and s.compiled.isupdate)
     
+    connection = property(lambda s:s._connection._branch())
+    
+    root_connection = property(lambda s:s._connection)
+    
     def __encode_param_keys(self, params):
         """apply string encoding to the keys of dictionary-based bind parameters"""
         if self.dialect.positional or self.dialect.supports_unicode_statements():
@@ -218,7 +222,7 @@ class DefaultExecutionContext(base.ExecutionContext):
         return re.match(r'SELECT', self.statement.lstrip(), re.I) is not None
 
     def create_cursor(self):
-        return self.connection.connection.cursor()
+        return self._connection.connection.cursor()
 
     def pre_execution(self):
         self.pre_exec()