From 8a7fc55abc9380773749135a28211d6cd5a78315 Mon Sep 17 00:00:00 2001
From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Mon, 27 Feb 2006 02:33:47 +0000
Subject: [PATCH] got mapper.using() to work, fixed push/pop mapper, custom
 session assignments

---
 lib/sqlalchemy/mapping/mapper.py      | 31 +++++++++++++++++++++++----
 lib/sqlalchemy/mapping/objectstore.py | 13 ++++++-----
 2 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/lib/sqlalchemy/mapping/mapper.py b/lib/sqlalchemy/mapping/mapper.py
index a4de28a2da..a7c56eb7d6 100644
--- a/lib/sqlalchemy/mapping/mapper.py
+++ b/lib/sqlalchemy/mapping/mapper.py
@@ -213,15 +213,22 @@ class Mapper(object):
             def init(self, *args, **kwargs):
                 nohist = kwargs.pop('_mapper_nohistory', False)
                 session = kwargs.pop('_sa_session', objectstore.get_session())
+                if not nohist:
+                    # register new with the correct session, before the object's 
+                    # constructor is called, since further assignments within the
+                    # constructor would otherwise bind it to whatever get_session() is.
+                    session.register_new(self)
                 if oldinit is not None:
                     try:
                         oldinit(self, *args, **kwargs)
                     except TypeError, msg:
                         # re-raise with the offending class name added to help in debugging
                         raise TypeError, '%s.%s' %(self.__class__.__name__, msg)
-                if not nohist:
-                    session.register_new(self)
-            self.class_.__init__ = init
+            # override oldinit, insuring that its not already one of our
+            # own modified inits
+            if oldinit is None or not hasattr(oldinit, '_sa_mapper_init'):
+                init._sa_mapper_init = True
+                self.class_.__init__ = init
         mapper_registry[self.class_] = self
         self.class_.c = self.c
         
@@ -300,7 +307,23 @@ class Mapper(object):
         mapper.__dict__.update(self.__dict__)
         mapper.props = self.props.copy()
         return mapper
-        
+    
+    def using(self, session):
+        """returns a proxying object to this mapper, which will execute methods on the mapper
+        within the context of the given session.  The session is placed as the "current" session
+        via the push_session/pop_session methods in the objectstore module."""
+        mapper = self
+        class Proxy(object):
+            def __getattr__(self, key):
+                def callit(*args, **kwargs):
+                    objectstore.push_session(session)
+                    try:
+                        return getattr(mapper, key)(*args, **kwargs)
+                    finally:
+                        objectstore.pop_session()
+                return callit
+        return Proxy()
+
     def options(self, *options):
         """uses this mapper as a prototype for a new mapper with different behavior.
         *options is a list of options directives, which include eagerload(), lazyload(), and noload()"""
diff --git a/lib/sqlalchemy/mapping/objectstore.py b/lib/sqlalchemy/mapping/objectstore.py
index 03a5523f73..df5f8d6b1d 100644
--- a/lib/sqlalchemy/mapping/objectstore.py
+++ b/lib/sqlalchemy/mapping/objectstore.py
@@ -369,7 +369,8 @@ class UnitOfWork(object):
         self.attributes.commit(obj)
         
     def register_new(self, obj):
-        self.new.append(obj)
+        if not self.new.contains(obj):
+            self.new.append(obj)
         
     def register_dirty(self, obj):
         if not self.dirty.contains(obj):
@@ -1059,17 +1060,15 @@ def get_session(obj=None):
 uow = get_session # deprecated
 
 def push_session(sess):
-    old = _sessions.get(thread.get_ident(), None)
+    old = get_session()
     sess._previous = old
-    _sessions[sess.hash_key] = sess
-    _sessions[thread.get_ident()] = sess
+    session_registry.set(sess)
     
 def pop_session():
-    sess = _sessions[thread.get_ident()]
+    sess = get_session()
     old = sess._previous
     sess._previous = None
-    _sessions[old.hash_key] = old
-    _sessions[thread.get_ident()] = old
+    session_registry.set(old)
     return old
     
 def using_session(sess, func):
-- 
2.47.2