Move Attach methods to context

To match EF

Author: John Campion

src/MongoFramework/IMongoDbContext.cs

@@ -19,5 +19,8 @@ public interface IMongoDbContext
 
 		void SaveChanges();
 		Task SaveChangesAsync(CancellationToken cancellationToken = default(CancellationToken));
+		
+		void Attach<TEntity>(TEntity entity) where TEntity : class;
+		void AttachRange<TEntity>(IEnumerable<TEntity> entities) where TEntity : class;
 	}
 }

src/MongoFramework/IMongoDbSet.cs

@@ -23,8 +23,6 @@ public interface IMongoDbSet<TEntity> : IMongoDbSet, IQueryable<TEntity> where T
 		TEntity Create();
 		void Add(TEntity entity);
 		void AddRange(IEnumerable<TEntity> entities);
-		void Attach(TEntity entity);
-		void AttachRange(IEnumerable<TEntity> entities);
 		void Update(TEntity entity);
 		void UpdateRange(IEnumerable<TEntity> entities);
 		void Remove(TEntity entity);

src/MongoFramework/IMongoDbTenantContext.cs

@@ -1,7 +1,11 @@
-namespace MongoFramework
+using System.Collections.Generic;
+
+namespace MongoFramework
 {
-	public interface IMongoDbTenantContext : IMongoDbContext
-	{
+	public interface IMongoDbTenantContext : IMongoDbContext
+	{
 		string TenantId { get; }
+		void CheckEntity(IHaveTenantId entity);
+		void CheckEntities(IEnumerable<IHaveTenantId> entity);
 	}
 }

src/MongoFramework/MongoDbContext.cs

@@ -9,6 +9,7 @@
 using System.Reflection;
 using System.Threading;
 using System.Threading.Tasks;
+using MongoFramework.Utilities;
 
 namespace MongoFramework
 {
@@ -107,6 +108,7 @@ public virtual async Task SaveChangesAsync(CancellationToken cancellationToken =
 			ChangeTracker.CommitChanges();
 			CommandStaging.CommitChanges();
 		}
+		
 		private static async Task InternalSaveChangesAsync<TEntity>(IMongoDbConnection connection, IEnumerable<IWriteCommand> commands, WriteModelOptions options, CancellationToken cancellationToken) where TEntity : class
 		{
 			await EntityIndexWriter.ApplyIndexingAsync<TEntity>(connection);
@@ -131,7 +133,30 @@ public IQueryable<TEntity> Query<TEntity>() where TEntity : class
 			var provider = new MongoFrameworkQueryProvider<TEntity>(Connection);
 			return new MongoFrameworkQueryable<TEntity>(provider);
 		}
+		
+		/// <summary>
+		/// Marks the entity as unchanged in the change tracker and starts tracking.
+		/// </summary>
+		/// <param name="entity"></param>
+		public virtual void Attach<TEntity>(TEntity entity) where TEntity : class
+		{
+			Check.NotNull(entity, nameof(entity));
+			ChangeTracker.SetEntityState(entity, EntityEntryState.NoChanges);
+		}
 
+		/// <summary>
+		/// Marks the collection of entities as unchanged in the change tracker and starts tracking.
+		/// </summary>
+		/// <param name="entities"></param>
+		public virtual void AttachRange<TEntity>(IEnumerable<TEntity> entities) where TEntity : class
+		{
+			Check.NotNull(entities, nameof(entities));
+			foreach (var entity in entities)
+			{
+				ChangeTracker.SetEntityState(entity, EntityEntryState.NoChanges);
+			}
+		}
+		
 		public void Dispose()
 		{
 			Dispose(true);

src/MongoFramework/MongoDbSet.cs

@@ -129,32 +129,8 @@ public virtual void AddRange(IEnumerable<TEntity> entities)
 			{
 				Context.ChangeTracker.SetEntityState(entity, EntityEntryState.Added);
 			}
-		}
-
-		/// <summary>
-		/// Marks the entity as unchanged in the change tracker and starts tracking.
-		/// </summary>
-		/// <param name="entity"></param>
-		public virtual void Attach(TEntity entity)
-		{
-			Check.NotNull(entity, nameof(entity));
-
-			Context.ChangeTracker.SetEntityState(entity, EntityEntryState.NoChanges);
 		}
-		/// <summary>
-		/// Marks the collection of entities as unchanged in the change tracker and starts tracking.
-		/// </summary>
-		/// <param name="entities"></param>
-		public virtual void AttachRange(IEnumerable<TEntity> entities)
-		{
-			Check.NotNull(entities, nameof(entities));
-
-			foreach (var entity in entities)
-			{
-				Context.ChangeTracker.SetEntityState(entity, EntityEntryState.NoChanges);
-			}
-		}
-
+		
 		/// <summary>
 		/// Marks the entity for updating.
 		/// </summary>

src/MongoFramework/MongoDbTenantContext.cs

@@ -1,6 +1,7 @@
-using MongoFramework.Infrastructure.Commands;
-using MongoFramework.Utilities;
-using System;
+using System.Collections.Generic;
+using MongoFramework.Infrastructure;
+using MongoFramework.Infrastructure.Commands;
+using MongoFramework.Utilities;
 
 namespace MongoFramework
 {
@@ -23,5 +24,52 @@ protected override WriteModelOptions GetWriteModelOptions()
 		{
 			return new WriteModelOptions { TenantId = TenantId };
 		}
+		
+		public virtual void CheckEntity(IHaveTenantId entity)
+		{
+			Check.NotNull(entity, nameof(entity));
+			
+			if (entity.TenantId != TenantId)
+			{
+				throw new MultiTenantException($"Entity type {entity.GetType().Name}, tenant ID does not match. Expected: {TenantId}, Entity has: {entity.TenantId}");
+			}
+		}
+
+		public virtual void CheckEntities(IEnumerable<IHaveTenantId> entities)
+		{
+			Check.NotNull(entities, nameof(entities));
+			
+			foreach (var entity in entities)
+			{
+				CheckEntity(entity);
+			}
+		}
+		
+		/// <summary>
+		/// Marks the entity as unchanged in the change tracker and starts tracking.
+		/// </summary>
+		/// <param name="entity"></param>
+		public override void Attach<TEntity>(TEntity entity) where TEntity : class
+		{
+			if (typeof(IHaveTenantId).IsAssignableFrom(typeof(TEntity)))
+			{
+				CheckEntity(entity as IHaveTenantId);
+			}
+			base.Attach(entity);
+		}
+
+		/// <summary>
+		/// Marks the collection of entities as unchanged in the change tracker and starts tracking.
+		/// </summary>
+		/// <param name="entities"></param>
+		public override void AttachRange<TEntity>(IEnumerable<TEntity> entities) where TEntity : class
+		{
+			if (typeof(IHaveTenantId).IsAssignableFrom(typeof(TEntity)))
+			{
+				CheckEntities(entities as IEnumerable<IHaveTenantId>);
+			}
+			base.AttachRange(entities);
+		}
+
 	}
 }

src/MongoFramework/MongoDbTenantSet.cs

@@ -27,28 +27,8 @@ public MongoDbTenantSet(IMongoDbContext context) : base(context)
 		{
 			Context = context as IMongoDbTenantContext ?? throw new ArgumentException("Context provided to a MongoDbTenantSet must be IMongoDbTenantContext",nameof(context));
 		}
-
-		protected virtual void CheckEntity(TEntity entity)
-		{
-			Check.NotNull(entity, nameof(entity));
-			
-			if (entity.TenantId != Context.TenantId)
-			{
-				throw new MultiTenantException($"Entity type {entity.GetType().Name}, tenant ID does not match. Expected: {Context.TenantId}, Entity has: {entity.TenantId}");
-			}
-		}
-
-		protected virtual void CheckEntities(IEnumerable<TEntity> entities)
-		{
-			Check.NotNull(entities, nameof(entities));
-			
-			foreach (var entity in entities)
-			{
-				CheckEntity(entity);
-			}
-		}
-
-				/// <summary>
+		
+		/// <summary>
 		///     Finds an entity with the given primary key value. If an entity with the given primary key value
 		///     is being tracked by the context, then it is returned immediately without making a request to the
 		///     database. Otherwise, a query is made to the database for an entity with the given primary key value
@@ -138,41 +118,29 @@ public override void AddRange(IEnumerable<TEntity> entities)
 				entity.TenantId = Context.TenantId;
 			}
 			base.AddRange(entities);
-		}
-
-		public override void Attach(TEntity entity)
-		{
-			CheckEntity(entity);
-			base.Attach(entity);
-		}
-
-		public override void AttachRange(IEnumerable<TEntity> entities)
-		{
-			CheckEntities(entities);
-			base.AttachRange(entities);
-		}
-
+		}
+		
 		public override void Update(TEntity entity)
 		{
-			CheckEntity(entity);
+			Context.CheckEntity(entity);
 			base.Update(entity);
 		}
 
 		public override void UpdateRange(IEnumerable<TEntity> entities)
 		{
-			CheckEntities(entities);
+			Context.CheckEntities(entities);
 			base.UpdateRange(entities);
 		}
 
 		public override void Remove(TEntity entity)
 		{
-			CheckEntity(entity);
+			Context.CheckEntity(entity);
 			base.Remove(entity);
 		}
 
 		public override void RemoveRange(IEnumerable<TEntity> entities)
 		{
-			CheckEntities(entities);
+			Context.CheckEntities(entities);
 			base.RemoveRange(entities);
 		}
 
@@ -189,23 +157,23 @@ protected override IQueryable<TEntity> GetQueryable(bool trackEntities)
 		{
 			var key = Context.TenantId;
 			var queryable = Context.Query<TEntity>().Where(c => c.TenantId == key);
-			if (trackEntities)
-			{
-				var provider = queryable.Provider as IMongoFrameworkQueryProvider<TEntity>;
-				provider.EntityProcessors.Add(new EntityTrackingProcessor<TEntity>(Context));
+			if (trackEntities)
+			{
+				var provider = queryable.Provider as IMongoFrameworkQueryProvider<TEntity>;
+				provider.EntityProcessors.Add(new EntityTrackingProcessor<TEntity>(Context));
 			}
 			return queryable;
 		}
 
 		public IQueryable<TEntity> GetSearchTextQueryable(string search)
 		{			
 			var key = Context.TenantId;
-			var queryable = Context.Query<TEntity>().WhereFilter(b => b.Text(search)).Where(c => c.TenantId == key);
-			var provider = queryable.Provider as IMongoFrameworkQueryProvider<TEntity>;
-			provider.EntityProcessors.Add(new EntityTrackingProcessor<TEntity>(Context));
+			var queryable = Context.Query<TEntity>().WhereFilter(b => b.Text(search)).Where(c => c.TenantId == key);
+			var provider = queryable.Provider as IMongoFrameworkQueryProvider<TEntity>;
+			provider.EntityProcessors.Add(new EntityTrackingProcessor<TEntity>(Context));
 			return queryable;
-		}
-
+		}
+
 		#endregion
 	}
 }