1 files changed, 33 insertions, 13 deletions
diff --git a/BackEnd/Timeline/Services/DataManager.cs b/BackEnd/Timeline/Services/DataManager.cs
index d447b0d5..b697630c 100644
--- a/BackEnd/Timeline/Services/DataManager.cs
+++ b/BackEnd/Timeline/Services/DataManager.cs
@@ -22,29 +22,30 @@ namespace Timeline.Services
         /// increases its ref count and returns a tag to the entry.
         /// </summary>
         /// <param name="data">The data. Can't be null.</param>
+        /// <param name="saveDatabaseChange">If true save database change. Otherwise it does not save database change.</param>
         /// <returns>The tag of the created entry.</returns>
         /// <exception cref="ArgumentNullException">Thrown when <paramref name="data"/> is null.</exception>
-        public Task<string> RetainEntry(byte[] data);
+        public Task<string> RetainEntry(byte[] data, bool saveDatabaseChange = true);
 
         /// <summary>
         /// Decrease the the ref count of the entry.
         /// Remove it if ref count is zero.
         /// </summary>
         /// <param name="tag">The tag of the entry.</param>
+        /// <param name="saveDatabaseChange">If true save database change. Otherwise it does not save database change.</param>
         /// <exception cref="ArgumentNullException">Thrown when <paramref name="tag"/> is null.</exception>
         /// <remarks>
         /// It's no-op if entry with tag does not exist.
         /// </remarks>
-        public Task FreeEntry(string tag);
+        public Task FreeEntry(string tag, bool saveDatabaseChange = true);
 
         /// <summary>
-        /// Retrieve the entry with given tag.
+        /// Retrieve the entry with given tag. If not exist, returns null.
         /// </summary>
         /// <param name="tag">The tag of the entry.</param>
-        /// <returns>The data of the entry.</returns>
+        /// <returns>The data of the entry. If not exist, returns null.</returns>
         /// <exception cref="ArgumentNullException">Thrown when <paramref name="tag"/> is null.</exception>
-        /// <exception cref="InvalidOperationException">Thrown when entry with given tag does not exist.</exception>
-        public Task<byte[]> GetEntry(string tag);
+        public Task<byte[]?> GetEntry(string tag);
     }
 
     public class DataManager : IDataManager
@@ -58,7 +59,7 @@ namespace Timeline.Services
             _eTagGenerator = eTagGenerator;
         }
 
-        public async Task<string> RetainEntry(byte[] data)
+        public async Task<string> RetainEntry(byte[] data, bool saveDatabaseChange = true)
         {
             if (data == null)
                 throw new ArgumentNullException(nameof(data));
@@ -81,11 +82,14 @@ namespace Timeline.Services
             {
                 entity.Ref += 1;
             }
-            await _database.SaveChangesAsync();
+
+            if (saveDatabaseChange)
+                await _database.SaveChangesAsync();
+
             return tag;
         }
 
-        public async Task FreeEntry(string tag)
+        public async Task FreeEntry(string tag, bool saveDatabaseChange)
         {
             if (tag == null)
                 throw new ArgumentNullException(nameof(tag));
@@ -102,21 +106,37 @@ namespace Timeline.Services
                 {
                     entity.Ref -= 1;
                 }
-                await _database.SaveChangesAsync();
+
+                if (saveDatabaseChange)
+                    await _database.SaveChangesAsync();
             }
         }
 
-        public async Task<byte[]> GetEntry(string tag)
+        public async Task<byte[]?> GetEntry(string tag)
         {
             if (tag == null)
                 throw new ArgumentNullException(nameof(tag));
 
             var entity = await _database.Data.Where(d => d.Tag == tag).Select(d => new { d.Data }).SingleOrDefaultAsync();
 
-            if (entity == null)
-                throw new InvalidOperationException(Resources.Services.DataManager.ExceptionEntryNotExist);
+            if (entity is null)
+                return null;
 
             return entity.Data;
         }
     }
+
+    public static class DataManagerExtensions
+    {
+        /// <summary>
+        /// Try to get an entry and throw <see cref="DatabaseCorruptedException"/> if not exist.
+        /// </summary>
+        public static async Task<byte[]> GetEntryAndCheck(this IDataManager dataManager, string tag, string notExistMessage)
+        {
+            var data = await dataManager.GetEntry(tag);
+            if (data is null)
+                throw new DatabaseCorruptedException($"Can't get data of tag {tag}. {notExistMessage}");
+            return data;
+        }
+    }
 }