[Storage] Add user agent for dataplane cmdlets (#12255)

Author: blueww

src/Storage/Storage.Management/ChangeLog.md

@@ -18,6 +18,7 @@
         - Additional information about change #1
 -->
 ## Upcoming Release
+* Fixed the issue that UserAgent not added for some dataplane cmdlets.
 * Supported create/update Storage account with MinimumTlsVersion and AllowBlobPublicAccess
     -  `New-AzStorageAccount`
     -  `Set-AzStorageAccount`

src/Storage/Storage/Blob/Cmdlet/GetAzureStorageContainerStoredAccessPolicy.cs

@@ -86,7 +86,7 @@ internal async Task GetAzureContainerStoredAccessPolicyAsync(long taskId, IStora
         internal async Task<SharedAccessBlobPolicies> GetPoliciesAsync(IStorageBlobManagement localChannel, string containerName, string policyName)
         {
             CloudBlobContainer container = localChannel.GetContainerReference(containerName);
-            BlobContainerPermissions blobContainerPermissions = await localChannel.GetContainerPermissionsAsync(container, null, null, null, CmdletCancellationToken).ConfigureAwait(false);
+            BlobContainerPermissions blobContainerPermissions = await localChannel.GetContainerPermissionsAsync(container, null, null, OperationContext, CmdletCancellationToken).ConfigureAwait(false);
             return blobContainerPermissions.SharedAccessPolicies;
         }
 

src/Storage/Storage/Blob/Cmdlet/NewAzureStorageContainerStoredAccessPolicy.cs

@@ -77,7 +77,7 @@ internal string CreateAzureContainerStoredAccessPolicy(IStorageBlobManagement lo
 
             //Get existing permissions
             CloudBlobContainer container = localChannel.GetContainerReference(containerName);
-            BlobContainerPermissions blobContainerPermissions = localChannel.GetContainerPermissions(container);
+            BlobContainerPermissions blobContainerPermissions = localChannel.GetContainerPermissions(container, null, null, OperationContext);
 
             //Add new policy
             if (blobContainerPermissions.SharedAccessPolicies.Keys.Contains(policyName))
@@ -90,7 +90,7 @@ internal string CreateAzureContainerStoredAccessPolicy(IStorageBlobManagement lo
             blobContainerPermissions.SharedAccessPolicies.Add(policyName, policy);
 
             //Set permissions back to container
-            localChannel.SetContainerPermissions(container, blobContainerPermissions);
+            localChannel.SetContainerPermissions(container, blobContainerPermissions, null, null, OperationContext);
             return policyName;
         }
 

src/Storage/Storage/Blob/Cmdlet/RemoveAzureStorageContainerStoredAccessPolicy.cs

@@ -69,7 +69,7 @@ internal bool RemoveAzureContainerStoredAccessPolicy(IStorageBlobManagement loca
 
             //Get existing permissions
             CloudBlobContainer container = localChannel.GetContainerReference(containerName);
-            BlobContainerPermissions blobContainerPermissions = localChannel.GetContainerPermissions(container);
+            BlobContainerPermissions blobContainerPermissions = localChannel.GetContainerPermissions(container, null, null, OperationContext);
 
             //remove the specified policy
             if (!blobContainerPermissions.SharedAccessPolicies.Keys.Contains(policyName))
@@ -80,7 +80,7 @@ internal bool RemoveAzureContainerStoredAccessPolicy(IStorageBlobManagement loca
             if (ShouldProcess(policyName, "Remove policy"))
             {
                 blobContainerPermissions.SharedAccessPolicies.Remove(policyName);
-                localChannel.SetContainerPermissions(container, blobContainerPermissions);
+                localChannel.SetContainerPermissions(container, blobContainerPermissions, null, null, OperationContext);
                 success = true;
             }
 

src/Storage/Storage/Blob/Cmdlet/SetAzureStorageContainerAcl.cs

@@ -87,7 +87,7 @@ internal async Task SetContainerAcl(long taskId, IStorageBlobManagement localCha
             BlobContainerPermissions permissions;
             try
             {
-                permissions = localChannel.GetContainerPermissions(container);
+                permissions = localChannel.GetContainerPermissions(container, null, requestOptions, OperationContext);
             }
             catch (StorageException e) when (e.IsNotFoundException())
             {

src/Storage/Storage/Blob/Cmdlet/SetAzureStorageContainerStoredAccessPolicy.cs

@@ -78,7 +78,7 @@ internal string SetAzureContainerStoredAccessPolicy(IStorageBlobManagement local
         {
             //Get existing permissions
             CloudBlobContainer container = localChannel.GetContainerReference(containerName);
-            BlobContainerPermissions blobContainerPermissions = localChannel.GetContainerPermissions(container);
+            BlobContainerPermissions blobContainerPermissions = localChannel.GetContainerPermissions(container, null, null, OperationContext);
 
             //Set the policy with new value
             if (!blobContainerPermissions.SharedAccessPolicies.Keys.Contains(policyName))
@@ -91,7 +91,7 @@ internal string SetAzureContainerStoredAccessPolicy(IStorageBlobManagement local
             blobContainerPermissions.SharedAccessPolicies[policyName] = policy;
 
             //Set permission back to container
-            localChannel.SetContainerPermissions(container, blobContainerPermissions);
+            localChannel.SetContainerPermissions(container, blobContainerPermissions, null, null, OperationContext);
             WriteObject(AccessPolicyHelper.ConstructPolicyOutputPSObject<SharedAccessBlobPolicy>(blobContainerPermissions.SharedAccessPolicies, policyName));
             return policyName;
         }

src/Storage/Storage/Common/CmdletOperationContext.cs

@@ -15,6 +15,7 @@
 namespace Microsoft.WindowsAzure.Commands.Storage.Common
 {
     using Microsoft.Azure.Storage;
+    using Microsoft.WindowsAzure.Commands.Common;
     using System;
     using System.Threading;
     using XTable = Microsoft.Azure.Cosmos.Table;
@@ -148,6 +149,9 @@ public static OperationContext GetStorageOperationContext(Action<string> outputW
                 }
             };
 
+            context.CustomUserAgent = ApiConstants.UserAgentHeaderValue;
+
+
             return context;
         }
 
@@ -207,6 +211,8 @@ public static XTable.OperationContext GetStorageTableOperationContext(Action<str
                 }
             };
 
+            context.CustomUserAgent = ApiConstants.UserAgentHeaderValue;
+
             return context;
         }
 

src/Storage/Storage/Common/Util.cs

@@ -108,10 +108,10 @@ private static CloudBlob GetBlobReferenceFromServer(
                 return null;
             }
 
-            return GetCorrespondingTypeBlobReference(blob);
+            return GetCorrespondingTypeBlobReference(blob, operationContext);
         }
 
-        public static CloudBlob GetCorrespondingTypeBlobReference(CloudBlob blob)
+        public static CloudBlob GetCorrespondingTypeBlobReference(CloudBlob blob, OperationContext operationContext)
         {
             CloudBlob targetBlob;
             switch (blob.Properties.BlobType)
@@ -135,7 +135,7 @@ public static CloudBlob GetCorrespondingTypeBlobReference(CloudBlob blob)
 
             try
             {
-                Task.Run(() => targetBlob.FetchAttributesAsync()).Wait();
+                Task.Run(() => targetBlob.FetchAttributesAsync(null, null, operationContext)).Wait();
             }
             catch (AggregateException e) when (e.InnerException is StorageException)
             {

src/Storage/Storage/File/Cmdlet/GetAzureStorageShareStoredAccessPolicy.cs

@@ -70,7 +70,7 @@ internal async Task GetAzureShareStoredAccessPolicyAsync(long taskId, IStorageFi
         internal async Task<SharedAccessFilePolicies> GetPoliciesAsync(IStorageFileManagement localChannel, string shareName, string policyName)
         {
             CloudFileShare share = localChannel.GetShareReference(shareName);
-            FileSharePermissions permissions = await localChannel.GetSharePermissionsAsync(share, null, null, null, CmdletCancellationToken).ConfigureAwait(false);
+            FileSharePermissions permissions = await localChannel.GetSharePermissionsAsync(share, null, null, OperationContext, CmdletCancellationToken).ConfigureAwait(false);
             return permissions.SharedAccessPolicies;
         }
 

src/Storage/Storage/File/Cmdlet/NewAzureStorageShareStoredAccessPolicy.cs

@@ -93,7 +93,7 @@ public override void ExecuteCmdlet()
             //Set permissions back to container
             try
             {
-                Task.Run(() => fileShare.SetPermissionsAsync(fileSharePermissions)).Wait();
+                Task.Run(() => fileShare.SetPermissionsAsync(fileSharePermissions, null, null, OperationContext)).Wait();
             }
             catch (AggregateException e) when (e.InnerException is StorageException)
             {

src/Storage/Storage/File/Cmdlet/RemoveAzureStorageShareStoredAccessPolicy.cs

@@ -63,7 +63,7 @@ internal bool RemoveAzureShareStoredAccessPolicy(IStorageFileManagement localCha
             if (ShouldProcess(policyName, "Remove policy"))
             {
                 permissions.SharedAccessPolicies.Remove(policyName);
-                localChannel.SetSharePermissions(share, permissions);
+                localChannel.SetSharePermissions(share, permissions, null, null, OperationContext);
                 success = true;
             }
 

src/Storage/Storage/File/Cmdlet/SetAzureStorageShareStoredAccessPolicy.cs

@@ -74,7 +74,7 @@ internal string SetAzureShareStoredAccessPolicy(IStorageFileManagement localChan
             permissions.SharedAccessPolicies[policyName] = policy;
 
             //Set permission back to share
-            localChannel.SetSharePermissions(share, permissions);
+            localChannel.SetSharePermissions(share, permissions, null, null, OperationContext);
             WriteObject(AccessPolicyHelper.ConstructPolicyOutputPSObject<SharedAccessFilePolicy>(permissions.SharedAccessPolicies, policyName));
             return policyName;
         }

src/Storage/Storage/File/Cmdlet/StartAzureStorageFileCopy.cs

@@ -345,7 +345,7 @@ private void StartCopyFromFile()
                 taskId,
                 destFile,
                 () => this.ConfirmOverwrite(sourceFile.SnapshotQualifiedUri.ToString(), destFile.SnapshotQualifiedUri.ToString()),
-                () => destFile.StartCopyAsync(sourceFile.GenerateCopySourceFile()));
+                () => destFile.StartCopyAsync(sourceFile.GenerateCopySourceFile(), null, null, this.RequestOptions, this.OperationContext));
 
             this.RunTask(taskGenerator);
         }

src/Storage/Storage/Model/Contract/StorageBlobManagement.cs

@@ -577,7 +577,7 @@ public async Task<CloudBlob> GetBlobReferenceFromServerAsync(CloudBlobContainer
                 CloudBlob blob = container.GetBlobReference(blobName);
                 await blob.FetchAttributesAsync(accessCondition, options, operationContext, cancellationToken).ConfigureAwait(false);
 
-                return Util.GetCorrespondingTypeBlobReference(blob);
+                return Util.GetCorrespondingTypeBlobReference(blob, operationContext);
             }
             catch (XSCL.StorageException e)
             {

src/Storage/Storage/Queue/Cmdlet/GetAzureStorageQueueStoredAccessPolicy.cs

@@ -84,7 +84,7 @@ internal async Task GetAzureQueueStoredAccessPolicyAsync(long taskId, IStorageQu
         internal async Task<SharedAccessQueuePolicies> GetPoliciesAsync(IStorageQueueManagement localChannel, string queueName, string policyName)
         {
             CloudQueue queue = localChannel.GetQueueReference(queueName);
-            QueuePermissions queuePermissions = await localChannel.GetPermissionsAsync(queue).ConfigureAwait(false);
+            QueuePermissions queuePermissions = await localChannel.GetPermissionsAsync(queue, null, OperationContext).ConfigureAwait(false);
             return queuePermissions.SharedAccessPolicies;
         }
 

src/Storage/Storage/Queue/Cmdlet/NewAzureStorageQueueStoredAccessPolicy.cs

@@ -76,7 +76,7 @@ internal string CreateAzureQueueStoredAccessPolicy(IStorageQueueManagement local
 
             //Get existing permissions
             CloudQueue queue = Channel.GetQueueReference(queueName);
-            QueuePermissions queuePermissions = localChannel.GetPermissions(queue);
+            QueuePermissions queuePermissions = localChannel.GetPermissions(queue, this.RequestOptions, this.OperationContext);
 
             //Add new policy
             if (queuePermissions.SharedAccessPolicies.Keys.Contains(policyName))
@@ -89,7 +89,7 @@ internal string CreateAzureQueueStoredAccessPolicy(IStorageQueueManagement local
             queuePermissions.SharedAccessPolicies.Add(policyName, policy);
 
             //Set permissions back to queue
-            localChannel.SetPermissions(queue, queuePermissions);
+            localChannel.SetPermissions(queue, queuePermissions, null, OperationContext);
             return policyName;
         }
 

src/Storage/Storage/Queue/Cmdlet/RemoveAzureStorageQueueStoredAccessPolicy.cs

@@ -67,7 +67,7 @@ internal bool RemoveAzureQueueStoredAccessPolicy(IStorageQueueManagement localCh
 
             //Get existing permissions
             CloudQueue queue = Channel.GetQueueReference(queueName);
-            QueuePermissions queuePermissions = localChannel.GetPermissions(queue);
+            QueuePermissions queuePermissions = localChannel.GetPermissions(queue, this.RequestOptions, this.OperationContext);
 
             //remove the specified policy
             if (!queuePermissions.SharedAccessPolicies.Keys.Contains(policyName))
@@ -78,7 +78,7 @@ internal bool RemoveAzureQueueStoredAccessPolicy(IStorageQueueManagement localCh
             if (ShouldProcess(policyName, "Remove policy"))
             {
                 queuePermissions.SharedAccessPolicies.Remove(policyName);
-                localChannel.SetPermissions(queue, queuePermissions);
+                localChannel.SetPermissions(queue, queuePermissions, null, OperationContext);
                 success = true;
             }
 

src/Storage/Storage/Queue/Cmdlet/SetAzureStorageQueueStoredAccessPolicy.cs

@@ -76,7 +76,7 @@ internal string SetAzureQueueStoredAccessPolicy(IStorageQueueManagement localCha
         {
             //Get existing permissions
             CloudQueue queue = Channel.GetQueueReference(queueName);
-            QueuePermissions queuePermissions = localChannel.GetPermissions(queue);
+            QueuePermissions queuePermissions = localChannel.GetPermissions(queue, this.RequestOptions, this.OperationContext);
 
             //Set the policy with new value
             if (!queuePermissions.SharedAccessPolicies.Keys.Contains(policyName))
@@ -90,7 +90,7 @@ internal string SetAzureQueueStoredAccessPolicy(IStorageQueueManagement localCha
 
             //Set permission back to queue
             WriteObject(AccessPolicyHelper.ConstructPolicyOutputPSObject<SharedAccessQueuePolicy>(queuePermissions.SharedAccessPolicies, policyName));
-            localChannel.SetPermissions(queue, queuePermissions);
+            localChannel.SetPermissions(queue, queuePermissions, null, OperationContext);
             return policyName;
         }
 

src/Storage/Storage/Table/Cmdlet/GetAzureStorageTableStoredAccessPolicy.cs

@@ -84,7 +84,7 @@ internal async Task GetAzureTableStoredAccessPolicyAsync(long taskId, IStorageTa
         internal async Task<SharedAccessTablePolicies> GetPoliciesAsync(IStorageTableManagement localChannel, string tableName, string policyName)
         {
             CloudTable table = localChannel.GetTableReference(tableName);
-            TablePermissions tablePermissions = await localChannel.GetTablePermissionsAsync(table).ConfigureAwait(false);
+            TablePermissions tablePermissions = await localChannel.GetTablePermissionsAsync(table, null, TableOperationContext).ConfigureAwait(false);
             return tablePermissions.SharedAccessPolicies;
         }
 

src/Storage/Storage/Table/Cmdlet/NewAzureStorageTableStoredAccessPolicy.cs

@@ -77,7 +77,7 @@ internal string CreateAzureTableStoredAccessPolicy(IStorageTableManagement local
 
             //Get existing permissions
             CloudTable table = localChannel.GetTableReference(tableName);
-            TablePermissions tablePermissions = localChannel.GetTablePermissions(table);
+            TablePermissions tablePermissions = localChannel.GetTablePermissions(table, null, TableOperationContext);
 
             //Add new policy
             if (tablePermissions.SharedAccessPolicies.Keys.Contains(policyName))
@@ -90,7 +90,7 @@ internal string CreateAzureTableStoredAccessPolicy(IStorageTableManagement local
             tablePermissions.SharedAccessPolicies.Add(policyName, policy);
 
             //Set permissions back to table
-            localChannel.SetTablePermissions(table, tablePermissions);
+            localChannel.SetTablePermissions(table, tablePermissions, null, TableOperationContext);
             return policyName;
         }
 

src/Storage/Storage/Table/Cmdlet/RemoveAzureStorageTableStoredAccessPolicy.cs

@@ -67,7 +67,7 @@ internal bool RemoveAzureTableStoredAccessPolicy(IStorageTableManagement localCh
 
             //Get existing permissions
             CloudTable table = localChannel.GetTableReference(tableName);
-            TablePermissions tablePermissions = localChannel.GetTablePermissions(table);
+            TablePermissions tablePermissions = localChannel.GetTablePermissions(table, this.RequestOptions, this.TableOperationContext);
 
             //remove the specified policy
             if (!tablePermissions.SharedAccessPolicies.Keys.Contains(policyName))
@@ -78,7 +78,7 @@ internal bool RemoveAzureTableStoredAccessPolicy(IStorageTableManagement localCh
             if (ShouldProcess(policyName, "Remove policy"))
             {
                 tablePermissions.SharedAccessPolicies.Remove(policyName);
-                localChannel.SetTablePermissions(table, tablePermissions);
+                localChannel.SetTablePermissions(table, tablePermissions, null, TableOperationContext);
                 success = true;
             }
 

src/Storage/Storage/Table/Cmdlet/SetAzureStorageTableStoredAccessPolicy.cs

@@ -79,7 +79,7 @@ internal string SetAzureTableStoredAccessPolicy(IStorageTableManagement localCha
 
             //Get existing permissions
             CloudTable table = localChannel.GetTableReference(Table);
-            TablePermissions tablePermissions = localChannel.GetTablePermissions(table);
+            TablePermissions tablePermissions = localChannel.GetTablePermissions(table, this.RequestOptions, this.TableOperationContext);
 
             //Set the policy with new value
             if (!tablePermissions.SharedAccessPolicies.Keys.Contains(policyName))
@@ -92,7 +92,7 @@ internal string SetAzureTableStoredAccessPolicy(IStorageTableManagement localCha
             tablePermissions.SharedAccessPolicies[policyName] = policy;
 
             //Set permission back to table
-            localChannel.SetTablePermissions(table, tablePermissions);
+            localChannel.SetTablePermissions(table, tablePermissions, null, TableOperationContext);
             WriteObject(AccessPolicyHelper.ConstructPolicyOutputPSObject<SharedAccessTablePolicy>(tablePermissions.SharedAccessPolicies, policyName));
             return policyName;
         }