From 4fd16b8e34d685239ea40fc3d8e9b02cc8ccf323 Mon Sep 17 00:00:00 2001
From: ndr_brt <andrea.bertagnolli@gmail.com>
Date: Wed, 17 Apr 2024 17:16:00 +0200
Subject: [PATCH] feat: add data map on TransformContext (#4120)

---
 .../edc/transform/TransformerContextImpl.java | 19 ++++++++
 .../transform/TransformerContextImplTest.java | 45 ++++++++++++++++++-
 .../edc/transform/spi/TransformerContext.java | 18 ++++++++
 3 files changed, 80 insertions(+), 2 deletions(-)

diff --git a/core/common/lib/transform-lib/src/main/java/org/eclipse/edc/transform/TransformerContextImpl.java b/core/common/lib/transform-lib/src/main/java/org/eclipse/edc/transform/TransformerContextImpl.java
index 0b4ff6d2a00..e14f96b9204 100644
--- a/core/common/lib/transform-lib/src/main/java/org/eclipse/edc/transform/TransformerContextImpl.java
+++ b/core/common/lib/transform-lib/src/main/java/org/eclipse/edc/transform/TransformerContextImpl.java
@@ -20,11 +20,16 @@
 import org.jetbrains.annotations.Nullable;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicReference;
 
 public class TransformerContextImpl implements TransformerContext {
     private final List<String> problems = new ArrayList<>();
     private final TypeTransformerRegistry registry;
+    private final Map<Class<?>, Map<String, AtomicReference<?>>> data = new HashMap<>();
 
     public TransformerContextImpl(TypeTransformerRegistry registry) {
         this.registry = registry;
@@ -70,4 +75,18 @@ public Class<?> typeAlias(String type, Class<?> defaultType) {
         return registry.typeAlias(type, defaultType);
     }
 
+    @Override
+    public void setData(Class<?> type, String key, Object value) {
+        data.computeIfAbsent(type, t -> new HashMap<>()).put(key, new AtomicReference<>(value));
+    }
+
+    @Override
+    public Object consumeData(Class<?> type, String key) {
+        return Optional.of(type)
+                .map(data::get)
+                .map(typeMap -> typeMap.get(key))
+                .map(reference -> reference.getAndSet(null))
+                .orElse(null);
+    }
+
 }
diff --git a/core/common/lib/transform-lib/src/test/java/org/eclipse/edc/transform/TransformerContextImplTest.java b/core/common/lib/transform-lib/src/test/java/org/eclipse/edc/transform/TransformerContextImplTest.java
index bddda886667..cb072dcd292 100644
--- a/core/common/lib/transform-lib/src/test/java/org/eclipse/edc/transform/TransformerContextImplTest.java
+++ b/core/common/lib/transform-lib/src/test/java/org/eclipse/edc/transform/TransformerContextImplTest.java
@@ -14,7 +14,9 @@
 
 package org.eclipse.edc.transform;
 
+import org.eclipse.edc.transform.spi.TransformerContext;
 import org.eclipse.edc.transform.spi.TypeTransformerRegistry;
+import org.junit.jupiter.api.Nested;
 import org.junit.jupiter.api.Test;
 
 import static org.assertj.core.api.Assertions.assertThat;
@@ -26,8 +28,8 @@
 
 class TransformerContextImplTest {
 
-    private final TypeTransformerRegistry registry = mock(TypeTransformerRegistry.class);
-    private final TransformerContextImpl context = new TransformerContextImpl(registry);
+    private final TypeTransformerRegistry registry = mock();
+    private final TransformerContext context = new TransformerContextImpl(registry);
 
     @Test
     void shouldReturnTransformedInput() {
@@ -55,4 +57,43 @@ void shouldNotTransform_whenInputIsNull() {
         assertThat(result).isNull();
         verifyNoInteractions(registry);
     }
+
+    @Nested
+    class Data {
+        @Test
+        void shouldClearData_whenConsumed() {
+            context.setData(Integer.class, "key", "value");
+
+            var data = context.consumeData(Integer.class, "key");
+
+            assertThat(data).isEqualTo("value");
+            assertThat(context.consumeData(Integer.class, "key")).isNull();
+        }
+
+        @Test
+        void shouldReturnNull_whenKeyDoesNotExist() {
+            context.setData(Integer.class, "key", "value");
+
+            var data = context.consumeData(Integer.class, "unexistent-key");
+
+            assertThat(data).isNull();
+        }
+
+        @Test
+        void shouldReturnNull_whenTypeDoesNotExist() {
+            var data = context.consumeData(Double.class, "any");
+
+            assertThat(data).isNull();
+        }
+
+        @Test
+        void shouldPermitMultipleDataPerType() {
+            context.setData(Integer.class, "key", "value");
+            context.setData(Integer.class, "another-key", "another-value");
+
+            var data = context.consumeData(Integer.class, "key");
+
+            assertThat(data).isEqualTo("value");
+        }
+    }
 }
diff --git a/spi/common/transform-spi/src/main/java/org/eclipse/edc/transform/spi/TransformerContext.java b/spi/common/transform-spi/src/main/java/org/eclipse/edc/transform/spi/TransformerContext.java
index 09032bdf30d..c4d0e19adfa 100644
--- a/spi/common/transform-spi/src/main/java/org/eclipse/edc/transform/spi/TransformerContext.java
+++ b/spi/common/transform-spi/src/main/java/org/eclipse/edc/transform/spi/TransformerContext.java
@@ -72,4 +72,22 @@ default Class<?> typeAlias(String type) {
     default Class<?> typeAlias(String type, Class<?> defaultType) {
         return defaultType;
     }
+
+    /**
+     * Set context data to be consumed by the transformer operating on type.
+     *
+     * @param type the type of the transformer that will consume the data.
+     * @param key the key.
+     * @param value the value.
+     */
+    void setData(Class<?> type, String key, Object value);
+
+    /**
+     * Consume context data. After consuming the data will be cleared.
+     *
+     * @param type the type of the transformer that consumes the data.
+     * @param key the key.
+     * @return the value.
+     */
+    Object consumeData(Class<?> type, String key);
 }