From 26f41dab75ef5b676daf451011cd9366617e6ba9 Mon Sep 17 00:00:00 2001
From: Greg Burd <gburd@onshape.com>
Date: Wed, 17 Jan 2018 12:38:33 -0500
Subject: [PATCH] Add notion of statement cache to UnitOfWork.  Ignore call to
 useKeyspace when session isn't valid.

---
 .../net/helenus/core/AbstractUnitOfWork.java  | 33 ++++++++++++++++++-
 .../net/helenus/core/SessionInitializer.java  |  6 ++--
 .../java/net/helenus/core/UnitOfWork.java     |  5 +++
 3 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/src/main/java/net/helenus/core/AbstractUnitOfWork.java b/src/main/java/net/helenus/core/AbstractUnitOfWork.java
index e9d439b..52134d2 100644
--- a/src/main/java/net/helenus/core/AbstractUnitOfWork.java
+++ b/src/main/java/net/helenus/core/AbstractUnitOfWork.java
@@ -24,6 +24,7 @@ import com.google.common.collect.TreeTraverser;
 import java.io.Serializable;
 import java.util.*;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.stream.Collectors;
@@ -48,6 +49,7 @@ public abstract class AbstractUnitOfWork<E extends Exception>
   private final HelenusSession session;
   private final AbstractUnitOfWork<E> parent;
   private final Table<String, String, Either<Object, List<Facet>>> cache = HashBasedTable.create();
+  private final Map<String, Object> statementCache = new ConcurrentHashMap<String, Object>();
   protected String purpose;
   protected List<String> nestedPurposes = new ArrayList<String>();
   protected String info;
@@ -203,10 +205,23 @@ public abstract class AbstractUnitOfWork<E extends Exception>
     }
   }
 
+  @Override
+  public Optional<Object> cacheLookup(String key) {
+      AbstractUnitOfWork self = this;
+      do {
+          Object result = self.statementCache.get(key);
+          if (result != null) {
+              return result == deleted ? Optional.ofNullable(null) : Optional.of(result);
+          }
+          self = self.parent;
+      } while (self != null);
+      return Optional.empty();
+  }
+
   @Override
   public Optional<Object> cacheLookup(List<Facet> facets) {
     String tableName = CacheUtil.schemaName(facets);
-    Optional<Object> result = Optional.empty();
+      Optional<Object> result = Optional.empty();
     for (Facet facet : facets) {
       if (!facet.fixed()) {
         String columnName = facet.name() + "==" + facet.value();
@@ -243,6 +258,16 @@ public abstract class AbstractUnitOfWork<E extends Exception>
     return result;
   }
 
+  @Override
+  public void cacheEvict(String key) {
+      statementCache.remove(key);
+  }
+
+  @Override
+  public void cacheDelete(String key) {
+      statementCache.replace(key, deleted);
+  }
+
   @Override
   public List<Facet> cacheEvict(List<Facet> facets) {
     Either<Object, List<Facet>> deletedObjectFacets = Either.right(facets);
@@ -279,6 +304,11 @@ public abstract class AbstractUnitOfWork<E extends Exception>
     return facets;
   }
 
+  @Override
+  public Object cacheUpdate(String key, Object value) {
+      return statementCache.replace(key, value);
+  }
+
   @Override
   public Object cacheUpdate(Object value, List<Facet> facets) {
     Object result = null;
@@ -378,6 +408,7 @@ public abstract class AbstractUnitOfWork<E extends Exception>
       } else {
 
         // Merge cache and statistics into parent if there is one.
+        parent.statementCache.putAll(statementCache);
         parent.mergeCache(cache);
         parent.addBatched(batch);
         if (purpose != null) {
diff --git a/src/main/java/net/helenus/core/SessionInitializer.java b/src/main/java/net/helenus/core/SessionInitializer.java
index 579a1c6..0d37d21 100644
--- a/src/main/java/net/helenus/core/SessionInitializer.java
+++ b/src/main/java/net/helenus/core/SessionInitializer.java
@@ -273,8 +273,10 @@ public final class SessionInitializer extends AbstractSessionOperations {
   }
 
   public SessionInitializer use(String keyspace) {
-    session.execute(SchemaUtil.use(keyspace, false));
-    this.usingKeyspace = keyspace;
+    if (session != null) {
+      session.execute(SchemaUtil.use(keyspace, false));
+      this.usingKeyspace = keyspace;
+    }
     return this;
   }
 
diff --git a/src/main/java/net/helenus/core/UnitOfWork.java b/src/main/java/net/helenus/core/UnitOfWork.java
index 3c41717..f5d70e6 100644
--- a/src/main/java/net/helenus/core/UnitOfWork.java
+++ b/src/main/java/net/helenus/core/UnitOfWork.java
@@ -61,12 +61,17 @@ public interface UnitOfWork<X extends Exception> extends AutoCloseable {
 
   void addFuture(CompletableFuture<?> future);
 
+  Optional<Object> cacheLookup(String key);
   Optional<Object> cacheLookup(List<Facet> facets);
 
+  Object cacheUpdate(String key, Object value);
   Object cacheUpdate(Object pojo, List<Facet> facets);
 
+  void cacheEvict(String key);
   List<Facet> cacheEvict(List<Facet> facets);
 
+  public void cacheDelete(String key);
+
   String getPurpose();
 
   UnitOfWork setPurpose(String purpose);