Transfer command

config/sample_config.cfg

@@ -5,8 +5,10 @@ current = docker
 bootstrap_hosts = localhost
 bootstrap_port = 9092
 security_protocol = LOCAL
+schema_registry = localhost:8081
 
 [Context.docker]
 bootstrap_hosts = kafka
 bootstrap_port = 9093
 security_protocol = PLAINTEXT
+schema_registry = schema-registry:8081

docker-compose.yml

@@ -38,4 +38,44 @@ services:
       KAFKA_GROUP_INITIAL_REBALANCE_DELAY_MS: 0
 
 
+  schema_registry_source:
+    image: confluentinc/cp-schema-registry:5.2.2
+    container_name: schema_registry_source
+    environment:
+      - SCHEMA_REGISTRY_KAFKASTORE_CONNECTION_URL=zookeeper:2181
+      - SCHEMA_REGISTRY_HOST_NAME=schema-registry
+      - SCHEMA_REGISTRY_LISTENERS=http://0.0.0.0:8081
+      - SCHEMA_REGISTRY_KAFKASTORE_BOOTSTRAP_SERVERS=PLAINTEXT://kafka:9092
+    links:
+      - zookeeper
+      - kafka
+    ports:
+      - 8081:8081
+
+  schema_registry_target:
+    image: confluentinc/cp-schema-registry:5.2.2
+    container_name: schema_registry_target
+    environment:
+      - SCHEMA_REGISTRY_KAFKASTORE_CONNECTION_URL=zookeeper:2181
+      - SCHEMA_REGISTRY_HOST_NAME=schema-registry
+      - SCHEMA_REGISTRY_LISTENERS=http://0.0.0.0:8082
+      - SCHEMA_REGISTRY_KAFKASTORE_BOOTSTRAP_SERVERS=PLAINTEXT://kafka:9092
+    links:
+      - zookeeper
+      - kafka
+    ports:
+      - 8082:8081
+
+  esque:
+    build: .
+    volumes:
+      - .:/esque
+    environment:
+      ESQUE_TEST_ENV: "ci"
+    depends_on:
+      - kafka
+    command: >
+      -c
+        "(until (kafkacat -b kafka:9093 -X debug=all -L); do sleep 5s; done) \
+        && pytest tests/ --integration"
 

esque/avromessage.py

@@ -0,0 +1,124 @@
+import json
+import pathlib
+import pickle
+import struct
+from io import BytesIO
+from typing import Optional, Tuple, Dict, BinaryIO, Iterable
+import itertools as it
+
+import fastavro
+from confluent_kafka.cimpl import Message
+from confluent_kafka.avro import loads as load_schema
+
+from esque.message import FileWriter, FileReader, KafkaMessage
+from esque.schemaregistry import SchemaRegistryClient
+
+
+class DecodedAvroMessage:
+    def __init__(
+        self,
+        key: Optional[Dict],
+        value: Optional[Dict],
+        key_schema_id: int,
+        value_schema_id: int,
+    ):
+        self.key = key
+        self.value = value
+        self.key_schema_id = key_schema_id
+        self.value_schema_id = value_schema_id
+
+
+class AvroFileWriter(FileWriter):
+    def __init__(
+        self, working_dir: pathlib.Path, schema_registry_client: SchemaRegistryClient
+    ):
+        self.working_dir = working_dir
+        self.schema_registry_client = schema_registry_client
+        self.current_key_schema_id = None
+        self.current_value_schema_id = None
+        self.schema_dir_name = None
+        self.schema_version = it.count(1)
+
+    def write_message_to_file(self, message: Message, file: BinaryIO):
+        key_schema_id, decoded_key = self.decode_bytes(message.key())
+        value_schema_id, decoded_value = self.decode_bytes(message.value())
+        decoded_message = DecodedAvroMessage(
+            decoded_key, decoded_value, key_schema_id, value_schema_id
+        )
+
+        if self.schema_changed(decoded_message) or self.schema_dir_name is None:
+            self.schema_dir_name = (
+                f"{next(self.schema_version):04}_{key_schema_id}_{value_schema_id}"
+            )
+            self.current_key_schema_id = key_schema_id
+            self.current_value_schema_id = value_schema_id
+            self._dump_schemata(key_schema_id, value_schema_id)
+
+        serializable_message = {
+            "key": decoded_key,
+            "value": decoded_value,
+            "schema_directory_name": self.schema_dir_name,
+        }
+        pickle.dump(serializable_message, file)
+
+    def _dump_schemata(self, key_schema_id, value_schema_id):
+        directory = self.working_dir / self.schema_dir_name
+        directory.mkdir()
+        (directory / "key_schema.avsc").write_text(
+            json.dumps(
+                self.schema_registry_client.get_schema_from_id(
+                    key_schema_id
+                ).original_schema
+            )
+        )
+        (directory / "value_schema.avsc").write_text(
+            json.dumps(
+                self.schema_registry_client.get_schema_from_id(
+                    value_schema_id
+                ).original_schema
+            )
+        )
+
+    def decode_bytes(self, raw_data: Optional[bytes]) -> Tuple[int, Optional[Dict]]:
+        if raw_data is None:
+            return -1, None
+
+        with BytesIO(raw_data) as fake_stream:
+            schema_id = extract_schema_id(fake_stream.read(5))
+            parsed_schema = self.schema_registry_client.get_schema_from_id(
+                schema_id
+            ).parsed_schema
+            record = fastavro.schemaless_reader(fake_stream, parsed_schema)
+        return schema_id, record
+
+    def schema_changed(self, decoded_message: DecodedAvroMessage) -> bool:
+        return (
+            self.current_value_schema_id != decoded_message.value_schema_id
+            and decoded_message.value is not None
+        ) or self.current_key_schema_id != decoded_message.key_schema_id
+
+
+class AvroFileReader(FileReader):
+    def __init__(self, working_dir: pathlib.Path):
+        self.working_dir = working_dir
+
+    def read_from_file(self, file: BinaryIO) -> Iterable[KafkaMessage]:
+        while True:
+            try:
+                record = pickle.load(file)
+            except EOFError:
+                return
+
+            schema_directory = self.working_dir / record["schema_directory_name"]
+
+            key_schema = load_schema((schema_directory / "key_schema.avsc").read_text())
+            value_schema = load_schema(
+                (schema_directory / "value_schema.avsc").read_text()
+            )
+
+            yield KafkaMessage(record["key"], record["value"], key_schema, value_schema)
+
+
+def extract_schema_id(message: bytes) -> int:
+    _, schema_id = struct.unpack(">bI", message[:5])
+    return schema_id

esque/cli/commands.py

@@ -1,18 +1,34 @@
+import pathlib
+import time
+from pathlib import Path
 from time import sleep
 
 import click
-from click import version_option
-
 import yaml
+from click import version_option
 
 from esque.__version__ import __version__
 from esque.broker import Broker
-from esque.cli.helpers import ensure_approval
+from esque.cli.helpers import ensure_approval, DeleteOnFinished
 from esque.cli.options import State, no_verify_option, pass_state
-from esque.cli.output import bold, pretty, pretty_topic_diffs, get_output_new_topics
-from esque.clients import Consumer, Producer
+from esque.cli.output import (
+    bold,
+    pretty,
+    pretty_topic_diffs,
+    get_output_new_topics,
+    blue_bold,
+    green_bold,
+)
+from esque.clients import (
+    FileConsumer,
+    FileProducer,
+    AvroFileProducer,
+    AvroFileConsumer,
+    PingConsumer,
+    PingProducer,
+)
 from esque.cluster import Cluster
-from esque.config import PING_TOPIC, Config
+from esque.config import PING_TOPIC, Config, PING_GROUP_ID
 from esque.consumergroup import ConsumerGroupController
 from esque.errors import (
     ConsumerGroupDoesNotExistException,
@@ -280,6 +296,130 @@ def get_topics(state, topic, o):
         click.echo(topic.name)
 
 
+@esque.command(
+    "transfer", help="Transfer messages of a topic from one environment to another."
+)
+@click.argument("topic", required=True)
+@click.option(
+    "-f",
+    "--from",
+    "from_context",
+    help="Source Context",
+    type=click.STRING,
+    required=True,
+)
+@click.option(
+    "-t",
+    "--to",
+    "to_context",
+    help="Destination context",
+    type=click.STRING,
+    required=True,
+)
+@click.option(
+    "-n", "--numbers", help="Number of messages", type=click.INT, required=True
+)
+@click.option("--last/--first", default=False)
+@click.option(
+    "-a",
+    "--avro",
+    help="Set this flag if the topic contains avro data",
+    default=False,
+    is_flag=True,
+)
+@pass_state
+def transfer(
+    state: State,
+    topic: str,
+    from_context: str,
+    to_context: str,
+    numbers: int,
+    last: bool,
+    avro: bool,
+):
+    current_timestamp_milliseconds = int(round(time.time() * 1000))
+    temp_name = topic + "_" + str(current_timestamp_milliseconds)
+    group_id = "group_for_" + temp_name
+    directory_name = "temp-file_" + temp_name
+    base_dir = Path(directory_name)
+    state.config.context_switch(from_context)
+
+    with DeleteOnFinished(base_dir) as working_dir:
+        number_consumed_messages = _consume_to_file(
+            working_dir, topic, group_id, from_context, numbers, avro, last
+        )
+
+        if number_consumed_messages == 0:
+            click.echo(
+                click.style(
+                    "Execution stopped, because no messages consumed.", fg="red"
+                )
+            )
+            return
+
+        click.echo(
+            "\nReady to produce to context "
+            + blue_bold(to_context)
+            + " and target topic "
+            + blue_bold(topic)
+        )
+
+        if not ensure_approval("Do you want to proceed?\n", no_verify=state.no_verify):
+            return
+
+        state.config.context_switch(to_context)
+        _produce_from_file(topic, to_context, working_dir, avro)
+
+
+def _produce_from_file(
+    topic: str, to_context: str, working_dir: pathlib.Path, avro: bool
+):
+    if avro:
+        producer = AvroFileProducer.create(working_dir)
+    else:
+        producer = FileProducer.create(working_dir)
+    click.echo(
+        "\nStart producing to topic "
+        + blue_bold(topic)
+        + " in target context "
+        + blue_bold(to_context)
+    )
+    number_produced_messages = producer.produce_from_file(topic)
+    click.echo(
+        green_bold(str(number_produced_messages))
+        + " messages successfully produced to context "
+        + green_bold(to_context)
+        + " and topic "
+        + green_bold(topic)
+        + "."
+    )
+
+
+def _consume_to_file(
+    working_dir: pathlib.Path,
+    topic: str,
+    group_id: str,
+    from_context: str,
+    numbers: int,
+    avro: bool,
+    last: bool,
+) -> int:
+    if avro:
+        consumer = AvroFileConsumer.create(group_id, topic, working_dir, last)
+    else:
+        consumer = FileConsumer.create(group_id, topic, working_dir, last)
+    click.echo(
+        "\nStart consuming from topic "
+        + blue_bold(topic)
+        + " in source context "
+        + blue_bold(from_context)
+    )
+    number_consumed_messages = consumer.consume_to_file(int(numbers))
+    click.echo(blue_bold(str(number_consumed_messages)) + " messages consumed.")
+
+    return number_consumed_messages
+
+
 @esque.command("ping", help="Tests the connection to the kafka cluster.")
 @click.option("-t", "--times", help="Number of pings.", default=10)
 @click.option("-w", "--wait", help="Seconds to wait between pings.", default=1)
@@ -293,13 +433,13 @@ def ping(state, times, wait):
         except TopicAlreadyExistsException:
             click.echo("Topic already exists.")
 
-        producer = Producer()
-        consumer = Consumer()
+        producer = PingProducer()
+        consumer = PingConsumer(PING_GROUP_ID, PING_TOPIC, True)
 
         click.echo(f"Ping with {state.cluster.bootstrap_servers}")
 
         for i in range(times):
-            producer.produce_ping()
+            producer.produce_ping(PING_TOPIC)
             _, delta = consumer.consume_ping()
             deltas.append(delta)
             click.echo(f"m_seq={i} time={delta:.2f}ms")

esque/cli/helpers.py

@@ -1,7 +1,23 @@
+import pathlib
+import shutil
+
 import click
 
 
 def ensure_approval(question: str, *, no_verify: bool = False) -> bool:
     if no_verify:
         return True
     return click.confirm(question)
+
+
+class DeleteOnFinished:
+    def __init__(self, dir_: pathlib.Path):
+        self._dir = dir_
+        self._dir.mkdir(parents=True)
+
+    def __enter__(self) -> pathlib.Path:
+        return self._dir
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self._dir.exists():
+            shutil.rmtree(self._dir)