HBASE-24032 [RSGroup] Assign created tables to respective rsgroup automatically instead of manual operations (#1318)

Signed-off-by: binlijin <binlijin@gmail.com>

hbase-rsgroup/src/main/java/org/apache/hadoop/hbase/rsgroup/RSGroupAdminEndpoint.java

@@ -33,6 +33,7 @@ import java.util.Set;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
+import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hbase.ClusterStatus;
 import org.apache.hadoop.hbase.Coprocessor;
 import org.apache.hadoop.hbase.CoprocessorEnvironment;
@@ -93,6 +94,7 @@ import org.apache.hadoop.hbase.security.UserProvider;
 import org.apache.hadoop.hbase.security.access.AccessChecker;
 import org.apache.hadoop.hbase.security.access.Permission;
 import org.apache.hadoop.hbase.zookeeper.ZooKeeperWatcher;
+import org.apache.hadoop.util.Shell;
 
 public class RSGroupAdminEndpoint extends RSGroupAdminService
     implements CoprocessorService, Coprocessor, MasterObserver {
@@ -107,6 +109,47 @@ public class RSGroupAdminEndpoint extends RSGroupAdminService
 
   /** Provider for mapping principal names to Users */
   private UserProvider userProvider;
+  /** Get rsgroup table mapping script */
+  private RSGroupMappingScript script;
+
+  // Package visibility for testing
+  static class RSGroupMappingScript {
+
+    static final String RS_GROUP_MAPPING_SCRIPT = "hbase.rsgroup.table.mapping.script";
+    static final String RS_GROUP_MAPPING_SCRIPT_TIMEOUT =
+      "hbase.rsgroup.table.mapping.script.timeout";
+
+    private Shell.ShellCommandExecutor rsgroupMappingScript;
+
+    RSGroupMappingScript(Configuration conf) {
+      String script = conf.get(RS_GROUP_MAPPING_SCRIPT);
+      if (script == null || script.isEmpty()) {
+        return;
+      }
+
+      rsgroupMappingScript = new Shell.ShellCommandExecutor(
+          new String[] { script, "", "" }, null, null,
+          conf.getLong(RS_GROUP_MAPPING_SCRIPT_TIMEOUT, 5000) // 5 seconds
+      );
+    }
+
+    String getRSGroup(String namespace, String tablename) {
+      if (rsgroupMappingScript == null) {
+        return RSGroupInfo.DEFAULT_GROUP;
+      }
+      String[] exec = rsgroupMappingScript.getExecString();
+      exec[1] = namespace;
+      exec[2] = tablename;
+      try {
+        rsgroupMappingScript.execute();
+      } catch (IOException e) {
+        LOG.error(e.getMessage() + " placing back to default rsgroup");
+        return RSGroupInfo.DEFAULT_GROUP;
+      }
+      return rsgroupMappingScript.getOutput().trim();
+    }
+
+  }
 
   @Override
   public void start(CoprocessorEnvironment env) throws IOException {
@@ -124,6 +167,7 @@ public class RSGroupAdminEndpoint extends RSGroupAdminService
 
     // set the user-provider.
     this.userProvider = UserProvider.instantiate(env.getConfiguration());
+    this.script = new RSGroupMappingScript(env.getConfiguration());
   }
 
   @Override
@@ -425,6 +469,16 @@ public class RSGroupAdminEndpoint extends RSGroupAdminService
         RSGroupInfo.DEFAULT_GROUP + "' for deploy of system table");
       groupName = RSGroupInfo.DEFAULT_GROUP;
     }
+
+    if (groupName == RSGroupInfo.DEFAULT_GROUP) {
+      TableName tableName = desc.getTableName();
+      groupName = script.getRSGroup(
+          tableName.getNamespaceAsString(),
+          tableName.getQualifierAsString()
+      );
+      LOG.info("rsgroup for " + tableName + " is " + groupName);
+    }
+
     RSGroupInfo rsGroupInfo = groupAdminServer.getRSGroupInfo(groupName);
     if (rsGroupInfo == null) {
       throw new ConstraintException("Default RSGroup (" + groupName + ") for this table's "

hbase-rsgroup/src/main/java/org/apache/hadoop/hbase/rsgroup/RSGroupInfoManagerImpl.java

@@ -815,7 +815,7 @@ public class RSGroupInfoManagerImpl implements RSGroupInfoManager, ServerListene
     Map<String,RSGroupInfo> newGroupMap = Maps.newHashMap(rsGroupMap);
     newGroupMap.put(srcGroupInfo.getName(), srcGroupInfo);
     newGroupMap.put(dstGroupInfo.getName(), dstGroupInfo);
-    flushConfig(newGroupMap);    
+    flushConfig(newGroupMap);
   }
 
   @Override

hbase-rsgroup/src/test/java/org/apache/hadoop/hbase/rsgroup/TestRSGroupMappingScript.java

@@ -0,0 +1,122 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.hadoop.hbase.rsgroup;
+
+import java.io.BufferedReader;
+import java.io.File;
+import java.io.FileOutputStream;
+import java.io.FileReader;
+import java.io.IOException;
+import java.io.PrintWriter;
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.apache.hadoop.hbase.HBaseTestingUtility;
+import org.apache.hadoop.hbase.TableName;
+import org.apache.hadoop.hbase.rsgroup.RSGroupAdminEndpoint.RSGroupMappingScript;
+import org.apache.hadoop.hbase.testclassification.SmallTests;
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+import org.junit.experimental.categories.Category;
+
+@Category({ SmallTests.class })
+public class TestRSGroupMappingScript {
+
+  private static final Log LOG = LogFactory.getLog(TestRSGroupMappingScript.class);
+
+  private static final HBaseTestingUtility UTIL = new HBaseTestingUtility();
+  private File script;
+
+  @BeforeClass
+  public static void setupScript() throws Exception {
+    String currentDir = new File("").getAbsolutePath();
+    UTIL.getConfiguration().set(
+        RSGroupMappingScript.RS_GROUP_MAPPING_SCRIPT,
+        currentDir + "/rsgroup_table_mapping.sh"
+    );
+  }
+
+  @Before
+  public void setup() throws Exception {
+    script = new File(UTIL.getConfiguration().get(RSGroupMappingScript.RS_GROUP_MAPPING_SCRIPT));
+    if (!script.createNewFile()) {
+      throw new IOException("Can't create script");
+    }
+
+    PrintWriter pw = new PrintWriter(new FileOutputStream(script));
+    try {
+      pw.println("#!/bin/bash");
+      pw.println("namespace=$1");
+      pw.println("tablename=$2");
+      pw.println("if [[ $namespace == test ]]; then");
+      pw.println("  echo test");
+      pw.println("elif [[ $tablename == *foo* ]]; then");
+      pw.println("  echo other");
+      pw.println("else");
+      pw.println("  echo default");
+      pw.println("fi");
+      pw.flush();
+    } finally {
+      pw.close();
+    }
+    boolean executable = script.setExecutable(true);
+    LOG.info("Created " + script  + ", executable=" + executable);
+    verifyScriptContent(script);
+  }
+
+  private void verifyScriptContent(File file) throws Exception {
+    BufferedReader reader = new BufferedReader(new FileReader(file));
+    String line;
+    while ((line = reader.readLine()) != null) {
+      LOG.info(line);
+    }
+  }
+
+  @Test
+  public void testScript() throws Exception {
+    RSGroupMappingScript script = new RSGroupMappingScript(UTIL.getConfiguration());
+    TableName testNamespace =
+      TableName.valueOf("test", "should_be_in_test");
+    String rsgroup = script.getRSGroup(
+      testNamespace.getNamespaceAsString(), testNamespace.getQualifierAsString()
+    );
+    Assert.assertEquals("test", rsgroup);
+
+    TableName otherName =
+      TableName.valueOf("whatever", "oh_foo_should_be_in_other");
+    rsgroup = script.getRSGroup(otherName.getNamespaceAsString(), otherName.getQualifierAsString());
+    Assert.assertEquals("other", rsgroup);
+
+    TableName defaultName =
+      TableName.valueOf("nono", "should_be_in_default");
+    rsgroup = script.getRSGroup(
+      defaultName.getNamespaceAsString(), defaultName.getQualifierAsString()
+    );
+    Assert.assertEquals("default", rsgroup);
+  }
+
+  @After
+  public void teardown() throws Exception {
+    if (script.exists()) {
+      script.delete();
+    }
+  }
+
+}