Fix InvalidOperationException in DBM propagation (#6233)

## Summary of changes

Now validates whether a `Connection` is `Open` before attempting to
issue SQL commands for DBM propagation.
Now copies over the `Transaction` field if there is one as we'd throw
`InvalidOperationException`
Avoids `Npgsql` `NotSupportedException` for DBM.

## Reason for change

Many, many errors related to not being in the `Transaction` scope

Many, many errors related to `Npgsql`

When using `Npgsql` `DbDataSource` (standard in v7.0)
the `.Connection` property throws a `NotSupportedException` as the
`DbDataSource` type doesn't have "Connections" as you'd normally expect.

The exact type that is getting hit here is `NpgsqlDataSourceCommand`


https://github.com/npgsql/npgsql/blob/e6c166ba51bc1632498c944981e648fa987b9c12/src/Npgsql/NpgsqlDataSourceCommand.cs#L62

https://www.npgsql.org/doc/basic-usage.html

## Implementation details

- Check for connection state
- Copy over `Transaction` to created command

## Test coverage

- Connection state is pretty difficult to replicate efficiently - any
feedback welcome there
- Added a `BeginTransaction` in the tests - there were many errors
before the fix applied, none afterwards.
- Will work on adding a proper sample application for `Npgsql`

## Other details
<!-- Fixes #{issue} -->

https://datadoghq.atlassian.net/browse/AIDM-434

<!-- ⚠️ Note: where possible, please obtain 2 approvals prior to
merging. Unless CODEOWNERS specifies otherwise, for external teams it is
typically best to have one review from a team member, and one review
from apm-dotnet. Trivial changes do not require 2 reviews. -->

tracer/src/Datadog.Trace.Trimming/build/Datadog.Trace.Trimming.xml

@@ -204,13 +204,15 @@
       <type fullname="System.Data.CommandType" />
       <type fullname="System.Data.Common.DbCommand" />
       <type fullname="System.Data.Common.DbConnectionStringBuilder" />
+      <type fullname="System.Data.ConnectionState" />
       <type fullname="System.Data.DbType" />
       <type fullname="System.Data.IDataParameter" />
       <type fullname="System.Data.IDataParameterCollection" />
       <type fullname="System.Data.IDataRecord" />
       <type fullname="System.Data.IDbCommand" />
       <type fullname="System.Data.IDbConnection" />
       <type fullname="System.Data.IDbDataParameter" />
+      <type fullname="System.Data.IDbTransaction" />
    </assembly>
    <assembly fullname="System.Data.SqlClient" />
    <assembly fullname="System.Data.SQLite" />

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/AdoNet/DbScopeFactory.cs

@@ -105,7 +105,7 @@ internal static class DbScopeFactory
                         var traceParentInjectedInComment = DatabaseMonitoringPropagator.PropagateDataViaComment(tracer.Settings.DbmPropagationMode, integrationId, command, tracer.DefaultServiceName, tagsFromConnectionString.DbName, tagsFromConnectionString.OutHost, scope.Span);
 
                         // try context injection only after comment injection, so that if it fails, we still have service level propagation
-                        var traceParentInjectedInContext = DatabaseMonitoringPropagator.PropagateDataViaContext(tracer.Settings.DbmPropagationMode, integrationId, command.Connection, scope.Span);
+                        var traceParentInjectedInContext = DatabaseMonitoringPropagator.PropagateDataViaContext(tracer.Settings.DbmPropagationMode, integrationId, command, scope.Span);
                         if (traceParentInjectedInComment || traceParentInjectedInContext)
                         {
                             tags.DbmTraceInjected = "true";

tracer/src/Datadog.Trace/DatabaseMonitoring/DatabaseMonitoringPropagator.cs

@@ -142,21 +142,37 @@ private static bool StartsWithHint(string commandText)
         /// Currently only working for MSSQL (uses an instruction that is specific to it)
         /// </summary>
         /// <returns>True if the traceparent information was set</returns>
-        internal static bool PropagateDataViaContext(DbmPropagationLevel propagationLevel, IntegrationId integrationId, IDbConnection? connection, Span span)
+        internal static bool PropagateDataViaContext(DbmPropagationLevel propagationLevel, IntegrationId integrationId, IDbCommand command, Span span)
         {
-            if (propagationLevel != DbmPropagationLevel.Full || integrationId != IntegrationId.SqlClient || connection == null)
+            if (propagationLevel != DbmPropagationLevel.Full || integrationId != IntegrationId.SqlClient)
             {
                 return false;
             }
 
+            // NOTE: For Npgsql command.Connection throws NotSupportedException for NpgsqlDataSourceCommand (v7.0+)
+            //       Since the feature isn't available for Npgsql we avoid this due to the integrationId check above
+            if (command.Connection == null)
+            {
+                return false;
+            }
+
+            if (command.Connection.State != ConnectionState.Open)
+            {
+                Log.Debug("PropagateDataViaContext did not have an Open connection, so it could not propagate Span data for DBM. Connection state was {ConnectionState}", command.Connection.State);
+
+                return false;
+            }
+
             var stopwatch = System.Diagnostics.Stopwatch.StartNew();
 
             const byte version = 0; // version can have a maximum value of 15 in the current format
             var sampled = SamplingPriorityValues.IsKeep(span.Context.TraceContext.GetOrMakeSamplingDecision());
             var contextValue = BuildContextValue(version, sampled, span.SpanId, span.TraceId128);
 
-            using (var injectionCommand = connection.CreateCommand())
+            using (var injectionCommand = command.Connection.CreateCommand())
             {
+                // if there is a Transaction we need to copy it or our ExecuteNonQuery will throw
+                injectionCommand.Transaction = command.Transaction;
                 injectionCommand.CommandText = SetContextCommand;
 
                 var parameter = injectionCommand.CreateParameter();

tracer/test/Datadog.Trace.ClrProfiler.IntegrationTests/AdoNet/MicrosoftDataSqliteTests.cs

@@ -8,14 +8,17 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Text.RegularExpressions;
 using System.Threading.Tasks;
 using Datadog.Trace.Configuration;
 using Datadog.Trace.TestHelpers;
+using VerifyXunit;
 using Xunit;
 using Xunit.Abstractions;
 
 namespace Datadog.Trace.ClrProfiler.IntegrationTests.AdoNet
 {
+    [UsesVerify]
     public class MicrosoftDataSqliteTests : TracingIntegrationTest
     {
         public MicrosoftDataSqliteTests(ITestOutputHelper output)
@@ -47,7 +50,7 @@ public async Task SubmitsTraces(string packageVersion, string metadataSchemaVers
                 return;
             }
 #endif
-            const int expectedSpanCount = 91;
+            const int expectedSpanCount = 105;
             const string dbType = "sqlite";
             const string expectedOperationName = dbType + ".query";
 
@@ -64,6 +67,22 @@ public async Task SubmitsTraces(string packageVersion, string metadataSchemaVers
             ValidateIntegrationSpans(spans, metadataSchemaVersion, expectedServiceName: clientSpanServiceName, isExternalSpan);
 
             telemetry.AssertIntegrationEnabled(IntegrationId.Sqlite);
+
+            var settings = VerifyHelper.GetSpanVerifierSettings();
+            settings.AddRegexScrubber(new Regex("Sqlite-Test-[a-zA-Z0-9]{32}"), "System-Data-SqlClient-Test-GUID");
+            settings.AddSimpleScrubber("out.host: localhost", "out.host: sqlserver");
+            settings.AddSimpleScrubber("out.host: (localdb)\\MSSQLLocalDB", "out.host: sqlserver");
+            settings.AddSimpleScrubber("out.host: sqledge_arm64", "out.host: sqlserver");
+            settings.AddSimpleScrubber("peer.service: localhost", "peer.service: sqlserver");
+            settings.AddSimpleScrubber("peer.service: (localdb)\\MSSQLLocalDB", "peer.service: sqlserver");
+            settings.AddSimpleScrubber("peer.service: sqledge_arm64", "peer.service: sqlserver");
+            settings.AddRegexScrubber(new Regex("dd.instrumentation.time_ms: \\d+.\\d+"), "dd.instrumentation.time_ms: 123.456");
+
+            var fileName = nameof(MicrosoftDataSqliteTests);
+
+            await VerifyHelper.VerifySpans(spans, settings)
+                  .DisableRequireUniquePrefix()
+                  .UseFileName($"{fileName}.Schema{metadataSchemaVersion.ToUpper()}");
         }
 
         [SkippableFact]

tracer/test/Datadog.Trace.ClrProfiler.IntegrationTests/AdoNet/SystemDataSqliteTests.cs

@@ -4,14 +4,17 @@
 // </copyright>
 
 using System.Linq;
+using System.Text.RegularExpressions;
 using System.Threading.Tasks;
 using Datadog.Trace.Configuration;
 using Datadog.Trace.TestHelpers;
+using VerifyXunit;
 using Xunit;
 using Xunit.Abstractions;
 
 namespace Datadog.Trace.ClrProfiler.IntegrationTests.AdoNet
 {
+    [UsesVerify]
     public class SystemDataSqliteTests : TracingIntegrationTest
     {
         public SystemDataSqliteTests(ITestOutputHelper output)
@@ -56,7 +59,7 @@ public async Task IntegrationDisabled()
 
         private async Task RunTest(string metadataSchemaVersion)
         {
-            const int expectedSpanCount = 91;
+            const int expectedSpanCount = 105;
             const string dbType = "sqlite";
             const string expectedOperationName = dbType + ".query";
 
@@ -72,6 +75,22 @@ private async Task RunTest(string metadataSchemaVersion)
             Assert.Equal(expectedSpanCount, spans.Count);
             ValidateIntegrationSpans(spans, metadataSchemaVersion, expectedServiceName: clientSpanServiceName, isExternalSpan);
             telemetry.AssertIntegrationEnabled(IntegrationId.Sqlite);
+
+            var settings = VerifyHelper.GetSpanVerifierSettings();
+            settings.AddRegexScrubber(new Regex("SQLite-Test-[a-zA-Z0-9]{32}"), "System-Data-SqlClient-Test-GUID");
+            settings.AddSimpleScrubber("out.host: localhost", "out.host: sqlserver");
+            settings.AddSimpleScrubber("out.host: (localdb)\\MSSQLLocalDB", "out.host: sqlserver");
+            settings.AddSimpleScrubber("out.host: sqledge_arm64", "out.host: sqlserver");
+            settings.AddSimpleScrubber("peer.service: localhost", "peer.service: sqlserver");
+            settings.AddSimpleScrubber("peer.service: (localdb)\\MSSQLLocalDB", "peer.service: sqlserver");
+            settings.AddSimpleScrubber("peer.service: sqledge_arm64", "peer.service: sqlserver");
+            settings.AddRegexScrubber(new Regex("dd.instrumentation.time_ms: \\d+.\\d+"), "dd.instrumentation.time_ms: 123.456");
+
+            var fileName = nameof(SystemDataSqliteTests);
+
+            await VerifyHelper.VerifySpans(spans, settings)
+                  .DisableRequireUniquePrefix()
+                  .UseFileName($"{fileName}.Schema{metadataSchemaVersion.ToUpper()}");
         }
     }
 }

tracer/test/Datadog.Trace.Tests/DatabaseMonitoring/DatabaseMonitoringPropagatorTests.cs

@@ -205,10 +205,12 @@ internal void ExpectedContextSet(string propagationMode, string integration, int
             var commandMock = new Mock<IDbCommand>();
             var parameterMock = new Mock<IDbDataParameter>();
             connectionMock.Setup(c => c.CreateCommand()).Returns(commandMock.Object);
+            connectionMock.SetupGet(c => c.State).Returns(ConnectionState.Open);
             commandMock.SetupSet(c => c.CommandText = It.IsAny<string>())
                        .Callback<string>(value => sql = value);
             commandMock.Setup(c => c.CreateParameter()).Returns(parameterMock.Object);
             commandMock.SetupGet(c => c.Parameters).Returns(Mock.Of<IDataParameterCollection>());
+            commandMock.SetupGet(c => c.Connection).Returns(connectionMock.Object);
             parameterMock.SetupSet(p => p.Value = It.IsAny<byte[]>())
                          .Callback<object>(value => context = (byte[])value);
 
@@ -217,7 +219,7 @@ internal void ExpectedContextSet(string propagationMode, string integration, int
                 var span = tracer.StartSpan("db.query", parent: SpanContext.None, serviceName: "pouet", traceId: new TraceId(Upper: 0xBABEBABEBABEBABE, Lower: 0xCAFECAFECAFECAFE), spanId: 0xBEEFBEEFBEEFBEEF);
                 span.Context.TraceContext.SetSamplingPriority(samplingPriority);
 
-                DatabaseMonitoringPropagator.PropagateDataViaContext(dbmPropagationLevel, integrationId, connectionMock.Object, span);
+                DatabaseMonitoringPropagator.PropagateDataViaContext(dbmPropagationLevel, integrationId, commandMock.Object, span);
 
                 if (shouldInject)
                 {