From 3ab0cb4f36d3f5fa07d4003b930917e533c7cba2 Mon Sep 17 00:00:00 2001
From: Marty T <120425148+tippmar-nr@users.noreply.github.com>
Date: Fri, 23 Aug 2024 15:40:06 -0500
Subject: [PATCH] fix: Preserve custom query name even if statement can't be
 parsed. (#2708) (#2709)

---
 .../Parsing/SqlParser.cs                      |  9 ++---
 .../Parsing/SqlParserTests.cs                 | 35 +++++++++++++++++++
 2 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/src/Agent/NewRelic/Agent/Extensions/NewRelic.Agent.Extensions/Parsing/SqlParser.cs b/src/Agent/NewRelic/Agent/Extensions/NewRelic.Agent.Extensions/Parsing/SqlParser.cs
index 4dd2c7eaa..83d290f08 100644
--- a/src/Agent/NewRelic/Agent/Extensions/NewRelic.Agent.Extensions/Parsing/SqlParser.cs
+++ b/src/Agent/NewRelic/Agent/Extensions/NewRelic.Agent.Extensions/Parsing/SqlParser.cs
@@ -239,7 +239,8 @@ private static ParseStatement CreateCompoundStatementParser(params ParseStatemen
                     }
                 }
 
-                return _nullParsedStatementStore.GetOrAdd(datastoreVendor, x => new ParsedSqlStatement(datastoreVendor, null, null));
+                // return a null statement with the model name if found, otherwise cache and return a null statement with no model name
+                return !string.IsNullOrEmpty(explicitModel) ? new ParsedSqlStatement(datastoreVendor, explicitModel, null) : _nullParsedStatementStore.GetOrAdd(datastoreVendor, x => new ParsedSqlStatement(datastoreVendor, null, null));
             };
         }
 
@@ -295,7 +296,7 @@ public virtual ParsedSqlStatement ParseStatement(DatastoreVendor vendor, Command
                         model = "ParseError";
                     }
                 }
-                
+
                 return CreateParsedDatabaseStatement(vendor, GetFullModel(model, explicitModel));
             }
 
@@ -323,7 +324,7 @@ public ParsedSqlStatement ParseStatement(DatastoreVendor vendor, CommandType com
                     string model = FromMatcher.Match(statement).Success
                         ? "(subquery)"
                         : "VARIABLE";
-                    
+
                     return new ParsedSqlStatement(vendor, GetFullModel(model, explicitModel), "select");
                 }
                 return null;
@@ -414,7 +415,7 @@ public static bool FixParameterizedSql(IDbCommand command)
                     Log.Debug("Not executing explain plan since DbType is Object.");
                     return false;
                 }
-                
+
                 // Parameter names can be supplied with the prefix @ or without
                 // if is supplied, remove the @ to the beginning of the param name
                 // This is to deal with regex issues, see below
diff --git a/tests/Agent/UnitTests/NewRelic.Agent.Extensions.Tests/Parsing/SqlParserTests.cs b/tests/Agent/UnitTests/NewRelic.Agent.Extensions.Tests/Parsing/SqlParserTests.cs
index ba5ea14fe..7bc8afeae 100644
--- a/tests/Agent/UnitTests/NewRelic.Agent.Extensions.Tests/Parsing/SqlParserTests.cs
+++ b/tests/Agent/UnitTests/NewRelic.Agent.Extensions.Tests/Parsing/SqlParserTests.cs
@@ -254,6 +254,41 @@ public void SqlParserTest_PullNameFromComment_ReplaceSlashes()
                 Assert.That(parsedDatabaseStatement.ToString(), Is.EqualTo("dude - [dudeservice|getalldudes]/select"));
             });
         }
+
+        [Test]
+        public void SqlParserTest_PullNameFromComment_IfStatementCannotBeParsed()
+        {
+            var parsedDatabaseStatement = SqlParser.GetParsedDatabaseStatement(DatastoreVendor.MSSQL, CommandType.Text,
+                """
+                /* NewRelicQueryName: MyCustomQueryName */
+                
+                -- This is a full-line comment
+                SELECT
+                    (
+                        SELECT COUNT(*)
+                        FROM
+                            SomeTable
+                        WHERE
+                            SomeColumn = @SomeValue
+                    ) +
+                    (
+                        SELECT COUNT(*)
+                        FROM
+                            SomeOtherTable
+                        WHERE
+                            SomeOtherColumn = @SomeOtherValue
+                    );
+                """);
+
+            Assert.That(parsedDatabaseStatement, Is.Not.Null);
+
+            Assert.Multiple(() =>
+            {
+                Assert.That(parsedDatabaseStatement.Model, Is.EqualTo("MyCustomQueryName"));
+
+                Assert.That(parsedDatabaseStatement.ToString(), Is.EqualTo("MyCustomQueryName/other"));
+            });
+        }
                 
 
         [Test]