Adding bucket_id support

src/main/java/io/appform/dropwizard/sharding/DBShardingBundleBase.java

@@ -40,6 +40,8 @@
 import io.appform.dropwizard.sharding.metrics.TransactionMetricManager;
 import io.appform.dropwizard.sharding.metrics.TransactionMetricObserver;
 import io.appform.dropwizard.sharding.observers.TransactionObserver;
+import io.appform.dropwizard.sharding.observers.bucket.BucketIdObserver;
+import io.appform.dropwizard.sharding.observers.bucket.BucketIdSaver;
 import io.appform.dropwizard.sharding.observers.internal.FilteringObserver;
 import io.appform.dropwizard.sharding.observers.internal.ListenerTriggeringObserver;
 import io.appform.dropwizard.sharding.observers.internal.TerminalTransactionObserver;
@@ -48,6 +50,7 @@
 import io.appform.dropwizard.sharding.sharding.ShardBlacklistingStore;
 import io.appform.dropwizard.sharding.sharding.ShardManager;
 import io.appform.dropwizard.sharding.sharding.impl.ConsistentHashBucketIdExtractor;
+import io.appform.dropwizard.sharding.utils.BucketCalculator;
 import io.appform.dropwizard.sharding.utils.ShardCalculator;
 import io.dropwizard.Configuration;
 import io.dropwizard.ConfiguredBundle;
@@ -79,7 +82,7 @@
  * Base for bundles. This cannot be used by clients. Use one of the derived classes.
  */
 @Slf4j
 public abstract class DBShardingBundleBase<T extends Configuration> implements ConfiguredBundle<T> {
 
     private static final String DEFAULT_NAMESPACE = "default";
     private static final String SHARD_ENV = "db.shards";
@@ -141,7 +144,7 @@
 
     public List<Class<?>> getInitialisedEntities() {
         if(this.initialisedEntities == null){
             throw new RuntimeException("DB sharding bundle is not initialised !");
         }
         return this.initialisedEntities;
     }
@@ -178,7 +181,7 @@
     public void run(T configuration, Environment environment) {
         val shardConfigurationListSize = getConfig(configuration).getShards().size();
         if (numShards != shardConfigurationListSize) {
             throw new RuntimeException(
                     "Shard count provided through environment does not match the size of the shard configuration list");
         }
         sessionFactories = shardBundles.stream().map(HibernateBundle::getSessionFactory).collect(Collectors.toList());
@@ -227,13 +230,13 @@
                 hibernateBundle.run(configuration, environment);
             } catch (Exception e) {
                 log.error("Error initializing db sharding bundle", e);
                 throw new RuntimeException(e);
             }
         });
     }
 
     @VisibleForTesting
     public void initBundles(Bootstrap bootstrap) {
         shardBundles.forEach(hibernameBundle -> initialize(bootstrap));
     }
 
@@ -253,10 +256,14 @@
     }
 
     private ShardingBundleOptions getShardingOptions(T configuration) {
         val shardingOptions = getConfig(configuration).getShardingOptions();
         return Objects.nonNull(shardingOptions) ? shardingOptions : new ShardingBundleOptions();
     }
 
+    public BucketCalculator<String> bucketCalculator() {
+        return new BucketCalculator<>(new ConsistentHashBucketIdExtractor<>(this.shardManager));
+    }
+
     public <EntityType, T extends Configuration>
     LookupDao<EntityType> createParentObjectDao(Class<EntityType> clazz) {
         return new LookupDao<>(this.sessionFactories, clazz,
@@ -293,7 +300,7 @@
                 rootObserver);
     }
 
     public <EntityType, T extends Configuration>
     CacheableLookupDao<EntityType> createParentObjectDao(
             Class<EntityType> clazz,
             BucketIdExtractor<String> bucketIdExtractor,
@@ -308,7 +315,7 @@
     }
 
 
     public <EntityType, T extends Configuration>
     RelationalDao<EntityType> createRelatedObjectDao(Class<EntityType> clazz) {
         return new RelationalDao<>(this.sessionFactories, clazz,
                 new ShardCalculator<>(this.shardManager,
@@ -334,7 +341,7 @@
     }
 
 
     public <EntityType, T extends Configuration>
     RelationalDao<EntityType> createRelatedObjectDao(
             Class<EntityType> clazz,
             BucketIdExtractor<String> bucketIdExtractor) {
@@ -346,7 +353,7 @@
                 rootObserver);
     }
 
     public <EntityType, T extends Configuration>
     CacheableRelationalDao<EntityType> createRelatedObjectDao(
             Class<EntityType> clazz,
             BucketIdExtractor<String> bucketIdExtractor,
@@ -369,7 +376,7 @@
                         new ConsistentHashBucketIdExtractor<>(this.shardManager)));
     }
 
     public <EntityType, DaoType extends AbstractDAO<EntityType>, T extends Configuration>
     WrapperDao<EntityType, DaoType> createWrapperDao(
             Class<DaoType> daoTypeClass,
             BucketIdExtractor<String> bucketIdExtractor) {
@@ -400,6 +407,7 @@
             }
             this.rootObserver = observer.setNext(rootObserver);
         }
+        rootObserver =  new BucketIdObserver(new BucketIdSaver(bucketCalculator()));
         rootObserver = new TransactionMetricObserver(new TransactionMetricManager(getMetricConfig(config),
                 metricRegistry)).setNext(rootObserver);
         rootObserver = new FilteringObserver(rootObserver).addFilters(filters);
@@ -408,11 +416,11 @@
         log.debug("Observer chain");
         rootObserver.visit(observer -> {
             log.debug(" Observer: {}", observer.getClass().getSimpleName());
             if (observer instanceof FilteringObserver) {
                 log.debug("  Filters:");
                 ((FilteringObserver) observer).getFilters().forEach(filter -> log.debug("    - {}", filter.getClass().getSimpleName()));
             }
             if (observer instanceof ListenerTriggeringObserver) {
                 log.debug("  Listeners:");
                 ((ListenerTriggeringObserver) observer).getListeners().forEach(filter -> log.debug("    - {}", filter.getClass().getSimpleName()));
             }

src/main/java/io/appform/dropwizard/sharding/dao/operations/OpContext.java

@@ -33,41 +33,41 @@
 
     <T, R> P visit(Get<T, R> opContext);
 
-    <T> P visit(GetAndUpdate<T> opContext);
+    <T> P visit(GetAndUpdate<T> opContext) throws Exception;
 
     <T, R> P visit(GetByLookupKey<T, R> opContext);
 
-    <T> P visit(GetAndUpdateByLookupKey<T> opContext);
+    <T> P visit(GetAndUpdateByLookupKey<T> opContext) throws Exception;
 
     <T> P visit(ReadOnlyForLookupDao<T> opContext);
 
     <T> P visit(ReadOnlyForRelationalDao<T> opContext);
 
-    <T> P visit(LockAndExecute<T> opContext);
+    <T> P visit(LockAndExecute<T> opContext) throws Exception;
 
-    P visit(UpdateByQuery opContext);
+    P visit(UpdateByQuery opContext) throws Exception;
 
-    <T> P visit(UpdateWithScroll<T> opContext);
+    <T> P visit(UpdateWithScroll<T> opContext) throws Exception;
 
-    <T> P visit(UpdateAll<T> opContext);
+    <T> P visit(UpdateAll<T> opContext) throws Exception;
 
-    <T> P visit(SelectAndUpdate<T> opContext);
+    <T> P visit(SelectAndUpdate<T> opContext) throws Exception;
 
     <T> P visit(RunInSession<T> opContext);
 
     <T> P visit(RunWithCriteria<T> opContext);
 
     P visit(DeleteByLookupKey opContext);
 
-    <U, V> P visit(Save<U, V> opContext);
+    <U, V> P visit(Save<U, V> opContext) throws Exception;
 
-    <T> P visit(SaveAll<T> opContext);
+    <T> P visit(SaveAll<T> opContext) throws Exception;
 
-    <T> P visit(CreateOrUpdateByLookupKey<T> opContext);
+    <T> P visit(CreateOrUpdateByLookupKey<T> opContext) throws Exception;
 
-    <T> P visit(CreateOrUpdate<T> opContext);
+    <T> P visit(CreateOrUpdate<T> opContext) throws Exception;
 
-    <T, U> P visit(CreateOrUpdateInLockedContext<T, U> opContext);
+    <T, U> P visit(CreateOrUpdateInLockedContext<T, U> opContext) throws Exception;
 
     <T, R> P visit(Select<T, R> opContext);
 

src/main/java/io/appform/dropwizard/sharding/observers/bucket/BucketIdObserver.java

@@ -0,0 +1,31 @@
+package io.appform.dropwizard.sharding.observers.bucket;
+
+import io.appform.dropwizard.sharding.execution.TransactionExecutionContext;
+import io.appform.dropwizard.sharding.observers.TransactionObserver;
+import lombok.extern.slf4j.Slf4j;
+
+import java.util.function.Supplier;
+
+@Slf4j
+public class BucketIdObserver extends TransactionObserver {
+    private BucketIdSaver bucketIdSaver;
+    private volatile boolean initDone;
+
+    public BucketIdObserver(BucketIdSaver bucketIdSaver) {
+        super(null);
+        this.bucketIdSaver = bucketIdSaver;
+        log.info("BucketId observer constructor called ");
+    }
+    public void init(BucketIdSaver bucketIdSaver) {
+        this.bucketIdSaver = bucketIdSaver;
+        this.initDone = true;
+        log.info("BucketId observer initialisation completed ");
+    }
+
+    @Override
+    public <T> T execute(TransactionExecutionContext context, Supplier<T> supplier) {
+        context.getOpContext().visit(this.bucketIdSaver);
+        return proceed(context, supplier);
+    }
+
+}