Open Telemetry : Adds query text in attribute

Microsoft.Azure.Cosmos/src/ChangeFeed/ChangeFeedIteratorCore.cs

@@ -229,7 +229,7 @@ public override async Task<ResponseMessage> ReadNextAsync(CancellationToken canc
                                                 operationType: OperationType.ReadFeed,
                                                 requestOptions: this.changeFeedRequestOptions,
                                                 task: (trace) => this.ReadNextInternalAsync(trace, cancellationToken),
-                                                openTelemetry: new (OpenTelemetryConstants.Operations.QueryChangeFeed, (response) => new OpenTelemetryResponse(responseMessage: response)),
+                                                openTelemetry: new (OpenTelemetryConstants.Operations.QueryChangeFeed, (response) => new OpenTelemetryResponse(responseMessage: response, querySpecFunc: () => this.changeFeedQuerySpec?.ToSqlQuerySpec())),
                                                 traceComponent: TraceComponent.ChangeFeed,
                                                 traceLevel: TraceLevel.Info);
         }

Microsoft.Azure.Cosmos/src/ChangeFeed/ChangeFeedQuerySpec.cs

@@ -54,5 +54,13 @@ internal bool ShouldSerializeQueryText()
         {
             return this.QueryText.Length > 0;
         }
+
+        /// <summary>
+        /// Converts to SQL Query Specs
+        /// </summary>
+        internal SqlQuerySpec ToSqlQuerySpec()
+        {
+            return new SqlQuerySpec(this.QueryText);
+        }
     }
 }

Microsoft.Azure.Cosmos/src/ChangeFeedProcessor/ChangeFeedEstimatorIterator.cs

@@ -122,7 +122,7 @@ public override Task<FeedResponse<ChangeFeedProcessorState>> ReadNextAsync(Cance
                 operationType: Documents.OperationType.ReadFeed,
                 requestOptions: null,
                 task: (trace) => this.ReadNextAsync(trace, cancellationToken),
-                openTelemetry: new (OpenTelemetryConstants.Operations.QueryChangeFeedEstimator, (response) => new OpenTelemetryResponse<ChangeFeedProcessorState>(responseMessage: response)),
+                openTelemetry: new (OpenTelemetryConstants.Operations.QueryChangeFeedEstimator, (response) => new OpenTelemetryResponse<ChangeFeedProcessorState>(responseMessage: response, querySpec: this.querySpec)),
                 traceComponent: TraceComponent.ChangeFeed,
                 traceLevel: TraceLevel.Info);
         }

Microsoft.Azure.Cosmos/src/CosmosClientTelemetryOptions.cs

@@ -46,5 +46,12 @@ public class CosmosClientTelemetryOptions
         /// These values decides whether to generate an <see cref="System.Diagnostics.Tracing.EventSource"/> with request diagnostics or not.
         /// </summary>
         public CosmosThresholdOptions CosmosThresholdOptions { get; set; } = new CosmosThresholdOptions();
+
+        /// <summary>
+        /// Enables printing query in Traces db.query.text attribute. By default, query is not printed.
+        /// Users have the option to enable printing parameterized or all queries, 
+        /// but has to beware that customer data may be shown when the later option is chosen. It's the user's responsibility to sanitize the queries if necessary.
+        /// </summary>
+        public QueryTextMode QueryTextMode { get; set; } = QueryTextMode.None;
     }
 }

Microsoft.Azure.Cosmos/src/Query/v3Query/QueryIterator.cs

@@ -43,7 +43,8 @@ private QueryIterator(
             RequestOptions requestOptions,
             CosmosClientContext clientContext,
             Guid correlatedActivityId,
-            ContainerInternal container)
+            ContainerInternal container,
+            SqlQuerySpec sqlQuerySpec)
         {
             this.cosmosQueryContext = cosmosQueryContext ?? throw new ArgumentNullException(nameof(cosmosQueryContext));
             this.queryPipelineStage = cosmosQueryExecutionContext ?? throw new ArgumentNullException(nameof(cosmosQueryExecutionContext));
@@ -53,6 +54,7 @@ private QueryIterator(
             this.hasMoreResults = true;
             this.correlatedActivityId = correlatedActivityId;
 
+            this.querySpec = sqlQuerySpec;
             this.container = container;
             this.operationName = OpenTelemetryConstants.Operations.QueryItems;
             this.operationType = Documents.OperationType.Query;
@@ -119,7 +121,8 @@ public static QueryIterator Create(
                         queryRequestOptions,
                         clientContext,
                         correlatedActivityId,
-                        containerCore);
+                        containerCore,
+                        sqlQuerySpec);
                 }
 
                 requestContinuationToken = tryParse.Result;
@@ -152,7 +155,8 @@ public static QueryIterator Create(
                 queryRequestOptions,
                 clientContext,
                 correlatedActivityId,
-                containerCore);
+                containerCore,
+                sqlQuerySpec);
         }
 
         public override bool HasMoreResults => this.hasMoreResults;

Microsoft.Azure.Cosmos/src/RequestOptions/ChangeFeedRequestOptions.cs

@@ -79,6 +79,13 @@ internal override void PopulateRequestOptions(RequestMessage request)
 #endif
         JsonSerializationFormatOptions JsonSerializationFormatOptions { get; set; }
 
+        /// <summary>
+        /// Enables printing query in Traces db.query.text attribute. By default, query is not printed.
+        /// Users have the option to enable printing parameterized or all queries, 
+        /// but has to beware that customer data may be shown when the later option is chosen. It's the user's responsibility to sanitize the queries if necessary.
+        /// </summary>
+        internal QueryTextMode? QueryTextMode { get; set; }
+
         internal ChangeFeedRequestOptions Clone()
         {
             return new ChangeFeedRequestOptions()

Microsoft.Azure.Cosmos/src/RequestOptions/QueryRequestOptions.cs

@@ -170,6 +170,13 @@ public ConsistencyLevel? ConsistencyLevel
         /// These options have no effect otherwise.
         /// </summary>
         public DedicatedGatewayRequestOptions DedicatedGatewayRequestOptions { get; set; }
+
+        /// <summary>
+        /// Enables printing query in Traces db.query.text attribute. By default, query is not printed.
+        /// Users have the option to enable printing parameterized or all queries, 
+        /// but has to beware that customer data may be shown when the later option is chosen. It's the user's responsibility to sanitize the queries if necessary.
+        /// </summary>
+        public QueryTextMode? QueryTextMode { get; set; } = null;
 
         internal CosmosElement CosmosElementContinuationToken { get; set; }
 

Microsoft.Azure.Cosmos/src/Resource/FeedIterators/FeedIterator.cs

@@ -7,6 +7,7 @@ namespace Microsoft.Azure.Cosmos
     using System;
     using System.Threading;
     using System.Threading.Tasks;
+    using Microsoft.Azure.Cosmos.Query.Core;
 
     /// <summary>
     /// Cosmos Result set iterator that keeps track of the continuation token when retrieving results form a query.
@@ -145,5 +146,10 @@ public void Dispose()
         /// Operation Type used for open telemetry traces, it will set optionally, otherwise default operation type will be used in traces
         /// </summary>
         internal Documents.OperationType? operationType;
+
+        /// <summary>
+        /// collect SQL query Specs for tracing
+        /// </summary>
+        internal SqlQuerySpec querySpec;
     }
 }

Microsoft.Azure.Cosmos/src/Resource/FeedIterators/FeedIteratorCore.cs

@@ -27,7 +27,6 @@ internal sealed class FeedIteratorCore : FeedIteratorInternal
         private readonly CosmosClientContext clientContext;
         private readonly string resourceLink;
         private readonly ResourceType resourceType;
-        private readonly SqlQuerySpec querySpec;
         private bool hasMoreResultsInternal;
 
         public FeedIteratorCore(

Microsoft.Azure.Cosmos/src/Resource/FeedIterators/FeedIteratorInlineCore.cs

@@ -28,6 +28,7 @@ internal FeedIteratorInlineCore(
             this.feedIteratorInternal = feedIteratorInternal;
             this.clientContext = clientContext;
 
+            this.querySpec = feedIteratorInternal.querySpec;
             this.container = feedIteratorInternal.container;
             this.databaseName = feedIteratorInternal.databaseName;
 
@@ -62,7 +63,7 @@ public override Task<ResponseMessage> ReadNextAsync(CancellationToken cancellati
                         task: (trace) => this.feedIteratorInternal.ReadNextAsync(trace, cancellationToken),
                         openTelemetry: new (this.operationName, (response) =>
                         {
-                            OpenTelemetryResponse openTelemetryResponse = new OpenTelemetryResponse(responseMessage: response);
+                            OpenTelemetryResponse openTelemetryResponse = new OpenTelemetryResponse(responseMessage: response, querySpecFunc: () => this.querySpec);
 
                             if (this.operationType.HasValue)
                             {

Microsoft.Azure.Cosmos/src/Resource/FeedIterators/FeedIteratorInlineCore{T}.cs

@@ -28,6 +28,7 @@ internal FeedIteratorInlineCore(
             this.feedIteratorInternal = feedIteratorInternal;
             this.clientContext = clientContext;
 
+            this.querySpec = feedIterator.querySpec;
             this.container = feedIteratorInternal.container;
             this.databaseName = feedIteratorInternal.databaseName;
 
@@ -62,7 +63,7 @@ public override Task<FeedResponse<T>> ReadNextAsync(CancellationToken cancellati
                         task: trace => this.feedIteratorInternal.ReadNextAsync(trace, cancellationToken),
                         openTelemetry: new (this.feedIteratorInternal.operationName, (response) =>
                         {
-                            OpenTelemetryResponse<T> openTelemetryResponse = new OpenTelemetryResponse<T>(responseMessage: response);
+                            OpenTelemetryResponse<T> openTelemetryResponse = new OpenTelemetryResponse<T>(responseMessage: response, querySpec: this.querySpec);
 
                             if (this.operationType.HasValue)
                             {

Microsoft.Azure.Cosmos/src/Resource/FeedIterators/FeedIterator{T}.cs

@@ -7,6 +7,7 @@ namespace Microsoft.Azure.Cosmos
     using System;
     using System.Threading;
     using System.Threading.Tasks;
+    using Microsoft.Azure.Cosmos.Query.Core;
 
     /// <summary>
     /// Cosmos Result set iterator that keeps track of the continuation token when retrieving results form a query.
@@ -133,5 +134,10 @@ public void Dispose()
         /// Operation Type used for open telemetry traces
         /// </summary>
         internal Documents.OperationType? operationType;
+
+        /// <summary>
+        /// collect SQL query Specs for tracing
+        /// </summary>
+        internal SqlQuerySpec querySpec;
     }
 }

Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryAttributeKeys.cs

@@ -137,6 +137,11 @@ internal sealed class OpenTelemetryAttributeKeys
         /// </summary>
         public const string CorrelatedActivityId = "db.cosmosdb.correlated_activity_id";
 
+        /// <summary>
+        /// Represents the Azure Cosmos DB SQL Query.
+        /// </summary>
+        public const string QueryText = "db.query.text";
+
         /// <summary>
         /// Represents the size of the batch operation.
         /// </summary>

Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryAttributes.cs

@@ -6,7 +6,7 @@ namespace Microsoft.Azure.Cosmos.Telemetry
 {
     using System;
     using System.Net;
-    using System.Security.AccessControl;
+    using Microsoft.Azure.Cosmos.Query.Core;
 
     internal class OpenTelemetryAttributes
     {
@@ -90,5 +90,10 @@ internal OpenTelemetryAttributes(RequestMessage requestMessage)
         /// Batch Size
         /// </summary>
         internal int? BatchSize { get; set; }
+
+        /// <summary>
+        /// Query Spec with Query Text and Parameters
+        /// </summary>
+        internal SqlQuerySpec QuerySpec { get; set; }
     }
 }

Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryCoreRecorder.cs

@@ -26,6 +26,7 @@ internal struct OpenTelemetryCoreRecorder : IDisposable
         private readonly OperationType operationType = OperationType.Invalid;
         private readonly string connectionModeCache = null;
 
+        private readonly QueryTextMode? queryTextMode = null;
         private OpenTelemetryAttributes response = null;
 
         /// <summary>
@@ -59,13 +60,15 @@ private OpenTelemetryCoreRecorder(
             string databaseName,
             OperationType operationType, 
             CosmosClientContext clientContext, 
-            CosmosThresholdOptions config)
+            CosmosThresholdOptions config,
+            QueryTextMode queryTextMode)
         {
             this.scope = scope;
             this.config = config;
 
             this.operationType = operationType;
             this.connectionModeCache = Enum.GetName(typeof(ConnectionMode), clientContext.ClientOptions.ConnectionMode);
+            this.queryTextMode = queryTextMode;
 
             if (scope.IsEnabled)
             {
@@ -107,7 +110,8 @@ public static OpenTelemetryCoreRecorder CreateOperationLevelParentActivity(
             string databaseName,
             Documents.OperationType operationType,
             CosmosClientContext clientContext,
-            CosmosThresholdOptions config)
+            CosmosThresholdOptions config, 
+            QueryTextMode queryTextMode)
         {
             return new OpenTelemetryCoreRecorder(
                         operationScope,
@@ -116,7 +120,8 @@ public static OpenTelemetryCoreRecorder CreateOperationLevelParentActivity(
                         databaseName,
                         operationType,
                         clientContext,
-                        config);
+                        config,
+                        queryTextMode);
         }
 
         public bool IsEnabled => this.scope.IsEnabled;
@@ -245,6 +250,15 @@ OperationType operationType
                     this.scope.AddAttribute(OpenTelemetryAttributeKeys.ActivityId, this.response.ActivityId);
                     this.scope.AddAttribute(OpenTelemetryAttributeKeys.CorrelatedActivityId, this.response.CorrelatedActivityId);
 
+                    if (this.response.QuerySpec is not null)
+                    {
+                        if (this.queryTextMode == QueryTextMode.All || 
+                            (this.queryTextMode == QueryTextMode.ParameterizedOnly && this.response.QuerySpec.ShouldSerializeParameters()))
+                        {
+                            this.scope.AddAttribute(OpenTelemetryAttributeKeys.QueryText, this.response.QuerySpec?.QueryText);
+                        }
+                    }
+
                     if (this.response.Diagnostics != null)
                     {
                         this.scope.AddAttribute(OpenTelemetryAttributeKeys.Region, ClientTelemetryHelper.GetContactedRegions(this.response.Diagnostics.GetContactedRegions()));

Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryRecorderFactory.cs

@@ -56,14 +56,17 @@ public static OpenTelemetryCoreRecorder CreateRecorder(Func<string> getOperation
                 {
                     scope.SetDisplayName($"{operationName} {containerName}");
 
+                    QueryTextMode queryTextMode = GetQueryTextMode(requestOptions, clientContext);
+
                     openTelemetryRecorder = OpenTelemetryCoreRecorder.CreateOperationLevelParentActivity(
                         operationScope: scope,
                         operationName: operationName,
                         containerName: containerName,
                         databaseName: databaseName,
                         operationType: operationType,
                         clientContext: clientContext,
-                        config: requestOptions?.CosmosThresholdOptions ?? clientContext.ClientOptions?.CosmosClientTelemetryOptions.CosmosThresholdOptions);
+                        config: requestOptions?.CosmosThresholdOptions ?? clientContext.ClientOptions?.CosmosClientTelemetryOptions.CosmosThresholdOptions,
+                        queryTextMode: queryTextMode);
                 }
 #if !INTERNAL
                 // If there are no listeners at operation level and no parent activity created.
@@ -83,5 +86,22 @@ public static OpenTelemetryCoreRecorder CreateRecorder(Func<string> getOperation
             }
             return openTelemetryRecorder;
         }
+
+        private static QueryTextMode GetQueryTextMode(RequestOptions requestOptions, CosmosClientContext clientContext)
+        {
+            QueryTextMode? queryTextMode = null;
+            if (requestOptions is QueryRequestOptions queryRequestOptions && queryRequestOptions.QueryTextMode.HasValue)
+            {
+                queryTextMode = queryRequestOptions.QueryTextMode.Value;
+            }
+            else if (requestOptions is ChangeFeedRequestOptions changeFeedRequestOptions && changeFeedRequestOptions.QueryTextMode.HasValue)
+            {
+                queryTextMode = changeFeedRequestOptions.QueryTextMode.Value;
+
+            }
+
+            queryTextMode ??= clientContext.ClientOptions?.CosmosClientTelemetryOptions?.QueryTextMode ?? QueryTextMode.None;
+            return queryTextMode.Value;
+        }
     }
 }

Microsoft.Azure.Cosmos/src/Telemetry/OpenTelemetry/OpenTelemetryResponse.cs

@@ -8,6 +8,8 @@ namespace Microsoft.Azure.Cosmos
     using System.IO;
     using System.Net;
     using Microsoft.Azure.Cosmos.Core.Trace;
+    using Microsoft.Azure.Cosmos.Query.Core;
+    using Microsoft.Azure.Documents;
     using Telemetry;
 
     internal sealed class OpenTelemetryResponse : OpenTelemetryAttributes
@@ -27,7 +29,7 @@ internal OpenTelemetryResponse(TransactionalBatchResponse responseMessage)
         {
         }
 
-        internal OpenTelemetryResponse(ResponseMessage responseMessage)
+        internal OpenTelemetryResponse(ResponseMessage responseMessage, Func<SqlQuerySpec> querySpecFunc = null)
            : this(
                   statusCode: responseMessage.StatusCode,
                   requestCharge: OpenTelemetryResponse.GetHeader(responseMessage)?.RequestCharge,
@@ -37,7 +39,8 @@ internal OpenTelemetryResponse(ResponseMessage responseMessage)
                   requestMessage: responseMessage.RequestMessage,
                   subStatusCode: OpenTelemetryResponse.GetHeader(responseMessage)?.SubStatusCode,
                   activityId: OpenTelemetryResponse.GetHeader(responseMessage)?.ActivityId,
-                  correlationId: OpenTelemetryResponse.GetHeader(responseMessage)?.CorrelatedActivityId)
+                  correlationId: OpenTelemetryResponse.GetHeader(responseMessage)?.CorrelatedActivityId,
+                  querySpecFunc: querySpecFunc)
         {
         }
 
@@ -51,7 +54,8 @@ private OpenTelemetryResponse(
             Documents.SubStatusCodes? subStatusCode,
             string activityId,
             string correlationId,
-            int? batchSize = null)
+            int? batchSize = null,
+            Func<SqlQuerySpec> querySpecFunc = null)
             : base(requestMessage)
         {
             this.StatusCode = statusCode;
@@ -63,6 +67,10 @@ private OpenTelemetryResponse(
             this.ActivityId = activityId;
             this.CorrelatedActivityId = correlationId;
             this.BatchSize = batchSize;
+            if (querySpecFunc != null)
+            {
+                this.QuerySpec = querySpecFunc();
+            }
         }
 
         private static string GetPayloadSize(ResponseMessage response)