├── .github └── workflows │ ├── kafka-utils-ci.yml │ └── pypi.yaml ├── .gitignore ├── .pre-commit-config.yaml ├── CHANGELOG.rst ├── LICENSE ├── MANIFEST.in ├── Makefile ├── README.md ├── docker-compose.yml ├── docker ├── itest │ ├── Dockerfile │ ├── download-kafka.sh │ └── run_tests.sh ├── kafka │ ├── Dockerfile │ ├── config.properties │ └── download-kafka.sh └── zookeeper │ └── Dockerfile ├── docs ├── Makefile └── source │ ├── conf.py │ ├── config.rst │ ├── index.rst │ ├── kafka_check.rst │ ├── kafka_cluster_manager.rst │ ├── kafka_consumer_manager.rst │ ├── kafka_corruption_check.rst │ └── kafka_rolling_restart.rst ├── kafka_utils ├── __init__.py ├── kafka_check │ ├── __init__.py │ ├── commands │ │ ├── __init__.py │ │ ├── command.py │ │ ├── min_isr.py │ │ ├── offline.py │ │ ├── replica_unavailability.py │ │ └── replication_factor.py │ ├── main.py │ ├── metadata_file.py │ └── status_code.py ├── kafka_cluster_manager │ ├── __init__.py │ ├── cluster_info │ │ ├── __init__.py │ │ ├── broker.py │ │ ├── cluster_balancer.py │ │ ├── cluster_topology.py │ │ ├── display.py │ │ ├── error.py │ │ ├── genetic_balancer.py │ │ ├── partition.py │ │ ├── partition_count_balancer.py │ │ ├── partition_measurer.py │ │ ├── replication_group_parser.py │ │ ├── rg.py │ │ ├── stats.py │ │ ├── topic.py │ │ └── util.py │ ├── cmds │ │ ├── __init__.py │ │ ├── command.py │ │ ├── decommission.py │ │ ├── preferred_replica_election.py │ │ ├── rebalance.py │ │ ├── replace.py │ │ ├── revoke_leadership.py │ │ ├── set_replication_factor.py │ │ ├── stats.py │ │ └── store_assignments.py │ └── main.py ├── kafka_consumer_manager │ ├── __init__.py │ ├── commands │ │ ├── __init__.py │ │ ├── copy_group.py │ │ ├── delete_group.py │ │ ├── list_groups.py │ │ ├── list_topics.py │ │ ├── offset_advance.py │ │ ├── offset_get.py │ │ ├── offset_manager.py │ │ ├── offset_restore.py │ │ ├── offset_rewind.py │ │ ├── offset_save.py │ │ ├── offset_set.py │ │ ├── offset_set_timestamp.py │ │ ├── offsets_for_timestamp.py │ │ ├── rename_group.py │ │ ├── unsubscribe_topics.py │ │ └── watermark_get.py │ ├── main.py │ └── util.py ├── kafka_corruption_check │ ├── __init__.py │ └── main.py ├── kafka_manual_throttle │ ├── __init__.py │ └── main.py ├── kafka_rolling_restart │ ├── __init__.py │ ├── main.py │ └── task.py ├── main.py └── util │ ├── __init__.py │ ├── client.py │ ├── config.py │ ├── error.py │ ├── metadata.py │ ├── monitoring.py │ ├── offsets.py │ ├── protocol.py │ ├── py.typed │ ├── serialization.py │ ├── ssh.py │ ├── utils.py │ ├── validation.py │ └── zookeeper.py ├── mypy.ini ├── requirements-dev.txt ├── requirements.txt ├── scripts ├── kafka-check ├── kafka-cluster-manager ├── kafka-consumer-manager ├── kafka-corruption-check ├── kafka-manual-throttle ├── kafka-rolling-restart └── kafka-utils ├── setup.py ├── tests ├── __init__.py ├── acceptance │ ├── change_topic_config.feature │ ├── config │ │ └── test.yaml │ ├── copy_group.feature │ ├── environment.py │ ├── kafka_consumer_manager.feature │ ├── list_groups.feature │ ├── list_topics.feature │ ├── min_isr.feature │ ├── offset_advance.feature │ ├── offset_get.feature │ ├── offset_restore.feature │ ├── offset_rewind.feature │ ├── offset_save.feature │ ├── offset_set.feature │ ├── rename_group.feature │ ├── replica_unavailability.feature │ ├── replication_factor.feature │ ├── steps │ │ ├── __init__.py │ │ ├── commit_fetch.py │ │ ├── common.py │ │ ├── config_update.py │ │ ├── copy_group.py │ │ ├── list_groups.py │ │ ├── list_topics.py │ │ ├── min_isr.py │ │ ├── offset_advance.py │ │ ├── offset_get.py │ │ ├── offset_restore.py │ │ ├── offset_rewind.py │ │ ├── offset_save.py │ │ ├── offset_set.py │ │ ├── rename_group.py │ │ ├── replica_unavailability.py │ │ ├── replication_factor.py │ │ ├── util.py │ │ └── watermark_get.py │ └── watermark_get.feature ├── kafka_check │ ├── __init__.py │ ├── test_is_first_broker.py │ ├── test_metadata_file.py │ ├── test_min_isr.py │ ├── test_offline.py │ ├── test_replica_unavailability.py │ └── test_replication_factor.py ├── kafka_cluster_manager │ ├── __init__.py │ ├── broker_test.py │ ├── cluster_balancer_test.py │ ├── cluster_topology_test.py │ ├── cmds │ │ ├── __init__.py │ │ ├── command_test.py │ │ ├── replace_broker_test.py │ │ └── set_replication_factor_test.py │ ├── conftest.py │ ├── decommission_test.py │ ├── genetic_balancer_test.py │ ├── helper.py │ ├── partition_count_balancer_test.py │ ├── partition_test.py │ ├── rg_test.py │ ├── stats_test.py │ ├── topic_test.py │ └── util_test.py ├── kafka_consumer_manager │ ├── __init__.py │ ├── test_copy_group.py │ ├── test_delete_group.py │ ├── test_list_groups.py │ ├── test_offset_advance.py │ ├── test_offset_get.py │ ├── test_offset_manager.py │ ├── test_offset_restore.py │ ├── test_offset_rewind.py │ ├── test_offset_save.py │ ├── test_offset_set.py │ ├── test_rename_group.py │ ├── test_unsubscribe_topics.py │ ├── test_utils.py │ └── test_watermark_get.py ├── kafka_corruption_check │ ├── __init__.py │ └── test_main.py ├── kafka_manual_throttle │ └── test_main.py ├── kafka_rolling_restart │ ├── __init__.py │ └── test_main.py └── util │ ├── config_test.py │ ├── metadata_test.py │ ├── test_monitoring.py │ ├── test_offsets.py │ ├── util_test.py │ ├── validation_test.py │ └── zookeeper_test.py ├── tox.ini └── tox_acceptance.ini /.github/workflows/kafka-utils-ci.yml: -------------------------------------------------------------------------------- 1 | # This workflow will install Python dependencies, run tests and lint with a variety of Python versions 2 | # For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions 3 | 4 | name: kafka-utils-ci 5 | on: [pull_request, release] 6 | 7 | jobs: 8 | tox: 9 | runs-on: ubuntu-22.04 10 | strategy: 11 | fail-fast: false 12 | matrix: 13 | toxenv: 14 | - py38-unittest 15 | - py38-kafka10-dockeritest 16 | - py38-kafka11-dockeritest 17 | env: 18 | PIP_INDEX_URL: https://pypi.python.org/simple 19 | steps: 20 | - uses: actions/checkout@v2 21 | - name: Set up Python 22 | uses: actions/setup-python@v2 23 | with: 24 | python-version: 3.8 25 | - name: Install dependencies 26 | run: | 27 | python -m pip install --upgrade pip 28 | pip install tox==3.2 29 | - name: Run tests 30 | run: tox -i https://pypi.python.org/simple -e ${{ matrix.toxenv }} 31 | -------------------------------------------------------------------------------- /.github/workflows/pypi.yaml: -------------------------------------------------------------------------------- 1 | name: Publish on PyPI 2 | 3 | on: 4 | push: 5 | tags: 6 | - v* 7 | 8 | jobs: 9 | publish: 10 | runs-on: ubuntu-latest 11 | 12 | steps: 13 | - uses: actions/checkout@v3 14 | - uses: actions/setup-python@v4 15 | with: 16 | python-version: 3.8 17 | 18 | - name: Install Python dependencies 19 | run: pip install wheel 20 | 21 | - name: Create a Wheel file and source distribution 22 | run: python setup.py sdist bdist_wheel 23 | 24 | - name: Publish distribution package to PyPI 25 | uses: pypa/gh-action-pypi-publish@v1.5.1 26 | with: 27 | user: __token__ 28 | password: ${{ secrets.PYPI_PASSWORD }} 29 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | 5 | # C extensions 6 | *.so 7 | 8 | # Pycharm 9 | .idea/ 10 | 11 | # Distribution / packaging 12 | .Python 13 | env/ 14 | build/ 15 | dist/ 16 | dist 17 | eggs/ 18 | sdist/ 19 | *.egg-info/ 20 | .installed.cfg 21 | *.egg 22 | *.swp 23 | 24 | # PyInstaller 25 | # Usually these files are written by a python script from a template 26 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 27 | *.manifest 28 | *.spec 29 | 30 | # Installer logs 31 | pip-log.txt 32 | pip-delete-this-directory.txt 33 | 34 | # Unit test / coverage reports 35 | .tox/ 36 | .coverage 37 | .cache 38 | nosetests.xml 39 | coverage.xml 40 | .pytest_cache/ 41 | 42 | # Sphinx documentation 43 | docs/_build/ 44 | 45 | # virtualenv 46 | virtualenv_run/ 47 | 48 | # VSCode 49 | .vscode/ 50 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | repos: 2 | - repo: https://github.com/pre-commit/pre-commit-hooks.git 3 | rev: v2.2.3 4 | hooks: 5 | - id: check-added-large-files 6 | - id: check-merge-conflict 7 | - id: trailing-whitespace 8 | - id: end-of-file-fixer 9 | - id: flake8 10 | args: ["--ignore=W504,E501"] 11 | - repo: https://github.com/pre-commit/mirrors-autopep8 12 | rev: v1.7.0 13 | hooks: 14 | - id: autopep8 15 | args: [--ignore=E501, --in-place] 16 | - repo: https://github.com/asottile/reorder_python_imports.git 17 | rev: v1.9.0 18 | hooks: 19 | - id: reorder-python-imports 20 | files: .*\.py$ 21 | - repo: https://github.com/asottile/pyupgrade 22 | rev: v3.1.0 23 | hooks: 24 | - id: pyupgrade 25 | args: ['--py37-plus'] 26 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | recursive-include kafka_utils *.py 2 | include README.md 3 | include LICENSE 4 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | ifeq ($(findstring .yelpcorp.com,$(shell hostname -f)), .yelpcorp.com) 2 | export PIP_INDEX_URL ?= https://pypi.yelpcorp.com/simple 3 | else 4 | export PIP_INDEX_URL ?= https://pypi.python.org/simple 5 | endif 6 | 7 | 8 | PACKAGE_VERSION := $(shell python setup.py --version) 9 | 10 | all: test 11 | 12 | clean: 13 | rm -rf build dist *.egg-info/ .tox/ virtualenv_run 14 | find . -name '*.pyc' -delete 15 | find . -name '__pycache__' -delete 16 | find . -name '*.deb' -delete 17 | find . -name '*.changes' -delete 18 | make -C docs clean 19 | 20 | test: 21 | tox -e py{37,38}-unittest 22 | 23 | acceptance: acceptance10 acceptance11 24 | 25 | acceptance10: 26 | tox -e py{37,38}-kafka10-dockeritest 27 | 28 | acceptance11: 29 | tox -e py{37,38}-kafka11-dockeritest 30 | 31 | coverage: 32 | tox -e coverage 33 | 34 | tag: 35 | git tag v${PACKAGE_VERSION} 36 | 37 | docs: 38 | tox -e docs 39 | 40 | .PHONY: all clean test coverage tag docs 41 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | **Deprecation Warning** 2 | 3 | Please note that this repo is not maintained in the open source community. The code and examples 4 | contained in this repository are for demonstration purposes only. 5 | 6 | You can read the latest from Yelp Engineering on our [tech blog](https://engineeringblog.yelp.com/). 7 | 8 | [![Build Status](https://github.com/Yelp/kafka-utils/workflows/kafka-utils-ci/badge.svg?branch=master)](https://github.com/Yelp/kafka-utils) 9 | 10 | # Kafka-Utils 11 | 12 | A suite of python tools to interact and manage Apache Kafka clusters. 13 | Kafka-Utils runs on python 3.7+. 14 | 15 | ## Configuration 16 | 17 | Kafka-Utils reads cluster configuration needed to access Kafka clusters from yaml files. Each cluster is identified by *type* and *name*. 18 | Multiple clusters of the same type should be listed in the same `type.yaml` file. 19 | The yaml files are read from `$KAFKA_DISCOVERY_DIR`, `$HOME/.kafka_discovery` and `/etc/kafka_discovery`, the former overrides the latter. 20 | 21 | 22 | Sample configuration for `sample_type` cluster at `/etc/kafka_discovery/sample_type.yaml` 23 | 24 | ```yaml 25 | --- 26 | clusters: 27 | cluster-1: 28 | broker_list: 29 | - "cluster-elb-1:9092" 30 | zookeeper: "11.11.11.111:2181,11.11.11.112:2181,11.11.11.113:2181/kafka-1" 31 | cluster-2: 32 | broker_list: 33 | - "cluster-elb-2:9092" 34 | zookeeper: "11.11.11.211:2181,11.11.11.212:2181,11.11.11.213:2181/kafka-2" 35 | local_config: 36 | cluster: cluster-1 37 | ``` 38 | 39 | ## Install 40 | 41 | From PyPI: 42 | ```shell 43 | $ pip install kafka-utils 44 | ``` 45 | 46 | 47 | ## Kafka-Utils command-line interface 48 | 49 | ### List all clusters 50 | 51 | ```shell 52 | $ kafka-utils 53 | cluster-type sample_type: 54 | cluster-name: cluster-1 55 | broker-list: cluster-elb-1:9092 56 | zookeeper: 11.11.11.111:2181,11.11.11.112:2181,11.11.11.113:2181/kafka-1 57 | cluster-name: cluster-2 58 | broker-list: cluster-elb-2:9092 59 | zookeeper: 11.11.11.211:2181,11.11.11.212:2181,11.11.11.213:2181/kafka-2 60 | ``` 61 | 62 | ### Get consumer offsets 63 | 64 | ```shell 65 | $ kafka-consumer-manager --cluster-type sample_type offset_get sample_consumer 66 | ``` 67 | 68 | ### Get consumer watermarks 69 | 70 | ```shell 71 | $ kafka-consumer-manager --cluster-type sample_type get_topic_watermark sample.topic 72 | 73 | ``` 74 | 75 | 76 | ### Rebalance cluster cluster1 of type sample_cluster 77 | 78 | ```shell 79 | $ kafka-cluster-manager --cluster-type sample_type --cluster-name cluster1 80 | --apply rebalance --brokers --leaders --max-partition-movements 10 81 | --max-leader-changes 15 82 | ``` 83 | 84 | ### Rolling-restart a cluster 85 | 86 | ```shell 87 | $ kafka-rolling-restart --cluster-type sample_type 88 | ``` 89 | 90 | ### Check in-sync replicas 91 | 92 | ```shell 93 | $ kafka-check --cluster-type sample_type min_isr 94 | ``` 95 | 96 | ### Check number of unavailable replicas 97 | 98 | ```shell 99 | $ kafka-check --cluster-type sample_type replica_unavailability 100 | ``` 101 | 102 | ## Documentation 103 | 104 | Read the documentation at [Read the Docs](http://kafka-utils.readthedocs.io/en/latest/). 105 | 106 | ## License 107 | 108 | Kafka-Utils is licensed under the Apache License, Version 2.0: http://www.apache.org/licenses/LICENSE-2.0 109 | 110 | ## Contributing 111 | 112 | Everyone is encouraged to contribute to Kafka-Utils by forking the 113 | [Github repository](http://github.com/Yelp/kafka-utils) and making a pull request or opening an issue. 114 | -------------------------------------------------------------------------------- /docker-compose.yml: -------------------------------------------------------------------------------- 1 | version: '2' 2 | 3 | services: 4 | itest: 5 | build: 6 | context: docker/itest 7 | args: 8 | - KAFKA_VERSION 9 | links: 10 | - kafka 11 | - zookeeper 12 | volumes: 13 | - .:/work 14 | command: echo "dummy command" 15 | 16 | kafka: 17 | build: 18 | context: docker/kafka 19 | args: 20 | - KAFKA_VERSION 21 | expose: 22 | - "9092" 23 | links: 24 | - zookeeper 25 | 26 | zookeeper: 27 | build: docker/zookeeper 28 | expose: 29 | - "2181" 30 | -------------------------------------------------------------------------------- /docker/itest/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM ubuntu:22.04 2 | MAINTAINER Team Data Streams Core 3 | ENV DEBIAN_FRONTEND noninteractive 4 | 5 | ARG KAFKA_VERSION 6 | # We need to install Java and Kafka in order to use Kafka CLI. The Kafka server 7 | # will never run in this container; the Kafka server will run in the "kafka" 8 | # container. 9 | 10 | # Install Java. 11 | RUN apt-get update && \ 12 | apt-get install -y \ 13 | software-properties-common \ 14 | openjdk-8-jdk 15 | 16 | # Install Kafka. 17 | RUN apt-get install -y \ 18 | unzip \ 19 | wget \ 20 | curl \ 21 | jq \ 22 | coreutils 23 | 24 | ENV JAVA_HOME "/usr/lib/jvm/java-8-openjdk-amd64/" 25 | ENV PATH "$JAVA_HOME/bin:$PATH" 26 | ENV SCALA_VERSION="2.11" 27 | ENV KAFKA_HOME /opt/kafka_${SCALA_VERSION}-${KAFKA_VERSION} 28 | COPY download-kafka.sh /tmp/download-kafka.sh 29 | RUN chmod 755 /tmp/download-kafka.sh 30 | RUN /tmp/download-kafka.sh && tar xfz /tmp/kafka_${SCALA_VERSION}-${KAFKA_VERSION}.tgz -C /opt && rm /tmp/kafka_${SCALA_VERSION}-${KAFKA_VERSION}.tgz 31 | ENV PATH="$PATH:$KAFKA_HOME/bin" 32 | 33 | RUN add-apt-repository ppa:deadsnakes/ppa 34 | 35 | # Install Python 36 | RUN apt-get update && apt-get install -y \ 37 | build-essential \ 38 | libffi-dev \ 39 | libssl-dev \ 40 | python3.7 \ 41 | python3.7-distutils \ 42 | python3.8 \ 43 | python3.8-distutils \ 44 | python-pkg-resources \ 45 | python-setuptools \ 46 | tox 47 | 48 | COPY run_tests.sh /scripts/run_tests.sh 49 | RUN chmod 755 /scripts/run_tests.sh 50 | 51 | WORKDIR /work 52 | -------------------------------------------------------------------------------- /docker/itest/download-kafka.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | 3 | url="https://archive.apache.org/dist/kafka/${KAFKA_VERSION}/kafka_${SCALA_VERSION}-${KAFKA_VERSION}.tgz" 4 | wget -q "${url}" -O "/tmp/kafka_${SCALA_VERSION}-${KAFKA_VERSION}.tgz" 5 | -------------------------------------------------------------------------------- /docker/itest/run_tests.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | set -e 4 | 5 | function do_at_exit { 6 | exit_status=$? 7 | rm -rf build/ dist/ kafka_utils.egg-info/ 8 | rm -rf .tox/log .tox/${ITEST_PYTHON_FACTOR}-acceptance .tox/dist_acceptance 9 | find . -name '*.pyc' -delete 10 | find . -name '__pycache__' -delete 11 | exit $exit_status 12 | } 13 | 14 | # Clean up artifacts from tests 15 | trap do_at_exit EXIT INT TERM 16 | 17 | tox -c tox_acceptance.ini -e ${ITEST_PYTHON_FACTOR}-acceptance 18 | -------------------------------------------------------------------------------- /docker/kafka/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM ubuntu:18.04 2 | MAINTAINER Team Data Streams Core 3 | ENV DEBIAN_FRONTEND noninteractive 4 | 5 | ARG KAFKA_VERSION 6 | 7 | # Install Kafka. 8 | RUN apt-get update && apt-get install -y \ 9 | unzip \ 10 | wget \ 11 | curl \ 12 | jq \ 13 | coreutils \ 14 | openjdk-8-jdk 15 | 16 | ENV JAVA_HOME "/usr/lib/jvm/java-8-openjdk-amd64/" 17 | ENV PATH "$PATH:$JAVA_HOME/bin" 18 | ENV SCALA_VERSION="2.11" 19 | ENV KAFKA_HOME /opt/kafka_${SCALA_VERSION}-${KAFKA_VERSION} 20 | COPY download-kafka.sh /tmp/download-kafka.sh 21 | RUN chmod 755 /tmp/download-kafka.sh 22 | RUN /tmp/download-kafka.sh && tar xfz /tmp/kafka_${SCALA_VERSION}-${KAFKA_VERSION}.tgz -C /opt && rm /tmp/kafka_${SCALA_VERSION}-${KAFKA_VERSION}.tgz 23 | ENV PATH="$PATH:$KAFKA_HOME/bin" 24 | 25 | COPY config.properties /server.properties 26 | 27 | CMD echo "Kafka starting" && kafka-server-start.sh /server.properties 28 | -------------------------------------------------------------------------------- /docker/kafka/download-kafka.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | 3 | url="https://archive.apache.org/dist/kafka/${KAFKA_VERSION}/kafka_${SCALA_VERSION}-${KAFKA_VERSION}.tgz" 4 | wget -q "${url}" -O "/tmp/kafka_${SCALA_VERSION}-${KAFKA_VERSION}.tgz" 5 | -------------------------------------------------------------------------------- /docker/zookeeper/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM ubuntu:18.04 2 | MAINTAINER Team Data Streams Core 3 | ENV DEBIAN_FRONTEND noninteractive 4 | 5 | RUN apt-get update && apt-get -y install zookeeper 6 | 7 | CMD /usr/share/zookeeper/bin/zkServer.sh start-foreground 8 | -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- 1 | # Makefile for Sphinx documentation 2 | # 3 | 4 | # You can set these variables from the command line. 5 | SPHINXOPTS = 6 | SPHINXBUILD = sphinx-build 7 | PAPER = 8 | BUILDDIR = build 9 | 10 | # Internal variables. 11 | PAPEROPT_a4 = -D latex_paper_size=a4 12 | PAPEROPT_letter = -D latex_paper_size=letter 13 | ALLSPHINXOPTS = -d $(BUILDDIR)/doctrees $(PAPEROPT_$(PAPER)) $(SPHINXOPTS) source 14 | 15 | .PHONY: help clean html dirhtml pickle json htmlhelp qthelp latex changes linkcheck doctest 16 | 17 | help: 18 | @echo "Please use \`make ' where is one of" 19 | @echo " html to make standalone HTML files" 20 | @echo " dirhtml to make HTML files named index.html in directories" 21 | @echo " pickle to make pickle files" 22 | @echo " json to make JSON files" 23 | @echo " htmlhelp to make HTML files and a HTML help project" 24 | @echo " qthelp to make HTML files and a qthelp project" 25 | @echo " latex to make LaTeX files, you can set PAPER=a4 or PAPER=letter" 26 | @echo " changes to make an overview of all changed/added/deprecated items" 27 | @echo " linkcheck to check all external links for integrity" 28 | @echo " doctest to run all doctests embedded in the documentation (if enabled)" 29 | 30 | clean: 31 | -rm -rf $(BUILDDIR)/* 32 | 33 | html: 34 | $(SPHINXBUILD) -b html $(ALLSPHINXOPTS) $(BUILDDIR)/html 35 | @echo 36 | @echo "Build finished. The HTML pages are in $(BUILDDIR)/html." 37 | 38 | dirhtml: 39 | $(SPHINXBUILD) -b dirhtml $(ALLSPHINXOPTS) $(BUILDDIR)/dirhtml 40 | @echo 41 | @echo "Build finished. The HTML pages are in $(BUILDDIR)/dirhtml." 42 | 43 | pickle: 44 | $(SPHINXBUILD) -b pickle $(ALLSPHINXOPTS) $(BUILDDIR)/pickle 45 | @echo 46 | @echo "Build finished; now you can process the pickle files." 47 | 48 | json: 49 | $(SPHINXBUILD) -b json $(ALLSPHINXOPTS) $(BUILDDIR)/json 50 | @echo 51 | @echo "Build finished; now you can process the JSON files." 52 | 53 | htmlhelp: 54 | $(SPHINXBUILD) -b htmlhelp $(ALLSPHINXOPTS) $(BUILDDIR)/htmlhelp 55 | @echo 56 | @echo "Build finished; now you can run HTML Help Workshop with the" \ 57 | ".hhp project file in $(BUILDDIR)/htmlhelp." 58 | 59 | qthelp: 60 | $(SPHINXBUILD) -b qthelp $(ALLSPHINXOPTS) $(BUILDDIR)/qthelp 61 | @echo 62 | @echo "Build finished; now you can run "qcollectiongenerator" with the" \ 63 | ".qhcp project file in $(BUILDDIR)/qthelp, like this:" 64 | @echo "# qcollectiongenerator $(BUILDDIR)/qthelp/Kafka-Utils.qhcp" 65 | @echo "To view the help file:" 66 | @echo "# assistant -collectionFile $(BUILDDIR)/qthelp/Kafka-Utils.qhc" 67 | 68 | latex: 69 | $(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(BUILDDIR)/latex 70 | @echo 71 | @echo "Build finished; the LaTeX files are in $(BUILDDIR)/latex." 72 | @echo "Run \`make all-pdf' or \`make all-ps' in that directory to" \ 73 | "run these through (pdf)latex." 74 | 75 | changes: 76 | $(SPHINXBUILD) -b changes $(ALLSPHINXOPTS) $(BUILDDIR)/changes 77 | @echo 78 | @echo "The overview file is in $(BUILDDIR)/changes." 79 | 80 | linkcheck: 81 | $(SPHINXBUILD) -b linkcheck $(ALLSPHINXOPTS) $(BUILDDIR)/linkcheck 82 | @echo 83 | @echo "Link check complete; look for any errors in the above output " \ 84 | "or in $(BUILDDIR)/linkcheck/output.txt." 85 | 86 | doctest: 87 | $(SPHINXBUILD) -b doctest $(ALLSPHINXOPTS) $(BUILDDIR)/doctest 88 | @echo "Testing of doctests in the sources finished, look at the " \ 89 | "results in $(BUILDDIR)/doctest/output.txt." 90 | -------------------------------------------------------------------------------- /docs/source/config.rst: -------------------------------------------------------------------------------- 1 | Configuration 2 | ############# 3 | 4 | Kafka-Utils reads the cluster configuration needed to access Kafka clusters from yaml files. 5 | Each cluster is identified by *type* and *name*. 6 | Multiple clusters of the same type should be listed in the same `type.yaml` file. 7 | The yaml files are read from :code:`$KAFKA_DISCOVERY_DIR`, :code:`$HOME/.kafka_discovery` and :code:`/etc/kafka_discovery`, 8 | the former overrides the latter. 9 | 10 | Sample configuration for :code:`sample_type` cluster at :code:`/etc/kafka_discovery/sample_type.yaml` 11 | 12 | .. code-block:: yaml 13 | 14 | --- 15 | clusters: 16 | cluster-1: 17 | broker_list: 18 | - "cluster-elb-1:9092" 19 | zookeeper: "11.11.11.111:2181,11.11.11.112:2181,11.11.11.113:2181/kafka-1" 20 | cluster-2: 21 | broker_list: 22 | - "cluster-elb-2:9092" 23 | zookeeper: "11.11.11.211:2181,11.11.11.212:2181,11.11.11.213:2181/kafka-2" 24 | local_config: 25 | cluster: cluster-1 26 | 27 | For example the kafka-cluster-manager command: 28 | 29 | .. code:: bash 30 | 31 | $ kafka-cluster-manager --cluster-type sample_type stats 32 | 33 | will pick up default cluster `cluster-1` from the local_config at /etc/kafka_discovery/sample_type.yaml to display 34 | statistics of default kafka-configuration. 35 | -------------------------------------------------------------------------------- /docs/source/index.rst: -------------------------------------------------------------------------------- 1 | Kafka-Utils v\ |version| 2 | ######################## 3 | 4 | Description 5 | *********** 6 | Kafka-Utils is a library containing tools to interact with kafka clusters and manage them. The tool provides utilities 7 | like listing of all the clusters, balancing the partition distribution across brokers and replication-groups, managing 8 | consumer groups, rolling-restart of the cluster, cluster healthchecks. 9 | 10 | For more information about Apache Kafka see the official `Kafka documentation`_. 11 | 12 | How to install 13 | ************** 14 | .. code-block:: bash 15 | 16 | $ pip install kafka-utils 17 | 18 | 19 | List available clusters. 20 | 21 | .. code-block:: bash 22 | 23 | $ kafka-utils 24 | Cluster type sample_type: 25 | Cluster name: cluster-1 26 | broker list: cluster-elb-1:9092 27 | zookeeper: 11.11.11.111:2181,11.11.11.112:2181,11.11.11.113:2181/kafka-1 28 | 29 | 30 | .. _Kafka documentation: http://kafka.apache.org/documentation.html#introduction 31 | 32 | .. toctree:: 33 | :maxdepth: -1 34 | 35 | config 36 | kafka_cluster_manager 37 | kafka_consumer_manager 38 | kafka_rolling_restart 39 | kafka_check 40 | kafka_corruption_check 41 | 42 | 43 | Indices and tables 44 | ================== 45 | 46 | * :ref:`genindex` 47 | * :ref:`modindex` 48 | * :ref:`search` 49 | -------------------------------------------------------------------------------- /docs/source/kafka_check.rst: -------------------------------------------------------------------------------- 1 | Kafka Check 2 | *********** 3 | 4 | The kafka-check command performs multiple checks on the health of the cluster. 5 | Each subcommand will run a different check. The tool can run on the broker 6 | itself or on any other machine, and it will check the health of the entire 7 | cluster. 8 | 9 | One possible way to deploy the tool is to install the kafka-utils package on 10 | every broker, and schedule kafka-check to run periodically on each machine 11 | with cron. Kafka-check provides two simple coordination mechanisms to make 12 | sure that the check only runs on a single broker per cluster. 13 | 14 | Coordination strategies: 15 | 16 | * First broker only: the script will only run on the broker with lowest broker id. 17 | * Controller only: the script will only run on the controller of the cluster. 18 | 19 | Coordination parameters: 20 | 21 | * :code:`--broker-id`: the id of the broker where the script is running. 22 | Set it to -1 if automatic broker ids are used. 23 | * :code:`--data-path DATA_PATH`: Path to the Kafka data folder, used in case of 24 | automatic broker ids to find the assigned id. 25 | * :code:`--controller-only`: if is specified, the script will only run on the 26 | controller. The execution on other brokers won't perform any check and it 27 | will always succeed. 28 | * :code:`--first-broker-only`: if specified, the command will only perform the 29 | check if broker_id is the lowest broker id in the cluster. If it is not the ' 30 | lowest, it will not perform any check and succeed immediately. 31 | 32 | Checking in-sync replicas 33 | ========================= 34 | The :code:`min_isr` subcommand checks if the number of in-sync replicas for a 35 | partition is equal or greater than the minimum number of in-sync replicas 36 | configured for the topic the partition belongs to. A topic specific 37 | :code:`min.insync.replicas` overrides the given default. 38 | 39 | The parameters for min_isr check are: 40 | 41 | * :code:`--default_min_isr DEFAULT_MIN_ISR`: Default min.isr value for cases without 42 | settings in Zookeeper for some topics. 43 | 44 | .. code-block:: bash 45 | 46 | $ kafka-check --cluster-type=sample_type min_isr 47 | OK: All replicas in sync. 48 | 49 | In case of min isr violations: 50 | 51 | .. code-block:: bash 52 | 53 | $ kafka-check --cluster-type=sample_type min_isr --default-min-isr 3 54 | 55 | isr=2 is lower than min_isr=3 for sample_topic:0 56 | CRITICAL: 1 partition(s) have the number of replicas in sync that is lower 57 | than the specified min ISR. 58 | 59 | Checking replicas available 60 | =========================== 61 | The :code:`replica_unavailability` subcommand checks if the number of replicas not 62 | available for communication is equal to zero. It will report the aggregated result 63 | of unavailable replicas of each broker if any. 64 | 65 | The parameters specific to replica_unavailability check are: 66 | 67 | .. code-block:: bash 68 | 69 | $ kafka-check --cluster-type=sample_type replica_unavailability 70 | OK: All replicas available for communication. 71 | 72 | In case of not first broker in the broker list in Zookeeper: 73 | 74 | .. code-block:: bash 75 | 76 | $ kafka-check --cluster-type=sample_type --broker-id 3 replica_unavailability --first-broker-only 77 | OK: Provided broker is not the first in broker-list. 78 | 79 | In case where some partitions replicas not available for communication. 80 | 81 | .. code-block:: bash 82 | 83 | $ kafka-check --cluster-type=sample_type replica_unavailability 84 | CRITICAL: 2 replica(s) unavailable for communication. 85 | 86 | Checking offline partitions 87 | =========================== 88 | The :code:`offline` subcommand checks if there are any offline partitions in the cluster. 89 | If any offline partition is found, it will terminate with an error, indicating the number 90 | of offline partitions. 91 | 92 | .. code-block:: bash 93 | 94 | $ kafka-check --cluster-type=sample_type offline 95 | CRITICAL: 64 offline partitions. 96 | -------------------------------------------------------------------------------- /docs/source/kafka_corruption_check.rst: -------------------------------------------------------------------------------- 1 | Corruption Check 2 | **************** 3 | 4 | The kafka-corruption-check script performs a check on the log files stored on 5 | the Kafka brokers. This tool finds all the log files modified in the specified 6 | time range and runs `DumpLogSegments 7 | `_ 8 | on them. The output is collected and filtered, and all information related to 9 | corrupted messages will be reported to the user. 10 | 11 | Even though this tool executes the log check with a low ionice priority, it can 12 | slow down the cluster given the high number of io operations required. Consider 13 | decreasing the batch size to reduce the additional load. 14 | 15 | Parameters 16 | ========== 17 | 18 | The parameters specific for kafka-corruption-check are: 19 | 20 | * ``--minutes N``: check the log files modified in the last ``N`` minutes. 21 | * ``--start-time START_TIME``: check the log files modified after 22 | ``START_TIME``. Example format: ``--start-time "2015-11-26 11:00:00"`` 23 | * ``--end-time END_TIME``: check the log files modified before ``END_TIME``. 24 | Example format: ``--end-time "2015-11-26 12:00:00"`` 25 | * ``--data-path``: the path to the data files on the Kafka broker. 26 | * ``--java-home``: the JAVA_HOME on the Kafka broker. 27 | * ``--batch-size BATCH_SIZE``: the number of files that will be checked 28 | in parallel on each broker. Default: 5. 29 | * ``--check-replicas``: if set it will also check the data on replicas. 30 | Default: false. 31 | * ``--verbose``: enable verbose output. 32 | 33 | Examples 34 | ======== 35 | 36 | Check all the files (leaders only) in the generic dev cluster and which were 37 | modified in the last 30 minutes: 38 | 39 | .. code-block:: bash 40 | 41 | $ kafka-corruption-check --cluster-type generic --cluster-name dev --data-path /var/kafka-logs --minutes 30 42 | Filtering leaders 43 | Broker: 0, leader of 9 over 13 files 44 | Broker: 1, leader of 4 over 11 files 45 | Starting 2 parallel processes 46 | Broker: broker0.example.org, 9 files to check 47 | Broker: broker1.example.org, 4 files to check 48 | Processes running: 49 | broker0.example.org: file 0 of 9 50 | broker0.example.org: file 5 of 9 51 | ERROR Host: broker0.example.org: /var/kafka-logs/test_topic-0/00000000000000003363.log 52 | ERROR Output: offset: 3371 position: 247 isvalid: false payloadsize: 22 magic: 0 compresscodec: NoCompressionCodec crc: 2230473982 53 | broker1.example.org: file 0 of 4 54 | 55 | In this example, one corrupted file was found in broker 0. 56 | 57 | Check all the files modified after the specified date, in both leaders and replicas: 58 | 59 | .. code-block:: bash 60 | 61 | $ kafka-corruption-check [...] --start-time "2015-11-26 11:00:00" --check-replicas 62 | 63 | Check all the files that were modified in the specified range: 64 | 65 | .. code-block:: bash 66 | 67 | $ kafka-corruption-check [...] --start-time "2015-11-26 11:00:00" --end-time "2015-11-26 12:00:00" 68 | -------------------------------------------------------------------------------- /kafka_utils/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2021 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | __version__ = '4.1.0' 15 | -------------------------------------------------------------------------------- /kafka_utils/kafka_check/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /kafka_utils/kafka_check/commands/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /kafka_utils/kafka_check/commands/offline.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from __future__ import annotations 15 | 16 | import itertools 17 | import sys 18 | from collections.abc import Collection 19 | from typing import Any 20 | from typing import cast 21 | from typing import Set 22 | from typing import Tuple 23 | 24 | from kafka_utils.kafka_check import status_code 25 | from kafka_utils.kafka_check.commands.command import KafkaCheckCmd 26 | from kafka_utils.util.metadata import get_topic_partition_with_error 27 | from kafka_utils.util.metadata import LEADER_NOT_AVAILABLE_ERROR 28 | 29 | 30 | class OfflineCmd(KafkaCheckCmd): 31 | 32 | def build_subparser(self, subparsers: Any) -> Any: 33 | subparser = subparsers.add_parser( 34 | 'offline', 35 | description='Check offline partitions on the specified broker', 36 | help='This subcommand will fail if there are any offline partitions ' 37 | 'in the cluster.' 38 | ) 39 | 40 | return subparser 41 | 42 | def run_command(self) -> tuple[int, dict[str, Any]]: 43 | """Checks the number of offline partitions""" 44 | offline = cast(Set[Tuple[str, int]], get_topic_partition_with_error( 45 | self.cluster_config, 46 | LEADER_NOT_AVAILABLE_ERROR, 47 | )) 48 | 49 | errcode = status_code.OK if not offline else status_code.CRITICAL 50 | out = _prepare_output(offline, self.args.verbose, self.args.head) 51 | return errcode, out 52 | 53 | 54 | def _prepare_output(partitions: Collection[tuple[str, int]], verbose: bool, head_limit: int) -> dict[str, Any]: 55 | """Returns dict with 'raw' and 'message' keys filled.""" 56 | out: dict[str, Any] = {} 57 | partitions_count = len(partitions) 58 | out['raw'] = { 59 | 'offline_count': partitions_count, 60 | } 61 | 62 | if head_limit != -1: 63 | partitions = list(itertools.islice(partitions, head_limit)) 64 | 65 | if partitions_count == 0: 66 | out['message'] = 'No offline partitions.' 67 | else: 68 | out['message'] = f"{partitions_count} offline partitions." 69 | if verbose: 70 | lines = ( 71 | f'{topic}:{partition}' 72 | for (topic, partition) in partitions 73 | ) 74 | title = f"Top {head_limit} partitions:\n" if head_limit != -1 else "Partitions:\n" 75 | out['verbose'] = title + "\n".join(lines) 76 | else: 77 | cmdline = sys.argv[:] 78 | cmdline.insert(1, '-v') 79 | out['message'] += '\nTo see all offline partitions run: ' + ' '.join(cmdline) 80 | 81 | if verbose: 82 | out['raw']['partitions'] = [ 83 | {'topic': topic, 'partition': partition} 84 | for (topic, partition) in partitions 85 | ] 86 | 87 | return out 88 | -------------------------------------------------------------------------------- /kafka_utils/kafka_check/metadata_file.py: -------------------------------------------------------------------------------- 1 | from __future__ import annotations 2 | 3 | from typing import TextIO 4 | 5 | 6 | def _parse_meta_properties_file(content: TextIO) -> int | None: 7 | for line in content: 8 | parts = line.rstrip().split("=") 9 | if len(parts) == 2 and parts[0] == "broker.id": 10 | return int(parts[1]) 11 | return None 12 | 13 | 14 | def _read_generated_broker_id(meta_properties_path: str) -> int: 15 | """reads broker_id from meta.properties file. 16 | 17 | :param string meta_properties_path: path for meta.properties file 18 | :returns int: broker_id from meta_properties_path 19 | """ 20 | try: 21 | with open(meta_properties_path) as f: 22 | broker_id = _parse_meta_properties_file(f) 23 | except OSError: 24 | raise OSError( 25 | f"Cannot open meta.properties file: {meta_properties_path}", 26 | ) 27 | except ValueError: 28 | raise ValueError("Broker id not valid") 29 | 30 | if broker_id is None: 31 | raise ValueError("Autogenerated broker id missing from data directory") 32 | 33 | return broker_id 34 | 35 | 36 | def get_broker_id(data_path: str) -> int: 37 | """This function will look into the data folder to get the automatically created 38 | broker_id. 39 | 40 | :param string data_path: the path to the kafka data folder 41 | :returns int: the real broker_id 42 | """ 43 | 44 | # Path to the meta.properties file. This is used to read the automatic broker id 45 | # if the given broker id is -1 46 | META_FILE_PATH = "{data_path}/meta.properties" 47 | 48 | if not data_path: 49 | raise ValueError("You need to specify the data_path if broker_id == -1") 50 | meta_properties_path = META_FILE_PATH.format(data_path=data_path) 51 | return _read_generated_broker_id(meta_properties_path) 52 | -------------------------------------------------------------------------------- /kafka_utils/kafka_check/status_code.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from __future__ import annotations 15 | 16 | import sys 17 | from typing import Any 18 | 19 | from typing_extensions import TypedDict 20 | 21 | from kafka_utils.util import print_json 22 | 23 | OK = 0 24 | WARNING = 1 25 | CRITICAL = 2 26 | 27 | STATUS_STRING = { 28 | OK: 'OK', 29 | WARNING: 'WARNING', 30 | CRITICAL: 'CRITICAL', 31 | } 32 | 33 | 34 | class TerminateMessageDict(TypedDict): 35 | message: str 36 | raw: str 37 | 38 | 39 | def prepare_terminate_message(string: str) -> TerminateMessageDict: 40 | return { 41 | 'message': string, 42 | 'raw': string, 43 | } 44 | 45 | 46 | def terminate(err_code: int, msg: dict[str, Any], json: bool) -> None: 47 | if json: 48 | output = { 49 | 'status': STATUS_STRING[err_code], 50 | 'data': msg['raw'], 51 | } 52 | print_json(output) 53 | else: 54 | print('{status}: {msg}'.format( 55 | status=STATUS_STRING[err_code], 56 | msg=msg['message'], 57 | )) 58 | if 'verbose' in msg: 59 | print(msg['verbose']) 60 | sys.exit(err_code) 61 | -------------------------------------------------------------------------------- /kafka_utils/kafka_cluster_manager/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /kafka_utils/kafka_cluster_manager/cluster_info/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /kafka_utils/kafka_cluster_manager/cluster_info/error.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from kafka_utils.util.error import KafkaToolError 15 | 16 | 17 | class InvalidBrokerIdError(KafkaToolError): 18 | """Raised when a broker id doesn't exist in the cluster.""" 19 | pass 20 | 21 | 22 | class InvalidPartitionError(KafkaToolError): 23 | """Raised when a partition tuple (topic, partition) doesn't exist in the cluster""" 24 | pass 25 | 26 | 27 | class InvalidReplicationFactorError(KafkaToolError): 28 | """Raised when an operation would result in an replication factor that is 29 | too small or too large for the cluster. 30 | """ 31 | pass 32 | 33 | 34 | class InvalidPartitionMeasurementError(KafkaToolError): 35 | """Raised when a partition is assigned a negative weight or size.""" 36 | pass 37 | 38 | 39 | class EmptyReplicationGroupError(KafkaToolError): 40 | """Raised when there are no brokers in a replication group.""" 41 | pass 42 | 43 | 44 | class BrokerDecommissionError(KafkaToolError): 45 | """Raised if it is not possible to move partition out 46 | from decommissioned brokers. 47 | """ 48 | pass 49 | 50 | 51 | class NotEligibleGroupError(KafkaToolError): 52 | """Raised when there are no brokers eligible to acquire a certain partition 53 | in a replication group. 54 | """ 55 | pass 56 | 57 | 58 | class RebalanceError(KafkaToolError): 59 | """Raised when a rebalance operation is not possible.""" 60 | pass 61 | -------------------------------------------------------------------------------- /kafka_utils/kafka_cluster_manager/cluster_info/partition_measurer.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from __future__ import annotations 15 | 16 | import argparse 17 | import itertools 18 | import shlex 19 | 20 | from kafka_utils.kafka_cluster_manager.cluster_info.broker import Broker 21 | from kafka_utils.util.config import ClusterConfig 22 | 23 | 24 | class PartitionMeasurer: 25 | """An interface used to gather metrics about a partition. 26 | 27 | :param cluster_config: ClusterConfig for the cluster. 28 | :param brokers: List of the cluster's brokers. 29 | :param assignment: The cluster's assignment. 30 | :param args: Namespace containing the command line arguments. 31 | """ 32 | 33 | def __init__( 34 | self, 35 | cluster_config: ClusterConfig, 36 | brokers: list[Broker], 37 | assignment: dict[tuple[str, int], list[int]], 38 | args: argparse.Namespace, 39 | ): 40 | self.cluster_config = cluster_config 41 | self.brokers = brokers 42 | self.assignment = assignment 43 | self.args = args 44 | if hasattr(args, 'measurer_args'): 45 | self.parse_args(list(itertools.chain.from_iterable( 46 | shlex.split(arg) for arg in args.measurer_args 47 | ))) 48 | else: 49 | self.parse_args([]) 50 | 51 | def parse_args(self, _measurer_args: list[str]) -> None: 52 | """Parse partition measurer command line arguments. 53 | 54 | :param _measurer_args: The list of arguments as strings. 55 | """ 56 | pass 57 | 58 | def get_weight(self, partition_name: tuple[str, int]) -> float: 59 | """Return a positive number representing the relative weight of this 60 | partition compared to the other partitions in the cluster. The weight 61 | is a measure of how much load this partition will place on any broker 62 | that it is assigned to. 63 | 64 | :param partition_name: A tuple with the topic id and partition id as the first and second elements respectively. 65 | """ 66 | raise NotImplementedError("Implement in subclass.") 67 | 68 | def get_size(self, partition_name: tuple[str, int]) -> float: 69 | """Return a positive number representing the size of this partition. 70 | The size is a measure of how expensive it is to move this partition 71 | from one broker to another. 72 | 73 | :param partition_name: A tuple with the topic id and partition id as the first and second elements respectively. 74 | """ 75 | raise NotImplementedError("Implement in subclass.") 76 | 77 | 78 | class UniformPartitionMeasurer(PartitionMeasurer): 79 | """An implementation of PartitionMeasurer that provides identital metrics 80 | for all partitions. 81 | """ 82 | 83 | def get_weight(self, partition_name: tuple[str, int]) -> float: 84 | return 1.0 85 | 86 | def get_size(self, partition_name: tuple[str, int]) -> float: 87 | return 1.0 88 | -------------------------------------------------------------------------------- /kafka_utils/kafka_cluster_manager/cluster_info/replication_group_parser.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from __future__ import annotations 15 | 16 | from kafka_utils.kafka_cluster_manager.cluster_info.broker import Broker 17 | 18 | 19 | class ReplicationGroupParser: 20 | """Base class for replication group parsers""" 21 | 22 | def get_replication_group(self, broker: Broker) -> str | None: 23 | """Implement the logic to extract the replication group id of a broker. 24 | 25 | This method is called with a broker object as argument. 26 | 27 | Example: 28 | 29 | .. code-block:: python 30 | 31 | class MyParser(ReplicationGroupParser): 32 | 33 | def get_replication_group(self, broker): 34 | return broker.metadata['host'].split('.', 2)[1] 35 | 36 | :param broker: py:class:`kafka_utils.kafka_cluster_manager.cluster_info.broker.Broker` 37 | :returns: a string representing the replication group name of the broker 38 | """ 39 | raise NotImplementedError("Implement in subclass") 40 | 41 | 42 | class DefaultReplicationGroupParser(ReplicationGroupParser): 43 | 44 | def get_replication_group(self, broker: Broker) -> str | None: 45 | """Default group is None. All brokers are considered in the same group. 46 | 47 | TODO: Support kafka 0.10 and rack tag. 48 | """ 49 | return None 50 | -------------------------------------------------------------------------------- /kafka_utils/kafka_cluster_manager/cluster_info/topic.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | """This class contains information for a topic object. 15 | 16 | Useful as part of reassignment project when deciding upon moving 17 | partitions of same topic over different brokers. 18 | """ 19 | from __future__ import annotations 20 | 21 | import logging 22 | from typing import TYPE_CHECKING 23 | 24 | if TYPE_CHECKING: 25 | from kafka_utils.kafka_cluster_manager.cluster_info.partition import Partition 26 | 27 | 28 | class Topic: 29 | """Information of a topic object. 30 | 31 | :params 32 | id: Name of the given topic 33 | replication_factor: replication factor of a given topic 34 | partitions: List of Partition objects 35 | """ 36 | 37 | def __init__(self, id: str, replication_factor: int = 0, partitions: set[Partition] | None = None) -> None: 38 | self._id = id 39 | self._replication_factor = replication_factor 40 | self._partitions = partitions or set() 41 | self.log = logging.getLogger(self.__class__.__name__) 42 | 43 | @property 44 | def id(self) -> str: 45 | return self._id 46 | 47 | @property 48 | def replication_factor(self) -> int: 49 | return self._replication_factor 50 | 51 | @property 52 | def partitions(self) -> set[Partition]: 53 | return self._partitions 54 | 55 | @property 56 | def weight(self) -> float: 57 | return sum( 58 | partition.weight * partition.replication_factor 59 | for partition in self._partitions 60 | ) 61 | 62 | def add_partition(self, partition: Partition) -> None: 63 | self._partitions.add(partition) 64 | 65 | def __str__(self) -> str: 66 | return f"{self._id}" 67 | 68 | def __repr__(self) -> str: 69 | return f"{self}" 70 | -------------------------------------------------------------------------------- /kafka_utils/kafka_cluster_manager/cmds/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /kafka_utils/kafka_cluster_manager/cmds/preferred_replica_election.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | import json 15 | import logging 16 | import sys 17 | 18 | from .command import ClusterManagerCmd 19 | 20 | 21 | class PreferredReplicaElectionCmd(ClusterManagerCmd): 22 | 23 | def __init__(self): 24 | super().__init__() 25 | self.log = logging.getLogger(self.__class__.__name__) 26 | 27 | def build_subparser(self, subparsers): 28 | subparser = subparsers.add_parser( 29 | 'preferred-replica-election', 30 | description='Generates json files for kafka-preferred-replica-election', 31 | help='This command is used to help generate kafka-preferred-replica-election' 32 | ' json files based on input from kafka-check.' 33 | ) 34 | subparser.add_argument( 35 | '--topic-partition-filter', 36 | type=str, 37 | required=True, 38 | help='Path to file containing a colon separated list of topic partitions to work on.' 39 | ' This is useful if you only want to operate on a limited set of topic partitions' 40 | ' Observe that kafka-check outputs colon separated lists, after ommitting the first $n' 41 | ' lines e.g. kafka-check offline | tail -n+$N', 42 | ) 43 | return subparser 44 | 45 | def run_command(self, cluster_topology, cluster_balancer): 46 | filter_set = self.get_topic_filter() 47 | output = {'partitions': []} 48 | # Validate the filter_set to check that all of the partitions provided are part of the cluster_topology 49 | for t_p in filter_set: 50 | if t_p not in cluster_topology.assignment: 51 | self.log.error(f'Topic partition {t_p} from filter not in cluster topology') 52 | sys.exit(1) 53 | topic, partition = t_p[0], t_p[1] 54 | output['partitions'].append({'topic': topic, 'partition': partition}) 55 | 56 | self.log.info('Preferred replica assignment json generated. Run kafka-preferred-replica-election using this json:') 57 | print(json.dumps(output)) 58 | 59 | if self.args.proposed_plan_file: 60 | self.write_json_plan(output, self.args.proposed_plan_file) 61 | self.log.info(f'Saved plan to file {self.args.proposed_plan_file}') 62 | -------------------------------------------------------------------------------- /kafka_utils/kafka_cluster_manager/cmds/revoke_leadership.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | import logging 15 | import sys 16 | 17 | from .command import ClusterManagerCmd 18 | from kafka_utils.util import positive_int 19 | from kafka_utils.util.validation import assignment_to_plan 20 | from kafka_utils.util.validation import validate_plan 21 | 22 | 23 | DEFAULT_MAX_LEADER_CHANGES = 5 24 | 25 | 26 | class RevokeLeadershipCmd(ClusterManagerCmd): 27 | 28 | def __init__(self): 29 | super().__init__() 30 | self.log = logging.getLogger(self.__class__.__name__) 31 | 32 | def build_subparser(self, subparsers): 33 | subparser = subparsers.add_parser( 34 | 'revoke-leadership', 35 | description='Re-assign leadership for all partitions on given brokers to other brokers', 36 | help='This command is used to move leadership for all partitions ' 37 | 'on given brokers to other brokers in balanced order. The generated plan include' 38 | ' only leadership changes.' 39 | ) 40 | subparser.add_argument( 41 | 'broker_ids', 42 | nargs='+', 43 | type=int, 44 | help='Broker ids of the brokers to revoke leadership for.', 45 | ) 46 | subparser.add_argument( 47 | '--max-leader-changes', 48 | type=positive_int, 49 | default=DEFAULT_MAX_LEADER_CHANGES, 50 | help='Maximum number of actions with leader-only changes.' 51 | ' DEFAULT: %(default)s', 52 | ) 53 | return subparser 54 | 55 | def run_command(self, cluster_topology, cluster_balancer): 56 | base_assignment = cluster_topology.assignment 57 | 58 | cluster_balancer.revoke_leadership(self.args.broker_ids) 59 | 60 | if not validate_plan( 61 | assignment_to_plan(cluster_topology.assignment), 62 | assignment_to_plan(base_assignment), 63 | ): 64 | self.log.error('Invalid assignment %s.', cluster_topology.assignment) 65 | print( 66 | f'Invalid assignment: {cluster_topology.assignment}', 67 | file=sys.stderr, 68 | ) 69 | sys.exit(1) 70 | 71 | # Reduce the proposed assignment based on max_leader_changes 72 | reduced_assignment = self.get_reduced_assignment( 73 | base_assignment, 74 | cluster_topology, 75 | 0, # Number of partition movements 76 | self.args.max_leader_changes, 77 | ) 78 | if reduced_assignment: 79 | self.process_assignment(reduced_assignment) 80 | else: 81 | msg = "Cluster already balanced. No more partitions as leaders in " \ 82 | "revoked-leadership brokers." 83 | self.log.info(msg) 84 | print(msg) 85 | -------------------------------------------------------------------------------- /kafka_utils/kafka_cluster_manager/cmds/stats.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | import json 15 | import logging 16 | 17 | from .command import ClusterManagerCmd 18 | from kafka_utils.kafka_cluster_manager.cluster_info.display import \ 19 | display_cluster_topology_stats 20 | from kafka_utils.util.validation import plan_to_assignment 21 | 22 | 23 | class StatsCmd(ClusterManagerCmd): 24 | 25 | def __init__(self): 26 | super().__init__() 27 | self.log = logging.getLogger(self.__class__.__name__) 28 | 29 | def build_subparser(self, subparsers): 30 | subparser = subparsers.add_parser( 31 | 'stats', 32 | description='Show imbalance statistics of cluster topology', 33 | help='This command is used to display imbalance statistics of current ' 34 | 'cluster-topology or cluster-topology after given assignment is ' 35 | 'applied.', 36 | ) 37 | subparser.add_argument( 38 | '--read-from-file', 39 | dest='plan_file_path', 40 | metavar='', 41 | type=str, 42 | help='Read the partition assignment from json file. Example format:' 43 | ' {"version": 1, "partitions": [{"topic": "foo", "partition": 1, ' 44 | '"replicas": [1,2,3]}]}', 45 | ) 46 | return subparser 47 | 48 | def run_command(self, cluster_topology, cluster_balancer): 49 | if self.args.plan_file_path: 50 | base_assignment = cluster_topology.assignment 51 | base_score = cluster_balancer.score() 52 | 53 | self.log.info( 54 | 'Integrating given assignment-plan in current cluster-topology.' 55 | ) 56 | cluster_topology.update_cluster_topology(self.get_assignment()) 57 | score = cluster_balancer.score() 58 | display_cluster_topology_stats(cluster_topology, base_assignment) 59 | if score is not None and base_score is not None: 60 | print('\nScore before: %f' % base_score) 61 | print('Score after: %f' % score) 62 | print('Score improvement %f' % (score - base_score)) 63 | else: 64 | score = cluster_balancer.score() 65 | display_cluster_topology_stats(cluster_topology) 66 | if score: 67 | print('\nScore: %f' % score) 68 | 69 | def get_assignment(self): 70 | """Parse the given json plan in dict format.""" 71 | try: 72 | plan = json.loads(open(self.args.plan_file_path).read()) 73 | return plan_to_assignment(plan) 74 | except OSError: 75 | self.log.exception( 76 | 'Given json file {file} not found.' 77 | .format(file=self.args.plan_file_path), 78 | ) 79 | raise 80 | except ValueError: 81 | self.log.exception( 82 | 'Given json file {file} could not be decoded.' 83 | .format(file=self.args.plan_file_path), 84 | ) 85 | raise 86 | except KeyError: 87 | self.log.exception( 88 | 'Given json file {file} could not be parsed in desired format.' 89 | .format(file=self.args.plan_file_path), 90 | ) 91 | raise 92 | -------------------------------------------------------------------------------- /kafka_utils/kafka_cluster_manager/cmds/store_assignments.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | import json 15 | import logging 16 | 17 | from .command import ClusterManagerCmd 18 | from kafka_utils.util.validation import assignment_to_plan 19 | 20 | 21 | class StoreAssignmentsCmd(ClusterManagerCmd): 22 | 23 | def __init__(self): 24 | super().__init__() 25 | self.log = logging.getLogger(self.__class__.__name__) 26 | 27 | def build_subparser(self, subparsers): 28 | subparser = subparsers.add_parser( 29 | 'store_assignments', 30 | description='Emit json encoding the current assignment of ' 31 | 'partitions to replicas.', 32 | help='''This command will not mutate the cluster\'s state. 33 | Output json is of this form: 34 | {"version":1,"partitions":[ 35 | {"topic": "foo1", "partition": 2, "replicas": [1, 2]}, 36 | {"topic": "foo1", "partition": 0, "replicas": [3, 4]}, 37 | {"topic": "foo2", "partition": 2, "replicas": [1, 2]}, 38 | {"topic": "foo2", "partition": 0, "replicas": [3, 4]}, 39 | {"topic": "foo1", "partition": 1, "replicas": [2, 3]}, 40 | {"topic": "foo2", "partition": 1, "replicas": [2, 3]}]}''' 41 | ) 42 | subparser.add_argument( 43 | '--json_out', 44 | type=str, 45 | help=('Path to json output file. ' 46 | 'Will output to stdout if not set. ' 47 | 'If file exists already, it will be clobbered.') 48 | ) 49 | return subparser 50 | 51 | def run_command(self, cluster_topology, _): 52 | plan_json = json.dumps(assignment_to_plan(cluster_topology.assignment)) 53 | if self.args.json_out: 54 | with open(self.args.json_out, 'w') as f: 55 | self.log.info( 56 | 'writing assignments as json to: %s', 57 | self.args.json_out, 58 | ) 59 | f.write(plan_json) 60 | else: 61 | self.log.info('writing assignments as json to stdout') 62 | print(plan_json) 63 | -------------------------------------------------------------------------------- /kafka_utils/kafka_consumer_manager/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | __version__ = "0.1.0" 15 | -------------------------------------------------------------------------------- /kafka_utils/kafka_consumer_manager/commands/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /kafka_utils/kafka_consumer_manager/commands/copy_group.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from __future__ import annotations 15 | 16 | import argparse 17 | import sys 18 | from typing import Any 19 | 20 | from .offset_manager import OffsetManagerBase 21 | from kafka_utils.util.client import KafkaToolClient 22 | from kafka_utils.util.config import ClusterConfig 23 | from kafka_utils.util.offsets import get_current_consumer_offsets 24 | from kafka_utils.util.offsets import set_consumer_offsets 25 | 26 | 27 | class CopyGroup(OffsetManagerBase): 28 | 29 | @classmethod 30 | def setup_subparser(cls, subparsers: Any) -> None: 31 | parser_copy_group = subparsers.add_parser( 32 | "copy_group", 33 | description="Copy specified consumer group details to a new group.", 34 | ) 35 | parser_copy_group.add_argument( 36 | 'source_groupid', 37 | help="Consumer Group to be copied.", 38 | ) 39 | parser_copy_group.add_argument( 40 | 'dest_groupid', 41 | help="New name for the consumer group being copied to.", 42 | ) 43 | parser_copy_group.add_argument( 44 | "--topic", 45 | help="Kafka topic whose offsets will be copied into destination group" 46 | " If no topic is specificed all topic offsets will be copied.", 47 | ) 48 | parser_copy_group.add_argument( 49 | "--partitions", 50 | nargs='+', 51 | type=int, 52 | help="List of partitions within the topic. If no partitions are " 53 | "specified, offsets from all partitions of the topic shall " 54 | "be copied.", 55 | ) 56 | parser_copy_group.set_defaults(command=cls.run) 57 | 58 | @classmethod 59 | def run(cls, args: argparse.Namespace, cluster_config: ClusterConfig) -> None: 60 | if args.source_groupid == args.dest_groupid: 61 | print( 62 | "Error: Source group ID and destination group ID are same.", 63 | file=sys.stderr, 64 | ) 65 | sys.exit(1) 66 | # Setup the Kafka client 67 | client = KafkaToolClient(cluster_config.broker_list) 68 | client.load_metadata_for_topics() 69 | source_topics = cls.preprocess_args( 70 | args.source_groupid, 71 | args.topic, 72 | args.partitions, 73 | cluster_config, 74 | client, 75 | use_admin_client=args.use_admin_client, 76 | ) 77 | 78 | cls.copy_group_kafka( 79 | client, 80 | source_topics, 81 | args.source_groupid, 82 | args.dest_groupid, 83 | ) 84 | 85 | @classmethod 86 | def copy_group_kafka(cls, client: KafkaToolClient, topics: dict[str, list[int]], source_group: str, destination_group: str) -> None: 87 | copied_offsets = get_current_consumer_offsets(client, source_group, topics) 88 | set_consumer_offsets(client, destination_group, copied_offsets) 89 | -------------------------------------------------------------------------------- /kafka_utils/kafka_consumer_manager/commands/delete_group.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from __future__ import annotations 15 | 16 | import argparse 17 | from typing import Any 18 | 19 | from .offset_manager import OffsetWriter 20 | from kafka_utils.util.client import KafkaToolClient 21 | from kafka_utils.util.config import ClusterConfig 22 | from kafka_utils.util.offsets import nullify_offsets 23 | from kafka_utils.util.offsets import set_consumer_offsets 24 | 25 | 26 | class DeleteGroup(OffsetWriter): 27 | 28 | @classmethod 29 | def setup_subparser(cls, subparsers: Any): 30 | parser_delete_group = subparsers.add_parser( 31 | "delete_group", 32 | description="Delete a consumer group by groupid. This " 33 | "tool shall delete all group offset metadata from Zookeeper.", 34 | add_help=False 35 | ) 36 | parser_delete_group.add_argument( 37 | "-h", "--help", action="help", 38 | help="Show this help message and exit." 39 | ) 40 | parser_delete_group.add_argument( 41 | 'groupid', 42 | help="Consumer Group IDs whose metadata shall be deleted." 43 | ) 44 | parser_delete_group.set_defaults(command=cls.run) 45 | 46 | @classmethod 47 | def run(cls, args: argparse.Namespace, cluster_config: ClusterConfig) -> None: 48 | # Setup the Kafka client 49 | client = KafkaToolClient(cluster_config.broker_list) 50 | client.load_metadata_for_topics() 51 | 52 | topics_dict = cls.preprocess_args( 53 | args.groupid, 54 | None, 55 | None, 56 | cluster_config, 57 | client, 58 | use_admin_client=args.use_admin_client, 59 | ) 60 | cls.delete_group_kafka(client, args.groupid, topics_dict) 61 | 62 | @classmethod 63 | def delete_group_kafka(cls, client: KafkaToolClient, group: str, topics: dict[str, dict[int, int]]) -> None: 64 | new_offsets = nullify_offsets(topics) 65 | set_consumer_offsets(client, group, new_offsets) 66 | -------------------------------------------------------------------------------- /kafka_utils/kafka_consumer_manager/commands/list_groups.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from __future__ import annotations 15 | 16 | import argparse 17 | from collections.abc import Collection 18 | from typing import Any 19 | 20 | from .offset_manager import OffsetManagerBase 21 | from kafka_utils.kafka_consumer_manager.util import get_kafka_group_reader 22 | from kafka_utils.util.config import ClusterConfig 23 | 24 | 25 | class ListGroups(OffsetManagerBase): 26 | 27 | @classmethod 28 | def setup_subparser(cls, subparsers: Any) -> None: 29 | parser_list_groups = subparsers.add_parser( 30 | "list_groups", 31 | description="List consumer groups.", 32 | add_help=False, 33 | ) 34 | parser_list_groups.set_defaults(command=cls.run) 35 | 36 | @classmethod 37 | def get_kafka_groups(cls, cluster_config: ClusterConfig, use_admin_client: bool = False) -> list[int]: 38 | '''Get the group_id of groups committed into Kafka.''' 39 | kafka_group_reader = get_kafka_group_reader(cluster_config, use_admin_client) 40 | groups_and_topics = kafka_group_reader.read_groups(list_only=True) 41 | return list(groups_and_topics.keys()) 42 | 43 | @classmethod 44 | def print_groups(cls, groups: Collection[int], cluster_config: ClusterConfig) -> None: 45 | print("Consumer Groups:") 46 | for groupid in groups: 47 | print(f"\t{groupid}") 48 | print( 49 | "{num_groups} groups found for cluster {cluster_name} " 50 | "of type {cluster_type}".format( 51 | num_groups=len(groups), 52 | cluster_name=cluster_config.name, 53 | cluster_type=cluster_config.type, 54 | ), 55 | ) 56 | 57 | @classmethod 58 | def run(cls, args: argparse.Namespace, cluster_config: ClusterConfig) -> None: 59 | groups = set() 60 | kafka_groups = cls.get_kafka_groups( 61 | cluster_config, 62 | args.use_admin_client, 63 | ) 64 | if kafka_groups: 65 | groups.update(kafka_groups) 66 | 67 | cls.print_groups(groups, cluster_config) 68 | -------------------------------------------------------------------------------- /kafka_utils/kafka_consumer_manager/commands/list_topics.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | import argparse 15 | import sys 16 | from typing import Any 17 | 18 | from .offset_manager import OffsetManagerBase 19 | from kafka_utils.util.client import KafkaToolClient 20 | from kafka_utils.util.config import ClusterConfig 21 | 22 | 23 | class ListTopics(OffsetManagerBase): 24 | 25 | @classmethod 26 | def setup_subparser(cls, subparsers: Any) -> None: 27 | parser_list_topics = subparsers.add_parser( 28 | "list_topics", 29 | description="List topics by consumer group.", 30 | add_help=False 31 | ) 32 | parser_list_topics.add_argument( 33 | "-h", "--help", action="help", 34 | help="Show this help message and exit." 35 | ) 36 | parser_list_topics.add_argument( 37 | 'groupid', 38 | help="Consumer Group ID whose topics shall be fetched." 39 | ) 40 | parser_list_topics.set_defaults(command=cls.run) 41 | 42 | @classmethod 43 | def run(cls, args: argparse.Namespace, cluster_config: ClusterConfig) -> None: 44 | # Setup the Kafka client 45 | client = KafkaToolClient(cluster_config.broker_list) 46 | client.load_metadata_for_topics() 47 | 48 | topics_dict = cls.preprocess_args( 49 | groupid=args.groupid, 50 | topic=None, 51 | partitions=None, 52 | cluster_config=cluster_config, 53 | client=client, 54 | fail_on_error=False, 55 | use_admin_client=args.use_admin_client, 56 | ) 57 | if not topics_dict: 58 | print("Consumer Group ID: {group} does not exist.".format( 59 | group=args.groupid, 60 | )) 61 | sys.exit(1) 62 | 63 | print(f"Consumer Group ID: {args.groupid}") 64 | for topic, partitions in topics_dict.items(): 65 | print(f"\tTopic: {topic}") 66 | print(f"\t\tPartitions: {partitions}") 67 | -------------------------------------------------------------------------------- /kafka_utils/kafka_consumer_manager/commands/offset_rewind.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | import sys 15 | 16 | from .offset_manager import OffsetWriter 17 | from kafka_utils.util.client import KafkaToolClient 18 | from kafka_utils.util.offsets import rewind_consumer_offsets 19 | 20 | 21 | class OffsetRewind(OffsetWriter): 22 | 23 | @classmethod 24 | def setup_subparser(cls, subparsers): 25 | parser_offset_rewind = subparsers.add_parser( 26 | "offset_rewind", 27 | description="Rewind consumer offsets for the specified consumer " 28 | "group to the earliest message in the topic partition", 29 | add_help=False 30 | ) 31 | parser_offset_rewind.add_argument( 32 | "-h", "--help", action="help", 33 | help="Show this help message and exit." 34 | ) 35 | parser_offset_rewind.add_argument( 36 | 'groupid', 37 | help="Consumer Group ID whose consumer offsets shall be rewinded." 38 | ) 39 | parser_offset_rewind.add_argument( 40 | "--topic", 41 | help="Kafka topic whose offsets shall be manipulated. If no topic is " 42 | "specified, offsets from all topics that the consumer is " 43 | "subscribed to, shall be rewinded." 44 | ) 45 | parser_offset_rewind.add_argument( 46 | "--partitions", nargs='+', type=int, 47 | help="List of partitions within the topic. If no partitions are " 48 | "specified, offsets from all partitions of the topic shall " 49 | "be rewinded." 50 | ) 51 | parser_offset_rewind.add_argument( 52 | '--force', 53 | action='store_true', 54 | help="Force the offset of the group to be committed even if " 55 | "it does not already exist." 56 | ) 57 | parser_offset_rewind.set_defaults(command=OffsetRewind.run) 58 | 59 | @classmethod 60 | def run(cls, args, cluster_config): 61 | # Setup the Kafka client 62 | client = KafkaToolClient(cluster_config.broker_list) 63 | client.load_metadata_for_topics() 64 | 65 | topics_dict = cls.preprocess_args( 66 | args.groupid, 67 | args.topic, 68 | args.partitions, 69 | cluster_config, 70 | client, 71 | force=args.force, 72 | use_admin_client=args.use_admin_client, 73 | ) 74 | try: 75 | rewind_consumer_offsets( 76 | client, 77 | args.groupid, 78 | topics_dict, 79 | ) 80 | except TypeError: 81 | print( 82 | "Error: Badly formatted input, please re-run command " 83 | "with --help option.", file=sys.stderr 84 | ) 85 | raise 86 | 87 | client.close() 88 | -------------------------------------------------------------------------------- /kafka_utils/kafka_consumer_manager/commands/rename_group.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | import sys 15 | 16 | from .offset_manager import OffsetManagerBase 17 | from kafka_utils.util.client import KafkaToolClient 18 | from kafka_utils.util.offsets import get_current_consumer_offsets 19 | from kafka_utils.util.offsets import nullify_offsets 20 | from kafka_utils.util.offsets import set_consumer_offsets 21 | 22 | 23 | class RenameGroup(OffsetManagerBase): 24 | 25 | @classmethod 26 | def setup_subparser(cls, subparsers): 27 | parser_rename_group = subparsers.add_parser( 28 | "rename_group", 29 | description="Rename specified consumer group ID to a new name. " 30 | "This tool shall migrate all offset metadata in Zookeeper.", 31 | add_help=False 32 | ) 33 | parser_rename_group.add_argument( 34 | "-h", "--help", action="help", 35 | help="Show this help message and exit." 36 | ) 37 | parser_rename_group.add_argument( 38 | 'old_groupid', 39 | help="Consumer Group ID to be renamed." 40 | ) 41 | parser_rename_group.add_argument( 42 | 'new_groupid', 43 | help="New name for the consumer group ID." 44 | ) 45 | parser_rename_group.set_defaults(command=cls.run) 46 | 47 | @classmethod 48 | def run(cls, args, cluster_config): 49 | if args.old_groupid == args.new_groupid: 50 | print( 51 | "Error: Old group ID and new group ID are the same.", 52 | file=sys.stderr, 53 | ) 54 | sys.exit(1) 55 | # Setup the Kafka client 56 | client = KafkaToolClient(cluster_config.broker_list) 57 | client.load_metadata_for_topics() 58 | 59 | topics_dict = cls.preprocess_args( 60 | groupid=args.old_groupid, 61 | topic=None, 62 | partitions=None, 63 | cluster_config=cluster_config, 64 | client=client, 65 | use_admin_client=args.use_admin_client, 66 | ) 67 | cls.rename_group( 68 | client, 69 | args.old_groupid, 70 | args.new_groupid, 71 | topics_dict, 72 | ) 73 | 74 | @classmethod 75 | def rename_group( 76 | cls, 77 | client, 78 | old_groupid, 79 | new_groupid, 80 | topics, 81 | ): 82 | copied_offsets = get_current_consumer_offsets( 83 | client, 84 | old_groupid, 85 | topics, 86 | ) 87 | set_consumer_offsets(client, new_groupid, copied_offsets) 88 | set_consumer_offsets( 89 | client, 90 | old_groupid, 91 | nullify_offsets(topics), 92 | ) 93 | -------------------------------------------------------------------------------- /kafka_utils/kafka_corruption_check/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Yelp/kafka-utils/def433ec4d07c60290d5dc937d3b4e5189eca9dc/kafka_utils/kafka_corruption_check/__init__.py -------------------------------------------------------------------------------- /kafka_utils/kafka_manual_throttle/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /kafka_utils/kafka_rolling_restart/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /kafka_utils/kafka_rolling_restart/task.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | import shlex 15 | 16 | 17 | class TaskFailedException(Exception): 18 | pass 19 | 20 | 21 | class Task: 22 | """Base class for implementing Task 23 | All the args passed can be accessed via self.args 24 | 25 | :param args: The program arguments 26 | """ 27 | 28 | def __init__(self, args): 29 | if args: 30 | self.args = self.parse_args(list( 31 | shlex.split(args) 32 | )) 33 | else: 34 | self.args = self.parse_args([]) 35 | 36 | def parse_args(self, args): 37 | """Parse args command line arguments. 38 | :param args: The list of arguments as strings. 39 | """ 40 | pass 41 | 42 | def run(self, host): 43 | """This contains the main logic of the task 44 | Please note an exception from this method will completely stop the restart 45 | :param host: the host on which precheck is executed on 46 | :type host: string 47 | """ 48 | raise NotImplementedError("Implemented in subclass") 49 | 50 | 51 | class PreStopTask(Task): 52 | """Class to be used for any pre stop checks""" 53 | 54 | 55 | class PostStopTask(Task): 56 | """Class to be used for any post stop checks""" 57 | -------------------------------------------------------------------------------- /kafka_utils/main.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | import argparse 15 | import logging 16 | 17 | from kafka_utils import __version__ 18 | from kafka_utils.util.config import iter_configurations 19 | 20 | 21 | logging.getLogger().addHandler(logging.NullHandler()) 22 | 23 | 24 | def parse_args(): 25 | """Parse the arguments.""" 26 | parser = argparse.ArgumentParser( 27 | description='Show available clusters.' 28 | ) 29 | parser.add_argument( 30 | '-v', 31 | '--version', 32 | action='version', 33 | version=f"%(prog)s {__version__}", 34 | ) 35 | parser.add_argument( 36 | '--discovery-base-path', 37 | dest='discovery_base_path', 38 | type=str, 39 | help='Path of the directory containing the .yaml config.' 40 | ' Default try: ' 41 | '$KAFKA_DISCOVERY_DIR, $HOME/.kafka_discovery, /etc/kafka_discovery', 42 | ) 43 | 44 | return parser.parse_args() 45 | 46 | 47 | def run(): 48 | args = parse_args() 49 | 50 | for config in iter_configurations(args.discovery_base_path): 51 | print(f"cluster-type {config.cluster_type}:") 52 | for cluster in config.get_all_clusters(): 53 | print( 54 | "\tcluster-name: {name}\n" 55 | "\tbroker-list: {b_list}\n" 56 | "\tzookeeper: {zk}".format( 57 | name=cluster.name, 58 | b_list=", ".join(cluster.broker_list), 59 | zk=cluster.zookeeper, 60 | ) 61 | ) 62 | -------------------------------------------------------------------------------- /kafka_utils/util/error.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from kafka.structs import OffsetCommitResponsePayload 15 | 16 | 17 | class KafkaToolError(Exception): 18 | """Base class for kafka tool exceptions.""" 19 | pass 20 | 21 | 22 | class ConfigurationError(KafkaToolError): 23 | """Error in configuration. For example: missing configuration file 24 | or misformatted configuration.""" 25 | pass 26 | 27 | 28 | class MissingConfigurationError(ConfigurationError): 29 | """Missing configuration file.""" 30 | pass 31 | 32 | 33 | class InvalidConfigurationError(ConfigurationError): 34 | """Invalid configuration file.""" 35 | pass 36 | 37 | 38 | class UnknownTopic(KafkaToolError): 39 | """Topic does not exist in kafka.""" 40 | pass 41 | 42 | 43 | class UnknownPartitions(KafkaToolError): 44 | """Partition doesn't exist in kafka.""" 45 | pass 46 | 47 | 48 | class OffsetCommitError(KafkaToolError): 49 | """Error during offset commit.""" 50 | 51 | def __init__(self, topic: str, partition: int, error: str) -> None: 52 | self.topic = topic 53 | self.partition = partition 54 | self.error = error 55 | 56 | def __eq__(self, other: object) -> bool: 57 | assert isinstance(other, (OffsetCommitError, OffsetCommitResponsePayload)) 58 | if all([ 59 | self.topic == other.topic, 60 | self.partition == other.partition, 61 | self.error == other.error, 62 | ]): 63 | return True 64 | return False 65 | 66 | 67 | class MaxConnectionAttemptsError(KafkaToolError): 68 | """Exceeded max connection attempts.""" 69 | pass 70 | -------------------------------------------------------------------------------- /kafka_utils/util/metadata.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from __future__ import annotations 15 | 16 | from collections import defaultdict 17 | 18 | from kafka.structs import PartitionMetadata 19 | 20 | from kafka_utils.util.client import KafkaToolClient 21 | from kafka_utils.util.config import ClusterConfig 22 | from kafka_utils.util.zookeeper import ZK 23 | 24 | 25 | LEADER_NOT_AVAILABLE_ERROR = 5 26 | REPLICA_NOT_AVAILABLE_ERROR = 9 27 | 28 | 29 | def get_topic_partition_metadata(hosts: list[str]) -> dict[str, dict[int, PartitionMetadata]]: 30 | """Returns topic-partition metadata from Kafka broker. 31 | 32 | kafka-python 1.3+ doesn't include partition metadata information in 33 | topic_partitions so we extract it from metadata ourselves. 34 | """ 35 | topic_partitions: dict[str, dict[int, PartitionMetadata]] = defaultdict(dict) 36 | 37 | kafka_client = KafkaToolClient(hosts, timeout=10) 38 | resp = kafka_client.send_metadata_request() 39 | 40 | for _, topic, partitions in resp.topics: 41 | for partition_error, partition, leader, replicas, isr in partitions: 42 | topic_partitions[topic][partition] = PartitionMetadata( 43 | topic, 44 | partition, 45 | leader, 46 | replicas, 47 | isr, 48 | partition_error, 49 | ) 50 | return topic_partitions 51 | 52 | 53 | def get_unavailable_brokers(zk: ZK, partition_metadata: PartitionMetadata) -> set[int]: 54 | """Returns the set of unavailable brokers from the difference of replica 55 | set of given partition to the set of available replicas. 56 | """ 57 | topic_data = zk.get_topics(partition_metadata.topic) 58 | topic = partition_metadata.topic 59 | partition = partition_metadata.partition 60 | expected_replicas = set(topic_data[topic]['partitions'][str(partition)]['replicas']) 61 | available_replicas = set(partition_metadata.replicas) 62 | return expected_replicas - available_replicas 63 | 64 | 65 | def get_topic_partition_with_error( 66 | cluster_config: ClusterConfig, 67 | error: int, 68 | fetch_unavailable_brokers: bool = False, 69 | ) -> set[tuple[str, int]] | tuple[set[tuple[str, int]], set[int]]: 70 | """Fetches the metadata from the cluster and returns the set of 71 | (topic, partition) tuples containing all the topic-partitions 72 | currently affected by the specified error. It also fetches unavailable-broker list 73 | if required.""" 74 | 75 | metadata = get_topic_partition_metadata(cluster_config.broker_list) 76 | affected_partitions = set() 77 | if fetch_unavailable_brokers: 78 | unavailable_brokers = set() 79 | with ZK(cluster_config) as zk: 80 | for partitions in metadata.values(): 81 | for partition_metadata in partitions.values(): 82 | if int(partition_metadata.error) == error: 83 | if fetch_unavailable_brokers: 84 | unavailable_brokers |= get_unavailable_brokers(zk, partition_metadata) 85 | affected_partitions.add((partition_metadata.topic, partition_metadata.partition)) 86 | 87 | if fetch_unavailable_brokers: 88 | return affected_partitions, unavailable_brokers 89 | else: 90 | return affected_partitions 91 | -------------------------------------------------------------------------------- /kafka_utils/util/protocol.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from __future__ import annotations 15 | 16 | import kafka.protocol.commit 17 | from kafka.protocol import KafkaProtocol 18 | from kafka.protocol.api import Response 19 | from kafka.protocol.commit import GroupCoordinatorRequest_v0 20 | from kafka.protocol.commit import OffsetCommitRequest_v2 21 | from kafka.structs import ConsumerMetadataResponse 22 | from kafka.structs import OffsetCommitRequestPayload 23 | from kafka.util import group_by_topic_and_partition 24 | 25 | 26 | class KafkaToolProtocol(KafkaProtocol): 27 | 28 | @classmethod 29 | def encode_offset_commit_request_kafka(cls, group: str, payloads: list[OffsetCommitRequestPayload]) -> OffsetCommitRequest_v2: 30 | """ 31 | Encode an OffsetCommitRequest struct 32 | Arguments: 33 | group: string, the consumer group you are committing offsets for 34 | payloads: list of OffsetCommitRequestPayload 35 | """ 36 | return OffsetCommitRequest_v2( 37 | consumer_group=group, 38 | consumer_group_generation_id=kafka.protocol.commit.OffsetCommitRequest[2].DEFAULT_GENERATION_ID, 39 | consumer_id='', 40 | retention_time=kafka.protocol.commit.OffsetCommitRequest[2].DEFAULT_RETENTION_TIME, 41 | topics=[( 42 | topic, 43 | [( 44 | partition, 45 | payload.offset, 46 | payload.metadata) 47 | for partition, payload in topic_payloads.items()]) 48 | for topic, topic_payloads in group_by_topic_and_partition(payloads).items()]) 49 | 50 | @classmethod 51 | def encode_consumer_metadata_request(cls, payloads: str) -> GroupCoordinatorRequest_v0: 52 | """ 53 | Encode a GroupCoordinatorRequest. Note that ConsumerMetadataRequest is 54 | renamed to GroupCoordinatorRequest in 0.9+. Interface is unchanged 55 | Arguments: 56 | payloads: string (consumer group) 57 | """ 58 | return GroupCoordinatorRequest_v0(payloads) 59 | 60 | @classmethod 61 | def decode_consumer_metadata_response(cls, response: Response) -> ConsumerMetadataResponse: 62 | """ 63 | Decode GroupCoordinatorResponse. Note that ConsumerMetadataResponse is 64 | renamed to GroupCoordinatorResponse in 0.9+ 65 | Arguments: 66 | response: response to decode 67 | """ 68 | return ConsumerMetadataResponse( 69 | response.error_code, 70 | response.coordinator_id, 71 | response.host, 72 | response.port, 73 | ) 74 | -------------------------------------------------------------------------------- /kafka_utils/util/py.typed: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Yelp/kafka-utils/def433ec4d07c60290d5dc937d3b4e5189eca9dc/kafka_utils/util/py.typed -------------------------------------------------------------------------------- /kafka_utils/util/serialization.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | import json 15 | from typing import Any 16 | 17 | 18 | def load_json(input_data: bytes) -> Any: 19 | data_str = input_data.decode() 20 | 21 | return json.loads(data_str) 22 | 23 | 24 | def dump_json(obj: Any) -> bytes: 25 | serialized = json.dumps(obj, sort_keys=True) 26 | 27 | return serialized.encode() 28 | -------------------------------------------------------------------------------- /kafka_utils/util/utils.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from __future__ import annotations 15 | 16 | import importlib 17 | import inspect 18 | import os 19 | import sys 20 | from types import ModuleType 21 | from typing import Collection 22 | 23 | 24 | def get_module(module_full_name: str) -> ModuleType: 25 | if ':' in module_full_name: 26 | path, module_name = module_full_name.rsplit(':', 1) 27 | if not os.path.isdir(path): 28 | print(f"{path} is not a valid directory", file=sys.stderr) 29 | sys.exit(1) 30 | sys.path.append(path) 31 | return importlib.import_module(module_name) 32 | else: 33 | return importlib.import_module(module_full_name) 34 | 35 | 36 | def child_class(class_types: Collection[type], base_class: type) -> type | None: 37 | """ 38 | Find the child-most class of `base_class`. 39 | 40 | Examples: 41 | class A: 42 | pass 43 | 44 | class B(A): 45 | pass 46 | 47 | class C(B): 48 | pass 49 | 50 | child_class([A, B, C], A) = C 51 | """ 52 | subclasses = set() 53 | for class_type in class_types: 54 | if class_type is base_class: 55 | continue 56 | if issubclass(class_type, base_class): 57 | subclasses.add(class_type) 58 | 59 | if len(subclasses) == 0: 60 | return None 61 | elif len(subclasses) == 1: 62 | return subclasses.pop() 63 | else: 64 | # If more than one class is a subclass of `base_class` 65 | # It is possible that one or more classes are subclasses of another 66 | # class (see example above). 67 | # Recursively find the child-most class. Break ties by returning any 68 | # child-most class. 69 | for c in subclasses: 70 | child = child_class(subclasses, c) 71 | if child is not None: 72 | return child 73 | return subclasses.pop() 74 | 75 | 76 | def dynamic_import(module_full_name: str, base_class: type) -> type | None: 77 | module = get_module(module_full_name) 78 | class_types = [ 79 | class_type 80 | for _, class_type in inspect.getmembers(module, inspect.isclass) 81 | ] 82 | return child_class(class_types, base_class) 83 | -------------------------------------------------------------------------------- /mypy.ini: -------------------------------------------------------------------------------- 1 | [mypy] 2 | python_version = 3.7 3 | pretty = true 4 | show_error_codes = true 5 | show_error_context = true 6 | show_column_numbers = true 7 | disallow_any_generics = true 8 | disallow_untyped_decorators = true 9 | warn_unused_ignores = true 10 | warn_redundant_casts = true 11 | 12 | [mypy-kafka_utils.util.*] 13 | disallow_untyped_defs = true 14 | disallow_incomplete_defs = true 15 | disallow_untyped_calls = true 16 | warn_unreachable = true 17 | 18 | [mypy-humanfriendly.*] 19 | ignore_missing_imports = true 20 | 21 | [mypy-kafka.*] 22 | ignore_missing_imports = true 23 | 24 | [mypy-kazoo.*] 25 | ignore_missing_imports = true 26 | 27 | [mypy-requests_futures.*] 28 | ignore_missing_imports = true 29 | -------------------------------------------------------------------------------- /requirements-dev.txt: -------------------------------------------------------------------------------- 1 | cryptography==3.0 2 | flake8 3 | importlib-metadata==3.7.2 4 | mock 5 | mypy 6 | pre-commit==1.14.4 7 | pytest==3.9.3 8 | setuptools==40.0.0 9 | tox-pip-extensions==1.2.1 10 | types-paramiko 11 | types-pytz 12 | types-PyYAML 13 | types-requests 14 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | PyStaticConfiguration==0.10.4 2 | PyYAML==5.1 3 | argcomplete==1.9.4 4 | boto==2.48.0 5 | cffi==1.12.2 6 | ecdsa==0.13.3 7 | enum34==1.1.6 8 | futures==3.2.0 9 | humanfriendly==4.8 10 | jsonschema==3.0.1 11 | kafka-python==1.4.7 12 | kazoo==2.6.1 13 | ndg-httpsclient==0.4.3 14 | paramiko==2.4.2 15 | pyOpenSSL==19.0.0 16 | pyasn1==0.4.5 17 | pytz==2018.4 18 | pycparser==2.18 19 | requests==2.21.0 20 | requests-futures==0.9.9 21 | setproctitle==1.1.10 22 | simplejson==3.16.0 23 | -------------------------------------------------------------------------------- /scripts/kafka-check: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | from kafka_utils.kafka_check.main import run 4 | 5 | if __name__ == '__main__': 6 | run() 7 | -------------------------------------------------------------------------------- /scripts/kafka-cluster-manager: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | from kafka_utils.kafka_cluster_manager.main import run 4 | 5 | if __name__ == '__main__': 6 | run() 7 | -------------------------------------------------------------------------------- /scripts/kafka-consumer-manager: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | from kafka_utils.kafka_consumer_manager.main import run 4 | 5 | if __name__ == '__main__': 6 | run() 7 | -------------------------------------------------------------------------------- /scripts/kafka-corruption-check: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | from kafka_utils.kafka_corruption_check.main import run 4 | 5 | if __name__ == '__main__': 6 | run() 7 | -------------------------------------------------------------------------------- /scripts/kafka-manual-throttle: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | from kafka_utils.kafka_manual_throttle.main import run 4 | 5 | if __name__ == '__main__': 6 | run() 7 | -------------------------------------------------------------------------------- /scripts/kafka-rolling-restart: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | from kafka_utils.kafka_rolling_restart.main import run 4 | 5 | if __name__ == '__main__': 6 | run() 7 | -------------------------------------------------------------------------------- /scripts/kafka-utils: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | from kafka_utils.main import run 4 | 5 | if __name__ == '__main__': 6 | run() 7 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | import os 15 | 16 | from setuptools import find_packages 17 | from setuptools import setup 18 | 19 | from kafka_utils import __version__ 20 | 21 | 22 | with open( 23 | os.path.join( 24 | os.path.abspath(os.path.dirname(__file__)), 25 | "README.md" 26 | ) 27 | ) as f: 28 | README = f.read() 29 | 30 | 31 | setup( 32 | name="kafka-utils", 33 | version=__version__, 34 | author="Team Data Streams Core", 35 | author_email="data-streams-core@yelp.com", 36 | description="Kafka management utils", 37 | packages=find_packages(exclude=["scripts*", "tests*"]), 38 | url="https://github.com/Yelp/kafka-utils", 39 | license="Apache License 2.0", 40 | long_description=README, 41 | long_description_content_type="text/markdown", 42 | keywords="apache kafka", 43 | scripts=[ 44 | "scripts/kafka-consumer-manager", 45 | "scripts/kafka-cluster-manager", 46 | "scripts/kafka-rolling-restart", 47 | "scripts/kafka-utils", 48 | "scripts/kafka-check", 49 | "scripts/kafka-corruption-check", 50 | ], 51 | python_requires='>=3.7', 52 | install_requires=[ 53 | "humanfriendly>=4.8", 54 | "kafka-python>=1.3.2,<1.5.0", 55 | "kazoo>=2.0,<3.0.0", 56 | "PyYAML>3.10", 57 | "pytz>=2014.1", 58 | "requests-futures>0.9.0", 59 | "paramiko>1.8.0,<3.0.0", 60 | "requests<3.0.0", 61 | "tenacity", 62 | "typing-extensions>=3.7.4", 63 | ], 64 | classifiers=[ 65 | "Development Status :: 4 - Beta", 66 | "License :: OSI Approved :: Apache Software License", 67 | "Programming Language :: Python", 68 | "Programming Language :: Python :: 3", 69 | "Programming Language :: Python :: 3.7", 70 | "Programming Language :: Python :: 3.8", 71 | "Environment :: Console", 72 | "Intended Audience :: Developers", 73 | "Intended Audience :: System Administrators", 74 | "Operating System :: POSIX", 75 | "Operating System :: MacOS :: MacOS X", 76 | ], 77 | package_data={ 78 | 'kafka_utils': ['util/py.typed'], 79 | }, 80 | ) 81 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/acceptance/change_topic_config.feature: -------------------------------------------------------------------------------- 1 | Feature: Change topic configuration using ZK 2 | 3 | @kafka10 4 | Scenario: Calling ZK to change topic level configs 5 | Given we have an existing kafka cluster with a topic 6 | when we set the configuration of the topic to 0 bytes 7 | then we produce to a kafka topic it should fail 8 | when we change the topic config in zk to 10000 bytes for kafka 10 9 | then we produce to a kafka topic it should succeed 10 | -------------------------------------------------------------------------------- /tests/acceptance/config/test.yaml: -------------------------------------------------------------------------------- 1 | --- 2 | clusters: 3 | test_cluster: 4 | broker_list: 5 | - "kafka:9092" 6 | zookeeper: "zookeeper:2181" 7 | local_config: 8 | cluster: test_cluster 9 | -------------------------------------------------------------------------------- /tests/acceptance/copy_group.feature: -------------------------------------------------------------------------------- 1 | Feature: kafka_consumer_manager copy_group subcommand 2 | 3 | Scenario: Calling the copy_group command 4 | Given we have an existing kafka cluster with a topic 5 | Given we have initialized kafka offsets storage 6 | Given we have a kafka consumer group 7 | when we call the copy_group command with a new groupid 8 | then the committed offsets in kafka for the new group will match the old group 9 | -------------------------------------------------------------------------------- /tests/acceptance/environment.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from steps.util import delete_topic 15 | from steps.util import list_topics 16 | 17 | 18 | def before_scenario(context, scenario): 19 | """Remove all topics from Kafka before each test""" 20 | topics = list_topics() 21 | for topic in topics.split('\n'): 22 | delete_topic(topic) 23 | -------------------------------------------------------------------------------- /tests/acceptance/kafka_consumer_manager.feature: -------------------------------------------------------------------------------- 1 | Feature: kafka_consumer_manager 2 | -------------------------------------------------------------------------------- /tests/acceptance/list_groups.feature: -------------------------------------------------------------------------------- 1 | Feature: kafka_consumer_manager list_groups subcommand 2 | 3 | Scenario: Calling the list_group command 4 | Given we have initialized kafka offsets storage 5 | Given we have a set of existing consumer groups 6 | when we call the list_groups command 7 | then the groups will be listed 8 | -------------------------------------------------------------------------------- /tests/acceptance/list_topics.feature: -------------------------------------------------------------------------------- 1 | Feature: kafka_consumer_manager list_topics subcommand 2 | 3 | Scenario: Calling the list_topics command 4 | Given we have an existing kafka cluster with a topic 5 | given we have a set of existing topics and a consumer group 6 | when we call the list_topics command 7 | then the topics will be listed 8 | -------------------------------------------------------------------------------- /tests/acceptance/min_isr.feature: -------------------------------------------------------------------------------- 1 | Feature: min_isr 2 | 3 | Scenario: Calling the min_isr command on empty cluster 4 | Given we have an existing kafka cluster 5 | when we call the min_isr command 6 | then OK min_isr will be printed 7 | 8 | Scenario: Calling the min_isr command on a cluster with isr greater or equal to min.isr for each topic 9 | Given we have an existing kafka cluster with a topic 10 | when we change min.isr settings for a topic to 1 11 | when we call the min_isr command 12 | then OK min_isr will be printed 13 | 14 | Scenario: Calling the min_isr command on a cluster with isr smaller than min.isr 15 | Given we have an existing kafka cluster with a topic 16 | when we change min.isr settings for a topic to 2 17 | when we call the min_isr command 18 | then CRITICAL min_isr will be printed 19 | -------------------------------------------------------------------------------- /tests/acceptance/offset_advance.feature: -------------------------------------------------------------------------------- 1 | Feature: kafka_consumer_manager offset_advance 2 | 3 | Scenario: Calling the offset_advance command 4 | Given we have an existing kafka cluster with a topic 5 | Given we have initialized kafka offsets storage 6 | when we produce some number of messages into the topic 7 | when we consume some number of messages from the topic 8 | when we call the offset_advance command and commit into kafka 9 | when we call the offset_get command 10 | then the latest message offsets will be shown 11 | -------------------------------------------------------------------------------- /tests/acceptance/offset_get.feature: -------------------------------------------------------------------------------- 1 | Feature: kafka_consumer_manager offset_get subcommand 2 | 3 | Scenario: Calling the offset_get command with json option 4 | Given we have an existing kafka cluster with a topic 5 | when we produce some number of messages into the topic 6 | when we consume some number of messages from the topic 7 | when we call the offset_get command with the json option 8 | then the correct json output will be shown 9 | 10 | Scenario: Committing offsets into Kafka and fetching offsets 11 | Given we have an existing kafka cluster with a topic 12 | Given we have initialized kafka offsets storage 13 | when we commit some offsets for a group into kafka 14 | when we fetch offsets for the group 15 | then the fetched offsets will match the committed offsets 16 | 17 | Scenario: Calling the offset_get command 18 | Given we have an existing kafka cluster with a topic 19 | Given we have initialized kafka offsets storage 20 | when we produce some number of messages into the topic 21 | when we consume some number of messages from the topic 22 | when we call the offset_set command and commit into kafka 23 | when we call the offset_get command 24 | then the offset that was committed into Kafka will be shown 25 | -------------------------------------------------------------------------------- /tests/acceptance/offset_restore.feature: -------------------------------------------------------------------------------- 1 | Feature: kafka_consumer_manager offset_restore subcommand 2 | 3 | # TODO: check if it passes 4 | Scenario: Calling the offset_restore command 5 | Given we have an existing kafka cluster with a topic 6 | Given we have initialized kafka offsets storage 7 | when we produce some number of messages into the topic 8 | when we consume some number of messages from the topic 9 | Given we have a json offsets file 10 | when we call the offset_restore command with the offsets file 11 | then the committed offsets will match the offsets file 12 | -------------------------------------------------------------------------------- /tests/acceptance/offset_rewind.feature: -------------------------------------------------------------------------------- 1 | Feature: kafka_consumer_manager offset_rewind subcommand 2 | 3 | Scenario: Calling the offset_rewind command 4 | Given we have an existing kafka cluster with a topic 5 | Given we have initialized kafka offsets storage 6 | Given we have a kafka consumer group 7 | when we call the offset_rewind command and commit into kafka 8 | when we call the offset_get command 9 | then consumer_group wont exist since it is rewind to low_offset 0 10 | -------------------------------------------------------------------------------- /tests/acceptance/offset_save.feature: -------------------------------------------------------------------------------- 1 | Feature: kafka_consumer_manager offset_save subcommand 2 | 3 | # This may fail TODO remove 4 | Scenario: Calling the offset_save command after consuming some messages 5 | Given we have an existing kafka cluster with a topic 6 | Given we have initialized kafka offsets storage 7 | when we produce some number of messages into the topic 8 | when we consume some number of messages from the topic 9 | when we call the offset_save command with an offsets file 10 | then the correct offsets will be saved into the given file 11 | 12 | Scenario: Calling offset_save after offset_restore 13 | Given we have an existing kafka cluster with a topic 14 | Given we have initialized kafka offsets storage 15 | when we produce some number of messages into the topic 16 | when we consume some number of messages from the topic 17 | Given we have a json offsets file 18 | when we call the offset_restore command with the offsets file 19 | when we call the offset_save command with an offsets file 20 | then the restored offsets will be saved into the given file 21 | -------------------------------------------------------------------------------- /tests/acceptance/offset_set.feature: -------------------------------------------------------------------------------- 1 | Feature: kafka_consumer_manager offset_set subcommand 2 | 3 | Scenario: Calling the offset_set command 4 | Given we have an existing kafka cluster with a topic 5 | when we produce some number of messages into the topic 6 | when we consume some number of messages from the topic 7 | when we call the offset_set command and commit into kafka 8 | then the committed offsets will match the specified offsets 9 | 10 | Scenario: Calling the offset_set command when the group doesn't exist 11 | Given we have an existing kafka cluster with a topic 12 | when we produce some number of messages into the topic 13 | when we call the offset_set command with a new groupid and the force option 14 | then the committed offsets will match the specified offsets 15 | -------------------------------------------------------------------------------- /tests/acceptance/rename_group.feature: -------------------------------------------------------------------------------- 1 | Feature: kafka_consumer_manager rename_group subcommand 2 | 3 | Scenario: Calling the rename_group command 4 | Given we have an existing kafka cluster with a topic 5 | when we produce some number of messages into the topic 6 | when we consume some number of messages from the topic 7 | when we call the rename_group command 8 | then the committed offsets in the new group will match the expected values 9 | then the group named has been changed 10 | -------------------------------------------------------------------------------- /tests/acceptance/replica_unavailability.feature: -------------------------------------------------------------------------------- 1 | Feature: replica_unavailability 2 | 3 | Scenario: Calling the replica_unavailability command on a cluster with unavailable replicas 4 | Given we have an existing kafka cluster with a topic 5 | when we call the replica_unavailability command 6 | then OK replica_unavailability will be printed 7 | -------------------------------------------------------------------------------- /tests/acceptance/replication_factor.feature: -------------------------------------------------------------------------------- 1 | Feature: replication_factor 2 | 3 | @kafka11 4 | Scenario: Calling the replication_factor check on empty cluster 5 | Given we have an existing kafka cluster 6 | when we call the replication_factor check 7 | then OK from replication_factor will be printed 8 | 9 | Scenario: Calling the replication_factor check on a cluster with proper replication factor for each topic 10 | Given we have an existing kafka cluster with a topic 11 | when we call the replication_factor check with adjusted min.isr 12 | then OK from replication_factor will be printed 13 | 14 | @kafka11 15 | Scenario: Calling the replication_factor check on a cluster with wrong replication factor for one of the topics 16 | Given we have an existing kafka cluster with a topic 17 | when we call the replication_factor check 18 | then CRITICAL from replication_factor will be printed 19 | -------------------------------------------------------------------------------- /tests/acceptance/steps/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/acceptance/steps/commit_fetch.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from behave import then 15 | from behave import when 16 | from steps.util import create_random_group_id 17 | from steps.util import get_cluster_config 18 | 19 | from kafka_utils.util.client import KafkaToolClient 20 | from kafka_utils.util.monitoring import get_current_consumer_offsets 21 | from kafka_utils.util.offsets import set_consumer_offsets 22 | 23 | 24 | TEST_OFFSET = 56 25 | 26 | 27 | def commit_offsets(offsets, group): 28 | # Setup the Kafka client 29 | config = get_cluster_config() 30 | client = KafkaToolClient(config.broker_list) 31 | set_consumer_offsets( 32 | client, 33 | group, 34 | offsets, 35 | ) 36 | client.close() 37 | 38 | 39 | def fetch_offsets(group, topics): 40 | # Setup the Kafka client 41 | config = get_cluster_config() 42 | client = KafkaToolClient(config.broker_list) 43 | offsets = get_current_consumer_offsets(client, group, topics, False) 44 | client.close() 45 | return offsets 46 | 47 | 48 | @when('we commit some offsets for a group into kafka') 49 | def step_impl4(context): 50 | context.offsets = {context.topic: {0: TEST_OFFSET}} 51 | context.group = create_random_group_id() 52 | commit_offsets(context.offsets, context.group) 53 | 54 | 55 | @when('we fetch offsets for the group') 56 | def step_impl4_3(context): 57 | topics = list(context.offsets.keys()) 58 | context.fetched_offsets = fetch_offsets(context.group, topics) 59 | 60 | 61 | @then('the fetched offsets will match the committed offsets') 62 | def step_impl5(context): 63 | assert context.fetched_offsets == context.offsets 64 | -------------------------------------------------------------------------------- /tests/acceptance/steps/common.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from behave import given 15 | from behave import when 16 | from steps.util import create_consumer_group 17 | from steps.util import create_random_group_id 18 | from steps.util import create_random_topic 19 | from steps.util import initialize_kafka_offsets_topic 20 | from steps.util import produce_example_msg 21 | 22 | PRODUCED_MSG_COUNT = 82 23 | CONSUMED_MSG_COUNT = 39 24 | 25 | 26 | @given('we have an existing kafka cluster with a topic') 27 | def step_impl1(context): 28 | context.topic = create_random_topic(1, 1) 29 | 30 | 31 | @given('we have a kafka consumer group') 32 | def step_impl2(context): 33 | context.group = create_random_group_id() 34 | context.client = create_consumer_group( 35 | context.topic, 36 | context.group, 37 | ) 38 | 39 | 40 | @when('we produce some number of messages into the topic') 41 | def step_impl3(context): 42 | produce_example_msg(context.topic, num_messages=PRODUCED_MSG_COUNT) 43 | context.msgs_produced = PRODUCED_MSG_COUNT 44 | 45 | 46 | @when('we consume some number of messages from the topic') 47 | def step_impl4(context): 48 | context.group = create_random_group_id() 49 | context.client = create_consumer_group( 50 | context.topic, 51 | context.group, 52 | num_messages=CONSUMED_MSG_COUNT, 53 | ) 54 | context.msgs_consumed = CONSUMED_MSG_COUNT 55 | 56 | 57 | @given('we have initialized kafka offsets storage') 58 | def step_impl5(context): 59 | 60 | initialize_kafka_offsets_topic() 61 | 62 | 63 | @given('we have an existing kafka cluster') 64 | def step_impl6(context): 65 | pass 66 | 67 | 68 | @given('we have an existing kafka cluster with multiple topics') 69 | def step_impl7(context): 70 | context.topic = [] 71 | context.topic.append(create_random_topic(1, 1, 'abcde')) 72 | context.topic.append(create_random_topic(1, 1, 'abcd')) 73 | -------------------------------------------------------------------------------- /tests/acceptance/steps/config_update.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | import time 15 | 16 | from behave import then 17 | from behave import when 18 | from kafka.errors import MessageSizeTooLargeError 19 | from steps.util import get_cluster_config 20 | from steps.util import produce_example_msg 21 | from steps.util import update_topic_config 22 | 23 | from kafka_utils.util.zookeeper import ZK 24 | 25 | 26 | @when('we set the configuration of the topic to 0 bytes') 27 | def step_impl1(context): 28 | context.output = update_topic_config( 29 | context.topic, 30 | 'max.message.bytes=0' 31 | ) 32 | 33 | 34 | @then('we produce to a kafka topic it should fail') 35 | def step_impl2(context): 36 | try: 37 | produce_example_msg(context.topic, num_messages=1) 38 | assert False, "Exception should not be raised" 39 | except MessageSizeTooLargeError as e: 40 | assert isinstance(e, MessageSizeTooLargeError) 41 | 42 | 43 | @when('we change the topic config in zk to 10000 bytes for kafka 10') 44 | def step_impl3(context): 45 | cluster_config = get_cluster_config() 46 | with ZK(cluster_config) as zk: 47 | current_config = zk.get_topic_config(context.topic) 48 | current_config['config']['max.message.bytes'] = '1000' 49 | zk.set_topic_config(context.topic, value=current_config) 50 | time.sleep(2) # sleeping for 2 seconds to ensure config is actually picked up 51 | 52 | 53 | @then('we produce to a kafka topic it should succeed') 54 | def step_impl5(context): 55 | try: 56 | produce_example_msg(context.topic, num_messages=1) 57 | except MessageSizeTooLargeError as e: 58 | assert False, "Exception should not be raised" 59 | assert isinstance(e, MessageSizeTooLargeError) 60 | -------------------------------------------------------------------------------- /tests/acceptance/steps/copy_group.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from behave import then 15 | from behave import when 16 | from steps.util import call_cmd 17 | 18 | from kafka_utils.util.offsets import get_current_consumer_offsets 19 | 20 | 21 | NEW_GROUP = 'new_group' 22 | 23 | 24 | def call_copy_group(old_group, new_group): 25 | cmd = ['kafka-consumer-manager', 26 | '--cluster-type', 'test', 27 | '--cluster-name', 'test_cluster', 28 | '--discovery-base-path', 'tests/acceptance/config', 29 | 'copy_group', 30 | old_group, 31 | new_group] 32 | return call_cmd(cmd) 33 | 34 | 35 | @when('we call the copy_group command with a new groupid') 36 | def step_impl3(context): 37 | call_copy_group(context.group, NEW_GROUP) 38 | 39 | 40 | @then('the committed offsets in kafka for the new group will match the old group') 41 | def step_impl4(context): 42 | old_group_offsets = get_current_consumer_offsets( 43 | context.client, 44 | context.group, 45 | [context.topic], 46 | ) 47 | new_group_offsets = get_current_consumer_offsets( 48 | context.client, 49 | NEW_GROUP, 50 | [context.topic], 51 | ) 52 | assert old_group_offsets == new_group_offsets 53 | -------------------------------------------------------------------------------- /tests/acceptance/steps/list_groups.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from behave import given 15 | from behave import then 16 | from behave import when 17 | from steps.util import call_cmd 18 | from steps.util import create_consumer_group 19 | from steps.util import create_random_group_id 20 | from steps.util import create_random_topic 21 | from steps.util import produce_example_msg 22 | 23 | 24 | @given('we have a set of existing consumer groups') 25 | def step_impl1(context): 26 | topic = create_random_topic(1, 1) 27 | produce_example_msg(topic) 28 | 29 | context.groups = [] 30 | for _ in range(3): 31 | group = create_random_group_id() 32 | context.groups.append(group) 33 | create_consumer_group(topic, group) 34 | 35 | 36 | def call_list_groups(): 37 | cmd = ['kafka-consumer-manager', 38 | '--cluster-type', 'test', 39 | '--cluster-name', 'test_cluster', 40 | '--discovery-base-path', 'tests/acceptance/config', 41 | 'list_groups'] 42 | return call_cmd(cmd) 43 | 44 | 45 | @when('we call the list_groups command') 46 | def step_impl3(context): 47 | context.output = call_list_groups() 48 | 49 | 50 | @then('the groups will be listed') 51 | def step_impl5(context): 52 | for group in context.groups: 53 | assert group in context.output 54 | -------------------------------------------------------------------------------- /tests/acceptance/steps/list_topics.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from behave import given 15 | from behave import then 16 | from behave import when 17 | from steps.util import call_cmd 18 | from steps.util import create_consumer_group 19 | from steps.util import create_random_topic 20 | from steps.util import produce_example_msg 21 | 22 | test_group = 'group1' 23 | test_topics = ['topic1', 'topic2', 'topic3'] 24 | 25 | 26 | @given('we have a set of existing topics and a consumer group') 27 | def step_impl1(context): 28 | for topic in test_topics: 29 | create_random_topic(1, 1, topic_name=topic) 30 | produce_example_msg(topic) 31 | create_consumer_group(topic, test_group) 32 | 33 | 34 | def call_list_topics(groupid): 35 | cmd = ['kafka-consumer-manager', 36 | '--cluster-type', 'test', 37 | '--cluster-name', 'test_cluster', 38 | '--discovery-base-path', 'tests/acceptance/config', 39 | 'list_topics', 40 | groupid] 41 | return call_cmd(cmd) 42 | 43 | 44 | @when('we call the list_topics command') 45 | def step_impl3(context): 46 | context.output = call_list_topics('group1') 47 | 48 | 49 | @then('the topics will be listed') 50 | def step_impl5(context): 51 | for topic in test_topics: 52 | assert topic in context.output 53 | -------------------------------------------------------------------------------- /tests/acceptance/steps/min_isr.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from behave import then 15 | from behave import when 16 | from steps.util import call_cmd 17 | from steps.util import get_cluster_config 18 | 19 | from kafka_utils.util.zookeeper import ZK 20 | 21 | 22 | ISR_CONF_NAME = 'min.insync.replicas' 23 | CONF_PATH = '/config/topics/' 24 | 25 | 26 | def call_min_isr(): 27 | cmd = ['kafka-check', 28 | '--cluster-type', 'test', 29 | '--cluster-name', 'test_cluster', 30 | '--discovery-base-path', 'tests/acceptance/config', 31 | 'min_isr'] 32 | return call_cmd(cmd) 33 | 34 | 35 | def set_min_isr(topic, min_isr): 36 | cluster_config = get_cluster_config() 37 | with ZK(cluster_config) as zk: 38 | config = zk.get_topic_config(topic) 39 | config['config'] = {ISR_CONF_NAME: str(min_isr)} 40 | zk.set_topic_config(topic, config) 41 | 42 | 43 | @when('we call the min_isr command') 44 | def step_impl2(context): 45 | context.min_isr_out = call_min_isr() 46 | 47 | 48 | @when('we change min.isr settings for a topic to 1') 49 | def step_impl3(context): 50 | set_min_isr(context.topic, 1) 51 | 52 | 53 | @when('we change min.isr settings for a topic to 2') 54 | def step_impl4(context): 55 | set_min_isr(context.topic, 2) 56 | 57 | 58 | @then('OK min_isr will be printed') 59 | def step_impl5(context): 60 | assert context.min_isr_out == 'OK: All replicas in sync.\n', context.min_isr_out 61 | 62 | 63 | @then('CRITICAL min_isr will be printed') 64 | def step_impl6(context): 65 | error_msg = ("CRITICAL: 1 partition(s) have the number of " 66 | "replicas in sync that is lower than the specified min ISR.\n") 67 | assert context.min_isr_out == error_msg, context.min_isr_out 68 | -------------------------------------------------------------------------------- /tests/acceptance/steps/offset_advance.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from behave import then 15 | from behave import when 16 | from steps.util import call_cmd 17 | from steps.util import get_cluster_config 18 | 19 | from kafka_utils.util.zookeeper import ZK 20 | 21 | 22 | def call_offset_advance(groupid, topic=None, force=False): 23 | cmd = ['kafka-consumer-manager', 24 | '--cluster-type', 'test', 25 | '--cluster-name', 'test_cluster', 26 | '--discovery-base-path', 'tests/acceptance/config', 27 | 'offset_advance', 28 | groupid] 29 | if topic: 30 | cmd.extend(['--topic', topic]) 31 | if force: 32 | cmd.extend(['--force']) 33 | return call_cmd(cmd) 34 | 35 | 36 | @when('we call the offset_advance command and commit into kafka') 37 | def step_impl3_2(context): 38 | call_offset_advance( 39 | context.group, 40 | topic=context.topic, 41 | force=True, 42 | ) 43 | 44 | 45 | @then('the committed offsets will match the latest message offsets') 46 | def step_impl4(context): 47 | cluster_config = get_cluster_config() 48 | with ZK(cluster_config) as zk: 49 | offsets = zk.get_group_offsets(context.group) 50 | assert offsets[context.topic]["0"] == context.msgs_produced 51 | 52 | 53 | @then('the latest message offsets will be shown') 54 | def step_impl5_2(context): 55 | offset = context.msgs_produced 56 | pattern = f'Current Offset: {offset}' 57 | assert pattern in context.output 58 | -------------------------------------------------------------------------------- /tests/acceptance/steps/offset_get.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from behave import then 15 | from behave import when 16 | from steps.util import call_offset_get 17 | from steps.util import load_json 18 | 19 | 20 | @when('we call the offset_get command') 21 | def step_impl4_2(context): 22 | context.output = call_offset_get(context.group) 23 | 24 | 25 | @when('we call the offset_get command with the json option') 26 | def step_impl4_5(context): 27 | context.output = call_offset_get(context.group, json=True) 28 | 29 | 30 | @then('the correct offset will be shown') 31 | def step_impl5(context): 32 | offset = context.msgs_consumed 33 | pattern = f'Current Offset: {offset}' 34 | assert pattern in context.output 35 | 36 | 37 | @then('the offset that was committed into Kafka will be shown') 38 | def step_impl5_2(context): 39 | offset = context.set_offset_kafka 40 | pattern = f'Current Offset: {offset}' 41 | assert pattern in context.output 42 | 43 | 44 | @then('the correct json output will be shown') 45 | def step_impl5_3(context): 46 | offset = context.msgs_consumed 47 | if context.msgs_produced > 0.0: 48 | percentage_distance = round((context.msgs_produced - offset) * 100.0 / context.msgs_produced, 2) 49 | else: 50 | percentage_distance = 0.0 51 | 52 | parsed_output = load_json(context.output) 53 | assert parsed_output == [ 54 | { 55 | "topic": context.topic, 56 | "partition": 0, 57 | "current": offset, 58 | "highmark": context.msgs_produced, 59 | "lowmark": 0, 60 | "offset_distance": context.msgs_produced - offset, 61 | "percentage_distance": percentage_distance 62 | } 63 | ] 64 | -------------------------------------------------------------------------------- /tests/acceptance/steps/offset_restore.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | import os 15 | import tempfile 16 | 17 | from behave import given 18 | from behave import then 19 | from behave import when 20 | from steps.util import call_cmd 21 | from steps.util import get_cluster_config 22 | 23 | from kafka_utils.util.client import KafkaToolClient 24 | from kafka_utils.util.offsets import get_current_consumer_offsets 25 | 26 | 27 | RESTORED_OFFSET = 55 28 | 29 | 30 | def create_restore_file(group, topic, offset): 31 | offset_restore_data = ''' 32 | {{ 33 | "groupid": "{group}", 34 | "offsets": {{ 35 | "{topic}": {{ 36 | "0": {offset} 37 | }} 38 | }} 39 | }} 40 | '''.format(group=group, topic=topic, offset=offset) 41 | 42 | f = tempfile.NamedTemporaryFile() 43 | 44 | offset_restore_data = offset_restore_data.encode() 45 | 46 | f.write(offset_restore_data) 47 | f.flush() 48 | return f 49 | 50 | 51 | def call_offset_restore(offsets_file): 52 | cmd = ['kafka-consumer-manager', 53 | '--cluster-type', 'test', 54 | '--cluster-name', 'test_cluster', 55 | '--discovery-base-path', 'tests/acceptance/config', 56 | 'offset_restore', 57 | offsets_file] 58 | return call_cmd(cmd) 59 | 60 | 61 | @given('we have a json offsets file') 62 | def step_impl2(context): 63 | context.restored_offset = RESTORED_OFFSET 64 | context.offsets_file = create_restore_file( 65 | context.group, 66 | context.topic, 67 | context.restored_offset, 68 | ) 69 | assert os.path.isfile(context.offsets_file.name) 70 | 71 | 72 | @when('we call the offset_restore command with the offsets file') 73 | def step_impl3_2(context): 74 | call_offset_restore(context.offsets_file.name) 75 | 76 | 77 | @then('the committed offsets will match the offsets file') 78 | def step_impl4(context): 79 | config = get_cluster_config() 80 | context.client = KafkaToolClient(config.broker_list) 81 | offsets = get_current_consumer_offsets( 82 | context.client, 83 | context.group, 84 | [context.topic], 85 | ) 86 | assert offsets[context.topic][0] == RESTORED_OFFSET 87 | -------------------------------------------------------------------------------- /tests/acceptance/steps/offset_rewind.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from behave import then 15 | from behave import when 16 | from steps.util import call_cmd 17 | from steps.util import get_cluster_config 18 | from steps.util import set_consumer_group_offset 19 | 20 | from kafka_utils.util.client import KafkaToolClient 21 | from kafka_utils.util.offsets import get_current_consumer_offsets 22 | 23 | 24 | def offsets_data(topic, offset): 25 | return '''{topic}.{partition}={offset}'''.format( 26 | topic=topic, 27 | partition='0', 28 | offset=offset, 29 | ) 30 | 31 | 32 | def call_offset_rewind(groupid, topic, force=False): 33 | cmd = ['kafka-consumer-manager', 34 | '--cluster-type', 'test', 35 | '--cluster-name', 'test_cluster', 36 | '--discovery-base-path', 'tests/acceptance/config', 37 | 'offset_rewind', 38 | groupid, 39 | '--topic', topic] 40 | if force: 41 | cmd.extend(['--force']) 42 | return call_cmd(cmd) 43 | 44 | 45 | @when('we set the offsets to a high number to emulate consumption') 46 | def step_impl2(context): 47 | set_consumer_group_offset( 48 | topic=context.topic, 49 | group=context.group, 50 | offset=100, 51 | ) 52 | 53 | 54 | @when('we call the offset_rewind command and commit into kafka') 55 | def step_impl3(context): 56 | call_offset_rewind(context.group, context.topic) 57 | 58 | 59 | @when('we call the offset_rewind command with a new groupid and the force option') 60 | def step_impl4(context): 61 | context.group = 'offset_rewind_created_group' 62 | call_offset_rewind( 63 | context.group, 64 | topic=context.topic, 65 | force=True, 66 | ) 67 | 68 | 69 | @then('the committed offsets will match the earliest message offsets') 70 | def step_impl5(context): 71 | config = get_cluster_config() 72 | context.client = KafkaToolClient(config.broker_list) 73 | offsets = get_current_consumer_offsets( 74 | context.client, 75 | context.group, 76 | [context.topic], 77 | ) 78 | assert offsets[context.topic][0] == 0 79 | 80 | 81 | @then('consumer_group wont exist since it is rewind to low_offset 0') 82 | def step_impl6(context): 83 | # Since using kafka offset storage, lowmark is 0, then a rewind to lowmark 84 | # will remove the consumer group from kafka 85 | assert "Current Offset" not in context.output 86 | 87 | 88 | @then('the earliest message offsets will be shown') 89 | def step_impl7(context): 90 | offset = 0 91 | pattern = f'Current Offset: {offset}' 92 | assert pattern in context.output 93 | -------------------------------------------------------------------------------- /tests/acceptance/steps/offset_save.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | import tempfile 15 | 16 | from behave import then 17 | from behave import when 18 | from steps.util import call_cmd 19 | from steps.util import load_json 20 | 21 | 22 | def create_saved_file(): 23 | return tempfile.NamedTemporaryFile() 24 | 25 | 26 | def call_offset_save(groupid, offsets_file): 27 | cmd = ['kafka-consumer-manager', 28 | '--cluster-type', 'test', 29 | '--cluster-name', 'test_cluster', 30 | '--discovery-base-path', 'tests/acceptance/config', 31 | 'offset_save', 32 | groupid, 33 | offsets_file] 34 | return call_cmd(cmd) 35 | 36 | 37 | @when('we call the offset_save command with an offsets file') 38 | def step_impl2_2(context): 39 | context.offsets_file = create_saved_file() 40 | call_offset_save(context.group, context.offsets_file.name) 41 | 42 | 43 | @then('the correct offsets will be saved into the given file') 44 | def step_impl3(context): 45 | offset = context.msgs_consumed 46 | 47 | data = load_json(context.offsets_file.read()) 48 | assert offset == data['offsets'][context.topic]['0'] 49 | context.offsets_file.close() 50 | 51 | 52 | @then('the restored offsets will be saved into the given file') 53 | def step_impl3_2(context): 54 | offset = context.restored_offset 55 | 56 | data = load_json(context.offsets_file.read()) 57 | assert offset == data['offsets'][context.topic]['0'] 58 | context.offsets_file.close() 59 | -------------------------------------------------------------------------------- /tests/acceptance/steps/offset_set.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from behave import then 15 | from behave import when 16 | from steps.util import call_cmd 17 | from steps.util import get_cluster_config 18 | 19 | from kafka_utils.util.client import KafkaToolClient 20 | from kafka_utils.util.offsets import get_current_consumer_offsets 21 | 22 | SET_OFFSET_KAFKA = 65 23 | 24 | 25 | def offsets_data(topic, offset): 26 | return '''{topic}.{partition}={offset}'''.format( 27 | topic=topic, 28 | partition='0', 29 | offset=offset, 30 | ) 31 | 32 | 33 | def call_offset_set(groupid, offsets_data, force=False): 34 | cmd = ['kafka-consumer-manager', 35 | '--cluster-type', 'test', 36 | '--cluster-name', 'test_cluster', 37 | '--discovery-base-path', 'tests/acceptance/config', 38 | 'offset_set', 39 | groupid, 40 | offsets_data] 41 | if force: 42 | cmd.extend(['--force']) 43 | return call_cmd(cmd) 44 | 45 | 46 | @when('we call the offset_set command and commit into kafka') 47 | def step_impl2_2(context): 48 | if not hasattr(context, 'group'): 49 | context.group = 'test_kafka_offset_group' 50 | context.offsets = offsets_data(context.topic, SET_OFFSET_KAFKA) 51 | context.set_offset_kafka = SET_OFFSET_KAFKA 52 | call_offset_set(context.group, context.offsets) 53 | 54 | 55 | @when('we call the offset_set command with a new groupid and the force option') 56 | def step_impl2_3(context): 57 | if not hasattr(context, 'group'): 58 | context.group = 'offset_set_created_group' 59 | context.offsets = offsets_data(context.topic, SET_OFFSET_KAFKA) 60 | context.set_offset_kafka = SET_OFFSET_KAFKA 61 | call_offset_set(context.group, context.offsets, force=True) 62 | 63 | 64 | @then('the committed offsets will match the specified offsets') 65 | def step_impl3(context): 66 | config = get_cluster_config() 67 | context.client = KafkaToolClient(config.broker_list) 68 | offsets = get_current_consumer_offsets( 69 | context.client, 70 | context.group, 71 | [context.topic], 72 | ) 73 | assert offsets[context.topic][0] == SET_OFFSET_KAFKA 74 | -------------------------------------------------------------------------------- /tests/acceptance/steps/rename_group.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from behave import then 15 | from behave import when 16 | from steps.util import call_cmd 17 | from steps.util import call_offset_get 18 | 19 | 20 | NEW_GROUP = 'new_group' 21 | 22 | 23 | def call_rename_group(old_group, new_group): 24 | cmd = ['kafka-consumer-manager', 25 | '--cluster-type', 'test', 26 | '--cluster-name', 'test_cluster', 27 | '--discovery-base-path', 'tests/acceptance/config', 28 | 'rename_group', 29 | old_group, 30 | new_group] 31 | return call_cmd(cmd) 32 | 33 | 34 | @when('we call the rename_group command') 35 | def step_impl2(context): 36 | call_rename_group(context.group, NEW_GROUP) 37 | 38 | 39 | @then('the committed offsets in the new group will match the expected values') 40 | def step_impl3(context): 41 | new_group = call_offset_get(NEW_GROUP, True) 42 | old_group = call_offset_get(context.group, True) 43 | assert "does not exist" in old_group 44 | assert context.topic in new_group 45 | 46 | 47 | @then('the group named has been changed') 48 | def step_impl4(context): 49 | new_groups = call_offset_get(NEW_GROUP) 50 | old_group = call_offset_get(context.group) 51 | assert "Offset Distance" in new_groups 52 | assert "Offset Distance" not in old_group 53 | -------------------------------------------------------------------------------- /tests/acceptance/steps/replica_unavailability.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from behave import then 15 | from behave import when 16 | from steps.util import call_cmd 17 | 18 | 19 | def call_replica_unavailability(): 20 | cmd = ['kafka-check', 21 | '--cluster-type', 'test', 22 | '--cluster-name', 'test_cluster', 23 | '--discovery-base-path', 'tests/acceptance/config', 24 | 'replica_unavailability'] 25 | return call_cmd(cmd) 26 | 27 | 28 | @when('we call the replica_unavailability command') 29 | def step_impl2(context): 30 | context.replica_unavailability_out = call_replica_unavailability() 31 | 32 | 33 | @then('OK replica_unavailability will be printed') 34 | def step_impl5(context): 35 | assert context.replica_unavailability_out == 'OK: All replicas available for communication.\n', context.replica_unavailability_out 36 | -------------------------------------------------------------------------------- /tests/acceptance/steps/replication_factor.py: -------------------------------------------------------------------------------- 1 | # Copyright 2018 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from behave import then 15 | from behave import when 16 | from steps.util import call_cmd 17 | 18 | 19 | def call_replication_factor_check(min_isr='1'): 20 | cmd = ['kafka-check', 21 | '--cluster-type', 'test', 22 | '--cluster-name', 'test_cluster', 23 | '--discovery-base-path', 'tests/acceptance/config', '-v', 24 | 'replication_factor', 25 | '--default-min-isr', min_isr] 26 | return call_cmd(cmd) 27 | 28 | 29 | @when('we call the replication_factor check') 30 | def step_impl2(context): 31 | context.replication_factor_out = call_replication_factor_check() 32 | 33 | 34 | @when('we call the replication_factor check with adjusted min.isr') 35 | def step_impl3(context): 36 | context.replication_factor_out = call_replication_factor_check('0') 37 | 38 | 39 | @then('OK from replication_factor will be printed') 40 | def step_impl5(context): 41 | msg = 'OK: All topics have proper replication factor.\n' 42 | assert context.replication_factor_out == msg, context.replication_factor_out 43 | 44 | 45 | @then('CRITICAL from replication_factor will be printed') 46 | def step_impl6(context): 47 | msg = ( 48 | "CRITICAL: 1 topic(s) have replication factor lower than specified min ISR + 1.\n" 49 | "Topics:\n" 50 | "replication_factor=1 is lower than min_isr=1 + 1 for {topic}\n" 51 | ).format(topic=context.topic) 52 | 53 | assert context.replication_factor_out == msg, context.replication_factor_out 54 | -------------------------------------------------------------------------------- /tests/acceptance/steps/watermark_get.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from behave import then 15 | from behave import when 16 | from steps.util import call_watermark_get 17 | 18 | 19 | @when('we call the watermark_get command') 20 | def step_impl4(context): 21 | context.output = call_watermark_get(context.topic) 22 | 23 | 24 | @then('the correct watermark will be shown') 25 | def step_impl5(context): 26 | highmark = context.msgs_produced 27 | highmark_pattern = f'High Watermark: {highmark}' 28 | lowmark_pattern = 'Low Watermark: 0' 29 | assert highmark_pattern in context.output 30 | assert lowmark_pattern in context.output 31 | 32 | 33 | @when('we call the watermark_get command with -r') 34 | def step_impl6(context): 35 | context.output = call_watermark_get('bc', regex=True) 36 | 37 | 38 | @when('we call the watermark_get command without -r') 39 | def step_impl7(context): 40 | context.output = call_watermark_get('abcd') 41 | 42 | 43 | @then('the correct two topics will be shown') 44 | def step_impl8(context): 45 | topic1 = context.topic[0] 46 | topic2 = context.topic[1] 47 | assert topic1 in context.output 48 | assert topic2 in context.output 49 | 50 | 51 | @then('the correct single topic will be shown') 52 | def step_impl9(context): 53 | topic1 = context.topic[0] 54 | topic2 = context.topic[1] 55 | assert topic1 not in context.output 56 | assert topic2 in context.output 57 | -------------------------------------------------------------------------------- /tests/acceptance/watermark_get.feature: -------------------------------------------------------------------------------- 1 | Feature: kafka_consumer_manager watermark_get subcommand 2 | 3 | Scenario: Calling the watermark_get command 4 | Given we have an existing kafka cluster with a topic 5 | when we produce some number of messages into the topic 6 | when we call the watermark_get command 7 | then the correct watermark will be shown 8 | 9 | Scenario: Calling the watermark_get command with -r 10 | Given we have an existing kafka cluster with multiple topics 11 | when we call the watermark_get command with -r 12 | then the correct two topics will be shown 13 | 14 | Scenario: Calling the watermark_get command without -r 15 | Given we have an existing kafka cluster with multiple topics 16 | when we call the watermark_get command without -r 17 | then the correct single topic will be shown 18 | -------------------------------------------------------------------------------- /tests/kafka_check/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/kafka_check/test_is_first_broker.py: -------------------------------------------------------------------------------- 1 | from kafka_utils.kafka_check.commands.command import is_first_broker 2 | 3 | 4 | def test_is_first_broker_empty_broker_list_return_false(): 5 | broker_id = 3434 6 | broker_ids = [] 7 | assert not is_first_broker(broker_ids, broker_id) 8 | 9 | 10 | def test_is_first_broker_returns_false_if_broker_id_not_exist_in_broker_id_list(): 11 | broker_id = 10 12 | broker_ids = [2, 3] 13 | assert not is_first_broker(broker_ids, broker_id) 14 | 15 | 16 | def test_is_first_broker_returns_true_if_broker_id_equal_min(): 17 | broker_ids = [1, 2, 3] 18 | broker_id = min(broker_ids) 19 | assert is_first_broker(broker_ids, broker_id) 20 | -------------------------------------------------------------------------------- /tests/kafka_check/test_metadata_file.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from kafka_utils.kafka_check.metadata_file import _parse_meta_properties_file 15 | 16 | 17 | META_PROPERTIES_CONTENT = [ 18 | "#\n", 19 | "#Thu May 12 17:59:12 BST 2016\n", 20 | "version=0\n", 21 | "broker.id=123\n", 22 | ] 23 | 24 | 25 | class TestParseMetaPropertiesFile: 26 | 27 | def test_parse_meta_properties_file(self): 28 | broker_id = _parse_meta_properties_file(META_PROPERTIES_CONTENT) 29 | 30 | assert broker_id == 123 31 | 32 | def test_parse_meta_properties_file_empty(self): 33 | broker_id = _parse_meta_properties_file([]) 34 | 35 | assert broker_id is None 36 | -------------------------------------------------------------------------------- /tests/kafka_check/test_offline.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from kafka_utils.kafka_check.commands.offline import _prepare_output 15 | 16 | 17 | OFFLINE_PARTITIONS = [ 18 | ('Topic0', 0), 19 | ('Topic0', 1), 20 | ('Topic13', 13), 21 | ] 22 | 23 | 24 | def test_prepare_output_ok_no_verbose(): 25 | expected = { 26 | 'message': "No offline partitions.", 27 | 'raw': { 28 | 'offline_count': 0, 29 | } 30 | } 31 | assert _prepare_output([], False, -1) == expected 32 | 33 | 34 | def test_prepare_output_ok_verbose(): 35 | expected = { 36 | 'message': "No offline partitions.", 37 | 'raw': { 38 | 'offline_count': 0, 39 | 'partitions': [], 40 | } 41 | } 42 | assert _prepare_output([], True, -1) == expected 43 | 44 | 45 | def test_prepare_output_critical_verbose(): 46 | expected = { 47 | 'message': "3 offline partitions.", 48 | 'verbose': ( 49 | "Partitions:\n" 50 | "Topic0:0\n" 51 | "Topic0:1\n" 52 | "Topic13:13" 53 | ), 54 | 'raw': { 55 | 'offline_count': 3, 56 | 'partitions': [ 57 | {'partition': 0, 'topic': 'Topic0'}, 58 | {'partition': 1, 'topic': 'Topic0'}, 59 | {'partition': 13, 'topic': 'Topic13'}, 60 | ], 61 | } 62 | } 63 | assert _prepare_output(OFFLINE_PARTITIONS, True, -1) == expected 64 | 65 | 66 | def test_prepare_output_critical_verbose_with_head(): 67 | expected = { 68 | 'message': "3 offline partitions.", 69 | 'verbose': ( 70 | "Top 2 partitions:\n" 71 | "Topic0:0\n" 72 | "Topic0:1" 73 | ), 74 | 'raw': { 75 | 'offline_count': 3, 76 | 'partitions': [ 77 | {'partition': 0, 'topic': 'Topic0'}, 78 | {'partition': 1, 'topic': 'Topic0'}, 79 | ], 80 | } 81 | } 82 | assert _prepare_output(OFFLINE_PARTITIONS, True, 2) == expected 83 | -------------------------------------------------------------------------------- /tests/kafka_check/test_replica_unavailability.py: -------------------------------------------------------------------------------- 1 | # Copyright 2019 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from kafka_utils.kafka_check.commands.replica_unavailability import _prepare_output 15 | 16 | 17 | UNAVAILABLE_REPLICAS = [ 18 | ('Topic0', 0), 19 | ('Topic0', 1), 20 | ('Topic13', 13), 21 | ] 22 | 23 | UNAVAILABLE_BROKERS = [123456, 987456] 24 | 25 | 26 | def test_prepare_output_ok_no_verbose(): 27 | expected = { 28 | 'message': "All replicas available for communication.", 29 | 'raw': { 30 | 'replica_unavailability_count': 0, 31 | } 32 | } 33 | assert _prepare_output([], [], False, -1) == expected 34 | 35 | 36 | def test_prepare_output_ok_verbose(): 37 | expected = { 38 | 'message': "All replicas available for communication.", 39 | 'raw': { 40 | 'replica_unavailability_count': 0, 41 | 'partitions': [], 42 | } 43 | } 44 | assert _prepare_output([], [], True, -1) == expected 45 | 46 | 47 | def test_prepare_output_critical_verbose(): 48 | expected = { 49 | 'message': "3 replicas unavailable for communication. Unavailable Brokers: 123456, 987456", 50 | 'verbose': ( 51 | "Partitions:\n" 52 | "Topic0:0\n" 53 | "Topic0:1\n" 54 | "Topic13:13" 55 | ), 56 | 'raw': { 57 | 'replica_unavailability_count': 3, 58 | 'partitions': [ 59 | {'partition': 0, 'topic': 'Topic0'}, 60 | {'partition': 1, 'topic': 'Topic0'}, 61 | {'partition': 13, 'topic': 'Topic13'}, 62 | ], 63 | } 64 | } 65 | assert _prepare_output(UNAVAILABLE_REPLICAS, UNAVAILABLE_BROKERS, True, -1) == expected 66 | 67 | 68 | def test_prepare_output_critical_verbose_with_head(): 69 | expected = { 70 | 'message': "3 replicas unavailable for communication. Unavailable Brokers: 123456, 987456", 71 | 'verbose': ( 72 | "Top 2 partitions:\n" 73 | "Topic0:0\n" 74 | "Topic0:1" 75 | ), 76 | 'raw': { 77 | 'replica_unavailability_count': 3, 78 | 'partitions': [ 79 | {'partition': 0, 'topic': 'Topic0'}, 80 | {'partition': 1, 'topic': 'Topic0'}, 81 | ], 82 | } 83 | } 84 | assert _prepare_output(UNAVAILABLE_REPLICAS, UNAVAILABLE_BROKERS, True, 2) == expected 85 | -------------------------------------------------------------------------------- /tests/kafka_cluster_manager/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/kafka_cluster_manager/cmds/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/kafka_cluster_manager/decommission_test.py: -------------------------------------------------------------------------------- 1 | from argparse import Namespace 2 | from unittest import mock 3 | 4 | import pytest 5 | 6 | from kafka_utils.kafka_cluster_manager.cluster_info \ 7 | .partition_count_balancer import PartitionCountBalancer 8 | from kafka_utils.kafka_cluster_manager.cmds import decommission 9 | from tests.kafka_cluster_manager.helper import broker_range 10 | 11 | 12 | @pytest.fixture 13 | def command_instance(): 14 | cmd = decommission.DecommissionCmd() 15 | cmd.args = mock.Mock(spec=Namespace) 16 | cmd.args.force_progress = False 17 | cmd.args.broker_ids = [] 18 | cmd.args.auto_max_movement_size = True 19 | cmd.args.max_partition_movements = 10 20 | cmd.args.max_leader_changes = 10 21 | return cmd 22 | 23 | 24 | def test_decommission_no_partitions_to_move(command_instance, create_cluster_topology): 25 | cluster_one_broker_empty = create_cluster_topology( 26 | assignment={('topic', 0): [0, 1]}, 27 | brokers=broker_range(3), 28 | ) 29 | command_instance.args.brokers_ids = [2] 30 | balancer = PartitionCountBalancer(cluster_one_broker_empty, command_instance.args) 31 | command_instance.run_command(cluster_one_broker_empty, balancer) 32 | -------------------------------------------------------------------------------- /tests/kafka_cluster_manager/helper.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from kafka_utils.kafka_cluster_manager.cluster_info.broker import Broker 15 | from kafka_utils.kafka_cluster_manager.cluster_info.partition import Partition 16 | 17 | 18 | def create_broker(broker_id, partitions): 19 | b = Broker(broker_id, partitions=set(partitions)) 20 | for p in partitions: 21 | p.add_replica(b) 22 | return b 23 | 24 | 25 | def create_and_attach_partition(topic, partition_id): 26 | partition = Partition(topic, partition_id) 27 | topic.add_partition(partition) 28 | return partition 29 | 30 | 31 | def broker_range(n): 32 | """Return list of brokers with broker ids ranging from 0 to n-1.""" 33 | return {str(x): {"host": "host%s" % x} for x in range(n)} 34 | -------------------------------------------------------------------------------- /tests/kafka_cluster_manager/topic_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from unittest.mock import Mock 15 | from unittest.mock import sentinel 16 | 17 | import pytest 18 | 19 | from kafka_utils.kafka_cluster_manager.cluster_info.partition import Partition 20 | from kafka_utils.kafka_cluster_manager.cluster_info.topic import Topic 21 | 22 | 23 | class TestTopic: 24 | 25 | @pytest.fixture 26 | def topic(self): 27 | return Topic('t0', 2, {sentinel.p1, sentinel.p2}) 28 | 29 | def test_id(self, topic): 30 | assert topic.id == 't0' 31 | 32 | def test_replication_factor(self, topic): 33 | assert topic.replication_factor == 2 34 | 35 | def test_partitions(self, topic): 36 | assert topic.partitions == {sentinel.p1, sentinel.p2} 37 | 38 | def test_add_partition(self): 39 | mock_partitions = { 40 | Mock( 41 | spec=Partition, 42 | replicas=[sentinel.r1, sentinel.r2], 43 | ), 44 | Mock( 45 | spec=Partition, 46 | replicas=[sentinel.r4, sentinel.r3], 47 | ), 48 | } 49 | topic = Topic('t0', 2, mock_partitions) 50 | new_partition = Mock(spec=Partition, replicas=[sentinel.r2]) 51 | topic.add_partition(new_partition) 52 | assert topic.partitions == mock_partitions | {new_partition} 53 | -------------------------------------------------------------------------------- /tests/kafka_cluster_manager/util_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | import pytest 15 | 16 | from kafka_utils.kafka_cluster_manager.cluster_info.util import compute_optimum 17 | from kafka_utils.kafka_cluster_manager.cluster_info.util import separate_groups 18 | 19 | 20 | def test_compute_optimum(): 21 | optimal, extra = compute_optimum(3, 10) 22 | 23 | assert optimal == 3 24 | assert extra == 1 25 | 26 | 27 | def test_compute_optimum_zero_groups(): 28 | with pytest.raises(ZeroDivisionError): 29 | optimal, extra = compute_optimum(0, 10) 30 | 31 | 32 | def test_compute_optimum_zero_elements(): 33 | optimal, extra = compute_optimum(10, 0) 34 | 35 | assert optimal == 0 36 | assert extra == 0 37 | 38 | 39 | def test_separate_groups_balanced(): 40 | groups = [4, 4, 4] 41 | total = 12 42 | 43 | overloaded, underloaded = separate_groups(groups, lambda x: x, total) 44 | 45 | assert not overloaded 46 | assert not underloaded 47 | 48 | 49 | def test_separate_groups_almost_balanced(): 50 | groups = [5, 5, 4] 51 | total = 14 52 | 53 | overloaded, underloaded = separate_groups(groups, lambda x: x, total) 54 | 55 | assert overloaded == [5, 5] 56 | assert underloaded == [4] 57 | 58 | 59 | def test_separate_groups_unbalanced(): 60 | groups = [4, 4, 3, 2] 61 | total = 13 62 | 63 | overloaded, underloaded = separate_groups(groups, lambda x: x, total) 64 | 65 | assert overloaded == [4, 4] 66 | assert underloaded == [2, 3] 67 | 68 | 69 | def test_separate_groups_balanced_greater_total(): 70 | groups = [4, 4, 4] 71 | total = 13 72 | 73 | overloaded, underloaded = separate_groups(groups, lambda x: x, total) 74 | 75 | assert underloaded == [4, 4, 4] 76 | 77 | 78 | def test_separate_groups_balanced_much_greater_total(): 79 | groups = [4, 4, 4] 80 | total = 20 81 | 82 | overloaded, underloaded = separate_groups(groups, lambda x: x, total) 83 | 84 | assert underloaded == [4, 4, 4] 85 | 86 | 87 | def test_separate_groups_balanced_smaller_total(): 88 | groups = [4, 4, 4] 89 | total = 6 90 | 91 | overloaded, underloaded = separate_groups(groups, lambda x: x, total) 92 | 93 | assert overloaded == [4, 4, 4] 94 | -------------------------------------------------------------------------------- /tests/kafka_consumer_manager/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/kafka_consumer_manager/test_copy_group.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | import contextlib 15 | from unittest import mock 16 | 17 | import pytest 18 | 19 | from kafka_utils.kafka_consumer_manager. \ 20 | commands.copy_group import CopyGroup 21 | 22 | 23 | @mock.patch( 24 | 'kafka_utils.kafka_consumer_manager.' 25 | 'commands.copy_group.KafkaToolClient', 26 | autospec=True, 27 | ) 28 | class TestCopyGroup: 29 | 30 | @contextlib.contextmanager 31 | def mock_kafka_info(self, topics_partitions): 32 | with mock.patch.object( 33 | CopyGroup, 34 | "preprocess_args", 35 | spec=CopyGroup.preprocess_args, 36 | return_value=topics_partitions, 37 | ) as mock_process_args, mock.patch( 38 | "kafka_utils.kafka_consumer_manager.util.prompt_user_input", 39 | autospec=True, 40 | ) as mock_user_confirm: 41 | yield mock_process_args, mock_user_confirm 42 | 43 | def test_run_with_kafka(self, mock_client): 44 | topics_partitions = { 45 | "topic1": [0, 1, 2], 46 | "topic2": [0, 1] 47 | } 48 | with self.mock_kafka_info( 49 | topics_partitions 50 | ) as (mock_process_args, mock_user_confirm),\ 51 | mock.patch('kafka_utils.kafka_consumer_manager.commands.' 52 | 'copy_group.get_current_consumer_offsets', 53 | autospec=True) as mock_get_current_consumer_offsets,\ 54 | mock.patch('kafka_utils.kafka_consumer_manager.commands.' 55 | 'copy_group.set_consumer_offsets', 56 | autospec=True) as mock_set_consumer_offsets: 57 | cluster_config = mock.Mock(zookeeper='some_ip') 58 | args = mock.Mock(source_groupid='old_group', dest_groupid='new_group') 59 | CopyGroup.run(args, cluster_config) 60 | assert mock_set_consumer_offsets.call_count == 1 61 | assert mock_get_current_consumer_offsets.call_count == 1 62 | 63 | def test_run_same_groupids(self, mock_client): 64 | topics_partitions = {} 65 | with self.mock_kafka_info( 66 | topics_partitions 67 | ) as (mock_process_args, mock_user_confirm): 68 | with pytest.raises(SystemExit) as ex: 69 | cluster_config = mock.Mock(zookeeper='some_ip') 70 | args = mock.Mock( 71 | source_groupid='my_group', 72 | dest_groupid='my_group', 73 | ) 74 | 75 | CopyGroup.run(args, cluster_config) 76 | 77 | assert ex.value.code == 1 78 | -------------------------------------------------------------------------------- /tests/kafka_consumer_manager/test_delete_group.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from unittest import mock 15 | 16 | import pytest 17 | 18 | from kafka_utils.kafka_consumer_manager.commands. \ 19 | delete_group import DeleteGroup 20 | 21 | 22 | class TestDeleteGroup: 23 | 24 | @pytest.yield_fixture 25 | def client(self): 26 | with mock.patch( 27 | 'kafka_utils.kafka_consumer_manager.' 28 | 'commands.delete_group.KafkaToolClient', 29 | autospec=True, 30 | ) as mock_client: 31 | yield mock_client 32 | 33 | @pytest.yield_fixture 34 | def offsets(self): 35 | yield {'topic1': {0: 100}} 36 | 37 | def test_delete_topic(self, client, offsets): 38 | new_offsets = {'topic1': {0: 0}} 39 | 40 | with mock.patch( 41 | 'kafka_utils.kafka_consumer_manager.' 42 | 'commands.delete_group.set_consumer_offsets', 43 | autospec=True, 44 | ) as mock_set: 45 | DeleteGroup.delete_group_kafka(client, 'some_group', offsets) 46 | 47 | assert mock_set.call_count == 1 48 | assert mock_set.call_args_list == [ 49 | mock.call( 50 | client, 51 | 'some_group', 52 | new_offsets, 53 | ), 54 | ] 55 | -------------------------------------------------------------------------------- /tests/kafka_consumer_manager/test_list_groups.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | import contextlib 15 | from unittest import mock 16 | 17 | from kafka_utils.kafka_consumer_manager. \ 18 | commands.list_groups import ListGroups 19 | 20 | 21 | class TestListGroups: 22 | 23 | @contextlib.contextmanager 24 | def mock_kafka_info(self): 25 | with mock.patch( 26 | "kafka_utils.kafka_consumer_manager." 27 | "commands.list_groups.get_kafka_group_reader", 28 | autospec=True 29 | ) as mock_kafka_reader: 30 | yield mock_kafka_reader 31 | 32 | def test_get_kafka_groups(self): 33 | with self.mock_kafka_info() as (mock_kafka_reader): 34 | expected_groups = { 35 | 'group_name': ['topic1', 'topic2'], 36 | 'group2': ['topic3', 'topic4'] 37 | } 38 | 39 | m = mock_kafka_reader.return_value 40 | m.read_groups.return_value = expected_groups 41 | 42 | cluster_config = mock.Mock(zookeeper='some_ip', type='some_cluster_type') 43 | cluster_config.configure_mock(name='some_cluster_name') 44 | 45 | result = ListGroups.get_kafka_groups(cluster_config) 46 | assert result == list(expected_groups.keys()) 47 | assert m.read_groups.call_count == 1 48 | 49 | @mock.patch("kafka_utils.kafka_consumer_manager.commands.list_groups.print", create=True) 50 | def test_print_groups(self, mock_print): 51 | groups = ['group1', 'group2', 'group3'] 52 | 53 | cluster_config = mock.Mock(zookeeper='some_ip', type='some_cluster_type') 54 | cluster_config.configure_mock(name='some_cluster_name') 55 | 56 | expected_print = [ 57 | mock.call("Consumer Groups:"), 58 | mock.call("\tgroup1"), 59 | mock.call("\tgroup2"), 60 | mock.call("\tgroup3"), 61 | mock.call("3 groups found for cluster some_cluster_name " 62 | "of type some_cluster_type"), 63 | ] 64 | 65 | ListGroups.print_groups(groups, cluster_config) 66 | assert mock_print.call_args_list == expected_print 67 | -------------------------------------------------------------------------------- /tests/kafka_consumer_manager/test_offset_advance.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from unittest import mock 15 | 16 | import pytest 17 | 18 | from kafka_utils.kafka_consumer_manager.commands. \ 19 | offset_advance import OffsetAdvance 20 | 21 | 22 | @mock.patch('kafka_utils.kafka_consumer_manager.' 23 | 'commands.offset_advance.KafkaToolClient') 24 | class TestOffsetAdvance: 25 | topics_partitions = { 26 | "topic1": [0, 1, 2], 27 | "topic2": [0, 1] 28 | } 29 | 30 | def test_run(self, mock_client): 31 | with mock.patch.object( 32 | OffsetAdvance, 33 | "preprocess_args", 34 | spec=OffsetAdvance.preprocess_args, 35 | return_value=self.topics_partitions, 36 | ), mock.patch( 37 | "kafka_utils.kafka_consumer_manager." 38 | "commands.offset_advance.advance_consumer_offsets", 39 | autospec=True 40 | ) as mock_advance: 41 | args = mock.Mock( 42 | groupid="some_group", 43 | topic=None, 44 | partitions=None, 45 | ) 46 | cluster_config = mock.Mock() 47 | OffsetAdvance.run(args, cluster_config) 48 | 49 | ordered_args, _ = mock_advance.call_args 50 | assert ordered_args[1] == args.groupid 51 | assert ordered_args[2] == self.topics_partitions 52 | mock_client.return_value.load_metadata_for_topics. \ 53 | assert_called_once_with() 54 | mock_client.return_value.close.assert_called_once_with() 55 | 56 | def test_run_type_error(self, mock_client): 57 | with mock.patch.object( 58 | OffsetAdvance, 59 | "preprocess_args", 60 | spec=OffsetAdvance.preprocess_args, 61 | return_value="some_string", 62 | ): 63 | args = mock.Mock( 64 | groupid="some_group", 65 | topic=None, 66 | partitions=None 67 | ) 68 | cluster_config = mock.Mock() 69 | with pytest.raises(TypeError): 70 | OffsetAdvance.run(args, cluster_config) 71 | -------------------------------------------------------------------------------- /tests/kafka_consumer_manager/test_offset_get.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from collections import OrderedDict 15 | from unittest import mock 16 | 17 | import pytest 18 | 19 | from kafka_utils.kafka_consumer_manager. \ 20 | commands.offset_get import OffsetGet 21 | from kafka_utils.util.monitoring import ConsumerPartitionOffsets 22 | 23 | 24 | class TestOffsetGet: 25 | 26 | @pytest.yield_fixture 27 | def client(self): 28 | with mock.patch( 29 | 'kafka_utils.kafka_consumer_manager.' 30 | 'commands.offset_get.KafkaToolClient', 31 | autospec=True, 32 | ) as mock_client: 33 | yield mock_client 34 | 35 | def test_get_offsets(self, client): 36 | consumer_group = 'group1' 37 | topics = {'topic1': {0: 100}} 38 | 39 | with mock.patch( 40 | 'kafka_utils.util.offsets._verify_topics_and_partitions', 41 | return_value=topics, 42 | autospec=True, 43 | ): 44 | OffsetGet.get_offsets( 45 | client, 46 | consumer_group, 47 | topics, 48 | ) 49 | 50 | assert client.load_metadata_for_topics.call_count == 1 51 | assert client.send_offset_fetch_request_kafka.call_count == 1 52 | 53 | def test_output_sorting(self): 54 | offsets = OrderedDict( 55 | [("topic2", [ConsumerPartitionOffsets("t2", 0, 0, 10, 0)]), 56 | ("topic1", [ConsumerPartitionOffsets("t1", 0, 5, 10, 0), 57 | ConsumerPartitionOffsets("t1", 1, 9, 10, 0)])]) 58 | 59 | sorted_dict = OffsetGet.sort_by_distance(offsets) 60 | assert list(sorted_dict.keys())[0] == "topic1" 61 | assert list(sorted_dict.keys())[1] == "topic2" 62 | 63 | def test_output_sorting_parcentage(self): 64 | offsets = OrderedDict( 65 | [("topic1", [ConsumerPartitionOffsets("t1", 0, 1, 10, 0), 66 | ConsumerPartitionOffsets("t1", 1, 2, 10, 0)]), 67 | ("topic2", [ConsumerPartitionOffsets("t2", 0, 900, 1000, 0)])]) 68 | 69 | sorted_dict = OffsetGet.sort_by_distance_percentage(offsets) 70 | assert list(sorted_dict.keys())[0] == "topic2" 71 | assert list(sorted_dict.keys())[1] == "topic1" 72 | -------------------------------------------------------------------------------- /tests/kafka_consumer_manager/test_offset_restore.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from unittest import mock 15 | 16 | import pytest 17 | 18 | from kafka_utils.kafka_consumer_manager. \ 19 | commands.offset_restore import OffsetRestore 20 | from kafka_utils.util.client import KafkaToolClient 21 | from kafka_utils.util.monitoring import ConsumerPartitionOffsets 22 | 23 | 24 | class TestOffsetRestore: 25 | 26 | topics_partitions = { 27 | "topic1": [0, 1, 2], 28 | "topic2": [0, 1, 2, 3], 29 | "topic3": [0, 1], 30 | } 31 | consumer_offsets_metadata = { 32 | 'topic1': 33 | [ 34 | ConsumerPartitionOffsets(topic='topic1', partition=0, current=20, highmark=655, lowmark=655), 35 | ConsumerPartitionOffsets(topic='topic1', partition=1, current=10, highmark=655, lowmark=655) 36 | ] 37 | } 38 | parsed_consumer_offsets = {'groupid': 'group1', 'offsets': {'topic1': {0: 10, 1: 20}}} 39 | new_consumer_offsets = {'topic1': {0: 10, 1: 20}} 40 | kafka_consumer_offsets = {'topic1': [ 41 | ConsumerPartitionOffsets(topic='topic1', partition=0, current=30, highmark=40, lowmark=10), 42 | ConsumerPartitionOffsets(topic='topic1', partition=1, current=20, highmark=40, lowmark=10), 43 | ]} 44 | 45 | @pytest.fixture 46 | def mock_kafka_client(self): 47 | mock_kafka_client = mock.MagicMock( 48 | spec=KafkaToolClient 49 | ) 50 | mock_kafka_client.get_partition_ids_for_topic. \ 51 | side_effect = self.topics_partitions 52 | return mock_kafka_client 53 | 54 | def test_build_new_offsets(self, mock_kafka_client): 55 | new_offsets = OffsetRestore.build_new_offsets( 56 | mock_kafka_client, 57 | {'topic1': {0: 10, 1: 20}}, 58 | {'topic1': [0, 1]}, 59 | self.kafka_consumer_offsets, 60 | ) 61 | 62 | assert new_offsets == self.new_consumer_offsets 63 | -------------------------------------------------------------------------------- /tests/kafka_consumer_manager/test_offset_rewind.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from unittest import mock 15 | 16 | import pytest 17 | 18 | from kafka_utils.kafka_consumer_manager. \ 19 | commands.offset_rewind import OffsetRewind 20 | 21 | 22 | class TestOffsetRewind: 23 | topics_partitions = { 24 | "topic1": [0, 1, 2], 25 | "topic2": [0, 1] 26 | } 27 | 28 | @mock.patch('kafka_utils.kafka_consumer_manager.' 29 | 'commands.offset_rewind.KafkaToolClient') 30 | def test_run(self, mock_client): 31 | with mock.patch.object( 32 | OffsetRewind, 33 | "preprocess_args", 34 | spec=OffsetRewind.preprocess_args, 35 | return_value=self.topics_partitions, 36 | ), mock.patch( 37 | "kafka_utils.kafka_consumer_manager." 38 | "commands.offset_rewind.rewind_consumer_offsets", 39 | autospec=True 40 | ) as mock_rewind: 41 | args = mock.Mock( 42 | groupid="some_group", 43 | topic=None, 44 | partitions=None 45 | ) 46 | cluster_config = mock.Mock() 47 | OffsetRewind.run(args, cluster_config) 48 | 49 | ordered_args, _ = mock_rewind.call_args 50 | assert ordered_args[1] == args.groupid 51 | assert ordered_args[2] == self.topics_partitions 52 | mock_client.return_value.load_metadata_for_topics. \ 53 | assert_called_once_with() 54 | mock_client.return_value.close.assert_called_once_with() 55 | 56 | @mock.patch('kafka_utils.kafka_consumer_manager' 57 | '.commands.offset_rewind.KafkaToolClient') 58 | def test_run_type_error(self, mock_client): 59 | with mock.patch.object( 60 | OffsetRewind, 61 | "preprocess_args", 62 | spec=OffsetRewind.preprocess_args, 63 | return_value="some_string", 64 | ): 65 | args = mock.Mock( 66 | groupid="some_group", 67 | topic=None, 68 | partitions=None 69 | ) 70 | cluster_config = mock.Mock() 71 | with pytest.raises(TypeError): 72 | OffsetRewind.run(args, cluster_config) 73 | -------------------------------------------------------------------------------- /tests/kafka_consumer_manager/test_offset_save.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from unittest import mock 15 | 16 | from kafka_utils.kafka_consumer_manager. \ 17 | commands.offset_save import OffsetSave 18 | from kafka_utils.util.monitoring import ConsumerPartitionOffsets 19 | 20 | 21 | class TestOffsetSave: 22 | topics_partitions = { 23 | "topic1": [0, 1, 2], 24 | "topic2": [0, 1, 2, 3], 25 | "topic3": [0, 1], 26 | } 27 | consumer_offsets_metadata = { 28 | 'topic1': 29 | [ 30 | ConsumerPartitionOffsets(topic='topic1', partition=0, current=10, highmark=655, lowmark=655), 31 | ConsumerPartitionOffsets(topic='topic1', partition=1, current=20, highmark=655, lowmark=655), 32 | ] 33 | } 34 | offset_data_file = {'groupid': 'group1', 'offsets': {'topic1': {0: 10, 1: 20}}} 35 | json_data = {'groupid': 'group1', 'offsets': {'topic1': {'0': 10, '1': 20}}} 36 | 37 | @mock.patch('kafka_utils.kafka_consumer_manager.' 38 | 'commands.offset_save.KafkaToolClient') 39 | def test_save_offsets(self, mock_client): 40 | with mock.patch.object( 41 | OffsetSave, 42 | "write_offsets_to_file", 43 | spec=OffsetSave.write_offsets_to_file, 44 | return_value=[], 45 | ) as mock_write_offsets: 46 | filename = 'offset_file' 47 | consumer_group = 'group1' 48 | OffsetSave.save_offsets( 49 | self.consumer_offsets_metadata, 50 | self.topics_partitions, 51 | filename, 52 | consumer_group, 53 | ) 54 | 55 | ordered_args, _ = mock_write_offsets.call_args 56 | assert ordered_args[0] == filename 57 | assert ordered_args[1] == self.offset_data_file 58 | 59 | @mock.patch('kafka_utils.kafka_consumer_manager.' 60 | 'commands.offset_save.KafkaToolClient') 61 | def test_run(self, mock_client): 62 | with mock.patch.object( 63 | OffsetSave, 64 | "preprocess_args", 65 | spec=OffsetSave.preprocess_args, 66 | return_value=self.topics_partitions, 67 | ), mock.patch( 68 | "kafka_utils.kafka_consumer_manager." 69 | "commands.offset_save.get_consumer_offsets_metadata", 70 | return_value=self.consumer_offsets_metadata, 71 | autospec=True, 72 | ), mock.patch.object( 73 | OffsetSave, 74 | "write_offsets_to_file", 75 | spec=OffsetSave.write_offsets_to_file, 76 | return_value=[], 77 | ) as mock_write_offsets: 78 | args = mock.Mock( 79 | groupid="group1", 80 | json_file="some_file", 81 | ) 82 | cluster_config = mock.Mock() 83 | OffsetSave.run(args, cluster_config) 84 | 85 | mock_client.return_value.load_metadata_for_topics. \ 86 | assert_called_once_with() 87 | mock_client.return_value.close.assert_called_once_with() 88 | ordered_args, _ = mock_write_offsets.call_args 89 | assert ordered_args[0] == "some_file" 90 | assert ordered_args[1] == self.offset_data_file 91 | -------------------------------------------------------------------------------- /tests/kafka_consumer_manager/test_rename_group.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | import contextlib 15 | from unittest import mock 16 | 17 | import pytest 18 | 19 | from kafka_utils.kafka_consumer_manager. \ 20 | commands.rename_group import RenameGroup 21 | 22 | 23 | @mock.patch( 24 | 'kafka_utils.kafka_consumer_manager.' 25 | 'commands.rename_group.KafkaToolClient', 26 | autospec=True, 27 | ) 28 | class TestRenameGroup: 29 | 30 | @contextlib.contextmanager 31 | def mock_kafka_info(self, topics_partitions): 32 | with mock.patch.object( 33 | RenameGroup, 34 | "preprocess_args", 35 | spec=RenameGroup.preprocess_args, 36 | return_value=topics_partitions, 37 | ) as mock_process_args, mock.patch( 38 | "kafka_utils.kafka_consumer_manager.util.prompt_user_input", 39 | autospec=True, 40 | ) as mock_user_confirm: 41 | yield mock_process_args, mock_user_confirm 42 | 43 | def test_run_with_kafka(self, mock_client): 44 | topics_partitions = { 45 | "topic1": [0, 1, 2], 46 | "topic2": [0, 1] 47 | } 48 | with self.mock_kafka_info( 49 | topics_partitions 50 | ) as (mock_process_args, mock_user_confirm),\ 51 | mock.patch('kafka_utils.kafka_consumer_manager.commands.' 52 | 'rename_group.get_current_consumer_offsets', 53 | autoapec=True) as mock_get_current_consumer_offsets,\ 54 | mock.patch('kafka_utils.kafka_consumer_manager.commands.' 55 | 'rename_group.set_consumer_offsets', 56 | autospec=True) as mock_set_consumer_offsets: 57 | cluster_config = mock.Mock(zookeeper='some_ip') 58 | args = mock.Mock(source_groupid='old_group', dest_groupid='new_group') 59 | RenameGroup.run(args, cluster_config) 60 | assert mock_set_consumer_offsets.call_count == 2 61 | assert mock_get_current_consumer_offsets.call_count == 1 62 | 63 | def test_run_same_groupids(self, mock_client): 64 | topics_partitions = {} 65 | with self.mock_kafka_info( 66 | topics_partitions 67 | ) as (mock_process_args, mock_user_confirm): 68 | with pytest.raises(SystemExit) as ex: 69 | cluster_config = mock.Mock(zookeeper='some_ip') 70 | args = mock.Mock( 71 | old_groupid='my_group', 72 | new_groupid='my_group', 73 | ) 74 | 75 | RenameGroup.run(args, cluster_config) 76 | 77 | ex.value.code == 1 78 | -------------------------------------------------------------------------------- /tests/kafka_consumer_manager/test_unsubscribe_topics.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from unittest import mock 15 | 16 | import pytest 17 | 18 | from kafka_utils.kafka_consumer_manager.commands. \ 19 | unsubscribe_topics import KafkaUnsubscriber 20 | from kafka_utils.kafka_consumer_manager.commands. \ 21 | unsubscribe_topics import UnsubscribeTopics 22 | 23 | 24 | class TestUnsubscribeTopics: 25 | 26 | @pytest.yield_fixture 27 | def client(self): 28 | with mock.patch( 29 | 'kafka_utils.kafka_consumer_manager.' 30 | 'commands.unsubscribe_topics.KafkaToolClient', 31 | autospec=True, 32 | ) as mock_client: 33 | yield mock_client 34 | 35 | @pytest.yield_fixture 36 | def kafka_unsubscriber(self): 37 | with mock.patch( 38 | 'kafka_utils.kafka_consumer_manager.' 39 | 'commands.unsubscribe_topics.KafkaUnsubscriber', 40 | autospec=True, 41 | ) as mock_kafka_unsub: 42 | yield mock_kafka_unsub 43 | 44 | topics_partitions = { 45 | 'topic1': [0, 1, 2], 46 | 'topic2': [0, 1], 47 | 'topic3': [0, 1, 2, 3], 48 | } 49 | 50 | cluster_config = mock.Mock(zookeeper='some_ip') 51 | 52 | def test_run_unsubscribe_kafka(self, client, kafka_unsubscriber): 53 | with mock.patch.object( 54 | UnsubscribeTopics, 55 | 'preprocess_args', 56 | spec=UnsubscribeTopics.preprocess_args, 57 | return_value=self.topics_partitions, 58 | ): 59 | args = mock.Mock(topic=None, topics=None) 60 | 61 | UnsubscribeTopics.run(args, self.cluster_config) 62 | 63 | kafka_obj = kafka_unsubscriber.return_value 64 | 65 | assert kafka_obj.unsubscribe_topics.call_count == 1 66 | 67 | def test_run_unsubscribe_only_topics(self, client, kafka_unsubscriber): 68 | with mock.patch.object( 69 | UnsubscribeTopics, 70 | 'preprocess_args', 71 | spec=UnsubscribeTopics.preprocess_args, 72 | return_value=self.topics_partitions, 73 | ): 74 | args = mock.Mock(topics=["topic1", "topic3"], topic=None, partitions=None) 75 | 76 | UnsubscribeTopics.run(args, self.cluster_config) 77 | 78 | def test_unsubscribe_topic(self, client): 79 | offsets = {'topic1': {0: 100}} 80 | new_offsets = {'topic1': {0: 0}} 81 | 82 | with mock.patch( 83 | 'kafka_utils.kafka_consumer_manager.' 84 | 'commands.unsubscribe_topics.get_current_consumer_offsets', 85 | autospec=True, 86 | return_value=offsets, 87 | ) as mock_get, mock.patch( 88 | 'kafka_utils.kafka_consumer_manager.' 89 | 'commands.unsubscribe_topics.set_consumer_offsets', 90 | autospec=True, 91 | ) as mock_set: 92 | unsubscriber = KafkaUnsubscriber(client) 93 | unsubscriber.delete_topic('some_group', 'topic1') 94 | 95 | assert mock_get.call_count == 1 96 | assert mock_set.call_count == 1 97 | assert mock_set.call_args_list == [ 98 | mock.call(client, 'some_group', new_offsets), 99 | ] 100 | -------------------------------------------------------------------------------- /tests/kafka_consumer_manager/test_watermark_get.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from unittest import mock 15 | 16 | import pytest 17 | 18 | from kafka_utils.kafka_consumer_manager. \ 19 | commands.watermark_get import WatermarkGet 20 | 21 | 22 | class TestGetWatermark: 23 | 24 | @pytest.yield_fixture 25 | def client(self): 26 | with mock.patch( 27 | 'kafka_utils.kafka_consumer_manager.' 28 | 'commands.watermark_get.KafkaToolClient', 29 | autospec=True, 30 | ) as mock_client: 31 | yield mock_client 32 | 33 | def test_get_watermark_for_topic(self, client): 34 | topics = '__consumer_offsets' 35 | client.topic_partitions = {} 36 | with mock.patch( 37 | 'kafka_utils.kafka_consumer_manager.commands.' 38 | 'watermark_get.get_watermark_for_topic', 39 | return_value={'test_topic': [1, 99, 3]}, 40 | autospec=True, 41 | ) as mock_get_watermark: 42 | WatermarkGet.get_watermarks( 43 | client, 44 | topics, 45 | exact=True 46 | ) 47 | assert mock_get_watermark.call_count == 1 48 | 49 | def test_get_watermark_for_regex(self, client): 50 | topics = '__consumer_*' 51 | client.topic_partitions = {} 52 | with mock.patch( 53 | 'kafka_utils.kafka_consumer_manager.commands.' 54 | 'watermark_get.get_watermark_for_regex', 55 | return_value={'__consumer_1': [1, 99, 3], 56 | '__consumer_2': [2, 100, 2]}, 57 | autospec=True, 58 | ) as mock_get_watermark: 59 | WatermarkGet.get_watermarks( 60 | client, 61 | topics, 62 | exact=False 63 | ) 64 | assert mock_get_watermark.call_count == 1 65 | -------------------------------------------------------------------------------- /tests/kafka_corruption_check/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Yelp/kafka-utils/def433ec4d07c60290d5dc937d3b4e5189eca9dc/tests/kafka_corruption_check/__init__.py -------------------------------------------------------------------------------- /tests/kafka_corruption_check/test_main.py: -------------------------------------------------------------------------------- 1 | from unittest import mock 2 | 3 | from kafka_utils.kafka_corruption_check import main 4 | 5 | 6 | FILE_LINE = "Dumping /path/to/file" 7 | CORRECT_LINE = "offset: 2413 position: 173 isvalid: true payloadsize: 3 magic: 0 compresscodec: NoCompressionCodec crc: 1879665293" 8 | INVALID_LINE = "offset: 2413 position: 173 isvalid: false payloadsize: 3 magic: 0 compresscodec: NoCompressionCodec crc: 1879665293" 9 | 10 | 11 | def mock_output(values): 12 | output = mock.Mock() 13 | output.readlines.return_value = values 14 | return output 15 | 16 | 17 | def test_find_files_cmd_min(): 18 | cmd = main.find_files_cmd("path", 60, None, None) 19 | assert cmd == 'find "path" -type f -name "*.log" -mmin -60' 20 | 21 | 22 | def test_find_files_cmd_start(): 23 | cmd = main.find_files_cmd("path", None, "START", None) 24 | assert cmd == 'find "path" -type f -name "*.log" -newermt "START"' 25 | 26 | 27 | def test_find_files_cmd_range(): 28 | cmd = main.find_files_cmd("path", None, "START", "END") 29 | assert cmd == 'find "path" -type f -name "*.log" -newermt "START" \\! -newermt "END"' 30 | 31 | 32 | @mock.patch( 33 | "kafka_utils.kafka_corruption_check." 34 | "main.print_line", 35 | ) 36 | def test_parse_output_correct(mock_print): 37 | main.parse_output("HOST", mock_output([CORRECT_LINE])) 38 | assert not mock_print.called 39 | 40 | 41 | @mock.patch( 42 | "kafka_utils.kafka_corruption_check." 43 | "main.print_line", 44 | ) 45 | def test_parse_output_invalid(mock_print): 46 | main.parse_output("HOST", mock_output([FILE_LINE, INVALID_LINE])) 47 | mock_print.assert_called_once_with("HOST", "/path/to/file", INVALID_LINE, "ERROR") 48 | 49 | 50 | @mock.patch( 51 | "kafka_utils.kafka_corruption_check." 52 | "main.get_partition_leaders", 53 | return_value={"t0-0": 1, "t0-1": 2}, 54 | ) 55 | def test_filter_leader_files(mock_get_partition): 56 | filtered = main.filter_leader_files(None, [(1, "host1", ["a/kafka-logs/t0-0/0123.log", 57 | "a/kafka-logs/t2-0/0123.log", 58 | "a/kafka-logs/t0-1/0123.log"]), 59 | (2, "host2", ["a/kafka-logs/t0-0/0123.log", 60 | "a/kafka-logs/t0-1/0123.log"])]) 61 | assert filtered == [(1, 'host1', ['a/kafka-logs/t0-0/0123.log', 62 | 'a/kafka-logs/t2-0/0123.log']), 63 | (2, 'host2', ['a/kafka-logs/t0-1/0123.log'])] 64 | -------------------------------------------------------------------------------- /tests/kafka_manual_throttle/test_main.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from unittest import mock 15 | 16 | import pytest 17 | 18 | from kafka_utils.kafka_manual_throttle import main 19 | from kafka_utils.util.zookeeper import ZK 20 | 21 | 22 | BROKER_ID = 0 23 | 24 | 25 | @pytest.fixture 26 | def mock_zk(): 27 | return mock.Mock(spec=ZK) 28 | 29 | 30 | def test_apply_throttles(mock_zk): 31 | mock_zk.get_broker_config.return_value = {} 32 | 33 | brokers = [0, 1, 2] 34 | main.apply_throttles(mock_zk, brokers, 42, 24) 35 | 36 | assert mock_zk.set_broker_config.call_count == len(brokers) 37 | 38 | expected_config = { 39 | "config": {"leader.replication.throttled.rate": "42", "follower.replication.throttled.rate": "24"}, 40 | } 41 | 42 | expected_set_calls = [ 43 | mock.call(0, expected_config), 44 | mock.call(1, expected_config), 45 | mock.call(2, expected_config), 46 | ] 47 | 48 | assert mock_zk.set_broker_config.call_args_list == expected_set_calls 49 | 50 | 51 | def test_clear_throttles(mock_zk): 52 | mock_zk.get_broker_config.return_value = {"config": {"leader.replication.throttled.rate": "42", "follower.replication.throttled.rate": "24"}} 53 | 54 | brokers = [0, 1, 2] 55 | main.clear_throttles(mock_zk, brokers) 56 | 57 | assert mock_zk.set_broker_config.call_count == len(brokers) 58 | 59 | expected_config = {"config": {}} 60 | expected_set_calls = [ 61 | mock.call(0, expected_config), 62 | mock.call(1, expected_config), 63 | mock.call(2, expected_config), 64 | ] 65 | 66 | assert mock_zk.set_broker_config.call_args_list == expected_set_calls 67 | 68 | 69 | @pytest.mark.parametrize( 70 | "current_config,leader,follower,expected_config", [ 71 | ({}, None, None, {}), 72 | ({}, 42, None, {"leader.replication.throttled.rate": 42}), 73 | ({}, 42, 24, {"leader.replication.throttled.rate": 42, "follower.replication.throttled.rate": 24}), 74 | ( 75 | {"leader.replication.throttled.rate": 42, "follower.replication.throttled.rate": 24}, 76 | None, 77 | 100, 78 | {"follower.replication.throttled.rate": 100}, 79 | ), 80 | ( 81 | {"other.config": "other"}, 82 | 42, 83 | 24, 84 | {"other.config": "other", "leader.replication.throttled.rate": 42, "follower.replication.throttled.rate": 24} 85 | ), 86 | ] 87 | ) 88 | def test_write_throttle(mock_zk, current_config, leader, follower, expected_config): 89 | mock_zk.get_broker_config.return_value = {"config": current_config} 90 | 91 | main.write_throttle(mock_zk, BROKER_ID, leader, follower) 92 | 93 | assert mock_zk.set_broker_config.call_count == 1 94 | 95 | expected_set_call = mock.call(BROKER_ID, {"config": expected_config}) 96 | assert mock_zk.set_broker_config.call_args_list == [expected_set_call] 97 | -------------------------------------------------------------------------------- /tests/kafka_rolling_restart/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | -------------------------------------------------------------------------------- /tests/util/util_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2016 Yelp Inc. 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | from argparse import ArgumentTypeError 15 | 16 | import pytest 17 | 18 | from kafka_utils.util import positive_float 19 | from kafka_utils.util import positive_int 20 | from kafka_utils.util import positive_nonzero_int 21 | from kafka_utils.util import tuple_alter 22 | from kafka_utils.util import tuple_remove 23 | from kafka_utils.util import tuple_replace 24 | from kafka_utils.util.utils import child_class 25 | 26 | 27 | def test_tuple_alter(): 28 | result = tuple_alter( 29 | (1, 2, 3, 2, 1), 30 | (2, lambda x: x + 1), 31 | (3, lambda x: x + 2), 32 | (4, lambda x: x + 3), 33 | ) 34 | 35 | assert result == (1, 2, 4, 4, 4) 36 | 37 | 38 | def test_tuple_remove(): 39 | assert tuple_remove((1, 2, 3, 2, 1), 1, 2, 3) == (2, 1) 40 | 41 | 42 | def test_tuple_replace(): 43 | result = tuple_replace( 44 | (1, 2, 3, 2, 1), 45 | (2, 1), 46 | (3, 2), 47 | (4, 3), 48 | ) 49 | 50 | assert result == (1, 2, 1, 2, 3) 51 | 52 | 53 | def test_positive_int_valid(): 54 | assert positive_int('123') == 123 55 | 56 | 57 | def test_positive_int_not_int(): 58 | with pytest.raises(ArgumentTypeError): 59 | positive_int('not_an_int') 60 | 61 | 62 | def test_positive_int_negative_int(): 63 | with pytest.raises(ArgumentTypeError): 64 | positive_int('-5') 65 | 66 | 67 | def test_positive_nonzero_int_valid(): 68 | assert positive_nonzero_int('123') == 123 69 | 70 | 71 | def test_positive_nonzero_int_not_int(): 72 | with pytest.raises(ArgumentTypeError): 73 | positive_nonzero_int('not_an_int') 74 | 75 | 76 | def test_positive_nonzero_int_zero(): 77 | with pytest.raises(ArgumentTypeError): 78 | positive_nonzero_int('0') 79 | 80 | 81 | def test_positive_float_valid(): 82 | assert positive_float('123.0') == 123.0 83 | 84 | 85 | def test_positive_float_not_float(): 86 | with pytest.raises(ArgumentTypeError): 87 | positive_float('not_a_float') 88 | 89 | 90 | def test_positive_float_negative_float(): 91 | with pytest.raises(ArgumentTypeError): 92 | positive_float('-1.45') 93 | 94 | 95 | def test_child_class(): 96 | class A: 97 | pass 98 | 99 | class B(A): 100 | pass 101 | 102 | class C(B): 103 | pass 104 | 105 | class D(A): 106 | pass 107 | 108 | assert child_class([A, B, C], A) == C 109 | assert child_class([A, B], A) == B 110 | assert child_class([A, B, D], A) in {B, D} 111 | -------------------------------------------------------------------------------- /tox.ini: -------------------------------------------------------------------------------- 1 | [tox] 2 | envlist = py{37,38}-unittest, py{37,38}-kafka{10,11}-dockeritest 3 | # the Github actions override the index server for our CI builds 4 | # The Makefile also sets PIP_INDEX_URL to the public PyPI when 5 | # running outside of Yelp. 6 | indexserver = 7 | default = https://pypi.yelpcorp.com/simple 8 | tox_pip_extensions_ext_pip_custom_platform = true 9 | tox_pip_extensions_ext_venv_update = true 10 | 11 | [testenv] 12 | deps = 13 | -rrequirements-dev.txt 14 | dockeritest: docker-compose==1.7.0 15 | whitelist_externals = /bin/bash 16 | passenv = ITEST_PYTHON_FACTOR KAFKA_VERSION ACCEPTANCE_TAGS 17 | setenv = 18 | py37: ITEST_PYTHON_FACTOR = py37 19 | py38: ITEST_PYTHON_FACTOR = py38 20 | kafka10: KAFKA_VERSION = 0.10.1.1 21 | kafka10: ACCEPTANCE_TAGS = ~kafka11 22 | kafka11: KAFKA_VERSION = 1.1.0 23 | kafka11: ACCEPTANCE_TAGS = 24 | commands = 25 | unittest: pre-commit install -f --install-hooks 26 | unittest: pre-commit run --all-files 27 | unittest: py.test -s {posargs} 28 | unittest: flake8 . 29 | unittest: mypy kafka_utils/ 30 | 31 | dockeritest: docker-compose rm --force 32 | dockeritest: docker-compose build 33 | dockeritest: /bin/bash -c " \ 34 | dockeritest: docker-compose run \ 35 | dockeritest: -e ITEST_PYTHON_FACTOR={env:ITEST_PYTHON_FACTOR} \ 36 | dockeritest: -e ACCEPTANCE_TAGS={env:ACCEPTANCE_TAGS} \ 37 | dockeritest: itest /scripts/run_tests.sh; exit_status=$?; \ 38 | dockeritest: docker-compose stop; \ 39 | dockeritest: docker network rm kafkautils_default; \ 40 | dockeritest: exit $exit_status" 41 | 42 | [testenv:coverage] 43 | deps = 44 | {[testenv]deps} 45 | coverage 46 | commands = 47 | coverage run --source kafka_utils/ -m pytest --strict {posargs} 48 | coverage report -m 49 | flake8 . 50 | 51 | [testenv:docs] 52 | basepython = python3.8 53 | deps = 54 | sphinx 55 | sphinx_rtd_theme 56 | changedir = docs 57 | commands = sphinx-build -b html -d build/doctrees source build/html 58 | 59 | [flake8] 60 | exclude = .svn,CVS,.bzr,.hg,.git,__pycache__,.tox,docs,virtualenv_run,.ropeproject 61 | ignore = E501,W605,W504 62 | 63 | [pytest] 64 | norecursedirs = .* virtualenv_run 65 | -------------------------------------------------------------------------------- /tox_acceptance.ini: -------------------------------------------------------------------------------- 1 | # this file is run via a Docker container 2 | # see docker/itest/run_tests.sh for more details 3 | [tox] 4 | envlist = py{37,38}-kafka{10,11}-dockeritest 5 | tox_pip_extensions_ext_pip_custom_platform = true 6 | tox_pip_extensions_ext_venv_update = true 7 | distdir = {toxworkdir}/dist_acceptance 8 | 9 | [testenv] 10 | deps = 11 | -rrequirements-dev.txt 12 | flake8 13 | mock 14 | acceptance: behave 15 | whitelist_externals = /bin/bash 16 | passenv = ITEST_PYTHON_FACTOR KAFKA_VERSION ACCEPTANCE_TAGS 17 | setenv = 18 | py37: ITEST_PYTHON_FACTOR = py37 19 | py38: ITEST_PYTHON_FACTOR = py38 20 | kafka10: KAFKA_VERSION = 0.10.1.1 21 | kafka10: ACCEPTANCE_TAGS = ~kafka11 22 | kafka11: KAFKA_VERSION = 1.1.0 23 | kafka11: ACCEPTANCE_TAGS = 24 | commands = 25 | acceptance: /bin/bash -c 'echo "Running acceptance tests using" $({envbindir}/python --version)' 26 | acceptance: /bin/bash -c 'env' 27 | acceptance: /bin/bash -c 'behave tests/acceptance --tags=$ACCEPTANCE_TAGS --no-capture' 28 | --------------------------------------------------------------------------------