Add download range to Get-AzureBatchNodeFileContent

src/ResourceManager/AzureBatch/ChangeLog.md

@@ -18,9 +18,10 @@
         - Additional information about change #1
 -->
 ## Current Release
+- Added new Get-AzureBatchJobPreparationAndReleaseTaskStatus cmdlet.
+- Added byte range start and end to Get-AzureBatchNodeFileContent parameters.
 
 ## Version 4.1.0
-- Added new Get-AzureBatchJobPreparationAndReleaseTaskStatus cmdlet.
 
 ## Version 3.1.0
 

src/ResourceManager/AzureBatch/Commands.Batch.Test/BatchTestHelpers.cs

@@ -313,6 +313,20 @@ public static RequestInterceptor CreateFakeGetFileAndPropertiesFromComputeNodeRe
             return interceptor;
         }
 
+        public static RequestInterceptor ExamineRequestInterceptor<T>(Action<T> assertAction) where T : class, IBatchRequest
+        {
+            RequestInterceptor interceptor = new RequestInterceptor(baseRequest =>
+            {
+                var request = baseRequest as T;
+
+                if (request != null)
+                {
+                    assertAction(request);
+                }
+            });
+            return interceptor;
+        }
+
         /// <summary>
         /// Builds a CertificateGetResponse object
         /// </summary>

src/ResourceManager/AzureBatch/Commands.Batch.Test/Files/GetBatchNodeFileContentCommandTests.cs

@@ -14,6 +14,7 @@
 
 using Microsoft.Azure.Batch;
 using Microsoft.Azure.Batch.Protocol;
+using Microsoft.Azure.Batch.Protocol.BatchRequests;
 using Microsoft.WindowsAzure.Commands.ScenarioTest;
 using Moq;
 using System;
@@ -114,5 +115,54 @@ public void GetBatchNodeFileByComputeNodeContentParametersTest()
                 cmdlet.ExecuteCmdlet();
             }
         }
+
+        [Theory]
+        [InlineData(null, 14L)]
+        [InlineData(7L, 14L)]
+        [InlineData(7L, null)]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void GetBatchNodeFileByteRangeSet_IsPopulatedInRequest(long? rangeStart, long? rangeEnd)
+        {
+            // Setup cmdlet without required parameters
+            BatchAccountContext context = BatchTestHelpers.CreateBatchContextWithKeys();
+            cmdlet.BatchContext = context;
+            cmdlet.JobId = null;
+            cmdlet.TaskId = null;
+            cmdlet.Name = null;
+            cmdlet.InputObject = null;
+            cmdlet.DestinationPath = null;
+            cmdlet.ByteRangeStart = rangeStart;
+            cmdlet.ByteRangeEnd = rangeEnd;
+
+            string fileName = "stdout.txt";
+            bool hit = false;
+            // Don't go to the service on a Get NodeFile call or Get NodeFile Properties call
+            RequestInterceptor interceptor = BatchTestHelpers.CreateFakeGetFileAndPropertiesFromTaskResponseInterceptor(cmdlet.Name);
+            RequestInterceptor examiner = BatchTestHelpers.ExamineRequestInterceptor<FileGetFromTaskBatchRequest>(req =>
+            {
+                hit = true;
+                Assert.Equal($"bytes={cmdlet.ByteRangeStart}-{cmdlet.ByteRangeEnd}", req.Options.OcpRange);
+            });
+
+            cmdlet.AdditionalBehaviors = new List<BatchClientBehavior> { examiner, interceptor };
+
+            using (MemoryStream memStream = new MemoryStream())
+            {
+                // Don't hit the file system during unit tests
+                cmdlet.DestinationStream = memStream;
+
+                Assert.Throws<ArgumentException>(() => cmdlet.ExecuteCmdlet());
+
+                // Fill required task details
+                cmdlet.JobId = "job-1";
+                cmdlet.TaskId = "task";
+                cmdlet.Name = fileName;
+
+                // Verify no exceptions occur
+                cmdlet.ExecuteCmdlet();
+
+                Assert.True(hit);
+            }
+        }
     }
 }

src/ResourceManager/AzureBatch/Commands.Batch/Files/GetBatchNodeFileContentCommand.cs

@@ -83,10 +83,44 @@ public class GetBatchNodeFileContentCommand : BatchObjectModelCmdletBase
         [ValidateNotNullOrEmpty]
         public Stream DestinationStream { get; set; }
 
+        [Parameter(HelpMessage = "The start of the byte range to be downloaded. If this is not specified and ByteRangeEnd is, this defaults to 0.")]
+        [ValidateNotNullOrEmpty]
+        [ValidateRange(0, long.MaxValue)]
+        public long? ByteRangeStart { get; set; }
+
+        [Parameter(HelpMessage = "The end of the byte range to be downloaded. If this is not specified and ByteRangeStart is, this defaults to the end of the file.")]
+        [ValidateNotNullOrEmpty]
+        [ValidateRange(0, long.MaxValue)]
+        public long? ByteRangeEnd { get; set; }
+
         public override void ExecuteCmdlet()
         {
-            DownloadNodeFileOptions options = new DownloadNodeFileOptions(this.BatchContext, this.JobId, this.TaskId, this.PoolId,
-                this.ComputeNodeId, this.Name, this.InputObject, this.DestinationPath, this.DestinationStream, this.AdditionalBehaviors);
+            if (this.ByteRangeEnd != null && this.ByteRangeStart == null)
+            {
+                this.ByteRangeStart = 0;
+            }
+
+            if (this.ByteRangeEnd == null && this.ByteRangeStart != null)
+            {
+                this.ByteRangeEnd = long.MaxValue;
+            }
+
+            var byteRange = this.ByteRangeStart != null && this.ByteRangeEnd != null
+                ? new DownloadNodeFileOptions.ByteRange(this.ByteRangeStart.Value, this.ByteRangeEnd.Value)
+                : null;
+
+            DownloadNodeFileOptions options = new DownloadNodeFileOptions(
+                this.BatchContext,
+                this.JobId,
+                this.TaskId,
+                this.PoolId,
+                this.ComputeNodeId,
+                this.Name,
+                this.InputObject,
+                this.DestinationPath,
+                this.DestinationStream,
+                byteRange,
+                this.AdditionalBehaviors);
 
             BatchClient.DownloadNodeFile(options);
         }

src/ResourceManager/AzureBatch/Commands.Batch/Models/BatchClient.Files.cs

@@ -17,6 +17,9 @@
 using System;
 using System.Collections.Generic;
 using System.IO;
+using System.Linq;
+using Protocol = Microsoft.Azure.Batch.Protocol;
+using BatchRequests = Microsoft.Azure.Batch.Protocol.BatchRequests;
 using NodeFile = Microsoft.Azure.Batch.NodeFile;
 
 namespace Microsoft.Azure.Commands.Batch.Models
@@ -219,12 +222,36 @@ public void DownloadNodeFile(DownloadNodeFileOptions options)
                     }
             }
 
-            DownloadNodeFileByInstance(nodeFile, options.DestinationPath, options.Stream, options.AdditionalBehaviors);
+            DownloadNodeFileByInstance(nodeFile, options.DestinationPath, options.Stream, options.Range, options.AdditionalBehaviors);
         }
 
         // Downloads the file represented by an NodeFile instance to the specified path.
-        private void DownloadNodeFileByInstance(NodeFile file, string destinationPath, Stream stream, IEnumerable<BatchClientBehavior> additionalBehaviors = null)
+        private void DownloadNodeFileByInstance(NodeFile file, string destinationPath, Stream stream, DownloadNodeFileOptions.ByteRange byteRange, IEnumerable<BatchClientBehavior> additionalBehaviors = null)
         {
+            // TODO: Update this to use the new built in support in the C# SDK when we update the C# SDK to 6.x or later
+            Protocol.RequestInterceptor interceptor = new Protocol.RequestInterceptor(baseRequest =>
+            {
+                var fromTaskRequest = baseRequest as BatchRequests.FileGetFromTaskBatchRequest;
+                if (fromTaskRequest != null && byteRange != null)
+                {
+                    fromTaskRequest.Options.OcpRange = $"bytes={byteRange.Start}-{byteRange.End}";
+                }
+
+                var fromNodeRequest = baseRequest as BatchRequests.FileGetFromComputeNodeBatchRequest;
+                if (fromNodeRequest != null && byteRange != null)
+                {
+                    fromNodeRequest.Options.OcpRange = $"bytes={byteRange.Start}-{byteRange.End}";
+                }
+            });
+
+            additionalBehaviors = additionalBehaviors != null ? new List<BatchClientBehavior> { interceptor }.Union(additionalBehaviors) :
+                new List<BatchClientBehavior>() { interceptor };
+
+            if (byteRange != null)
+            {
+                WriteVerbose(string.Format(Resources.DownloadingNodeFileByteRange, byteRange.Start, byteRange.End));
+            }
+
             if (stream != null)
             {
                 // Don't dispose supplied Stream

src/ResourceManager/AzureBatch/Commands.Batch/Models/DownloadNodeFileOptions.cs

@@ -22,8 +22,30 @@ namespace Microsoft.Azure.Commands.Batch.Models
 {
     public class DownloadNodeFileOptions : NodeFileOperationParameters
     {
-        public DownloadNodeFileOptions(BatchAccountContext context, string jobId, string taskId, string poolId, string computeNodeId, string nodeFileName,
-            PSNodeFile nodeFile, string destinationPath, Stream stream, IEnumerable<BatchClientBehavior> additionalBehaviors = null)
+        public class ByteRange
+        {
+            public ByteRange(long start, long end)
+            {
+                this.Start = start;
+                this.End = end;
+            }
+
+            public long Start { get; }
+            public long End { get; }
+        }
+
+        public DownloadNodeFileOptions(
+            BatchAccountContext context,
+            string jobId,
+            string taskId,
+            string poolId,
+            string computeNodeId,
+            string nodeFileName,
+            PSNodeFile nodeFile,
+            string destinationPath,
+            Stream stream,
+            ByteRange byteRange,
+            IEnumerable<BatchClientBehavior> additionalBehaviors = null)
             : base(context, jobId, taskId, poolId, computeNodeId, nodeFileName, nodeFile, additionalBehaviors)
         {
             if (string.IsNullOrWhiteSpace(destinationPath) && stream == null)
@@ -32,6 +54,7 @@ public DownloadNodeFileOptions(BatchAccountContext context, string jobId, string
             }
 
             this.DestinationPath = destinationPath;
+            this.Range = byteRange;
             this.Stream = stream;
         }
 
@@ -40,6 +63,11 @@ public DownloadNodeFileOptions(BatchAccountContext context, string jobId, string
         /// </summary>
         public string DestinationPath { get; set; }
 
+        /// <summary>
+        /// The byte range to be downloaded.
+        /// </summary>
+        public ByteRange Range { get; }
+
         /// <summary>
         /// The Stream into which the node file data will be written. This stream will not be closed or rewound by this call.
         /// </summary>

src/ResourceManager/AzureBatch/Commands.Batch/Properties/Resources.resx

@@ -484,4 +484,7 @@
   <data name="GetJobPreparationAndReleaseStatus" xml:space="preserve">
     <value>Getting job preparation and release status for job "{0}"</value>
   </data>
+  <data name="DownloadingNodeFileByteRange" xml:space="preserve">
+    <value>Downloading byte range {0} to {1}</value>
+  </data>
 </root>

src/ResourceManager/AzureBatch/Commands.Batch/help/Get-AzureBatchNodeFileContent.md

@@ -15,37 +15,39 @@ Gets a Batch node file.
 ### Task_Id_Path
 ```
 Get-AzureBatchNodeFileContent -JobId <String> -TaskId <String> [-Name] <String> -DestinationPath <String>
- -BatchContext <BatchAccountContext> [<CommonParameters>]
+ [-ByteRangeStart <Int64>] [-ByteRangeEnd <Int64>] -BatchContext <BatchAccountContext> [<CommonParameters>]
 ```
 
 ### Task_Id_Stream
 ```
 Get-AzureBatchNodeFileContent -JobId <String> -TaskId <String> [-Name] <String> -DestinationStream <Stream>
- -BatchContext <BatchAccountContext> [<CommonParameters>]
+ [-ByteRangeStart <Int64>] [-ByteRangeEnd <Int64>] -BatchContext <BatchAccountContext> [<CommonParameters>]
 ```
 
 ### ComputeNode_Id_Path
 ```
 Get-AzureBatchNodeFileContent [-PoolId] <String> [-ComputeNodeId] <String> [-Name] <String>
- -DestinationPath <String> -BatchContext <BatchAccountContext> [<CommonParameters>]
+ -DestinationPath <String> [-ByteRangeStart <Int64>] [-ByteRangeEnd <Int64>]
+ -BatchContext <BatchAccountContext> [<CommonParameters>]
 ```
 
 ### ComputeNode_Id_Stream
 ```
 Get-AzureBatchNodeFileContent [-PoolId] <String> [-ComputeNodeId] <String> [-Name] <String>
- -DestinationStream <Stream> -BatchContext <BatchAccountContext> [<CommonParameters>]
+ -DestinationStream <Stream> [-ByteRangeStart <Int64>] [-ByteRangeEnd <Int64>]
+ -BatchContext <BatchAccountContext> [<CommonParameters>]
 ```
 
 ### InputObject_Path
 ```
-Get-AzureBatchNodeFileContent [[-InputObject] <PSNodeFile>] -DestinationPath <String>
- -BatchContext <BatchAccountContext> [<CommonParameters>]
+Get-AzureBatchNodeFileContent [[-InputObject] <PSNodeFile>] -DestinationPath <String> [-ByteRangeStart <Int64>]
+ [-ByteRangeEnd <Int64>] -BatchContext <BatchAccountContext> [<CommonParameters>]
 ```
 
 ### InputObject_Stream
 ```
 Get-AzureBatchNodeFileContent [[-InputObject] <PSNodeFile>] -DestinationStream <Stream>
- -BatchContext <BatchAccountContext> [<CommonParameters>]
+ [-ByteRangeStart <Int64>] [-ByteRangeEnd <Int64>] -BatchContext <BatchAccountContext> [<CommonParameters>]
 ```
 
 ## DESCRIPTION
@@ -130,6 +132,34 @@ Accept pipeline input: True (ByValue)
 Accept wildcard characters: False
 ```
 
+### -ByteRangeEnd
+The end of the byte range to be downloaded.
+```yaml
+Type: Int64
+Parameter Sets: (All)
+Aliases: 
+
+Required: False
+Position: Named
+Default value: None
+Accept pipeline input: False
+Accept wildcard characters: False
+```
+
+### -ByteRangeStart
+The start of the byte range to be downloaded.
+```yaml
+Type: Int64
+Parameter Sets: (All)
+Aliases: 
+
+Required: False
+Position: Named
+Default value: None
+Accept pipeline input: False
+Accept wildcard characters: False
+```
+
 ### -ComputeNodeId
 Specifies the ID of the compute node that contains the node file that this cmdlet returns.