lib/neo4j/active_base.rb in neo4j-8.0.11 vs lib/neo4j/active_base.rb in neo4j-8.0.12

@@ -13,11 +13,11 @@
       def on_establish_session(&block)
         @establish_session_block = block
       end
 
       def establish_session
-        @establish_session_block.call if @establish_session_block
+        make_session_wrap!(@establish_session_block.call) if @establish_session_block
       end
 
       def current_transaction_or_session
         current_transaction || current_session
       end
@@ -26,11 +26,11 @@
         current_transaction_or_session.query(*args)
       end
 
       # Should support setting session via config options
       def current_session=(session)
-        SessionRegistry.current_session = session
+        SessionRegistry.current_session = make_session_wrap!(session)
       end
 
       def current_adaptor=(adaptor)
         self.current_session = Neo4j::Core::CypherSession.new(adaptor)
       end
@@ -74,9 +74,14 @@
 
       private
 
       def validate_model_schema!
         Neo4j::ModelSchema.validate_model_schema! unless Neo4j::Migrations.currently_running_migrations
+      end
+
+      def make_session_wrap!(session)
+        session.adaptor.instance_variable_get('@options')[:wrap_level] = :proc
+        session
       end
     end
   end
 end