NHibernate + MiniProfiler

gistfile1.cs

using System;
using System.Collections.Generic;
using System.Data;
using System.Data.Common;
using System.Data.SqlClient;
using System.Linq;
using System.Runtime.Remoting.Messaging;
using System.Runtime.Remoting.Proxies;
using System.Text;
using System.Web;
using Castle.DynamicProxy;
using MvcMiniProfiler;
using MvcMiniProfiler.Data;
using NHibernate.AdoNet;
using NHibernate.AdoNet.Util;
using NHibernate.Driver;
using NHibernate.Exceptions;
using NHibernate.Util;

namespace Cygnus.Infrastructure
{
    public class ProfiledSql2008ClientDriver : Sql2008ClientDriver, IEmbeddedBatcherFactoryProvider
    {
        ProxyGenerator pg = new ProxyGenerator();
        public override IDbCommand CreateCommand()
        {
            IDbCommand command = base.CreateCommand();
            
            if (MiniProfiler.Current != null)
                command = pg.CreateInterfaceProxyWithTarget(command, new ProfileInterceptor());

            return command;
        }

        System.Type IEmbeddedBatcherFactoryProvider.BatcherFactoryClass
        {
            get { return typeof(ProfiledSqlClientBatchingBatcherFactory); }
        }

    }

    public class ProfiledSqlClientBatchingBatcherFactory : SqlClientBatchingBatcherFactory 
    {
        public override NHibernate.Engine.IBatcher CreateBatcher(ConnectionManager connectionManager, NHibernate.IInterceptor interceptor)
        {
            return new ProfiledSqlClientBatchingBatcher(connectionManager, interceptor);
        }
    }
    public class ProfiledSqlClientBatchingBatcher : AbstractBatcher
    {
        private int batchSize;
        private int totalExpectedRowsAffected;
        private SqlClientSqlCommandSet currentBatch;
        private StringBuilder currentBatchCommandsLog;
        private readonly int defaultTimeout;

        public ProfiledSqlClientBatchingBatcher(ConnectionManager connectionManager, NHibernate.IInterceptor interceptor)
            : base(connectionManager, interceptor)
        {
            batchSize = Factory.Settings.AdoBatchSize;
            defaultTimeout = PropertiesHelper.GetInt32(NHibernate.Cfg.Environment.CommandTimeout, NHibernate.Cfg.Environment.Properties, -1);

            currentBatch = CreateConfiguredBatch();
            //we always create this, because we need to deal with a scenario in which
            //the user change the logging configuration at runtime. Trying to put this
            //behind an if(log.IsDebugEnabled) will cause a null reference exception 
            //at that point.
            currentBatchCommandsLog = new StringBuilder().AppendLine("Batch commands:");
        }

        public override int BatchSize
        {
            get { return batchSize; }
            set { batchSize = value; }
        }

        protected override int CountOfStatementsInCurrentBatch
        {
            get { return currentBatch.CountOfCommands; }
        }

        public override void AddToBatch(IExpectation expectation)
        {
            totalExpectedRowsAffected += expectation.ExpectedRowCount;
            IDbCommand batchUpdate = CurrentCommand;
            Driver.AdjustCommand(batchUpdate);
            string lineWithParameters = null;
            var sqlStatementLogger = Factory.Settings.SqlStatementLogger;
            if (sqlStatementLogger.IsDebugEnabled || log.IsDebugEnabled)
            {
                lineWithParameters = sqlStatementLogger.GetCommandLineWithParameters(batchUpdate);
                var formatStyle = sqlStatementLogger.DetermineActualStyle(FormatStyle.Basic);
                lineWithParameters = formatStyle.Formatter.Format(lineWithParameters);
                currentBatchCommandsLog.Append("command ")
                    .Append(currentBatch.CountOfCommands)
                    .Append(":")
                    .AppendLine(lineWithParameters);
            }
            if (log.IsDebugEnabled)
            {
                log.Debug("Adding to batch:" + lineWithParameters);
            }
            if (batchUpdate is System.Data.SqlClient.SqlCommand)
                currentBatch.Append((System.Data.SqlClient.SqlCommand)batchUpdate);
            else
            {
                var sqlCommand = new System.Data.SqlClient.SqlCommand(
                batchUpdate.CommandText,
                (SqlConnection)batchUpdate.Connection,
                (SqlTransaction)batchUpdate.Transaction);
                foreach (SqlParameter p in batchUpdate.Parameters)
                {
                    sqlCommand.Parameters.Add(
                    new SqlParameter(
                    p.ParameterName,
                    p.SqlDbType,
                    p.Size,
                    p.Direction,
                    p.IsNullable,
                    p.Precision,
                    p.Scale,
                    p.SourceColumn,
                    p.SourceVersion,
                    p.Value));
                }
                currentBatch.Append(sqlCommand);
            }

            if (currentBatch.CountOfCommands >= batchSize)
            {
                ExecuteBatchWithTiming(batchUpdate);
            }
        }

        protected override void DoExecuteBatch(IDbCommand ps)
        {
            log.DebugFormat("Executing batch");
            CheckReaders();
            Prepare(currentBatch.BatchCommand);
            if (Factory.Settings.SqlStatementLogger.IsDebugEnabled)
            {
                Factory.Settings.SqlStatementLogger.LogBatchCommand(currentBatchCommandsLog.ToString());
                currentBatchCommandsLog = new StringBuilder().AppendLine("Batch commands:");
            }

            int rowsAffected;
            try
            {
                rowsAffected = currentBatch.ExecuteNonQuery();
            }
            catch (DbException e)
            {
                throw ADOExceptionHelper.Convert(Factory.SQLExceptionConverter, e, "could not execute batch command.");
            }

            Expectations.VerifyOutcomeBatched(totalExpectedRowsAffected, rowsAffected);

            currentBatch.Dispose();
            totalExpectedRowsAffected = 0;
            currentBatch = CreateConfiguredBatch();
        }

        private SqlClientSqlCommandSet CreateConfiguredBatch()
        {
            var result = new SqlClientSqlCommandSet();
            if (defaultTimeout > 0)
            {
                try
                {
                    result.CommandTimeout = defaultTimeout;
                }
                catch (Exception e)
                {
                    if (log.IsWarnEnabled)
                    {
                        log.Warn(e.ToString());
                    }
                }
            }

            return result;
        }
    }
    public class ProfileInterceptor : IInterceptor
    {
        public void Intercept(IInvocation invocation)
        {
            var profiler = MiniProfiler.Current as IDbProfiler;

            var executeType = GetExecuteType(invocation);
            var cmd = (DbCommand) invocation.InvocationTarget;
            if (executeType != ExecuteType.None)
                profiler.ExecuteStart(cmd, executeType);
            invocation.Proceed();
            object returnValue = invocation.ReturnValue;

            if (executeType == ExecuteType.Reader)
                returnValue = new ProfiledDbDataReader((DbDataReader)returnValue, cmd.Connection, profiler);

            //IMessage returnMessage = new ReturnMessage(returnValue, invocation.Arguments, invocation.Arguments.Count(), methodMessage.LogicalCallContext, methodMessage);

            if (executeType == ExecuteType.Reader)
                profiler.ExecuteFinish(cmd, executeType, (DbDataReader)returnValue);
            else if (executeType != ExecuteType.None)
                profiler.ExecuteFinish(cmd, executeType);

        }
        private static ExecuteType GetExecuteType(IInvocation invocation)
        {
            switch (invocation.Method.Name)
            {
                case "ExecuteNonQuery":
                    return ExecuteType.NonQuery;
                case "ExecuteReader":
                    return ExecuteType.Reader;
                case "ExecuteScalar":
                    return ExecuteType.Scalar;
                default:
                    return ExecuteType.None;
            }
        }

    }
}