added tests for partitioned collections

IntegrationTests/CreatePartitionedTests.cs

@@ -0,0 +1,68 @@
+using IntegrationTests.Infrastructure;
+using MongoDbGenericRepository.Models;
+using NUnit.Framework;
+using System.Collections.Generic;
+using System.Threading.Tasks;
+
+namespace IntegrationTests
+{
+    public class CreateTestsPartitionedDocument : PartitionedDocument
+    {
+        public CreateTestsPartitionedDocument() : base("TestPartitionKey")
+        {
+            Version = 1;
+        }
+        public string SomeContent { get; set; }
+    }
+
+    public class CreatePartitionedTests : BaseMongoDbRepositoryTests<CreateTestsPartitionedDocument>
+    {
+        [Test]
+        public void AddOne()
+        {
+            // Arrange
+            var document = new CreateTestsPartitionedDocument();
+            // Act
+            SUT.AddOne(document);
+            // Assert
+            long count = SUT.Count<CreateTestsPartitionedDocument>(e => e.Id == document.Id, PartitionKey);
+            Assert.AreEqual(1, count);
+        }
+
+        [Test]
+        public async Task AddOneAsync()
+        {
+            // Arrange
+            var document = new CreateTestsPartitionedDocument();
+            // Act
+            await SUT.AddOneAsync(document);
+            // Assert
+            long count = SUT.Count<CreateTestsPartitionedDocument>(e => e.Id == document.Id, PartitionKey);
+            Assert.AreEqual(1, count);
+        }
+
+        [Test]
+        public void AddMany()
+        {
+            // Arrange
+            var documents = new List<CreateTestsPartitionedDocument> { new CreateTestsPartitionedDocument(), new CreateTestsPartitionedDocument() };
+            // Act
+            SUT.AddMany(documents);
+            // Assert
+            long count = SUT.Count<CreateTestsPartitionedDocument>(e => e.Id == documents[0].Id || e.Id == documents[1].Id, PartitionKey);
+            Assert.AreEqual(2, count);
+        }
+
+        [Test]
+        public async Task AddManyAsync()
+        {
+            // Arrange
+            var documents = new List<CreateTestsPartitionedDocument> { new CreateTestsPartitionedDocument(), new CreateTestsPartitionedDocument() };
+            // Act
+            await SUT.AddManyAsync(documents);
+            // Assert
+            long count = SUT.Count<CreateTestsPartitionedDocument>(e => e.Id == documents[0].Id || e.Id == documents[1].Id, PartitionKey);
+            Assert.AreEqual(2, count);
+        }
+    }
+}

IntegrationTests/CreateTests.cs

@@ -2,41 +2,31 @@
 using MongoDbGenericRepository.Models;
 using NUnit.Framework;
 using System.Collections.Generic;
-using System.Configuration;
 using System.Threading.Tasks;
 
 namespace IntegrationTests
 {
-    public class CreateTests
+    public class CreateTestsDocument : Document
     {
-        /// <summary>
+        public CreateTestsDocument()
-        /// SUT: System Under Test
-        /// </summary>
-        private static ITestRepository SUT { get; set; }
-
-        [OneTimeSetUp]
-        public void Init()
         {
-            var connectionString = ConfigurationManager.ConnectionStrings["MongoDbTests"].ConnectionString;
+            Version = 2;
-            SUT = new TestRepository(connectionString, "MongoDbTests");
-        }
-
-        [OneTimeTearDown]
-        public void Cleanup()
-        {
-            // We drop the collection at the end of each test session.
-            SUT.DropTestCollection<InsertTestsDocument>();
         }
+        public string SomeContent { get; set; }
+    }
 
+    [TestFixture]
+    public class CreateTests : BaseMongoDbRepositoryTests<CreateTestsDocument>
+    {
         [Test]
         public void AddOne()
         {
             // Arrange
-            var document = new InsertTestsDocument();
+            var document = new CreateTestsDocument();
             // Act
             SUT.AddOne(document);
             // Assert
-            long count = SUT.Count<InsertTestsDocument>(e => e.Id == document.Id);
+            long count = SUT.Count<CreateTestsDocument>(e => e.Id == document.Id);
             Assert.AreEqual(1, count);
         }
 
@@ -44,11 +34,11 @@ namespace IntegrationTests
         public async Task AddOneAsync()
         {
             // Arrange
-            var document = new InsertTestsDocument();
+            var document = new CreateTestsDocument();
             // Act
             await SUT.AddOneAsync(document);
             // Assert
-            long count = SUT.Count<InsertTestsDocument>(e => e.Id == document.Id);
+            long count = SUT.Count<CreateTestsDocument>(e => e.Id == document.Id);
             Assert.AreEqual(1, count);
         }
 
@@ -56,11 +46,11 @@ namespace IntegrationTests
         public void AddMany()
         {
             // Arrange
-            var documents = new List<InsertTestsDocument> { new InsertTestsDocument(), new InsertTestsDocument() };
+            var documents = new List<CreateTestsDocument> { new CreateTestsDocument(), new CreateTestsDocument() };
             // Act
             SUT.AddMany(documents);
             // Assert
-            long count = SUT.Count<InsertTestsDocument>(e => e.Id == documents[0].Id || e.Id == documents[1].Id);
+            long count = SUT.Count<CreateTestsDocument>(e => e.Id == documents[0].Id || e.Id == documents[1].Id);
             Assert.AreEqual(2, count);
         }
 
@@ -68,26 +58,12 @@ namespace IntegrationTests
         public async Task AddManyAsync()
         {
             // Arrange
-            var documents = new List<InsertTestsDocument> { new InsertTestsDocument(), new InsertTestsDocument() };
+            var documents = new List<CreateTestsDocument> { new CreateTestsDocument(), new CreateTestsDocument() };
             // Act
             await SUT.AddManyAsync(documents);
             // Assert
-            long count = SUT.Count<InsertTestsDocument>(e => e.Id == documents[0].Id || e.Id == documents[1].Id);
+            long count = SUT.Count<CreateTestsDocument>(e => e.Id == documents[0].Id || e.Id == documents[1].Id);
             Assert.AreEqual(2, count);
         }
-
-
-        #region Utils
-
-        private class InsertTestsDocument : Document
-        {
-            public InsertTestsDocument()
-            {
-                Version = 2;
-            }
-            public string SomeContent { get; set; }
-        }
-
-        #endregion
     }
 }

IntegrationTests/Infrastructure/BaseMongoDbRepositoryTests.cs

@@ -0,0 +1,51 @@
+using MongoDbGenericRepository.Models;
+using NUnit.Framework;
+using System.Configuration;
+
+namespace IntegrationTests.Infrastructure
+{
+    public class BaseMongoDbRepositoryTests<T> where T : Document, new()
+    {
+        public T GetDocumentInstance()
+        {
+            return new T();
+        }
+
+        public BaseMongoDbRepositoryTests()
+        {
+            var type = GetDocumentInstance();
+            if (type is IPartitionedDocument)
+            {
+                PartitionKey = ((IPartitionedDocument)type).PartitionKey;
+            }
+        }
+
+        public string PartitionKey { get; set; }
+
+        /// <summary>
+        /// SUT: System Under Test
+        /// </summary>
+        protected static ITestRepository SUT { get; set; }
+
+        [OneTimeSetUp]
+        public void Init()
+        {
+            var connectionString = ConfigurationManager.ConnectionStrings["MongoDbTests"].ConnectionString;
+            SUT = new TestRepository(connectionString, "MongoDbTests");
+        }
+
+        [OneTimeTearDown]
+        public void Cleanup()
+        {
+            // We drop the collection at the end of each test session.
+            if (!string.IsNullOrEmpty(PartitionKey))
+            {
+                SUT.DropTestCollection<T>(PartitionKey);
+            }
+            else
+            {
+                SUT.DropTestCollection<T>();
+            }
+        }
+    }
+}

IntegrationTests/Infrastructure/ITestRepository.cs

@@ -5,5 +5,6 @@ namespace IntegrationTests
     public interface ITestRepository : IBaseMongoRepository
     {
         void DropTestCollection<TDocument>();
+        void DropTestCollection<TDocument>(string partitionKey);
     }
 }

IntegrationTests/Infrastructure/TestRepository.cs

@@ -13,5 +13,10 @@ namespace IntegrationTests.Infrastructure
         {
             _mongoDbContext.DropCollection<TDocument>();
         }
+
+        public void DropTestCollection<TDocument>(string partitionKey)
+        {
+            _mongoDbContext.DropCollection<TDocument>(partitionKey);
+        }
     }
 }

IntegrationTests/IntegrationTests.csproj

@@ -56,10 +56,13 @@
     <Reference Include="System.Xml" />
   </ItemGroup>
   <ItemGroup>
+    <Compile Include="Infrastructure\BaseMongoDbRepositoryTests.cs" />
+    <Compile Include="CreatePartitionedTests.cs" />
     <Compile Include="Infrastructure\ITestRepository.cs" />
     <Compile Include="Infrastructure\TestRepository.cs" />
     <Compile Include="CreateTests.cs" />
     <Compile Include="Properties\AssemblyInfo.cs" />
+    <Compile Include="ReadTests.cs" />
     <Compile Include="UpdateTests.cs" />
   </ItemGroup>
   <ItemGroup>

IntegrationTests/ReadTests.cs

@@ -0,0 +1,19 @@
+using MongoDbGenericRepository.Models;
+using NUnit.Framework;
+
+namespace IntegrationTests
+{
+    public class ReadTestsDocument : Document
+    {
+        public ReadTestsDocument()
+        {
+            Version = 2;
+        }
+        public string SomeContent { get; set; }
+    }
+
+    [TestFixture]
+    public class ReadTests
+    {
+    }
+}

IntegrationTests/UpdateTests.cs

@@ -1,10 +1,4 @@
-using System;
+namespace IntegrationTests
-using System.Collections.Generic;
-using System.Linq;
-using System.Text;
-using System.Threading.Tasks;
-
-namespace IntegrationTests
 {
     class UpdateTests
     {

MongoDbGenericRepository/IMongoDbContext.cs

@@ -9,20 +9,25 @@ namespace MongoDbGenericRepository
         /// The private GetCollection method
         /// </summary>
         /// <typeparam name="TDocument"></typeparam>
-        /// <returns></returns>
         IMongoCollection<TDocument> GetCollection<TDocument>();
 
         /// <summary>
-        /// The private GetCollection method
+        /// Returns a collection for a document type that has a partition key.
         /// </summary>
         /// <typeparam name="TDocument"></typeparam>
-        /// <returns></returns>
+        /// <param name="partitionKey">The value of the partition key.</param>
-        IMongoCollection<TDocument> GetCollection<TDocument>(TDocument document) where TDocument : IDocument;
+        IMongoCollection<TDocument> GetCollection<TDocument>(string partitionKey) where TDocument : IDocument;
 
         /// <summary>
         /// Drops a collection, use very carefully.
         /// </summary>
         /// <typeparam name="TDocument"></typeparam>
         void DropCollection<TDocument>();
+
+        /// <summary>
+        /// Drops a collection having a partitionkey, use very carefully.
+        /// </summary>
+        /// <typeparam name="TDocument"></typeparam>
+        void DropCollection<TDocument>(string partitionKey);
     }
 }

MongoDbGenericRepository/Models/Document.cs

@@ -23,13 +23,6 @@ namespace MongoDbGenericRepository.Models
         [BsonId]
         public Guid Id { get; set; }
 
-        /// <summary>
-        /// The name of the property used for partitioning the collection
-        /// This will not be inserted into the collection.
-        /// This partition key will be prepended to the collection name to create a new collection.
-        /// </summary>
-        public string PartitionKey { get; set; }
-
         /// <summary>
         /// The datetime in UTC at which the document was added.
         /// </summary>
@@ -40,4 +33,18 @@ namespace MongoDbGenericRepository.Models
         /// </summary>
         public int Version { get; set; }
     }
+
+    public class PartitionedDocument : Document, IPartitionedDocument
+    {
+        public PartitionedDocument(string partitionKey)
+        {
+            PartitionKey = partitionKey;
+        }
+        /// <summary>
+        /// The name of the property used for partitioning the collection
+        /// This will not be inserted into the collection.
+        /// This partition key will be prepended to the collection name to create a new collection.
+        /// </summary>
+        public string PartitionKey { get; set; }
+    }
 }

MongoDbGenericRepository/Models/IDocument.cs

@@ -2,11 +2,25 @@
 
 namespace MongoDbGenericRepository.Models
 {
+    /// <summary>
+    /// This class represents a basic document that can be stored in MongoDb.
+    /// Your document must implement this class in order for the MongoDbRepository to handle them.
+    /// </summary>
     public interface IDocument
     {
         DateTime AddedAtUtc { get; set; }
         Guid Id { get; set; }
-        string PartitionKey { get; set; }
         int Version { get; set; }
     }
+
+    /// <summary>
+    /// This class represents a document that can be inserted in a collection that can be partitioned.
+    /// The partition key allows for the creation of different collections having the same document schema.
+    /// This can be useful if you are planning to build a Software as a Service (SaaS) Platform, or if you want to reduce indexing.
+    /// You could for example insert Logs in different collections based on the week and year they where created, or their Log category/source.
+    /// </summary>
+    public interface IPartitionedDocument : IDocument
+    {
+        string PartitionKey { get; set; }
+    }
 }

MongoDbGenericRepository/MongoDbContext.cs

@@ -35,13 +35,13 @@ namespace MongoDbGenericRepository
         }
 
         /// <summary>
-        /// The private GetCollection method
+        /// Returns a collection for a document type that has a partition key.
         /// </summary>
         /// <typeparam name="TDocument"></typeparam>
-        /// <returns></returns>
+        /// <param name="partitionKey">The value of the partition key.</param>
-        public IMongoCollection<TDocument> GetCollection<TDocument>(TDocument document) where TDocument : IDocument
+        public IMongoCollection<TDocument> GetCollection<TDocument>(string partitionKey) where TDocument : IDocument
         {
-            return _database.GetCollection<TDocument>(PluralizePartitioned(document));
+            return _database.GetCollection<TDocument>(partitionKey +"-"+ Pluralize<TDocument>());
         }
 
         /// <summary>
@@ -53,14 +53,18 @@ namespace MongoDbGenericRepository
             _database.DropCollection(Pluralize<TDocument>());
         }
 
+        /// <summary>
+        /// Drops a collection having a partitionkey, use very carefully.
+        /// </summary>
+        /// <typeparam name="TDocument"></typeparam>
+        public void DropCollection<TDocument>(string partitionKey)
+        {
+            _database.DropCollection(partitionKey + "-" + Pluralize<TDocument>());
+        }
+
         private string Pluralize<TDocument>()
         {
             return typeof(TDocument).Name.ToLower() + "s";
         }
-
-        private string PluralizePartitioned<TDocument>(TDocument document) where TDocument : IDocument
-        {
-            return document.PartitionKey + typeof(TDocument).Name.ToLower() + "s";
-        }
     }
 }

MongoDbGenericRepository/MongoDbRepository.cs

@@ -10,6 +10,41 @@ namespace MongoDbGenericRepository
 {
     public interface IBaseMongoRepository
     {
+        #region Create
+
+        /// <summary>
+        /// Asynchronously adds a document to the collection.
+        /// </summary>
+        /// <typeparam name="TDocument"></typeparam>
+        /// <param name="document">The document you want to add.</param>
+        Task AddOneAsync<TDocument>(TDocument document) where TDocument : IDocument;
+
+        /// <summary>
+        /// Adds a document to the collection.
+        /// Populates the Id and AddedAtUtc fields if necessary.
+        /// </summary>
+        /// <typeparam name="TDocument"></typeparam>
+        /// <param name="document">The document you want to add.</param>
+        void AddOne<TDocument>(TDocument document) where TDocument : IDocument;
+
+        /// <summary>
+        /// Asynchronously adds a list of documents to the collection.
+        /// Populates the Id and AddedAtUtc fields if necessary.
+        /// </summary>
+        /// <typeparam name="TDocument"></typeparam>
+        /// <param name="document">The document you want to add.</param>
+        Task AddManyAsync<TDocument>(IEnumerable<TDocument> documents) where TDocument : IDocument;
+
+        /// <summary>
+        /// Adds a list of documents to the collection.
+        /// Populates the Id and AddedAtUtc fields if necessary.
+        /// </summary>
+        /// <typeparam name="TDocument"></typeparam>
+        /// <param name="document">The document you want to add.</param>
+        void AddMany<TDocument>(IEnumerable<TDocument> documents) where TDocument : IDocument;
+
+        #endregion
+
         #region Read
 
         /// <summary>
@@ -80,46 +115,16 @@ namespace MongoDbGenericRepository
         /// </summary>
         /// <typeparam name="TDocument"></typeparam>
         /// <param name="filter">A LINQ expression filter.</param>
-        long Count<TDocument>(Expression<Func<TDocument, bool>> filter) where TDocument : IDocument;
+        long Count<TDocument>(Expression<Func<TDocument, bool>> filter, string partitionKey = null) where TDocument : IDocument;
 
         #endregion Get
 
-        #region Create
-
-        /// <summary>
-        /// Asynchronously adds a document to the collection.
-        /// </summary>
-        /// <typeparam name="TDocument"></typeparam>
-        /// <param name="document">The document you want to add.</param>
-        Task AddOneAsync<TDocument>(TDocument document) where TDocument : IDocument;
-
-        /// <summary>
-        /// Adds a document to the collection.
-        /// Populates the Id and AddedAtUtc fields if necessary.
-        /// </summary>
-        /// <typeparam name="TDocument"></typeparam>
-        /// <param name="document">The document you want to add.</param>
-        void AddOne<TDocument>(TDocument document) where TDocument : IDocument;
-
-        /// <summary>
-        /// Asynchronously adds a list of documents to the collection.
-        /// Populates the Id and AddedAtUtc fields if necessary.
-        /// </summary>
-        /// <typeparam name="TDocument"></typeparam>
-        /// <param name="document">The document you want to add.</param>
-        Task AddManyAsync<TDocument>(IEnumerable<TDocument> documents) where TDocument : IDocument;
-
-        /// <summary>
-        /// Adds a list of documents to the collection.
-        /// Populates the Id and AddedAtUtc fields if necessary.
-        /// </summary>
-        /// <typeparam name="TDocument"></typeparam>
-        /// <param name="document">The document you want to add.</param>
-        void AddMany<TDocument>(IEnumerable<TDocument> documents) where TDocument : IDocument;
-
-        #endregion
     }
 
+    /// <summary>
+    /// The base Repository, it is meant to be inherited from by your custom custom MongoRepository implementation.
+    /// Its constructor must be given a connection string and a database name.
+    /// </summary>
     public abstract class BaseMongoRepository : IBaseMongoRepository
     {
         public string ConnectionString { get; set; }
@@ -137,6 +142,68 @@ namespace MongoDbGenericRepository
 
         protected IMongoDbContext _mongoDbContext = null;
 
+        #region Create
+
+        /// <summary>
+        /// Asynchronously adds a document to the collection.
+        /// Populates the Id and AddedAtUtc fields if necessary.
+        /// </summary>
+        /// <typeparam name="TDocument"></typeparam>
+        /// <param name="document">The document you want to add.</param>
+        public async Task AddOneAsync<TDocument>(TDocument document) where TDocument : IDocument
+        {
+            FormatDocument(document);
+            await HandlePartitioned(document).InsertOneAsync(document);
+        }
+
+        /// <summary>
+        /// Adds a document to the collection.
+        /// Populates the Id and AddedAtUtc fields if necessary.
+        /// </summary>
+        /// <typeparam name="TDocument"></typeparam>
+        /// <param name="document">The document you want to add.</param>
+        public void AddOne<TDocument>(TDocument document) where TDocument : IDocument
+        {
+            FormatDocument(document);
+            HandlePartitioned(document).InsertOne(document);
+        }
+
+        /// <summary>
+        /// Asynchronously adds a list of documents to the collection.
+        /// Populates the Id and AddedAtUtc fields if necessary.
+        /// </summary>
+        /// <typeparam name="TDocument"></typeparam>
+        /// <param name="document">The document you want to add.</param>
+        public async Task AddManyAsync<TDocument>(IEnumerable<TDocument> documents) where TDocument : IDocument
+        {
+            if (!documents.Any())
+            {
+                return;
+            }
+            foreach (var doc in documents)
+            {
+                FormatDocument(doc);
+            }
+            await HandlePartitioned(documents.FirstOrDefault()).InsertManyAsync(documents);
+        }
+
+        /// <summary>
+        /// Adds a list of documents to the collection.
+        /// Populates the Id and AddedAtUtc fields if necessary.
+        /// </summary>
+        /// <typeparam name="TDocument"></typeparam>
+        /// <param name="document">The document you want to add.</param>
+        public void AddMany<TDocument>(IEnumerable<TDocument> documents) where TDocument : IDocument
+        {
+            foreach (var document in documents)
+            {
+                FormatDocument(document);
+            }
+            HandlePartitioned(documents.FirstOrDefault()).InsertMany(documents);
+        }
+
+        #endregion Create
+
         #region Read
 
         /// <summary>
@@ -227,6 +294,7 @@ namespace MongoDbGenericRepository
         /// </summary>
         /// <typeparam name="TDocument"></typeparam>
         /// <param name="filter">A LINQ expression filter.</param>
+        /// <param name="partitionKey">An optional partitionKey</param>
         public async Task<long> CountAsync<TDocument>(Expression<Func<TDocument, bool>> filter) where TDocument : IDocument
         {
             return await GetCollection<TDocument>().CountAsync(filter);
@@ -237,9 +305,10 @@ namespace MongoDbGenericRepository
         /// </summary>
         /// <typeparam name="TDocument"></typeparam>
         /// <param name="filter">A LINQ expression filter.</param>
-        public long Count<TDocument>(Expression<Func<TDocument, bool>> filter) where TDocument : IDocument
+        /// <param name="partitionKey">An optional partitionKey</param>
+        public long Count<TDocument>(Expression<Func<TDocument, bool>> filter, string partitionKey = null) where TDocument : IDocument
         {
-            return GetCollection<TDocument>().Find(filter).Count();
+            return HandlePartitioned<TDocument>(partitionKey).Find(filter).Count();
         }
 
         /// <summary>
@@ -272,78 +341,6 @@ namespace MongoDbGenericRepository
 
         #endregion Get
 
-        #region Create
-
-
-        private void FormatDocument<TDocument>(TDocument document) where TDocument : IDocument
-        {
-            if (document.Id == default(Guid))
-            {
-                document.Id = Guid.NewGuid();
-            }
-
-            if (document.AddedAtUtc == default(DateTime))
-            {
-                document.AddedAtUtc = DateTime.UtcNow;
-            }
-        }
-
-        /// <summary>
-        /// Asynchronously adds a document to the collection.
-        /// Populates the Id and AddedAtUtc fields if necessary.
-        /// </summary>
-        /// <typeparam name="TDocument"></typeparam>
-        /// <param name="document">The document you want to add.</param>
-        public async Task AddOneAsync<TDocument>(TDocument document) where TDocument : IDocument
-        {
-            FormatDocument(document);
-            await GetCollection<TDocument>().InsertOneAsync(document);
-        }
-
-        /// <summary>
-        /// Adds a document to the collection.
-        /// Populates the Id and AddedAtUtc fields if necessary.
-        /// </summary>
-        /// <typeparam name="TDocument"></typeparam>
-        /// <param name="document">The document you want to add.</param>
-        public void AddOne<TDocument>(TDocument document) where TDocument : IDocument
-        {
-            FormatDocument(document);
-            GetCollection<TDocument>().InsertOne(document);
-        }
-
-        /// <summary>
-        /// Asynchronously adds a list of documents to the collection.
-        /// Populates the Id and AddedAtUtc fields if necessary.
-        /// </summary>
-        /// <typeparam name="TDocument"></typeparam>
-        /// <param name="document">The document you want to add.</param>
-        public async Task AddManyAsync<TDocument>(IEnumerable<TDocument> documents) where TDocument : IDocument
-        {
-            foreach (var document in documents)
-            {
-                FormatDocument(document);
-            }
-            await GetCollection<TDocument>().InsertManyAsync(documents);
-        }
-
-        /// <summary>
-        /// Adds a list of documents to the collection.
-        /// Populates the Id and AddedAtUtc fields if necessary.
-        /// </summary>
-        /// <typeparam name="TDocument"></typeparam>
-        /// <param name="document">The document you want to add.</param>
-        public void AddMany<TDocument>(IEnumerable<TDocument> documents) where TDocument : IDocument
-        {
-            foreach (var document in documents)
-            {
-                FormatDocument(document);
-            }
-            GetCollection<TDocument>().InsertMany(documents);
-        }
-
-        #endregion Create
-
         #region Delete
 
         /// <summary>
@@ -509,13 +506,14 @@ namespace MongoDbGenericRepository
         #endregion Find And Update
 
         /// <summary>
-        /// The private GetCollection method
+        /// 
         /// </summary>
         /// <typeparam name="TDocument"></typeparam>
+        /// <param name="partitionKey"></param>
         /// <returns></returns>
-        private IMongoCollection<TDocument> GetCollection<TDocument>(TDocument document) where TDocument : IDocument
+        private IMongoCollection<TDocument> GetCollection<TDocument>(string partitionKey) where TDocument : IDocument
         {
-            return _mongoDbContext.GetCollection<TDocument>(document);
+            return _mongoDbContext.GetCollection<TDocument>(partitionKey);
         }
 
         /// <summary>
@@ -527,5 +525,40 @@ namespace MongoDbGenericRepository
         {
             return _mongoDbContext.GetCollection<TDocument>();
         }
+
+        private IMongoCollection<TDocument> HandlePartitioned<TDocument>(TDocument document) where TDocument : IDocument
+        {
+            if (document is IPartitionedDocument)
+            {
+                return GetCollection<TDocument>(((IPartitionedDocument)document).PartitionKey);
+            }
+            return GetCollection<TDocument>();
+        }
+
+        private IMongoCollection<TDocument> HandlePartitioned<TDocument>(string partitionKey) where TDocument : IDocument
+        {
+            if (!string.IsNullOrEmpty(partitionKey))
+            {
+                return GetCollection<TDocument>(partitionKey);
+            }
+            return GetCollection<TDocument>();
+        }
+
+        private void FormatDocument<TDocument>(TDocument document) where TDocument : IDocument
+        {
+            if (document == null)
+            {
+                throw new ArgumentNullException(nameof(document));
+            }
+            if (document.Id == default(Guid))
+            {
+                document.Id = Guid.NewGuid();
+            }
+
+            if (document.AddedAtUtc == default(DateTime))
+            {
+                document.AddedAtUtc = DateTime.UtcNow;
+            }
+        }
     }
 }