From 4ae649eccc1a5be1872919509ce2b1fc53db1839 Mon Sep 17 00:00:00 2001
From: Castix <vincenzo.castiglia@axant.it>
Date: Fri, 18 Mar 2022 13:10:44 +0100
Subject: [PATCH] helper to replace session - from #45

---
 .gitignore                         |  2 ++
 ming/odm/mapper.py                 | 10 ++++++++++
 ming/tests/odm/test_declarative.py | 26 ++++++++++++++++++++++++++
 setup.cfg                          |  3 +++
 4 files changed, 41 insertions(+)

diff --git a/.gitignore b/.gitignore
index 83e8a08..bece568 100644
--- a/.gitignore
+++ b/.gitignore
@@ -23,3 +23,5 @@ parsetab_*
 .tox
 .eggs/
 venv/
+.\#*
+\#*\#
\ No newline at end of file
diff --git a/ming/odm/mapper.py b/ming/odm/mapper.py
index 6e89a22..8ae80fa 100644
--- a/ming/odm/mapper.py
+++ b/ming/odm/mapper.py
@@ -61,6 +61,16 @@ def __init__(self, mapped_class, collection, session, **kwargs):
             raise TypeError('Unknown kwd args: %r' % kwargs)
         self._instrument_class(properties, include_properties, exclude_properties)
 
+    @classmethod
+    def replace_session(cls, session):
+        for _mapper in cls.all_mappers():
+            _mapper.session = session
+            _mapper.mapped_class.query.session = session
+            _mapper.mapped_class.__mongometa__.session = session
+            _mapper._compiled = False
+            _mapper.compile()
+            _mapper.session.ensure_indexes(_mapper.collection)
+
     def __repr__(self):
         return '<Mapper %s:%s>' % (
             self.mapped_class.__name__, self.collection.m.collection_name)
diff --git a/ming/tests/odm/test_declarative.py b/ming/tests/odm/test_declarative.py
index 9cbd9f3..876ae75 100644
--- a/ming/tests/odm/test_declarative.py
+++ b/ming/tests/odm/test_declarative.py
@@ -1,6 +1,7 @@
 import sys
 from collections import defaultdict
 from unittest import TestCase, SkipTest
+from unittest.mock import MagicMock
 
 from ming import schema as S
 from ming import create_datastore
@@ -811,3 +812,28 @@ def test_hook_base(self):
                          [
                              {'_id': doc._id, 'a': doc.a}
                          ])
+
+
+class TestReplacingSession(TestCase):
+
+    def setUp(self):
+        Mapper._mapper_by_classname.clear()
+        self.datastore = create_datastore('mim:///test_db')
+        self.session = ODMSession(bind=self.datastore)
+        class Basic(MappedClass):
+            class __mongometa__:
+                name = 'hook'
+                session = self.session
+            _id = FieldProperty(S.ObjectId)
+            a = FieldProperty(int)
+        Mapper.compile_all()
+        self.Basic = Basic
+        self.session.remove(self.Basic)
+
+    def test_hook_base(self):
+        assert id(self.Basic.query.session) == id(self.session)
+        session2 = MagicMock()
+        new_session = ODMSession(bind=session2)
+        Mapper.replace_session(new_session)
+        assert id(self.Basic.query.session) == id(new_session)
+        assert id(self.session) != id(new_session)
diff --git a/setup.cfg b/setup.cfg
index ebde9bc..1be30d6 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,3 +1,6 @@
+[pylint]
+# disabling protected-access because of mongodb _id property
+disable = protected-access
 
 [nosetests]
 detailed-errors=1