Allow setting a custom placement policy per table through yb-admin_cl…

src/yb/client/client-internal.cc

@@ -1344,6 +1344,9 @@ void GetTableSchemaRpc::Finished(const Status& status) {
       if (resp_.has_index_info()) {
         info_->index_info.emplace(resp_.index_info());
       }
+      if (resp_.has_replication_info()) {
+        info_->replication_info.emplace(resp_.replication_info());
+      }
       CHECK_GT(info_->table_id.size(), 0) << "Running against a too-old master";
       info_->colocated = resp_.colocated();
     }

src/yb/client/table.cc

@@ -131,6 +131,10 @@ const bool YBTable::colocated() const {
   return info_.colocated;
 }
 
+const boost::optional<master::ReplicationInfoPB>& YBTable::replication_info() const {
+  return info_.replication_info;
+}
+
 std::string YBTable::ToString() const {
   return strings::Substitute(
       "$0 $1 IndexInfo: $2 IndexMap $3", (IsIndex() ? "Index Table" : "Normal Table"), id(),

src/yb/client/table.h

@@ -47,6 +47,7 @@ struct YBTableInfo {
   boost::optional<IndexInfo> index_info;
   YBTableType table_type;
   bool colocated;
+  boost::optional<master::ReplicationInfoPB> replication_info;
 };
 
 // A YBTable represents a table on a particular cluster. It holds the current
@@ -101,6 +102,9 @@ class YBTable : public std::enable_shared_from_this<YBTable> {
   // Is the table colocated?
   const bool colocated() const;
 
+  // Returns the replication info for the table.
+  const boost::optional<master::ReplicationInfoPB>& replication_info() const;
+
   std::string ToString() const;
   //------------------------------------------------------------------------------------------------
   // CQL support

src/yb/client/table_alterer.cc

@@ -59,6 +59,11 @@ YBTableAlterer* YBTableAlterer::SetTableProperties(const TableProperties& table_
   return this;
 }
 
+YBTableAlterer* YBTableAlterer::replication_info(const master::ReplicationInfoPB& ri) {
+  replication_info_.emplace(ri);
+  return this;
+}
+
 YBTableAlterer* YBTableAlterer::SetWalRetentionSecs(const uint32_t wal_retention_secs) {
   wal_retention_secs_ = wal_retention_secs;
   return this;
@@ -97,7 +102,8 @@ Status YBTableAlterer::ToRequest(master::AlterTableRequestPB* req) {
     return status_;
   }
 
-  if (!rename_to_ && steps_.empty() && !table_properties_ && !wal_retention_secs_) {
+  if (!rename_to_ && steps_.empty() && !table_properties_ && !wal_retention_secs_ &&
+      !replication_info_) {
     return STATUS(InvalidArgument, "No alter steps provided");
   }
 
@@ -169,6 +175,11 @@ Status YBTableAlterer::ToRequest(master::AlterTableRequestPB* req) {
     req->set_wal_retention_secs(*wal_retention_secs_);
   }
 
+  if (replication_info_) {
+    // TODO: Maybe add checks for the sanity of the replication_info.
+    req->mutable_replication_info()->CopyFrom(replication_info_.get());
+  }
+
   return Status::OK();
 }
 

src/yb/client/table_alterer.h

@@ -62,7 +62,7 @@ class YBTableAlterer {
 
   YBTableAlterer* SetWalRetentionSecs(const uint32_t wal_retention_secs);
 
-      // Set the timeout for the operation. This includes any waiting
+  // Set the timeout for the operation. This includes any waiting
   // after the alter has been submitted (i.e if the alter is slow
   // to be performed on a large table, it may time out and then
   // later be successful).
@@ -73,6 +73,9 @@ class YBTableAlterer {
   // If not provided, defaults to true.
   YBTableAlterer* wait(bool wait);
 
+  // Set replication info for the table.
+  YBTableAlterer* replication_info(const master::ReplicationInfoPB& ri);
+
   // Alters the table.
   //
   // The return value may indicate an error in the alter operation, or a
@@ -111,6 +114,8 @@ class YBTableAlterer {
 
   boost::optional<uint32_t> wal_retention_secs_;
 
+  boost::optional<master::ReplicationInfoPB> replication_info_;
+
   DISALLOW_COPY_AND_ASSIGN(YBTableAlterer);
 };
 

src/yb/integration-tests/CMakeLists.txt

@@ -127,6 +127,7 @@ ADD_YB_TEST(log_version-test)
 ADD_YB_TEST(load_balancer-test)
 ADD_YB_TEST(load_balancer_multi_table-test)
 ADD_YB_TEST(load_balancer_respect_affinity-test)
+ADD_YB_TEST(load_balancer_placement_policy-test)
 
 set(YB_TEST_LINK_LIBS_SAVED ${YB_TEST_LINK_LIBS})
 set(YB_TEST_LINK_LIBS ${YB_TEST_LINK_LIBS} cassandra cql_test_util)