Include WorkflowInstanceId and TaskId in Payload

Signed-off-by: Victor Chang <[email protected]>

src/Api/Storage/Payload.cs

@@ -64,6 +64,10 @@ public enum PayloadState
 
         public string CorrelationId { get; init; }
 
+        public string? WorkflowInstanceId { get; init; }
+
+        public string? TaskId { get; init; }
+
         public int Count { get => Files.Count; }
 
         public bool HasTimedOut { get => ElapsedTime().TotalSeconds >= Timeout; }
@@ -80,13 +84,20 @@ public TimeSpan Elapsed
         public int FilesFailedToUpload { get => Files.Count(p => p.IsUploadFailed); }
 
         public Payload(string key, string correlationId, uint timeout)
+            : this(key, correlationId, null, null, timeout)
+        {
+        }
+
+        public Payload(string key, string correlationId, string? workflowInstanceId, string? taskId, uint timeout)
         {
             Guard.Against.NullOrWhiteSpace(key, nameof(key));
 
             Files = new List<FileStorageMetadata>();
             _lastReceived = new Stopwatch();
 
             CorrelationId = correlationId;
+            WorkflowInstanceId = workflowInstanceId;
+            TaskId = taskId;
             MachineName = Environment.MachineName;
             DateTimeCreated = DateTime.UtcNow;
             PayloadId = Guid.NewGuid();

src/Database/EntityFramework/Migrations/20230802003305_R4_0.4.0.cs → src/Database/EntityFramework/Migrations/20230808233742_R4_0.4.0.cs

@@ -8,6 +8,18 @@ public partial class R4_040 : Migration
     {
         protected override void Up(MigrationBuilder migrationBuilder)
         {
+            migrationBuilder.AddColumn<string>(
+                name: "TaskId",
+                table: "Payloads",
+                type: "TEXT",
+                nullable: true);
+
+            migrationBuilder.AddColumn<string>(
+                name: "WorkflowInstanceId",
+                table: "Payloads",
+                type: "TEXT",
+                nullable: true);
+
             migrationBuilder.AddColumn<string>(
                 name: "PluginAssemblies",
                 table: "MonaiApplicationEntities",
@@ -18,6 +30,14 @@ protected override void Up(MigrationBuilder migrationBuilder)
 
         protected override void Down(MigrationBuilder migrationBuilder)
         {
+            migrationBuilder.DropColumn(
+                name: "TaskId",
+                table: "Payloads");
+
+            migrationBuilder.DropColumn(
+                name: "WorkflowInstanceId",
+                table: "Payloads");
+
             migrationBuilder.DropColumn(
                 name: "PluginAssemblies",
                 table: "MonaiApplicationEntities");

src/Database/EntityFramework/Migrations/InformaticsGatewayContextModelSnapshot.cs

@@ -15,7 +15,7 @@ partial class InformaticsGatewayContextModelSnapshot : ModelSnapshot
         protected override void BuildModel(ModelBuilder modelBuilder)
         {
 #pragma warning disable 612, 618
-            modelBuilder.HasAnnotation("ProductVersion", "6.0.15");
+            modelBuilder.HasAnnotation("ProductVersion", "6.0.21");
 
             modelBuilder.Entity("Monai.Deploy.InformaticsGateway.Api.DestinationApplicationEntity", b =>
                 {
@@ -273,9 +273,15 @@ protected override void BuildModel(ModelBuilder modelBuilder)
                     b.Property<int>("State")
                         .HasColumnType("INTEGER");
 
+                    b.Property<string>("TaskId")
+                        .HasColumnType("TEXT");
+
                     b.Property<uint>("Timeout")
                         .HasColumnType("INTEGER");
 
+                    b.Property<string>("WorkflowInstanceId")
+                        .HasColumnType("TEXT");
+
                     b.HasKey("PayloadId");
 
                     b.HasIndex(new[] { "CorrelationId", "PayloadId" }, "idx_payload_ids")

src/Database/EntityFramework/Test/PayloadRepositoryTest.cs

@@ -61,7 +61,7 @@ public PayloadRepositoryTest(SqliteDatabaseFixture databaseFixture)
         [Fact]
         public async Task GivenAPayload_WhenAddingToDatabase_ExpectItToBeSaved()
         {
-            var payload = new Payload(Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), 5);
+            var payload = new Payload(Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), 5);
             payload.Add(new DicomFileStorageMetadata(Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), Guid.NewGuid().ToString()));
             payload.State = Payload.PayloadState.Move;
 
@@ -75,6 +75,8 @@ public async Task GivenAPayload_WhenAddingToDatabase_ExpectItToBeSaved()
             Assert.Equal(payload.Count, actual!.Count);
             Assert.Equal(payload.RetryCount, actual!.RetryCount);
             Assert.Equal(payload.CorrelationId, actual!.CorrelationId);
+            Assert.Equal(payload.WorkflowInstanceId, actual!.WorkflowInstanceId);
+            Assert.Equal(payload.TaskId, actual!.TaskId);
             Assert.Equal(payload.CalledAeTitle, actual!.CalledAeTitle);
             Assert.Equal(payload.CallingAeTitle, actual!.CallingAeTitle);
             Assert.Equal(payload.Timeout, actual!.Timeout);
@@ -84,7 +86,7 @@ public async Task GivenAPayload_WhenAddingToDatabase_ExpectItToBeSaved()
         [Fact]
         public async Task GivenAPayload_WhenRemoveIsCalled_ExpectItToDeleted()
         {
-            var payload = new Payload(Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), 5);
+            var payload = new Payload(Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), 5);
             payload.Add(new DicomFileStorageMetadata(Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), Guid.NewGuid().ToString()));
             payload.State = Payload.PayloadState.Move;
 

src/Database/MongoDB/Integration.Test/PayloadRepositoryTest.cs

@@ -64,7 +64,7 @@ public PayloadRepositoryTest(MongoDatabaseFixture databaseFixture)
         [Fact]
         public async Task GivenAPayload_WhenAddingToDatabase_ExpectItToBeSaved()
         {
-            var payload = new Payload(Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), 5);
+            var payload = new Payload(Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), 5);
             payload.Add(new DicomFileStorageMetadata(Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), Guid.NewGuid().ToString()));
             payload.State = Payload.PayloadState.Move;
 
@@ -80,6 +80,8 @@ public async Task GivenAPayload_WhenAddingToDatabase_ExpectItToBeSaved()
             Assert.Equal(payload.Count, actual!.Count);
             Assert.Equal(payload.RetryCount, actual!.RetryCount);
             Assert.Equal(payload.CorrelationId, actual!.CorrelationId);
+            Assert.Equal(payload.WorkflowInstanceId, actual!.WorkflowInstanceId);
+            Assert.Equal(payload.TaskId, actual!.TaskId);
             Assert.Equal(payload.CalledAeTitle, actual!.CalledAeTitle);
             Assert.Equal(payload.CallingAeTitle, actual!.CallingAeTitle);
             Assert.Equal(payload.Timeout, actual!.Timeout);
@@ -119,7 +121,7 @@ public async Task GivenDestinationApplicationEntitiesInTheDatabase_WhenToListIsC
         [Fact]
         public async Task GivenAPayload_WhenUpdateIsCalled_ExpectItToSaved()
         {
-            var payload = new Payload(Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), 5);
+            var payload = new Payload(Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), 5);
             payload.Add(new DicomFileStorageMetadata(Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), Guid.NewGuid().ToString()));
 
             var store = new PayloadRepository(_serviceScopeFactory.Object, _logger.Object, _options, _databaseFixture.Options);
@@ -139,6 +141,8 @@ public async Task GivenAPayload_WhenUpdateIsCalled_ExpectItToSaved()
             Assert.Equal(updated.Count, actual!.Count);
             Assert.Equal(updated.RetryCount, actual!.RetryCount);
             Assert.Equal(updated.CorrelationId, actual!.CorrelationId);
+            Assert.Equal(payload.WorkflowInstanceId, actual!.WorkflowInstanceId);
+            Assert.Equal(payload.TaskId, actual!.TaskId);
             Assert.Equal(updated.CalledAeTitle, actual!.CalledAeTitle);
             Assert.Equal(updated.CallingAeTitle, actual!.CallingAeTitle);
             Assert.Equal(updated.Timeout, actual!.Timeout);

src/InformaticsGateway/Services/Connectors/PayloadAssembler.cs

@@ -24,10 +24,8 @@
 using DotNext.Threading;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Logging;
-using Microsoft.Extensions.Options;
 using Monai.Deploy.InformaticsGateway.Api;
 using Monai.Deploy.InformaticsGateway.Api.Storage;
-using Monai.Deploy.InformaticsGateway.Configuration;
 using Monai.Deploy.InformaticsGateway.Database.Api.Repositories;
 using Monai.Deploy.InformaticsGateway.Logging;
 
@@ -40,7 +38,6 @@ namespace Monai.Deploy.InformaticsGateway.Services.Connectors
     internal sealed partial class PayloadAssembler : IPayloadAssembler, IDisposable
     {
         internal const int DEFAULT_TIMEOUT = 5;
-        private readonly IOptions<InformaticsGatewayConfiguration> _options;
         private readonly ILogger<PayloadAssembler> _logger;
         private readonly IServiceScopeFactory _serviceScopeFactory;
 
@@ -50,11 +47,9 @@ internal sealed partial class PayloadAssembler : IPayloadAssembler, IDisposable
         private readonly System.Timers.Timer _timer;
 
         public PayloadAssembler(
-            IOptions<InformaticsGatewayConfiguration> options,
             ILogger<PayloadAssembler> logger,
             IServiceScopeFactory serviceScopeFactory)
         {
-            _options = options ?? throw new ArgumentNullException(nameof(options));
             _logger = logger ?? throw new ArgumentNullException(nameof(logger));
             _serviceScopeFactory = serviceScopeFactory ?? throw new ArgumentNullException(nameof(serviceScopeFactory));
 
@@ -103,7 +98,7 @@ public async Task<Guid> Queue(string bucket, FileStorageMetadata file, uint time
 
             using var _ = _logger.BeginScope(new LoggingDataDictionary<string, object>() { { "CorrelationId", file.CorrelationId } });
 
-            var payload = await CreateOrGetPayload(bucket, file.CorrelationId, timeout).ConfigureAwait(false);
+            var payload = await CreateOrGetPayload(bucket, file.CorrelationId, file.WorkflowInstanceId, file.TaskId, timeout).ConfigureAwait(false);
             payload.Add(file);
             _logger.FileAddedToBucket(payload.Key, payload.Count);
             return payload.PayloadId;
@@ -127,7 +122,7 @@ private async void OnTimedEvent(Object source, System.Timers.ElapsedEventArgs e)
                 await _intializedTask.ConfigureAwait(false);
 
                 _timer.Enabled = false;
-                if (_payloads.Count > 0)
+                if (!_payloads.IsEmpty)
                 {
                     _logger.BucketsActive(_payloads.Count);
                 }
@@ -198,13 +193,13 @@ private async Task QueueBucketForNotification(string key, Payload payload)
             }
         }
 
-        private async Task<Payload> CreateOrGetPayload(string key, string correationId, uint timeout)
+        private async Task<Payload> CreateOrGetPayload(string key, string correlationId, string? workflowInstanceId, string? taskId, uint timeout)
         {
             return await _payloads.GetOrAdd(key, x => new AsyncLazy<Payload>(async () =>
             {
                 var scope = _serviceScopeFactory.CreateScope();
                 var repository = scope.ServiceProvider.GetRequiredService<IPayloadRepository>();
-                var newPayload = new Payload(key, correationId, timeout);
+                var newPayload = new Payload(key, correlationId, workflowInstanceId, taskId, timeout);
                 await repository.AddAsync(newPayload).ConfigureAwait(false);
                 _logger.BucketCreated(key, timeout);
                 return newPayload;

src/InformaticsGateway/Test/Services/Common/ExternalAppPluginTest.cs

@@ -108,6 +108,7 @@ public async Task ExternalAppPlugin_Should_Replace_StudyUid_Plus_SaveData()
             var dataset = new DicomDataset
             {
                 { DicomTag.PatientID, "PID" },
+                { DicomTag.AccessionNumber, "AccesssionNumber" },
                 { DicomTag.StudyInstanceUID, DicomUIDGenerator.GenerateDerivedFromUUID() },
                 { DicomTag.SeriesInstanceUID, DicomUIDGenerator.GenerateDerivedFromUUID() },
                 { DicomTag.SOPInstanceUID, DicomUIDGenerator.GenerateDerivedFromUUID() },

src/InformaticsGateway/Test/Services/Connectors/PayloadAssemblerTest.cs

@@ -67,15 +67,14 @@ public PayloadAssemblerTest()
         [Fact]
         public void GivenAPayloadAssembler_WhenInitialized_ExpectParametersToBeValidated()
         {
-            Assert.Throws<ArgumentNullException>(() => new PayloadAssembler(null, null, null));
-            Assert.Throws<ArgumentNullException>(() => new PayloadAssembler(_options, null, null));
-            Assert.Throws<ArgumentNullException>(() => new PayloadAssembler(_options, _logger.Object, null));
+            Assert.Throws<ArgumentNullException>(() => new PayloadAssembler(null, null));
+            Assert.Throws<ArgumentNullException>(() => new PayloadAssembler(_logger.Object, null));
         }
 
         [RetryFact(10, 200)]
         public async Task GivenAFileStorageMetadata_WhenQueueingWihtoutSpecifyingATimeout_ExpectDefaultTimeoutToBeUsed()
         {
-            var payloadAssembler = new PayloadAssembler(_options, _logger.Object, _serviceScopeFactory.Object);
+            var payloadAssembler = new PayloadAssembler(_logger.Object, _serviceScopeFactory.Object);
 
             _ = Assert.ThrowsAsync<OperationCanceledException>(async () => await Task.Run(() => payloadAssembler.Dequeue(_cancellationTokenSource.Token)));
 
@@ -91,7 +90,7 @@ public async Task GivenFileStorageMetadataInTheDatabase_AtServiceStartup_ExpectP
         {
             _repository.Setup(p => p.RemovePendingPayloadsAsync(It.IsAny<CancellationToken>()));
 
-            var payloadAssembler = new PayloadAssembler(_options, _logger.Object, _serviceScopeFactory.Object);
+            var payloadAssembler = new PayloadAssembler(_logger.Object, _serviceScopeFactory.Object);
             await Task.Delay(250);
             payloadAssembler.Dispose();
             _cancellationTokenSource.Cancel();
@@ -102,7 +101,7 @@ public async Task GivenFileStorageMetadataInTheDatabase_AtServiceStartup_ExpectP
         [RetryFact(10, 200)]
         public async Task GivenAPayloadAssembler_WhenDisposed_ExpectResourceToBeCleanedUp()
         {
-            var payloadAssembler = new PayloadAssembler(_options, _logger.Object, _serviceScopeFactory.Object);
+            var payloadAssembler = new PayloadAssembler(_logger.Object, _serviceScopeFactory.Object);
 
             _ = Assert.ThrowsAsync<OperationCanceledException>(async () => await Task.Run(() => payloadAssembler.Dequeue(_cancellationTokenSource.Token)));
 
@@ -118,7 +117,7 @@ public async Task GivenAPayloadAssembler_WhenDisposed_ExpectResourceToBeCleanedU
         [RetryFact(10, 200)]
         public async Task GivenAPayloadThatHasNotCompleteUploads_WhenProcessedByTimedEvent_ExpectToBeRemovedFromQueue()
         {
-            var payloadAssembler = new PayloadAssembler(_options, _logger.Object, _serviceScopeFactory.Object);
+            var payloadAssembler = new PayloadAssembler(_logger.Object, _serviceScopeFactory.Object);
 
             var file1 = new TestStorageInfo(Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), "file1", ".txt");
             var file2 = new TestStorageInfo(Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), "file1", ".txt");
@@ -138,7 +137,7 @@ public async Task GivenAPayloadThatHasNotCompleteUploads_WhenProcessedByTimedEve
         [RetryFact(10, 200)]
         public async Task GivenAPayloadThatHasCompletedUploads_WhenProcessedByTimedEvent_ExpectToBeAddedToQueue()
         {
-            var payloadAssembler = new PayloadAssembler(_options, _logger.Object, _serviceScopeFactory.Object);
+            var payloadAssembler = new PayloadAssembler(_logger.Object, _serviceScopeFactory.Object);
 
             var file = new TestStorageInfo(Guid.NewGuid().ToString(), Guid.NewGuid().ToString(), "file1", ".txt");
             file.File.SetUploaded("bucket");