Update document model async operations

aws · Oct 30, 2024 · 2946318 · 2946318
1 parent 8491f2e
commit 2946318
Show file tree

Hide file tree

Showing 6 changed files with 214 additions and 69 deletions.
diff --git a/sdk/src/Services/DynamoDBv2/Custom/DocumentModel/_async/DocumentBatchGet.Async.cs b/sdk/src/Services/DynamoDBv2/Custom/DocumentModel/_async/DocumentBatchGet.Async.cs
@@ -14,6 +14,7 @@
  */
 #pragma warning disable 1574
 
+using Amazon.Runtime.Telemetry.Tracing;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -33,9 +34,13 @@ public partial interface IDocumentBatchGet
     public partial class DocumentBatchGet : IDocumentBatchGet
     {
         /// <inheritdoc/>
-        public Task ExecuteAsync(CancellationToken cancellationToken = default(CancellationToken))
+        public async Task ExecuteAsync(CancellationToken cancellationToken = default(CancellationToken))
         {
-            return ExecuteHelperAsync(cancellationToken);
+            var operationName = DynamoDBTelemetry.ExtractOperationName(nameof(DocumentBatchGet), nameof(ExecuteAsync));
+            using (DynamoDBTelemetry.CreateSpan(TracerProvider, operationName, spanKind: SpanKind.CLIENT))
+            {
+                await ExecuteHelperAsync(cancellationToken).ConfigureAwait(false);
+            }
         }
     }
 
@@ -53,9 +58,13 @@ public partial interface IMultiTableDocumentBatchGet
     public partial class MultiTableDocumentBatchGet : IMultiTableDocumentBatchGet
     {
         /// <inheritdoc/>
-        public Task ExecuteAsync(CancellationToken cancellationToken = default(CancellationToken))
+        public async Task ExecuteAsync(CancellationToken cancellationToken = default(CancellationToken))
         {
-            return ExecuteHelperAsync(cancellationToken);
+            var operationName = DynamoDBTelemetry.ExtractOperationName(nameof(MultiTableDocumentBatchGet), nameof(ExecuteAsync));
+            using (DynamoDBTelemetry.CreateSpan(TracerProvider, operationName, spanKind: SpanKind.CLIENT))
+            {
+                await ExecuteHelperAsync(cancellationToken).ConfigureAwait(false);
+            }
         }
     }
 }
diff --git a/sdk/src/Services/DynamoDBv2/Custom/DocumentModel/_async/DocumentBatchWrite.Async.cs b/sdk/src/Services/DynamoDBv2/Custom/DocumentModel/_async/DocumentBatchWrite.Async.cs
@@ -35,9 +35,15 @@ public partial interface IDocumentBatchWrite
     public partial class DocumentBatchWrite : IDocumentBatchWrite
     {
         /// <inheritdoc/>
-        public Task ExecuteAsync(CancellationToken cancellationToken = default(CancellationToken))
+        public async Task ExecuteAsync(CancellationToken cancellationToken = default(CancellationToken))
         {
-            return ExecuteHelperAsync(cancellationToken);
+            {
+                var operationName = DynamoDBTelemetry.ExtractOperationName(nameof(DocumentBatchWrite), nameof(ExecuteAsync));
+                using (DynamoDBTelemetry.CreateSpan(TracerProvider, operationName, spanKind: SpanKind.CLIENT))
+                {
+                    await ExecuteHelperAsync(cancellationToken).ConfigureAwait(false);
+                }
+            }
         }
     }
 
@@ -57,9 +63,15 @@ public partial interface IMultiTableDocumentBatchWrite
     public partial class MultiTableDocumentBatchWrite : IMultiTableDocumentBatchWrite
     {
         /// <inheritdoc/>
-        public Task ExecuteAsync(CancellationToken cancellationToken = default(CancellationToken))
+        public async Task ExecuteAsync(CancellationToken cancellationToken = default(CancellationToken))
         {
-            return ExecuteHelperAsync(cancellationToken);
+            {
+                var operationName = DynamoDBTelemetry.ExtractOperationName(nameof(MultiTableDocumentBatchWrite), nameof(ExecuteAsync));
+                using (DynamoDBTelemetry.CreateSpan(TracerProvider, operationName, spanKind: SpanKind.CLIENT))
+                {
+                    await ExecuteHelperAsync(cancellationToken).ConfigureAwait(false);
+                }
+            }
         }
     }
 }
diff --git a/sdk/src/Services/DynamoDBv2/Custom/DocumentModel/_async/DocumentTransactGet.Async.cs b/sdk/src/Services/DynamoDBv2/Custom/DocumentModel/_async/DocumentTransactGet.Async.cs
@@ -13,6 +13,7 @@
  * permissions and limitations under the License.
  */
 
+using Amazon.Runtime.Telemetry.Tracing;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -36,9 +37,13 @@ public partial interface IDocumentTransactGet
     public partial class DocumentTransactGet : IDocumentTransactGet
     {
         /// <inheritdoc/>
-        public Task ExecuteAsync(CancellationToken cancellationToken = default(CancellationToken))
+        public async Task ExecuteAsync(CancellationToken cancellationToken = default(CancellationToken))
         {
-            return ExecuteHelperAsync(cancellationToken);
+            var operationName = DynamoDBTelemetry.ExtractOperationName(nameof(DocumentTransactGet), nameof(ExecuteAsync));
+            using (DynamoDBTelemetry.CreateSpan(TracerProvider, operationName, spanKind: SpanKind.CLIENT))
+            {
+                await ExecuteHelperAsync(cancellationToken).ConfigureAwait(false);
+            }
         }
     }
 
@@ -56,9 +61,13 @@ public partial interface IMultiTableDocumentTransactGet
     public partial class MultiTableDocumentTransactGet : IMultiTableDocumentTransactGet
     {
         /// <inheritdoc/>
-        public Task ExecuteAsync(CancellationToken cancellationToken = default(CancellationToken))
+        public async Task ExecuteAsync(CancellationToken cancellationToken = default(CancellationToken))
         {
-            return ExecuteHelperAsync(cancellationToken);
+            var operationName = DynamoDBTelemetry.ExtractOperationName(nameof(MultiTableDocumentTransactGet), nameof(ExecuteAsync));
+            using (DynamoDBTelemetry.CreateSpan(TracerProvider, operationName, spanKind: SpanKind.CLIENT))
+            {
+                await ExecuteHelperAsync(cancellationToken).ConfigureAwait(false);
+            }
         }
     }
 }
diff --git a/sdk/src/Services/DynamoDBv2/Custom/DocumentModel/_async/DocumentTransactWrite.Async.cs b/sdk/src/Services/DynamoDBv2/Custom/DocumentModel/_async/DocumentTransactWrite.Async.cs
@@ -13,6 +13,7 @@
  * permissions and limitations under the License.
  */
 
+using Amazon.Runtime.Telemetry.Tracing;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -31,9 +32,13 @@ public partial interface IDocumentTransactWrite
     public partial class DocumentTransactWrite : IDocumentTransactWrite
     {
         /// <inheritdoc/>
-        public Task ExecuteAsync(CancellationToken cancellationToken = default(CancellationToken))
+        public async Task ExecuteAsync(CancellationToken cancellationToken = default(CancellationToken))
         {
-            return ExecuteHelperAsync(cancellationToken);
+            var operationName = DynamoDBTelemetry.ExtractOperationName(nameof(DocumentTransactWrite), nameof(ExecuteAsync));
+            using (DynamoDBTelemetry.CreateSpan(TracerProvider, operationName, spanKind: SpanKind.CLIENT))
+            {
+                await ExecuteHelperAsync(cancellationToken).ConfigureAwait(false);
+            }
         }
     }
 
@@ -50,9 +55,13 @@ public partial interface IMultiTableDocumentTransactWrite
     public partial class MultiTableDocumentTransactWrite : IMultiTableDocumentTransactWrite
     {
         /// <inheritdoc/>
-        public Task ExecuteAsync(CancellationToken cancellationToken = default(CancellationToken))
+        public async Task ExecuteAsync(CancellationToken cancellationToken = default(CancellationToken))
         {
-            return ExecuteHelperAsync(cancellationToken);
+            var operationName = DynamoDBTelemetry.ExtractOperationName(nameof(MultiTableDocumentTransactWrite), nameof(ExecuteAsync));
+            using (DynamoDBTelemetry.CreateSpan(TracerProvider, operationName, spanKind: SpanKind.CLIENT))
+            {
+                await ExecuteHelperAsync(cancellationToken).ConfigureAwait(false);
+            }
         }
     }
 }
diff --git a/sdk/src/Services/DynamoDBv2/Custom/DocumentModel/_async/Search.Async.cs b/sdk/src/Services/DynamoDBv2/Custom/DocumentModel/_async/Search.Async.cs
@@ -14,6 +14,7 @@
  */
 #pragma warning disable 1574
 
+using Amazon.Runtime.Telemetry.Tracing;
 using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
@@ -46,15 +47,23 @@ public partial interface ISearch
     public partial class Search : ISearch
     {
         /// <inheritdoc/>
-        public Task<List<Document>> GetNextSetAsync(CancellationToken cancellationToken = default(CancellationToken))
+        public async Task<List<Document>> GetNextSetAsync(CancellationToken cancellationToken = default(CancellationToken))
         {
-            return GetNextSetHelperAsync(cancellationToken);
+            var operationName = DynamoDBTelemetry.ExtractOperationName(nameof(Search), nameof(GetNextSetAsync));
+            using (DynamoDBTelemetry.CreateSpan(TracerProvider, operationName, spanKind: SpanKind.CLIENT))
+            {
+                return await GetNextSetHelperAsync(cancellationToken).ConfigureAwait(false);
+            }
         }
 
         /// <inheritdoc/>
-        public Task<List<Document>> GetRemainingAsync(CancellationToken cancellationToken = default(CancellationToken))
+        public async Task<List<Document>> GetRemainingAsync(CancellationToken cancellationToken = default(CancellationToken))
         {
-            return GetRemainingHelperAsync(cancellationToken);
+            var operationName = DynamoDBTelemetry.ExtractOperationName(nameof(Search), nameof(GetRemainingAsync));
+            using (DynamoDBTelemetry.CreateSpan(TracerProvider, operationName, spanKind: SpanKind.CLIENT))
+            {
+                return await GetRemainingHelperAsync(cancellationToken).ConfigureAwait(false);
+            }
         }
     }
 }