JAVA-2262: Allow user to store custom state in MapperContext

olim7t · olim7t · commit bc409445ccd2 · 2019-06-21T15:13:58.000-07:00
diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/mapper/QueryProviderIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/mapper/QueryProviderIT.java
@@ -41,6 +41,7 @@
 import com.datastax.oss.driver.api.testinfra.session.SessionRule;
 import com.datastax.oss.driver.categories.ParallelizableTests;
 import java.util.Objects;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.junit.BeforeClass;
 import org.junit.ClassRule;
 import org.junit.Test;
@@ -56,6 +57,10 @@ public class QueryProviderIT {
   private static SessionRule<CqlSession> sessionRule = SessionRule.builder(ccm).build();
   @ClassRule public static TestRule chain = RuleChain.outerRule(ccm).around(sessionRule);
 
+  // Dummy counter to exercize the "custom state" feature: it gets incremented each time the query
+  // provider is called.
+  private static AtomicInteger executionCount = new AtomicInteger();
+
   private static SensorDao dao;
 
   @BeforeClass
@@ -70,7 +75,10 @@ public static void setup() {
             .setExecutionProfile(sessionRule.slowProfile())
             .build());
 
-    SensorMapper mapper = new QueryProviderIT_SensorMapperBuilder(session).build();
+    SensorMapper mapper =
+        new QueryProviderIT_SensorMapperBuilder(session)
+            .withCustomState("executionCount", executionCount)
+            .build();
     dao = mapper.sensorDao(sessionRule.keyspace());
   }
 
@@ -85,11 +93,15 @@ public void should_invoke_query_provider() {
     dao.save(readingFeb1);
     dao.save(readingJan31);
 
+    assertThat(executionCount.get()).isEqualTo(0);
+
     assertThat(dao.findSlice(1, null, null).all())
         .containsExactly(readingFeb3, readingFeb2, readingFeb1, readingJan31);
     assertThat(dao.findSlice(1, 2, null).all())
         .containsExactly(readingFeb3, readingFeb2, readingFeb1);
     assertThat(dao.findSlice(1, 2, 3).all()).containsExactly(readingFeb3);
+
+    assertThat(executionCount.get()).isEqualTo(3);
   }
 
   @Mapper
@@ -109,12 +121,14 @@ public interface SensorDao {
 
   public static class FindSliceProvider {
     private final CqlSession session;
+    private final AtomicInteger executionCount;
     private final EntityHelper<SensorReading> sensorReadingHelper;
     private final Select selectStart;
 
     public FindSliceProvider(
         MapperContext context, EntityHelper<SensorReading> sensorReadingHelper) {
       this.session = context.getSession();
+      this.executionCount = ((AtomicInteger) context.getCustomState().get("executionCount"));
       this.sensorReadingHelper = sensorReadingHelper;
       this.selectStart =
           sensorReadingHelper.selectStart().whereColumn("id").isEqualTo(bindMarker());
@@ -125,6 +139,8 @@ public PagingIterable<SensorReading> findSlice(int id, Integer month, Integer da
         throw new IllegalArgumentException("Can't specify day if month is null");
       }
 
+      executionCount.incrementAndGet();
+
       Select select = this.selectStart;
       if (month != null) {
         select = select.whereColumn("month").isEqualTo(bindMarker());
diff --git a/mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/mapper/MapperBuilderGenerator.java b/mapper-processor/src/main/java/com/datastax/oss/driver/internal/mapper/processor/mapper/MapperBuilderGenerator.java
@@ -70,7 +70,8 @@ protected JavaFile.Builder getContents() {
                     .addModifiers(Modifier.PUBLIC)
                     .addAnnotation(Override.class)
                     .returns(ClassName.get(interfaceElement))
-                    .addStatement("$1T context = new $1T(session)", DefaultMapperContext.class)
+                    .addStatement(
+                        "$1T context = new $1T(session, customState)", DefaultMapperContext.class)
                     .addStatement(
                         "return new $T(context)",
                         GeneratedNames.mapperImplementation(interfaceElement))
diff --git a/mapper-runtime/src/main/java/com/datastax/oss/driver/api/mapper/MapperBuilder.java b/mapper-runtime/src/main/java/com/datastax/oss/driver/api/mapper/MapperBuilder.java
@@ -17,6 +17,11 @@
 
 import com.datastax.oss.driver.api.core.CqlSession;
 import com.datastax.oss.driver.api.mapper.annotations.Mapper;
+import com.datastax.oss.driver.api.mapper.annotations.QueryProvider;
+import edu.umd.cs.findbugs.annotations.NonNull;
+import edu.umd.cs.findbugs.annotations.Nullable;
+import java.util.HashMap;
+import java.util.Map;
 
 /**
  * Builds an instance of a {@link Mapper}-annotated interface wrapping a {@link CqlSession}.
@@ -27,9 +32,27 @@
 public abstract class MapperBuilder<MapperT> {
 
   protected final CqlSession session;
+  protected Map<Object, Object> customState;
 
   protected MapperBuilder(CqlSession session) {
     this.session = session;
+    this.customState = new HashMap<>();
+  }
+
+  /**
+   * Stores custom state that will be propagated to {@link MapperContext#getCustomState()}.
+   *
+   * <p>This is intended mainly for {@link QueryProvider} methods: since provider classes are
+   * instantiated directly by the generated mapper code, they have no way to access non-static state
+   * from the rest of your application. This method allows you to pass that state while building the
+   * mapper, and access it later at runtime.
+   *
+   * <p>Note that this state will be accessed concurrently, it should be thread-safe.
+   */
+  @NonNull
+  public MapperBuilder<MapperT> withCustomState(@Nullable Object key, @Nullable Object value) {
+    customState.put(key, value);
+    return this;
   }
 
   public abstract MapperT build();
diff --git a/mapper-runtime/src/main/java/com/datastax/oss/driver/api/mapper/MapperContext.java b/mapper-runtime/src/main/java/com/datastax/oss/driver/api/mapper/MapperContext.java
@@ -20,6 +20,7 @@
 import com.datastax.oss.driver.api.mapper.entity.naming.NameConverter;
 import edu.umd.cs.findbugs.annotations.NonNull;
 import edu.umd.cs.findbugs.annotations.Nullable;
+import java.util.Map;
 
 /**
  * A runtime context that gets passed from the mapper to DAO components to share global resources
@@ -52,4 +53,13 @@ public interface MapperContext {
    */
   @NonNull
   NameConverter getNameConverter(Class<? extends NameConverter> converterClass);
+
+  /**
+   * Retrieves any custom state that was set while building the mapper with {@link
+   * MapperBuilder#withCustomState(Object, Object)}.
+   *
+   * <p>The returned map is immutable. If no state was set on the builder, it will be empty.
+   */
+  @NonNull
+  Map<Object, Object> getCustomState();
 }
diff --git a/mapper-runtime/src/main/java/com/datastax/oss/driver/api/mapper/annotations/QueryProvider.java b/mapper-runtime/src/main/java/com/datastax/oss/driver/api/mapper/annotations/QueryProvider.java
@@ -15,6 +15,7 @@
  */
 package com.datastax.oss.driver.api.mapper.annotations;
 
+import com.datastax.oss.driver.api.mapper.MapperBuilder;
 import com.datastax.oss.driver.api.mapper.MapperContext;
 import com.datastax.oss.driver.api.mapper.entity.EntityHelper;
 import java.lang.annotation.ElementType;
@@ -52,6 +53,8 @@
  *
  * <p>The parameters and return type are completely free-form, as long as they match those of the
  * provider method.
+ *
+ * @see MapperBuilder#withCustomState(Object, Object)
  */
 @Target(ElementType.METHOD)
 @Retention(RetentionPolicy.CLASS)
diff --git a/mapper-runtime/src/main/java/com/datastax/oss/driver/internal/mapper/DefaultMapperContext.java b/mapper-runtime/src/main/java/com/datastax/oss/driver/internal/mapper/DefaultMapperContext.java
@@ -19,9 +19,11 @@
 import com.datastax.oss.driver.api.core.CqlSession;
 import com.datastax.oss.driver.api.mapper.MapperContext;
 import com.datastax.oss.driver.api.mapper.entity.naming.NameConverter;
+import com.datastax.oss.protocol.internal.util.collection.NullAllowingImmutableMap;
 import edu.umd.cs.findbugs.annotations.NonNull;
 import edu.umd.cs.findbugs.annotations.Nullable;
 import java.lang.reflect.InvocationTargetException;
+import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
@@ -32,28 +34,38 @@ public class DefaultMapperContext implements MapperContext {
   private final CqlIdentifier keyspaceId;
   private final CqlIdentifier tableId;
   private final ConcurrentMap<Class<? extends NameConverter>, NameConverter> nameConverterCache;
+  private final Map<Object, Object> customState;
 
-  public DefaultMapperContext(@NonNull CqlSession session) {
-    this(session, null, null, new ConcurrentHashMap<>());
+  public DefaultMapperContext(
+      @NonNull CqlSession session, @NonNull Map<Object, Object> customState) {
+    this(
+        session,
+        null,
+        null,
+        new ConcurrentHashMap<>(),
+        NullAllowingImmutableMap.copyOf(customState));
   }
 
   private DefaultMapperContext(
       CqlSession session,
       CqlIdentifier keyspaceId,
       CqlIdentifier tableId,
-      ConcurrentMap<Class<? extends NameConverter>, NameConverter> nameConverterCache) {
+      ConcurrentMap<Class<? extends NameConverter>, NameConverter> nameConverterCache,
+      Map<Object, Object> customState) {
     this.session = session;
     this.keyspaceId = keyspaceId;
     this.tableId = tableId;
     this.nameConverterCache = nameConverterCache;
+    this.customState = customState;
   }
 
   public DefaultMapperContext withKeyspaceAndTable(
       @Nullable CqlIdentifier newKeyspaceId, @Nullable CqlIdentifier newTableId) {
     return (Objects.equals(newKeyspaceId, this.keyspaceId)
             && Objects.equals(newTableId, this.tableId))
         ? this
-        : new DefaultMapperContext(session, newKeyspaceId, newTableId, nameConverterCache);
+        : new DefaultMapperContext(
+            session, newKeyspaceId, newTableId, nameConverterCache, customState);
   }
 
   @NonNull
@@ -81,6 +93,12 @@ public NameConverter getNameConverter(Class<? extends NameConverter> converterCl
         converterClass, DefaultMapperContext::buildNameConverter);
   }
 
+  @NonNull
+  @Override
+  public Map<Object, Object> getCustomState() {
+    return customState;
+  }
+
   private static NameConverter buildNameConverter(Class<? extends NameConverter> converterClass) {
     try {
       return converterClass.getDeclaredConstructor().newInstance();
