diff --git a/src/Core/Resolvers/MsSqlQueryExecutor.cs b/src/Core/Resolvers/MsSqlQueryExecutor.cs index 368e5d6b00..8052d8b787 100644 --- a/src/Core/Resolvers/MsSqlQueryExecutor.cs +++ b/src/Core/Resolvers/MsSqlQueryExecutor.cs @@ -3,6 +3,7 @@ using System.Data; using System.Data.Common; +using System.Diagnostics; using System.Net; using System.Security.Claims; using System.Text; @@ -520,6 +521,49 @@ public override string GetSessionParamsQuery(HttpContext? httpContext, IDictiona sessionMapQuery = sessionMapQuery.Append(statementToSetReadOnlyParam); } + // Add OpenTelemetry correlation values for observability. + // These allow correlating database queries with distributed traces. + Activity? currentActivity = Activity.Current; + if (currentActivity is not null) + { + string traceIdParamName = $"{SESSION_PARAM_NAME}{counter.Next()}"; + parameters.Add(traceIdParamName, new(currentActivity.TraceId.ToString())); + sessionMapQuery.Append($"EXEC sp_set_session_context 'dab.trace_id', {traceIdParamName}, @read_only = 0;"); + + string spanIdParamName = $"{SESSION_PARAM_NAME}{counter.Next()}"; + parameters.Add(spanIdParamName, new(currentActivity.SpanId.ToString())); + sessionMapQuery.Append($"EXEC sp_set_session_context 'dab.span_id', {spanIdParamName}, @read_only = 0;"); + } + + // Add OBO-specific observability values when user-delegated auth is enabled. + // These values are for observability/auditing only and MUST NOT be used for authorization decisions. + // For row-level security, use database policies with the user's actual claims. + if (_dataSourceUserDelegatedAuth.ContainsKey(dataSourceName)) + { + // Set auth type indicator for OBO requests + string authTypeParamName = $"{SESSION_PARAM_NAME}{counter.Next()}"; + parameters.Add(authTypeParamName, new("obo")); + sessionMapQuery.Append($"EXEC sp_set_session_context 'dab.auth_type', {authTypeParamName}, @read_only = 0;"); + + // Set user identifier (oid preferred, fallback to sub) for auditing/observability + string? userId = httpContext.User.FindFirst("oid")?.Value ?? httpContext.User.FindFirst("sub")?.Value; + if (!string.IsNullOrWhiteSpace(userId)) + { + string userIdParamName = $"{SESSION_PARAM_NAME}{counter.Next()}"; + parameters.Add(userIdParamName, new(userId)); + sessionMapQuery.Append($"EXEC sp_set_session_context 'dab.user_id', {userIdParamName}, @read_only = 0;"); + } + + // Set tenant identifier for auditing/observability + string? tenantId = httpContext.User.FindFirst("tid")?.Value; + if (!string.IsNullOrWhiteSpace(tenantId)) + { + string tenantIdParamName = $"{SESSION_PARAM_NAME}{counter.Next()}"; + parameters.Add(tenantIdParamName, new(tenantId)); + sessionMapQuery.Append($"EXEC sp_set_session_context 'dab.tenant_id', {tenantIdParamName}, @read_only = 0;"); + } + } + return sessionMapQuery.ToString(); } diff --git a/src/Core/Resolvers/OboSqlTokenProvider.cs b/src/Core/Resolvers/OboSqlTokenProvider.cs index a4bbf05030..f3528261dd 100644 --- a/src/Core/Resolvers/OboSqlTokenProvider.cs +++ b/src/Core/Resolvers/OboSqlTokenProvider.cs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +using System.Diagnostics; using System.Net; using System.Security.Claims; using System.Security.Cryptography; @@ -66,13 +67,19 @@ public OboSqlTokenProvider( { if (principal is null) { - _logger.LogWarning("Cannot acquire OBO token: ClaimsPrincipal is null."); + _logger.LogWarning( + "{EventType}: Cannot acquire OBO token - ClaimsPrincipal is null (traceId: {TraceId}).", + "OboValidationFailed", + Activity.Current?.TraceId.ToString() ?? "none"); return null; } if (string.IsNullOrWhiteSpace(incomingJwtAssertion)) { - _logger.LogWarning("Cannot acquire OBO token: Incoming JWT assertion is null or empty."); + _logger.LogWarning( + "{EventType}: Cannot acquire OBO token - Incoming JWT assertion is null or empty (traceId: {TraceId}).", + "OboValidationFailed", + Activity.Current?.TraceId.ToString() ?? "none"); return null; } @@ -80,7 +87,10 @@ public OboSqlTokenProvider( string? subjectId = ExtractSubjectId(principal); if (string.IsNullOrWhiteSpace(subjectId)) { - _logger.LogWarning("Cannot acquire OBO token: Neither 'oid' nor 'sub' claim found in token."); + _logger.LogWarning( + "{EventType}: Cannot acquire OBO token - Neither 'oid' nor 'sub' claim found in token (traceId: {TraceId}).", + "OboValidationFailed", + Activity.Current?.TraceId.ToString() ?? "none"); throw new DataApiBuilderException( message: DataApiBuilderException.OBO_IDENTITY_CLAIMS_MISSING, statusCode: HttpStatusCode.Unauthorized, @@ -90,7 +100,10 @@ public OboSqlTokenProvider( string? tenantId = principal.FindFirst("tid")?.Value; if (string.IsNullOrWhiteSpace(tenantId)) { - _logger.LogWarning("Cannot acquire OBO token: 'tid' (tenant id) claim not found or empty in token."); + _logger.LogWarning( + "{EventType}: Cannot acquire OBO token - 'tid' (tenant id) claim not found or empty in token (traceId: {TraceId}).", + "OboValidationFailed", + Activity.Current?.TraceId.ToString() ?? "none"); throw new DataApiBuilderException( message: DataApiBuilderException.OBO_TENANT_CLAIM_MISSING, statusCode: HttpStatusCode.Unauthorized, @@ -115,9 +128,11 @@ public OboSqlTokenProvider( { wasCacheMiss = true; _logger.LogInformation( - "OBO token cache MISS for subject {SubjectId} (tenant: {TenantId}). Acquiring new token from Azure AD.", + "{EventType}: OBO token cache MISS for subject {SubjectId} (tenant: {TenantId}, traceId: {TraceId}). Acquiring new token from Azure AD.", + "OboTokenCacheMiss", subjectId, - tenantId); + tenantId, + Activity.Current?.TraceId.ToString() ?? "none"); AuthenticationResult result = await _msalClient.AcquireTokenOnBehalfOfAsync( scopes, @@ -144,8 +159,10 @@ public OboSqlTokenProvider( ctx.Options.SetSkipDistributedCache(true, true); _logger.LogInformation( - "OBO token ACQUIRED for subject {SubjectId}. Expires: {ExpiresOn}, Cache TTL: {CacheDuration}.", + "{EventType}: OBO token ACQUIRED for subject {SubjectId} (traceId: {TraceId}). Expires: {ExpiresOn}, Cache TTL: {CacheDuration}.", + "OboTokenAcquired", subjectId, + Activity.Current?.TraceId.ToString() ?? "none", result.ExpiresOn, cacheDuration); @@ -155,7 +172,11 @@ public OboSqlTokenProvider( if (!string.IsNullOrEmpty(accessToken) && !wasCacheMiss) { - _logger.LogInformation("OBO token cache HIT for subject {SubjectId}.", subjectId); + _logger.LogInformation( + "{EventType}: OBO token cache HIT for subject {SubjectId} (traceId: {TraceId}).", + "OboTokenCacheHit", + subjectId, + Activity.Current?.TraceId.ToString() ?? "none"); } return accessToken; @@ -164,8 +185,10 @@ public OboSqlTokenProvider( { _logger.LogError( ex, - "Failed to acquire OBO token for subject {SubjectId}. Error: {ErrorCode} - {Message}", + "{EventType}: Failed to acquire OBO token for subject {SubjectId} (traceId: {TraceId}). Error: {ErrorCode} - {Message}", + "OboTokenAcquisitionFailed", subjectId, + Activity.Current?.TraceId.ToString() ?? "none", ex.ErrorCode, ex.Message); throw new DataApiBuilderException( diff --git a/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs b/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs index 778144c6d2..8b357b22df 100644 --- a/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs +++ b/src/Service.Tests/UnitTests/SqlQueryExecutorUnitTests.cs @@ -1014,6 +1014,206 @@ private static Mock CreateHttpContextAccessorWithAuthentic #endregion + /// + /// Validates that GetSessionParamsQuery includes all observability values: + /// - OpenTelemetry correlation values (dab.trace_id, dab.span_id) when an Activity is present + /// - OBO observability values (dab.auth_type, dab.user_id, dab.tenant_id) when user-delegated auth is enabled + /// + [TestMethod, TestCategory(TestCategory.MSSQL)] + public void GetSessionParamsQuery_IncludesAllObservabilityValues_WhenActivityAndOboEnabled() + { + // Arrange + TestHelper.SetupDatabaseEnvironment(TestCategory.MSSQL); + + // Create runtime config with user-delegated-auth enabled and set-session-context + RuntimeConfig runtimeConfig = new( + Schema: "UnitTestSchema", + DataSource: new DataSource( + DatabaseType: DatabaseType.MSSQL, + ConnectionString: "Server=localhost;Database=TestDb;", + Options: new Dictionary { { "set-session-context", true } }) + { + UserDelegatedAuth = new UserDelegatedAuthOptions( + Enabled: true, + Provider: "EntraId", + DatabaseAudience: "https://database.windows.net/") + }, + Runtime: new( + Rest: new(), + GraphQL: new(), + Mcp: new(), + Host: new(Cors: null, Authentication: null)), + Entities: new(new Dictionary())); + + MockFileSystem fileSystem = new(); + fileSystem.AddFile(FileSystemRuntimeConfigLoader.DEFAULT_CONFIG_FILE_NAME, new MockFileData(runtimeConfig.ToJson())); + FileSystemRuntimeConfigLoader loader = new(fileSystem); + RuntimeConfigProvider runtimeConfigProvider = new(loader); + + Mock>> queryExecutorLogger = new(); + Mock httpContextAccessor = new(); + DbExceptionParser dbExceptionParser = new MsSqlDbExceptionParser(runtimeConfigProvider); + + MsSqlQueryExecutor msSqlQueryExecutor = new( + runtimeConfigProvider, + dbExceptionParser, + queryExecutorLogger.Object, + httpContextAccessor.Object); + + // Create a mock HttpContext with OBO-specific claims (oid, tid, sub) + Mock mockContext = new(); + Mock mockRequest = new(); + Mock mockHeaders = new(); + + mockHeaders.Setup(h => h["Authorization"]).Returns("Bearer test-token"); + mockRequest.Setup(r => r.Headers).Returns(mockHeaders.Object); + mockContext.Setup(c => c.Request).Returns(mockRequest.Object); + + var identity = new System.Security.Claims.ClaimsIdentity( + new[] + { + new System.Security.Claims.Claim("oid", "00000000-0000-0000-0000-000000000001"), + new System.Security.Claims.Claim("tid", "11111111-1111-1111-1111-111111111111"), + new System.Security.Claims.Claim("sub", "test-subject") + }, + "TestAuth"); + var principal = new System.Security.Claims.ClaimsPrincipal(identity); + mockContext.Setup(c => c.User).Returns(principal); + + Dictionary parameters = new(); + + // Act - Create an Activity to simulate OpenTelemetry tracing + using ActivitySource activitySource = new("TestActivitySource"); + using ActivityListener listener = new() + { + ShouldListenTo = _ => true, + Sample = (ref ActivityCreationOptions _) => ActivitySamplingResult.AllData + }; + ActivitySource.AddActivityListener(listener); + + using Activity testActivity = activitySource.StartActivity("TestOperation")!; + Assert.IsNotNull(testActivity, "Activity should be created for test"); + + string sessionParamsQuery = msSqlQueryExecutor.GetSessionParamsQuery( + mockContext.Object, + parameters, + runtimeConfigProvider.GetConfig().DefaultDataSourceName); + + // Assert + Assert.IsFalse(string.IsNullOrEmpty(sessionParamsQuery), "Session params query should not be empty"); + + // Verify OpenTelemetry correlation values are included + Assert.IsTrue( + sessionParamsQuery.Contains("'dab.trace_id'"), + "Session params query should include dab.trace_id"); + Assert.IsTrue( + sessionParamsQuery.Contains("'dab.span_id'"), + "Session params query should include dab.span_id"); + + // Verify the correlation values are in the parameters + Assert.IsTrue( + parameters.Values.Any(p => p.Value?.ToString() == testActivity.TraceId.ToString()), + $"Parameters should contain trace_id value: {testActivity.TraceId}"); + Assert.IsTrue( + parameters.Values.Any(p => p.Value?.ToString() == testActivity.SpanId.ToString()), + $"Parameters should contain span_id value: {testActivity.SpanId}"); + + // Verify OBO-specific observability values are included + Assert.IsTrue( + sessionParamsQuery.Contains("'dab.auth_type'"), + "Session params query should include dab.auth_type for OBO"); + Assert.IsTrue( + sessionParamsQuery.Contains("'dab.user_id'"), + "Session params query should include dab.user_id for OBO"); + Assert.IsTrue( + sessionParamsQuery.Contains("'dab.tenant_id'"), + "Session params query should include dab.tenant_id for OBO"); + + // Verify the OBO parameter values are correct + Assert.IsTrue( + parameters.Values.Any(p => p.Value?.ToString() == "obo"), + "Parameters should contain auth_type value: obo"); + Assert.IsTrue( + parameters.Values.Any(p => p.Value?.ToString() == "00000000-0000-0000-0000-000000000001"), + "Parameters should contain user_id value (oid)"); + Assert.IsTrue( + parameters.Values.Any(p => p.Value?.ToString() == "11111111-1111-1111-1111-111111111111"), + "Parameters should contain tenant_id value"); + } + + /// + /// Validates that GetSessionParamsQuery does NOT include correlation values + /// when no Activity is present. + /// + [TestMethod, TestCategory(TestCategory.MSSQL)] + public void GetSessionParamsQuery_ExcludesCorrelationIds_WhenNoActivity() + { + // Arrange + TestHelper.SetupDatabaseEnvironment(TestCategory.MSSQL); + RuntimeConfig runtimeConfig = new( + Schema: "UnitTestSchema", + DataSource: new DataSource( + DatabaseType: DatabaseType.MSSQL, + ConnectionString: "Server=localhost;Database=TestDb;", + Options: new Dictionary { { "set-session-context", true } }), + Runtime: new( + Rest: new(), + GraphQL: new(), + Mcp: new(), + Host: new(Cors: null, Authentication: null)), + Entities: new(new Dictionary())); + + MockFileSystem fileSystem = new(); + fileSystem.AddFile(FileSystemRuntimeConfigLoader.DEFAULT_CONFIG_FILE_NAME, new MockFileData(runtimeConfig.ToJson())); + FileSystemRuntimeConfigLoader loader = new(fileSystem); + RuntimeConfigProvider runtimeConfigProvider = new(loader); + + Mock>> queryExecutorLogger = new(); + Mock httpContextAccessor = new(); + DbExceptionParser dbExceptionParser = new MsSqlDbExceptionParser(runtimeConfigProvider); + + MsSqlQueryExecutor msSqlQueryExecutor = new( + runtimeConfigProvider, + dbExceptionParser, + queryExecutorLogger.Object, + httpContextAccessor.Object); + + // Create a mock HttpContext with a simple authenticated user + Mock mockContext = new(); + Mock mockRequest = new(); + Mock mockHeaders = new(); + + mockHeaders.Setup(h => h["Authorization"]).Returns(string.Empty); + mockRequest.Setup(r => r.Headers).Returns(mockHeaders.Object); + mockContext.Setup(c => c.Request).Returns(mockRequest.Object); + + var identity = new System.Security.Claims.ClaimsIdentity( + new[] { new System.Security.Claims.Claim("sub", "test-user") }, + "TestAuth"); + var principal = new System.Security.Claims.ClaimsPrincipal(identity); + mockContext.Setup(c => c.User).Returns(principal); + + Dictionary parameters = new(); + + // Act - Ensure no Activity is present (Activity.Current should be null) + // We don't start any activity here + string sessionParamsQuery = msSqlQueryExecutor.GetSessionParamsQuery( + mockContext.Object, + parameters, + runtimeConfigProvider.GetConfig().DefaultDataSourceName); + + // Assert + Assert.IsFalse(string.IsNullOrEmpty(sessionParamsQuery), "Session params query should not be empty (has user claims)"); + + // Verify trace_id and span_id are NOT included when no Activity + Assert.IsFalse( + sessionParamsQuery.Contains("'dab.trace_id'"), + "Session params query should NOT include dab.trace_id when no Activity present"); + Assert.IsFalse( + sessionParamsQuery.Contains("'dab.span_id'"), + "Session params query should NOT include dab.span_id when no Activity present"); + } + [TestCleanup] public void CleanupAfterEachTest() {