From 073c32f01be86c21def3f9401a168905ae8109f4 Mon Sep 17 00:00:00 2001
From: Parag Jain <pjain11@yahoo-inc.com>
Date: Mon, 24 Aug 2015 13:45:35 -0500
Subject: [PATCH] Poll rules immediately after change

---
 .../metadata/SQLMetadataRuleManager.java      |  8 +++--
 .../metadata/SQLMetadataRuleManagerTest.java  | 32 +++++++++++++++----
 2 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/server/src/main/java/io/druid/metadata/SQLMetadataRuleManager.java b/server/src/main/java/io/druid/metadata/SQLMetadataRuleManager.java
index d4d9e4779d6..b17b5ae870b 100644
--- a/server/src/main/java/io/druid/metadata/SQLMetadataRuleManager.java
+++ b/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;
   }
 
diff --git a/server/src/test/java/io/druid/metadata/SQLMetadataRuleManagerTest.java b/server/src/test/java/io/druid/metadata/SQLMetadataRuleManagerTest.java
index 7061cfef9d3..eefc4919c6d 100644
--- a/server/src/test/java/io/druid/metadata/SQLMetadataRuleManagerTest.java
+++ b/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
   {