From 306bb1767ee1396d7b5ff3df64b956eebc0c40cc Mon Sep 17 00:00:00 2001
From: Tim Cools <tim@soloco.be>
Date: Thu, 10 Mar 2016 22:38:32 +0100
Subject: [PATCH 1/2] Add document loaded/addedForStorage interception methods

---
 .../Services/DirtyTrackingIdentityMapTests.cs |  30 ++--
 .../Services/IdentityMapTests.cs              |  22 +--
 .../Using_DocumentSessionListener_Tests.cs    | 154 +++++++++++++++++-
 src/Marten/DocumentStore.cs                   |   4 +-
 src/Marten/IDocumentSessionListener.cs        |  23 ++-
 .../Services/DirtyTrackingIdentityMap.cs      |  24 ++-
 src/Marten/Services/IdentityMap.cs            |  40 +++--
 7 files changed, 243 insertions(+), 54 deletions(-)

diff --git a/src/Marten.Testing/Services/DirtyTrackingIdentityMapTests.cs b/src/Marten.Testing/Services/DirtyTrackingIdentityMapTests.cs
index 8f14af7c05..3a09bd065a 100644
--- a/src/Marten.Testing/Services/DirtyTrackingIdentityMapTests.cs
+++ b/src/Marten.Testing/Services/DirtyTrackingIdentityMapTests.cs
@@ -18,7 +18,7 @@ public void get_value_on_first_request()
 
             var serializer = new JilSerializer();
 
-            var map = new DirtyTrackingIdentityMap(serializer);
+            var map = new DirtyTrackingIdentityMap(serializer, null);
 
             var target2 = map.Get<Target>(target.Id, serializer.ToJson(target));
 
@@ -34,7 +34,7 @@ public void get_with_concrete_type()
 
             var json = serializer.ToJson(camaro);
 
-            var map = new DirtyTrackingIdentityMap(serializer);
+            var map = new DirtyTrackingIdentityMap(serializer, null);
 
             map.Get<NulloIdentityMapTests.Car>(camaro.Id, typeof(NulloIdentityMapTests.Camaro), json)
                 .ShouldBeOfType<NulloIdentityMapTests.Camaro>()
@@ -50,7 +50,7 @@ public void get_value_on_subsequent_requests()
 
             var serializer = new JilSerializer();
 
-            var map = new DirtyTrackingIdentityMap(serializer);
+            var map = new DirtyTrackingIdentityMap(serializer, null);
 
             var target2 = map.Get<Target>(target.Id, serializer.ToJson(target));
             var target3 = map.Get<Target>(target.Id, serializer.ToJson(target));
@@ -74,7 +74,7 @@ public void get_value_on_first_request_with_lazy_json()
 
             var serializer = new JilSerializer();
 
-            var map = new DirtyTrackingIdentityMap(serializer);
+            var map = new DirtyTrackingIdentityMap(serializer, null);
 
             var json = serializer.ToJson(target);
             var clonedTarget = serializer.FromJson<Target>(json);
@@ -96,7 +96,7 @@ public void get_value_on_subsequent_requests_with_lazy_json()
 
             var serializer = new JilSerializer();
 
-            var map = new DirtyTrackingIdentityMap(serializer);
+            var map = new DirtyTrackingIdentityMap(serializer, null);
 
             var target2 = map.Get<Target>(target.Id, () => new FetchResult<Target>(target, serializer.ToJson(target)));
             var target3 = map.Get<Target>(target.Id, () => new FetchResult<Target>(target, serializer.ToJson(target)));
@@ -123,7 +123,7 @@ public void detect_changes_with_no_changes()
 
             var serializer = new JilSerializer();
 
-            var map = new DirtyTrackingIdentityMap(serializer);
+            var map = new DirtyTrackingIdentityMap(serializer, null);
 
 
             var a1 = map.Get<Target>(a.Id, serializer.ToJson(a));
@@ -147,7 +147,7 @@ public void detect_changes_with_multiple_dirties()
 
             var serializer = new JilSerializer();
 
-            var map = new DirtyTrackingIdentityMap(serializer);
+            var map = new DirtyTrackingIdentityMap(serializer, null);
 
 
             var a1 = map.Get<Target>(a.Id, serializer.ToJson(a));
@@ -176,7 +176,7 @@ public void detect_changes_then_clear_the_changes()
 
             var serializer = new JilSerializer();
 
-            var map = new DirtyTrackingIdentityMap(serializer);
+            var map = new DirtyTrackingIdentityMap(serializer, null);
 
 
             var a1 = map.Get<Target>(a.Id, serializer.ToJson(a));
@@ -206,7 +206,7 @@ public void remove_item()
 
             var serializer = new JilSerializer();
 
-            var map = new DirtyTrackingIdentityMap(serializer);
+            var map = new DirtyTrackingIdentityMap(serializer, null);
 
             var target3 = map.Get<Target>(target.Id, serializer.ToJson(target));
 
@@ -225,7 +225,7 @@ public void store()
             var target = Target.Random();
             var serializer = new JilSerializer();
 
-            var map = new DirtyTrackingIdentityMap(serializer);
+            var map = new DirtyTrackingIdentityMap(serializer, null);
 
             map.Store(target.Id, target);
 
@@ -238,7 +238,7 @@ public void get_with_miss_in_database()
         {
             var serializer = new JilSerializer();
 
-            var map = new DirtyTrackingIdentityMap(serializer);
+            var map = new DirtyTrackingIdentityMap(serializer, null);
             map.Get<Target>(Guid.NewGuid(), () => null).ShouldBeNull();
         }
 
@@ -248,7 +248,7 @@ public void has_positive()
             var target = Target.Random();
             var serializer = new JilSerializer();
 
-            var map = new DirtyTrackingIdentityMap(serializer);
+            var map = new DirtyTrackingIdentityMap(serializer, null);
 
             map.Store(target.Id, target);
 
@@ -261,7 +261,7 @@ public void has_negative()
         {
             var serializer = new JilSerializer();
 
-            var map = new DirtyTrackingIdentityMap(serializer);
+            var map = new DirtyTrackingIdentityMap(serializer, null);
             map.Has<Target>(Guid.NewGuid()).ShouldBeFalse();
         }
 
@@ -271,13 +271,11 @@ public void retrieve()
             var target = Target.Random();
             var serializer = new JilSerializer();
 
-            var map = new DirtyTrackingIdentityMap(serializer);
+            var map = new DirtyTrackingIdentityMap(serializer, null);
 
             map.Store(target.Id, target);
 
             map.Retrieve<Target>(target.Id).ShouldBeTheSameAs(target);
-
         }
-
     }
 }
\ No newline at end of file
diff --git a/src/Marten.Testing/Services/IdentityMapTests.cs b/src/Marten.Testing/Services/IdentityMapTests.cs
index 23931a374f..c832848f59 100644
--- a/src/Marten.Testing/Services/IdentityMapTests.cs
+++ b/src/Marten.Testing/Services/IdentityMapTests.cs
@@ -16,7 +16,7 @@ public void get_value_on_first_request()
 
             var serializer = new JilSerializer();
 
-            var map = new IdentityMap(serializer);
+            var map = new IdentityMap(serializer, null);
 
             var target2 = map.Get<Target>(target.Id, serializer.ToJson(target));
 
@@ -32,7 +32,7 @@ public void get_with_concrete_type()
 
             var json = serializer.ToJson(camaro);
 
-            var map = new IdentityMap(serializer);
+            var map = new IdentityMap(serializer, null);
 
             map.Get<NulloIdentityMapTests.Car>(camaro.Id, typeof(NulloIdentityMapTests.Camaro), json)
                 .ShouldBeOfType<NulloIdentityMapTests.Camaro>()
@@ -48,7 +48,7 @@ public void get_value_on_subsequent_requests()
 
             var serializer = new JilSerializer();
 
-            var map = new IdentityMap(serializer);
+            var map = new IdentityMap(serializer, null);
 
             var target2 = map.Get<Target>(target.Id, serializer.ToJson(target));
             var target3 = map.Get<Target>(target.Id, serializer.ToJson(target));
@@ -74,7 +74,7 @@ public void remove_item()
 
             var serializer = new JilSerializer();
 
-            var map = new IdentityMap(serializer);
+            var map = new IdentityMap(serializer, null);
 
             var target3 = map.Get<Target>(target.Id, serializer.ToJson(target));
 
@@ -94,7 +94,7 @@ public void get_value_on_first_request_with_lazy_json()
 
             var serializer = new JilSerializer();
 
-            var map = new IdentityMap(serializer);
+            var map = new IdentityMap(serializer, null);
 
             var json = serializer.ToJson(target);
             var clonedTarget = serializer.FromJson<Target>(json);
@@ -116,7 +116,7 @@ public void get_value_on_subsequent_requests_with_lazy_json()
 
             var serializer = new JilSerializer();
 
-            var map = new IdentityMap(serializer);
+            var map = new IdentityMap(serializer, null);
 
             var target2 = map.Get<Target>(target.Id, () => new FetchResult<Target>(target, serializer.ToJson(target)));
             var target3 = map.Get<Target>(target.Id, () => new FetchResult<Target>(target, serializer.ToJson(target)));
@@ -139,7 +139,7 @@ public void store()
             var target = Target.Random();
             var serializer = new JilSerializer();
 
-            var map = new IdentityMap(serializer);
+            var map = new IdentityMap(serializer, null);
 
             map.Store(target.Id, target);
 
@@ -152,7 +152,7 @@ public void get_with_miss_in_database()
         {
             var serializer = new JilSerializer();
 
-            var map = new IdentityMap(serializer);
+            var map = new IdentityMap(serializer, null);
             map.Get<Target>(Guid.NewGuid(), () => null).ShouldBeNull();
         }
 
@@ -162,7 +162,7 @@ public void has_positive()
             var target = Target.Random();
             var serializer = new JilSerializer();
 
-            var map = new IdentityMap(serializer);
+            var map = new IdentityMap(serializer, null);
 
             map.Store(target.Id, target);
 
@@ -175,7 +175,7 @@ public void has_negative()
         {
             var serializer = new JilSerializer();
 
-            var map = new IdentityMap(serializer);
+            var map = new IdentityMap(serializer, null);
             map.Has<Target>(Guid.NewGuid()).ShouldBeFalse();
         }
 
@@ -185,7 +185,7 @@ public void retrieve()
             var target = Target.Random();
             var serializer = new JilSerializer();
 
-            var map = new IdentityMap(serializer);
+            var map = new IdentityMap(serializer, null);
 
             map.Store(target.Id, target);
 
diff --git a/src/Marten.Testing/Using_DocumentSessionListener_Tests.cs b/src/Marten.Testing/Using_DocumentSessionListener_Tests.cs
index bb89c651f9..122154c5c0 100644
--- a/src/Marten.Testing/Using_DocumentSessionListener_Tests.cs
+++ b/src/Marten.Testing/Using_DocumentSessionListener_Tests.cs
@@ -1,5 +1,10 @@
-using System.Threading.Tasks;
+using System;
+using System.Collections.Generic;
+using System.Diagnostics;
+using System.Linq;
+using System.Threading.Tasks;
 using Marten.Testing.Documents;
+using Shouldly;
 using Xunit;
 
 namespace Marten.Testing
@@ -65,10 +70,157 @@ public async Task call_listener_events_on_synchronous_session_saves_async()
                 }
             }
         }
+
+        [Fact]
+        public void call_listener_events_on_document_store()
+        {
+            var stub1 = new StubDocumentSessionListener();
+            var stub2 = new StubDocumentSessionListener();
+
+            using (var store = DocumentStore.For(_ =>
+            {
+                _.Connection(ConnectionSource.ConnectionString);
+                _.AutoCreateSchemaObjects = AutoCreate.All;
+
+                _.Listeners.Add(stub1);
+                _.Listeners.Add(stub2);
+            }))
+            {
+                using (var session = store.OpenSession())
+                {
+                    var user1 = new User { Id = Guid.NewGuid() };
+                    var user2 = new User { Id = Guid.NewGuid() };
+
+                    session.Store(user1, user2);
+
+                    stub1.StoredDocuments.ShouldContainKeyAndValue(user1.Id, user1);
+                    stub1.StoredDocuments.ShouldContainKeyAndValue(user2.Id, user2);
+
+                    stub2.StoredDocuments.ShouldContainKeyAndValue(user1.Id, user1);
+                    stub2.StoredDocuments.ShouldContainKeyAndValue(user2.Id, user2);
+                }
+            }
+        }
+
+        [Fact]
+        public void call_listener_events_on_document_store_objects()
+        {
+            var stub1 = new StubDocumentSessionListener();
+            var stub2 = new StubDocumentSessionListener();
+
+            using (var store = DocumentStore.For(_ =>
+            {
+                _.Connection(ConnectionSource.ConnectionString);
+                _.AutoCreateSchemaObjects = AutoCreate.All;
+
+                _.Listeners.Add(stub1);
+                _.Listeners.Add(stub2);
+            }))
+            {
+                using (var session = store.OpenSession())
+                {
+                    var user1 = new User { Id = Guid.NewGuid() };
+                    var user2 = new User { Id = Guid.NewGuid() };
+
+                    session.StoreObjects(new [] { user1, user2 });
+
+                    stub1.StoredDocuments.ShouldContainKeyAndValue(user1.Id, user1);
+                    stub1.StoredDocuments.ShouldContainKeyAndValue(user2.Id, user2);
+
+                    stub2.StoredDocuments.ShouldContainKeyAndValue(user1.Id, user1);
+                    stub2.StoredDocuments.ShouldContainKeyAndValue(user2.Id, user2);
+                }
+            }
+        }
+
+        [Fact]
+        public void call_listener_events_on_document_load()
+        {
+            var stub1 = new StubDocumentSessionListener();
+            var stub2 = new StubDocumentSessionListener();
+
+            using (var store = DocumentStore.For(_ =>
+            {
+                _.Connection(ConnectionSource.ConnectionString);
+                _.AutoCreateSchemaObjects = AutoCreate.All;
+
+                _.Listeners.Add(stub1);
+                _.Listeners.Add(stub2);
+            }))
+            {
+                var user1 = new User { Id = Guid.NewGuid() };
+                var user2 = new User { Id = Guid.NewGuid() };
+
+                using (var session = store.OpenSession())
+                {
+                    session.StoreObjects(new[] { user1, user2 });
+                    session.SaveChanges();
+                }
+
+                using (var session = store.OpenSession())
+                {
+                    var user = session.Load<User>(user1.Id);
+
+                    stub1.LoadedDocuments.ShouldContainKeyAndValue(user1.Id, user);
+                    stub2.LoadedDocuments.ShouldContainKeyAndValue(user1.Id, user);
+                }
+            }
+        }
+
+        [Fact]
+        public void call_listener_events_on_document_query()
+        {
+            var stub1 = new StubDocumentSessionListener();
+            var stub2 = new StubDocumentSessionListener();
+
+            using (var store = DocumentStore.For(_ =>
+            {
+                _.Connection(ConnectionSource.ConnectionString);
+                _.AutoCreateSchemaObjects = AutoCreate.All;
+
+                _.Listeners.Add(stub1);
+                _.Listeners.Add(stub2);
+            }))
+            {
+                var user1 = new User { Id = Guid.NewGuid() };
+                var user2 = new User { Id = Guid.NewGuid() };
+
+                using (var session = store.OpenSession())
+                {
+                    session.StoreObjects(new[] { user1, user2 });
+                    session.SaveChanges();
+                }
+
+                using (var session = store.OpenSession())
+                {
+                    var users = session.Query<User>().ToList();
+
+                    stub1.LoadedDocuments.ShouldContainKeyAndValue(user1.Id, users.FirstOrDefault(where => where.Id == user1.Id));
+                    stub1.LoadedDocuments.ShouldContainKeyAndValue(user2.Id, users.FirstOrDefault(where => where.Id == user2.Id));
+
+                    stub2.LoadedDocuments.ShouldContainKeyAndValue(user1.Id, users.FirstOrDefault(where => where.Id == user1.Id));
+                    stub2.LoadedDocuments.ShouldContainKeyAndValue(user2.Id, users.FirstOrDefault(where => where.Id == user2.Id));
+                }
+            }
+        }
     }
 
     public class StubDocumentSessionListener : DocumentSessionListenerBase
     {
+        public override void DocumentLoaded(object id, object document)
+        {
+            LoadedDocuments.Add(id, document);
+        }
+
+        public IDictionary<object, object> LoadedDocuments { get; } = new Dictionary<object, object>();
+
+        public override void DocumentAddedForStorage(object id, object document)
+        {
+            StoredDocuments.Add(id, document);
+        }
+
+        public IDictionary<object, object> StoredDocuments { get; } = new Dictionary<object, object>();
+
         public override void BeforeSaveChanges(IDocumentSession session)
         {
             SaveChangesSession = session;
diff --git a/src/Marten/DocumentStore.cs b/src/Marten/DocumentStore.cs
index b0d254a290..a9316271ce 100644
--- a/src/Marten/DocumentStore.cs
+++ b/src/Marten/DocumentStore.cs
@@ -226,10 +226,10 @@ private IIdentityMap createMap(DocumentTracking tracking)
                     return new NulloIdentityMap(_serializer);
 
                 case DocumentTracking.IdentityOnly:
-                    return new IdentityMap(_serializer);
+                    return new IdentityMap(_serializer, _options.Listeners);
 
                 case DocumentTracking.DirtyTracking:
-                    return new DirtyTrackingIdentityMap(_serializer);
+                    return new DirtyTrackingIdentityMap(_serializer, _options.Listeners);
 
                 default:
                     throw new ArgumentOutOfRangeException(nameof(tracking));
diff --git a/src/Marten/IDocumentSessionListener.cs b/src/Marten/IDocumentSessionListener.cs
index cb990ca72d..6158105a99 100644
--- a/src/Marten/IDocumentSessionListener.cs
+++ b/src/Marten/IDocumentSessionListener.cs
@@ -22,6 +22,17 @@ public interface IDocumentSessionListener
         void AfterCommit(IDocumentSession session);
 
         Task AfterCommitAsync(IDocumentSession session);
+
+
+        /// <summary>
+        /// After a document is loaded
+        /// </summary>
+        void DocumentLoaded(object id, object document);
+
+        /// <summary>
+        /// After a document is added for storage
+        /// </summary>
+        void DocumentAddedForStorage(object id, object document);
     }
 
     public abstract class DocumentSessionListenerBase : IDocumentSessionListener
@@ -31,7 +42,7 @@ public virtual void BeforeSaveChanges(IDocumentSession session)
             // Nothing
         }
 
-        public async virtual Task BeforeSaveChangesAsync(IDocumentSession session)
+        public virtual async Task BeforeSaveChangesAsync(IDocumentSession session)
         {
             await Task.Run(() => BeforeSaveChanges(session));
         }
@@ -45,5 +56,15 @@ public virtual async Task AfterCommitAsync(IDocumentSession session)
         {
             await Task.Run(() => AfterCommit(session));
         }
+
+        public virtual void DocumentLoaded(object id, object document)
+        {
+            // Nothing
+        }
+
+        public virtual void DocumentAddedForStorage(object id, object document)
+        {
+            // Nothing
+        }
     }
 }
\ No newline at end of file
diff --git a/src/Marten/Services/DirtyTrackingIdentityMap.cs b/src/Marten/Services/DirtyTrackingIdentityMap.cs
index ef67ef5d98..f5dea9a0e3 100644
--- a/src/Marten/Services/DirtyTrackingIdentityMap.cs
+++ b/src/Marten/Services/DirtyTrackingIdentityMap.cs
@@ -11,14 +11,15 @@ namespace Marten.Services
     public class DirtyTrackingIdentityMap : IIdentityMap, IDocumentTracker
     {
         private readonly ISerializer _serializer;
+        private readonly IEnumerable<IDocumentSessionListener> _listeners;
 
         private readonly Cache<Type, ConcurrentDictionary<int, TrackedEntity>> _objects
             = new Cache<Type, ConcurrentDictionary<int, TrackedEntity>>(_ => new ConcurrentDictionary<int, TrackedEntity>());
 
-
-        public DirtyTrackingIdentityMap(ISerializer serializer)
+        public DirtyTrackingIdentityMap(ISerializer serializer, IEnumerable<IDocumentSessionListener> listeners)
         {
             _serializer = serializer;
+            _listeners = listeners?.Any() == true ? listeners : null;
         }
 
         public T Get<T>(object id, Func<FetchResult<T>> result) where T : class
@@ -27,14 +28,14 @@ public T Get<T>(object id, Func<FetchResult<T>> result) where T : class
             {
                 var fetchResult = result();
 
+                _listeners?.Each(listener => listener.DocumentLoaded(id, fetchResult?.Document));
+
                 return new TrackedEntity(id, typeof(T), fetchResult?.Document, fetchResult?.Json, _serializer);
             }).Document as T;
         }
 
         public async Task<T> GetAsync<T>(object id, Func<CancellationToken, Task<FetchResult<T>>> result, CancellationToken token = default(CancellationToken)) where T : class
         {
-
-
             var dict = _objects[typeof(T)];
             var hashCode = id.GetHashCode();
 
@@ -46,6 +47,8 @@ public T Get<T>(object id, Func<FetchResult<T>> result) where T : class
             var fetchResult = await result(token).ConfigureAwait(false);
             if (fetchResult == null) return null;
 
+            _listeners?.Each(listener => listener.DocumentLoaded(id, fetchResult.Document));
+
             dict[hashCode] = new TrackedEntity(id, typeof(T), fetchResult.Document, fetchResult.Json, _serializer);
 
             return fetchResult?.Document;
@@ -53,12 +56,19 @@ public T Get<T>(object id, Func<FetchResult<T>> result) where T : class
 
         public T Get<T>(object id, string json) where T : class
         {
-            return _objects[typeof(T)].GetOrAdd(id.GetHashCode(), _ => new TrackedEntity(id, _serializer, typeof(T), json)).Document.As<T>();
+            return Get<T>(id, typeof (T), json);
         }
 
         public T Get<T>(object id, Type concreteType, string json) where T : class
         {
-            return _objects[typeof(T)].GetOrAdd(id.GetHashCode(), _ => new TrackedEntity(id, _serializer, concreteType, json)).Document.As<T>();
+            return _objects[typeof(T)].GetOrAdd(id.GetHashCode(), _ =>
+            {
+                var trackedEntity = new TrackedEntity(id, _serializer, concreteType, json);
+
+                _listeners?.Each(listener => listener.DocumentLoaded(id, trackedEntity.Document));
+
+                return trackedEntity;
+            }).Document.As<T>();
         }
 
         public void Remove<T>(object id)
@@ -81,6 +91,8 @@ public void Store<T>(object id, T entity)
                 }
             }
 
+            _listeners?.Each(listener => listener.DocumentAddedForStorage(id, entity));
+
             dictionary.AddOrUpdate(hashCode, new TrackedEntity(id, _serializer, typeof(T), entity), (i, e) => e);
         }
 
diff --git a/src/Marten/Services/IdentityMap.cs b/src/Marten/Services/IdentityMap.cs
index 05351053e6..8ead4808a1 100644
--- a/src/Marten/Services/IdentityMap.cs
+++ b/src/Marten/Services/IdentityMap.cs
@@ -1,5 +1,7 @@
 using System;
 using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Baseline;
@@ -9,18 +11,25 @@ namespace Marten.Services
     public class IdentityMap : IIdentityMap
     {
         private readonly ISerializer _serializer;
+        private readonly IEnumerable<IDocumentSessionListener> _listeners;
 
         private readonly Cache<Type, ConcurrentDictionary<object, object>> _objects
             = new Cache<Type, ConcurrentDictionary<object, object>>(_ => new ConcurrentDictionary<object, object>());
 
-        public IdentityMap(ISerializer serializer)
+        public IdentityMap(ISerializer serializer, IEnumerable<IDocumentSessionListener> listeners)
         {
             _serializer = serializer;
+            _listeners = listeners?.Any() == true ? listeners : null;
         }
 
         public T Get<T>(object id, Func<FetchResult<T>> result) where T : class
         {
-            return _objects[typeof(T)].GetOrAdd(id, _ => result()?.Document).As<T>();
+            return _objects[typeof(T)].GetOrAdd(id, _ =>
+            {
+                var document = result()?.Document;
+                _listeners?.Each(listener => listener.DocumentLoaded(id, document));
+                return document;
+            }).As<T>();
         }
 
         public async Task<T> GetAsync<T>(object id, Func<CancellationToken, Task<FetchResult<T>>> result, CancellationToken token = default(CancellationToken)) where T : class
@@ -37,22 +46,27 @@ public T Get<T>(object id, Func<FetchResult<T>> result) where T : class
 
             dict[id] = fetchResult.Document;
 
+            _listeners?.Each(listener => listener.DocumentLoaded(id, fetchResult.Document));
+
             return fetchResult.Document;
         }
 
         public T Get<T>(object id, string json) where T : class
         {
-            return _objects[typeof(T)].GetOrAdd(id, _ =>
-            {
-                return Deserialize<T>(json);
-            }).As<T>();
+            return Get<T>(id, typeof (T), json);
         }
 
         public T Get<T>(object id, Type concreteType, string json) where T : class
         {
             return (T)_objects[typeof(T)].GetOrAdd(id, _ =>
             {
-                return json.IsEmpty() ? null : _serializer.FromJson(concreteType, json);
+                if (json.IsEmpty()) return null;
+
+                var document = _serializer.FromJson(concreteType, json);
+
+                _listeners?.Each(listener => listener.DocumentLoaded(id, document));
+
+                return document;
             });
         }
 
@@ -75,6 +89,8 @@ public void Store<T>(object id, T entity)
                 }
             }
 
+            _listeners?.Each(listener => listener.DocumentAddedForStorage(id, entity));
+
             dictionary.AddOrUpdate(id, entity, (i, e) => e);
         }
 
@@ -89,15 +105,5 @@ public T Retrieve<T>(object id) where T : class
             var dict = _objects[typeof(T)];
             return dict.ContainsKey(id) ? dict[id] as T : null;
         }
-
-        private T Deserialize<T>(string text) where T : class
-        {
-            if (text.IsEmpty())
-            {
-                return null;
-            }
-
-            return _serializer.FromJson<T>(text);
-        }
     }
 }
\ No newline at end of file

From 94206e14e109c61426c243e0094e1267d55c79a4 Mon Sep 17 00:00:00 2001
From: Tim Cools <tim@soloco.be>
Date: Thu, 10 Mar 2016 22:53:58 +0100
Subject: [PATCH 2/2] Add document loaded/addedForStorage interception dirty
 tracking session tests

---
 .../Using_DocumentSessionListener_Tests.cs    | 134 ++++++++++++++++++
 1 file changed, 134 insertions(+)

diff --git a/src/Marten.Testing/Using_DocumentSessionListener_Tests.cs b/src/Marten.Testing/Using_DocumentSessionListener_Tests.cs
index 122154c5c0..8358add9e2 100644
--- a/src/Marten.Testing/Using_DocumentSessionListener_Tests.cs
+++ b/src/Marten.Testing/Using_DocumentSessionListener_Tests.cs
@@ -203,6 +203,140 @@ public void call_listener_events_on_document_query()
                 }
             }
         }
+
+        [Fact]
+        public void call_listener_events_on_document_store_and_dirty_tracking_session()
+        {
+            var stub1 = new StubDocumentSessionListener();
+            var stub2 = new StubDocumentSessionListener();
+
+            using (var store = DocumentStore.For(_ =>
+            {
+                _.Connection(ConnectionSource.ConnectionString);
+                _.AutoCreateSchemaObjects = AutoCreate.All;
+
+                _.Listeners.Add(stub1);
+                _.Listeners.Add(stub2);
+            }))
+            {
+                using (var session = store.DirtyTrackedSession())
+                {
+                    var user1 = new User { Id = Guid.NewGuid() };
+                    var user2 = new User { Id = Guid.NewGuid() };
+
+                    session.Store(user1, user2);
+
+                    stub1.StoredDocuments.ShouldContainKeyAndValue(user1.Id, user1);
+                    stub1.StoredDocuments.ShouldContainKeyAndValue(user2.Id, user2);
+
+                    stub2.StoredDocuments.ShouldContainKeyAndValue(user1.Id, user1);
+                    stub2.StoredDocuments.ShouldContainKeyAndValue(user2.Id, user2);
+                }
+            }
+        }
+
+        [Fact]
+        public void call_listener_events_on_document_store_objects_and_dirty_tracking_session()
+        {
+            var stub1 = new StubDocumentSessionListener();
+            var stub2 = new StubDocumentSessionListener();
+
+            using (var store = DocumentStore.For(_ =>
+            {
+                _.Connection(ConnectionSource.ConnectionString);
+                _.AutoCreateSchemaObjects = AutoCreate.All;
+
+                _.Listeners.Add(stub1);
+                _.Listeners.Add(stub2);
+            }))
+            {
+                using (var session = store.DirtyTrackedSession())
+                {
+                    var user1 = new User { Id = Guid.NewGuid() };
+                    var user2 = new User { Id = Guid.NewGuid() };
+
+                    session.StoreObjects(new[] { user1, user2 });
+
+                    stub1.StoredDocuments.ShouldContainKeyAndValue(user1.Id, user1);
+                    stub1.StoredDocuments.ShouldContainKeyAndValue(user2.Id, user2);
+
+                    stub2.StoredDocuments.ShouldContainKeyAndValue(user1.Id, user1);
+                    stub2.StoredDocuments.ShouldContainKeyAndValue(user2.Id, user2);
+                }
+            }
+        }
+
+        [Fact]
+        public void call_listener_events_on_document_load_and_dirty_tracking_session()
+        {
+            var stub1 = new StubDocumentSessionListener();
+            var stub2 = new StubDocumentSessionListener();
+
+            using (var store = DocumentStore.For(_ =>
+            {
+                _.Connection(ConnectionSource.ConnectionString);
+                _.AutoCreateSchemaObjects = AutoCreate.All;
+
+                _.Listeners.Add(stub1);
+                _.Listeners.Add(stub2);
+            }))
+            {
+                var user1 = new User { Id = Guid.NewGuid() };
+                var user2 = new User { Id = Guid.NewGuid() };
+
+                using (var session = store.OpenSession())
+                {
+                    session.StoreObjects(new[] { user1, user2 });
+                    session.SaveChanges();
+                }
+
+                using (var session = store.DirtyTrackedSession())
+                {
+                    var user = session.Load<User>(user1.Id);
+
+                    stub1.LoadedDocuments.ShouldContainKeyAndValue(user1.Id, user);
+                    stub2.LoadedDocuments.ShouldContainKeyAndValue(user1.Id, user);
+                }
+            }
+        }
+
+        [Fact]
+        public void call_listener_events_on_document_query_and_dirty_tracking_session()
+        {
+            var stub1 = new StubDocumentSessionListener();
+            var stub2 = new StubDocumentSessionListener();
+
+            using (var store = DocumentStore.For(_ =>
+            {
+                _.Connection(ConnectionSource.ConnectionString);
+                _.AutoCreateSchemaObjects = AutoCreate.All;
+
+                _.Listeners.Add(stub1);
+                _.Listeners.Add(stub2);
+            }))
+            {
+                var user1 = new User { Id = Guid.NewGuid() };
+                var user2 = new User { Id = Guid.NewGuid() };
+
+                using (var session = store.OpenSession())
+                {
+                    session.StoreObjects(new[] { user1, user2 });
+                    session.SaveChanges();
+                }
+
+                using (var session = store.DirtyTrackedSession())
+                {
+                    var users = session.Query<User>().ToList();
+
+                    stub1.LoadedDocuments.ShouldContainKeyAndValue(user1.Id, users.FirstOrDefault(where => where.Id == user1.Id));
+                    stub1.LoadedDocuments.ShouldContainKeyAndValue(user2.Id, users.FirstOrDefault(where => where.Id == user2.Id));
+
+                    stub2.LoadedDocuments.ShouldContainKeyAndValue(user1.Id, users.FirstOrDefault(where => where.Id == user1.Id));
+                    stub2.LoadedDocuments.ShouldContainKeyAndValue(user2.Id, users.FirstOrDefault(where => where.Id == user2.Id));
+                }
+            }
+        }
+
     }
 
     public class StubDocumentSessionListener : DocumentSessionListenerBase