Merge pull request #1661 from pjain1/poll_rules

Poll rules immediately after change

server/src/main/java/io/druid/metadata/SQLMetadataRuleManager.java

@@ -32,7 +32,6 @@ import com.google.inject.Inject;
 import com.metamx.common.Pair;
 import com.metamx.common.lifecycle.LifecycleStart;
 import com.metamx.common.lifecycle.LifecycleStop;
-import com.metamx.common.logger.Logger;
 import com.metamx.emitter.EmittingLogger;
 import io.druid.audit.AuditEntry;
 import io.druid.audit.AuditInfo;
@@ -377,7 +376,12 @@ public class SQLMetadataRuleManager implements MetadataRuleManager
         return false;
       }
     }
-
+    try {
+      poll();
+    }
+    catch (Exception e) {
+      log.error(e, String.format("Exception while polling for rules after overriding the rule for %s", dataSource));
+    }
     return true;
   }
 

server/src/test/java/io/druid/metadata/SQLMetadataRuleManagerTest.java

@@ -17,11 +17,9 @@
 
 package io.druid.metadata;
 
-import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.common.base.Suppliers;
 import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Lists;
 import io.druid.audit.AuditEntry;
 import io.druid.audit.AuditInfo;
 import io.druid.audit.AuditManager;
@@ -29,18 +27,14 @@ import io.druid.client.DruidServer;
 import io.druid.jackson.DefaultObjectMapper;
 import io.druid.server.audit.SQLAuditManager;
 import io.druid.server.audit.SQLAuditManagerConfig;
-import io.druid.server.coordinator.rules.ForeverDropRule;
-import io.druid.server.coordinator.rules.ForeverLoadRule;
 import io.druid.server.coordinator.rules.IntervalLoadRule;
-import io.druid.server.coordinator.rules.PeriodLoadRule;
 import io.druid.server.coordinator.rules.Rule;
 import io.druid.server.metrics.NoopServiceEmitter;
-import junit.framework.Assert;
-import org.joda.time.DateTime;
 import org.joda.time.Interval;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
+import org.junit.Assert;
 import org.skife.jdbi.v2.Handle;
 import org.skife.jdbi.v2.tweak.HandleCallback;
 
@@ -83,6 +77,30 @@ public class SQLMetadataRuleManagerTest
     );
   }
 
+  @Test
+  public void testRuleInsert()
+  {
+    List<Rule> rules = Arrays.<Rule>asList(
+        new IntervalLoadRule(
+            new Interval("2015-01-01/2015-02-01"), ImmutableMap.<String, Integer>of(
+            DruidServer.DEFAULT_TIER,
+            DruidServer.DEFAULT_NUM_REPLICANTS
+        )
+        )
+    );
+    AuditInfo auditInfo = new AuditInfo("test_author", "test_comment", "127.0.0.1");
+    ruleManager.overrideRule(
+        "test_dataSource",
+        rules,
+        auditInfo
+    );
+    // New rule should be be reflected in the in memory rules map immediately after being set by user
+    Map<String, List<Rule>> allRules = ruleManager.getAllRules();
+    Assert.assertEquals(1, allRules.size());
+    Assert.assertEquals(1, allRules.get("test_dataSource").size());
+    Assert.assertEquals(rules.get(0), allRules.get("test_dataSource").get(0));
+  }
+
   @Test
   public void testAuditEntryCreated() throws Exception
   {