#!/bin/bash
# This hacky script will read the migration JSON file you provide as argument 1 and execute the reassignemnt up to the partition you specify as argument 2.
# This is useful if you don't want to migrate all paritions at once but not have to re-create the reassignment plan everytime
# and keep the reassignments consistent across the whole topic.
# Will work with any properly formatted kafka parition reassignment JSON file.
# Recommended to use with https://github.com/SiftScience/kafka-assigner to give you more control over the reassignemnt.


# will migrate up to the partition you specify
# Example:
# ./reassign-partition.sh ~/my-reassignemnt-plan.json 5
# will migrate partitions 0 through 5.
# ./reassign-partition.sh ~/my-reassignment-plan.json 0
# will only migrate the first partition

# The script will wait until Kafka verifies that the reassignment is complete before exiting.
# this means you can chain it like so:
# ./reassign-partition.sh ~/my-reassignemnt-plan.json 2 && sleep 120 && ./reassign-partition.sh ~/my-reassignemnt-plan.json 5

REASSIGN_FILE=$1
PARTITION=$2

if [[ -z ${REASSIGN_FILE} ]] || [[ ! -f ${REASSIGN_FILE} ]] ;then
  echo "Reassignment plan required as first argument. Needs to be an existing file."
  exit 1
fi

if [[ -z ${PARTITION} ]];then
  echo "Partition required as second argument."
  exit 1
fi


MAX_PART=$(jq '.partitions | length' $REASSIGN_FILE)

if [[ "$PARTITION" -ge "$MAX_PART" ]]; then
  echo "Reassignment plan only has ${MAX_PART} partitions. 0-$(($MAX_PART-1))"
  exit 1
fi

if [[ "$PARTITION" -lt 0 ]]; then
  PARTITION=0
else
  PARTITION=$(($PARTITION+1))
fi

FILTERED_REASSIGN_FILE="${REASSIGN_FILE}-FILTERED"
FILTERED=$(jq "del(.partitions[$PARTITION:] )" $REASSIGN_FILE)

echo $FILTERED > $FILTERED_REASSIGN_FILE

/opt/kafka/bin/kafka-reassign-partitions.sh --zookeeper zookeeper-node.example.com  --reassignment-json-file $FILTERED_REASSIGN_FILE --execute

until /opt/kafka/bin/kafka-reassign-partitions.sh --zookeeper zookeeper-node.example.com  --reassignment-json-file $FILTERED_REASSIGN_FILE --verify | tee /dev/tty | grep -m 1 ",$((${PARTITION}-1))] completed successfully"; do
sleep 1
done