Merge pull request #65 from SeriousM/sql_output_enhancement

Add select statement for OUTPUT parameters.
MiniProfiler · May 29, 2014 · 5f84308 · 5f84308
2 parents d052401 + e703305
commit 5f84308
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 9 deletions.
diff --git a/StackExchange.Profiling.Tests/SqlFormatterTest.cs b/StackExchange.Profiling.Tests/SqlFormatterTest.cs
@@ -244,11 +244,11 @@ public void StoredProcedureCallWithNormalAndReturnParameter()
         }
 
         [Test]
-        public void StoredProcedureCallWithOutputParameter()
+        public void StoredProcedureCallWithOneOutputParameter()
         {
             // arrange
             _commandText = "dbo.SOMEPROCEDURE";
-            const string expectedOutput = "DECLARE @x int = 123;\r\n\r\nEXEC dbo.SOMEPROCEDURE @x = @x OUTPUT;";
+            const string expectedOutput = "DECLARE @x int = 123;\r\n\r\nEXEC dbo.SOMEPROCEDURE @x = @x OUTPUT;\r\nSELECT @x AS x;";
             CreateDbCommand(CommandType.StoredProcedure);
             // note: since the sql-OUTPUT parameters can be read within the procedure, we need to support setting the value
             AddDbParameter<int>("x", 123, ParameterDirection.Output);
@@ -260,12 +260,48 @@ public void StoredProcedureCallWithOutputParameter()
             Assert.AreEqual(expectedOutput, actualOutput);
         }
 
+        [Test]
+        public void StoredProcedureCallWithTwoOutputParameter()
+        {
+            // arrange
+            _commandText = "dbo.SOMEPROCEDURE";
+            const string expectedOutput = "DECLARE @x int = 123,\r\n        @y int = 123;\r\n\r\nEXEC dbo.SOMEPROCEDURE @x = @x OUTPUT, @y = @y OUTPUT;\r\nSELECT @x AS x, @y AS y;";
+            CreateDbCommand(CommandType.StoredProcedure);
+            // note: since the sql-OUTPUT parameters can be read within the procedure, we need to support setting the value
+            AddDbParameter<int>("x", 123, ParameterDirection.Output);
+            AddDbParameter<int>("y", 123, ParameterDirection.Output);
+
+            // act
+            var actualOutput = GenerateOutput();
+
+            // assert
+            Assert.AreEqual(expectedOutput, actualOutput);
+        }
+
+        [Test]
+        public void StoredProcedureCallWithOneOutputParameterAndOneReturnParameter()
+        {
+            // arrange
+            _commandText = "dbo.SOMEPROCEDURE";
+            string expectedOutput = "DECLARE @x int = 123,\r\n        @retval int;\r\n\r\nEXEC @retval = dbo.SOMEPROCEDURE @x = @x OUTPUT;\r\nSELECT @retval AS ReturnValue, @x AS x;";
+            CreateDbCommand(CommandType.StoredProcedure);
+            // note: since the sql-OUTPUT parameters can be read within the procedure, we need to support setting the value
+            AddDbParameter<int>("x", 123, ParameterDirection.Output);
+            AddDbParameter<int>("retval", null, ParameterDirection.ReturnValue);
+
+            // act
+            var actualOutput = GenerateOutput();
+
+            // assert
+            Assert.AreEqual(expectedOutput, actualOutput);
+        }
+
         [Test]
         public void StoredProcedureCallWithInOutputParameter()
         {
             // arrange
             _commandText = "dbo.SOMEPROCEDURE";
-            const string expectedOutput = "DECLARE @x int = 123;\r\n\r\nEXEC dbo.SOMEPROCEDURE @x = @x OUTPUT;";
+            const string expectedOutput = "DECLARE @x int = 123;\r\n\r\nEXEC dbo.SOMEPROCEDURE @x = @x OUTPUT;\r\nSELECT @x AS x;";
             CreateDbCommand(CommandType.StoredProcedure);
             // note: since the sql-OUTPUT parameters can be read within the procedure, we need to support setting the value
             AddDbParameter<int>("x", 123, ParameterDirection.InputOutput);

diff --git a/StackExchange.Profiling/SqlFormatters/SqlServerFormatter.cs b/StackExchange.Profiling/SqlFormatters/SqlServerFormatter.cs
@@ -72,7 +72,7 @@ public virtual string FormatSql(string commandText, List<SqlTimingParameter> par
                     .AppendLine();
             }
 
-            // only treat 'StoredProcedure' differently since 'Text' may contain 'TableDirect' and 'StoredProcedure'
+            // only treat 'StoredProcedure' differently since 'Text' may contain 'TableDirect' or 'StoredProcedure'
             if (command != null && command.CommandType == CommandType.StoredProcedure)
             {
                 GenerateStoreProcedureCall(commandText, parameters, buffer);
@@ -101,13 +101,32 @@ private void GenerateStoreProcedureCall(string commandText, List<SqlTimingParame
             GenerateStoredProcedureParameters(buffer, parameters);
             buffer.Append(";");
 
-            if (returnValueParameter != null)
-            {
-                buffer.AppendLine().Append("SELECT @").Append(returnValueParameter.Name).Append(" AS ReturnValue;");
-            }
+	        GenerateSelectStatement(buffer, parameters);
         }
 
-        private static SqlTimingParameter GetReturnValueParameter(List<SqlTimingParameter> parameters)
+	    private void GenerateSelectStatement(StringBuilder buffer, List<SqlTimingParameter> parameters)
+	    {
+		    if (parameters == null) return;
+
+		    var parametersToSelect = parameters.Where(
+			    x => x.Direction == ParameterDirection.InputOutput.ToString() ||
+			         x.Direction == ParameterDirection.Output.ToString())
+					 .Select(x => x.Name.TrimStart('@'))
+					 .Select(x => "@" + x + " AS " + x)
+					 .ToList();
+
+		    var returnValueParameter = parameters.SingleOrDefault(x => x.Direction == ParameterDirection.ReturnValue.ToString());
+			if (returnValueParameter != null)
+			{
+				parametersToSelect.Insert(0, "@" + returnValueParameter.Name + " AS ReturnValue");
+			}
+
+		    if (!parametersToSelect.Any()) return;
+
+			buffer.AppendLine().Append("SELECT ").Append(string.Join(", ", parametersToSelect)).Append(";");
+	    }
+
+	    private static SqlTimingParameter GetReturnValueParameter(List<SqlTimingParameter> parameters)
         {
             if (parameters == null || !parameters.Any()) return null;
             return parameters.FirstOrDefault(x => x.Direction == ParameterDirection.ReturnValue.ToString());