KENGINE-283 Topic ID Mismatch Issue - Exception added to disallow Alt…

core/src/main/scala/kafka/admin/TopicCommand.scala

@@ -21,7 +21,7 @@ import java.util
 import java.util.{Collections, Properties}
 
 import joptsimple._
-import kafka.common.AdminCommandFailedException
+import kafka.common.{AddPartitionNotAllowedDueToTopicIDMismatchException, AdminCommandFailedException}
 import kafka.log.LogConfig
 import kafka.server.ConfigType
 import kafka.utils.Implicits._
@@ -385,22 +385,31 @@ object TopicCommand extends Logging {
           println(s"Updated config for topic $topic.")
         }
 
-        if(tp.hasPartitions) {
-          if (Topic.isInternal(topic)) {
-            throw new IllegalArgumentException(s"The number of partitions for the internal topic $topic cannot be changed.")
-          }
-          println("WARNING: If partitions are increased for a topic that has a key, the partition " +
-            "logic or ordering of the messages will be affected")
-          val existingAssignment = zkClient.getFullReplicaAssignmentForTopics(immutable.Set(topic)).map {
-            case (topicPartition, assignment) => topicPartition.partition -> assignment
+        try {
+          if (tp.hasPartitions) {
+            if (Topic.isInternal(topic)) {
+              throw new IllegalArgumentException(s"The number of partitions for the internal topic $topic cannot be changed.")
+            }
+            println("WARNING: If partitions are increased for a topic that has a key, the partition " +
+              "logic or ordering of the messages will be affected")
+
+            val existingAssignment = zkClient.getFullReplicaAssignmentForTopicsInAlterTopic(immutable.Set(topic)).map {
+              case (topicPartition, assignment) => topicPartition.partition -> assignment
+            }
+            if (existingAssignment.isEmpty)
+              throw new InvalidTopicException(s"The topic $topic does not exist")
+
+            val newAssignment = tp.replicaAssignment.getOrElse(Map()).drop(existingAssignment.size)
+
+            val allBrokers = adminZkClient.getBrokerMetadatas()
+            val partitions: Integer = tp.partitions.getOrElse(1)
+            adminZkClient.addPartitions(topic, existingAssignment, allBrokers, partitions, Option(newAssignment).filter(_.nonEmpty))
+            println("Adding partitions succeeded!")
           }
-          if (existingAssignment.isEmpty)
-            throw new InvalidTopicException(s"The topic $topic does not exist")
-          val newAssignment = tp.replicaAssignment.getOrElse(Map()).drop(existingAssignment.size)
-          val allBrokers = adminZkClient.getBrokerMetadatas()
-          val partitions: Integer = tp.partitions.getOrElse(1)
-          adminZkClient.addPartitions(topic, existingAssignment, allBrokers, partitions, Option(newAssignment).filter(_.nonEmpty))
-          println("Adding partitions succeeded!")
+        } catch {
+          case _: AddPartitionNotAllowedDueToTopicIDMismatchException =>
+            println("WARNING:- Operation not allowed in this version")
+            throw new AddPartitionNotAllowedDueToTopicIDMismatchException("The TopicZNode already contains a TopicID. The AlterTopic operation will result in metadata inconsistencies.")
         }
       }
     }

core/src/main/scala/kafka/common/AddPartitionNotAllowedDueToTopicIDMismatchException.scala

@@ -0,0 +1,5 @@
+package kafka.common
+
+class AddPartitionNotAllowedDueToTopicIDMismatchException(message: String) extends RuntimeException(message) {
+  def this() = this(null)
+}

core/src/main/scala/kafka/zk/KafkaZkClient.scala

@@ -21,6 +21,7 @@ import java.util.Properties
 import com.yammer.metrics.core.MetricName
 import kafka.api.LeaderAndIsr
 import kafka.cluster.Broker
+import kafka.common.AddPartitionNotAllowedDueToTopicIDMismatchException
 import kafka.controller.{KafkaController, LeaderIsrAndControllerEpoch, ReplicaAssignment}
 import kafka.log.LogConfig
 import kafka.metrics.KafkaMetricsGroup
@@ -585,6 +586,31 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
     getFullReplicaAssignmentForTopics(topics).mapValues(_.replicas).toMap
   }
 
+  /**
+   * Gets the replica assignments for the given topics, is only called within AlterTopic
+   * because this is the only method where topic ID can be rewritten in the assignments
+   * If a topicID(topic_id) exists, it will be rewritten without one on doing the AlterTopic operation
+   * Exception is thrown in order to stop the operation from happening altogether
+   *
+   * @param topics the topics whose partitions we wish to get the assignments for.
+   * @return the full replica assignment for each partition from the given topics.
+   */
+  def getFullReplicaAssignmentForTopicsInAlterTopic(topics: Set[String]): Map[TopicPartition, ReplicaAssignment] = {
+    val getDataRequests = topics.map(topic => GetDataRequest(TopicZNode.path(topic), ctx = Some(topic)))
+    val getDataResponses = retryRequestsUntilConnected(getDataRequests.toSeq)
+    getDataResponses.flatMap { getDataResponse =>
+      val topic = getDataResponse.ctx.get.asInstanceOf[String]
+      if (TopicZNode.checkForTopicIDInJSON(getDataResponse.data) == Some(true)) {
+        throw new AddPartitionNotAllowedDueToTopicIDMismatchException("The TopicZNode already contains a TopicID. The AlterTopic operation will result in metadata inconsistencies.")
+      }
+      getDataResponse.resultCode match {
+        case Code.OK => TopicZNode.decode(topic, getDataResponse.data)
+        case Code.NONODE => Map.empty[TopicPartition, ReplicaAssignment]
+        case _ => throw getDataResponse.resultException.get
+      }
+    }.toMap
+  }
+
   /**
     * Gets the replica assignments for the given topics.
     * @param topics the topics whose partitions we wish to get the assignments for.
@@ -1673,7 +1699,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
 
   def secure: Boolean = isSecure
 
-  private[zk] def retryRequestUntilConnected[Req <: AsyncRequest](request: Req, expectedControllerZkVersion: Int = ZkVersion.MatchAnyVersion): Req#Response = {
+  private[kafka] def retryRequestUntilConnected[Req <: AsyncRequest](request: Req, expectedControllerZkVersion: Int = ZkVersion.MatchAnyVersion): Req#Response = {
     retryRequestsUntilConnected(Seq(request), expectedControllerZkVersion).head
   }
 
@@ -1688,7 +1714,7 @@ class KafkaZkClient private[zk] (zooKeeperClient: ZooKeeperClient, isSecure: Boo
     }
   }
 
-  private def retryRequestsUntilConnected[Req <: AsyncRequest](requests: Seq[Req]): Seq[Req#Response] = {
+  private[kafka] def retryRequestsUntilConnected[Req <: AsyncRequest](requests: Seq[Req]): Seq[Req#Response] = {
     val remainingRequests = new mutable.ArrayBuffer(requests.size) ++= requests
     val responses = new mutable.ArrayBuffer[Req#Response]
     while (remainingRequests.nonEmpty) {

core/src/main/scala/kafka/zk/ZkData.scala

@@ -259,6 +259,19 @@ object TopicZNode {
       "removing_replicas" -> removingReplicasAssignmentJson.asJava
     ).asJava)
   }
+
+  def checkForTopicIDInJSON(bytes: Array[Byte]): Option[Boolean] = {
+      Json.parseBytes(bytes).map { js =>
+      val assignmentJson = js.asJsonObject
+      val topicId = assignmentJson.get("topic_id").map(_.to[String])
+
+        if (topicId != None) {
+          return Option(true)
+        }
+        else false
+      }
+  }
+
   def decode(topic: String, bytes: Array[Byte]): Map[TopicPartition, ReplicaAssignment] = {
     def getReplicas(replicasJsonOpt: Option[JsonObject], partition: String): Seq[Int] = {
       replicasJsonOpt match {

core/src/test/scala/unit/kafka/admin/TopicCommandTest.scala

@@ -17,9 +17,12 @@
 package kafka.admin
 
 import kafka.admin.TopicCommand.{TopicCommandOptions, ZookeeperTopicService}
+import kafka.common.AddPartitionNotAllowedDueToTopicIDMismatchException
+
 import kafka.server.ConfigType
-import kafka.utils.{Logging, TestUtils}
-import kafka.zk.{ConfigEntityChangeNotificationZNode, DeleteTopicsTopicZNode, ZooKeeperTestHarness}
+import kafka.utils.{Json, Logging, TestUtils}
+import kafka.zk.{ConfigEntityChangeNotificationZNode, DeleteTopicsTopicZNode, TopicZNode, ZkVersion, ZooKeeperTestHarness}
+import kafka.zookeeper.{GetDataRequest, SetDataRequest}
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.config.{ConfigException, ConfigResource}
 import org.apache.kafka.common.errors.{InvalidPartitionsException, InvalidReplicationFactorException, TopicExistsException}
@@ -29,7 +32,11 @@ import org.junit.rules.TestName
 import org.junit.{After, Before, Rule, Test}
 import org.scalatest.Assertions.intercept
 
+import scala.collection.{Map, mutable, Seq}
 import scala.util.Random
+import scala.jdk.CollectionConverters._
+import java.util
+import java.util.UUID
 
 class TopicCommandTest extends ZooKeeperTestHarness with Logging with RackAwareTest {
 
@@ -597,4 +604,62 @@ class TopicCommandTest extends ZooKeeperTestHarness with Logging with RackAwareT
     expectAlterInternalTopicPartitionCountFailed(Topic.GROUP_METADATA_TOPIC_NAME)
     expectAlterInternalTopicPartitionCountFailed(Topic.TRANSACTION_STATE_TOPIC_NAME)
   }
+
+  // Test to check if Alter Topic is blocked when topic_id exists in order to avoid topic ID mismatch issues
+  // If topic_id is present in the TopicZNode due to a newer controller version setting it in processTopicIDs, we want to prevent
+  // rewriting of the TopicZNode without a topic ID in the current older client by blocking the Alter Topic command when used with the Zookeeper flag.
+  // Note:- topic_id is different from confluent_topic_id, existence of confluent_topic_id doesn't cause any problems
+  @Test
+  def testAlterTopicWithTopicIDPresent(): Unit = {
+    def topicZNodeEncodeForTest(topicId: Option[UUID]) : Array[Byte] = {
+      val replicaAssignmentJson = mutable.Map[String, util.List[Int]]()
+      val addingReplicasAssignmentJson = mutable.Map[String, util.List[Int]]()
+      val removingReplicasAssignmentJson = mutable.Map[String, util.List[Int]]()
+      val observersAssignment = mutable.Map.empty[String, util.List[Int]]
+      val targetObserversAssignment = mutable.Map.empty[String, util.List[Int]]
+      val topicAssignment = mutable.Map(
+        "version" -> 3,
+        "partitions" -> replicaAssignmentJson.asJava,
+        "adding_replicas" -> addingReplicasAssignmentJson.asJava,
+        "removing_replicas" -> removingReplicasAssignmentJson.asJava,
+        "observers" -> observersAssignment.asJava,
+        "target_observers" -> targetObserversAssignment.asJava,
+        "topic_id" -> topicId.get.toString
+      ).asJava
+      Json.encodeAsBytes(topicAssignment)
+    }
+    val brokers = List(0)
+    TestUtils.createBrokersInZk(zkClient, brokers)
+
+    adminZkClient.createTopic(testTopicName, 1, 1)
+
+    val encodedData = topicZNodeEncodeForTest(Some(UUID.randomUUID))
+    val setDataRequest = SetDataRequest(TopicZNode.path(testTopicName), encodedData, ZkVersion.MatchAnyVersion)
+    zkClient.retryRequestUntilConnected(setDataRequest)
+
+    def expectAlterTopicFailedIfTopicIDPresent(topic: String): Unit = {
+      try {
+        topicService.alterTopic(new TopicCommandOptions(
+          Array("--topic", topic, "--partitions", "2")))
+        fail("Should have thrown an AddPartitionNotAllowedDueToTopicIDMismatchException")
+      } catch {
+        case _: AddPartitionNotAllowedDueToTopicIDMismatchException => // expected
+      }
+    }
+
+    // We want to make sure that the topic_id is still present after the request fails and no other overwriting occurred
+    def checkIfTopicIDIsStillPresent(topic: String): Seq[Boolean] = {
+      val getDataRequests = GetDataRequest(TopicZNode.path(topic), ctx = Some(topic))
+      val getDataResponses = zkClient.retryRequestsUntilConnected(Seq(getDataRequests))
+      getDataResponses.flatMap { getDataResponse =>
+        if (TopicZNode.checkForTopicIDInJSON(getDataResponse.data) == Some(true)) {
+          return Seq(true)
+        }
+        else Seq(false)
+      }
+    }
+
+    expectAlterTopicFailedIfTopicIDPresent(testTopicName)
+    assertEquals(checkIfTopicIDIsStillPresent(testTopicName),Seq(true))
+  }
 }