receive_answers.py

import subprocess
import os
import time
import ujson
import logging
import subprocess
import argparse

from kafka import KafkaConsumer
from pathlib import Path
from datetime import datetime


import subprocess

parser=argparse.ArgumentParser()
parser.add_argument("--query_name")
parser.add_argument("--save_path")
args=parser.parse_args()

query_name = args.query_name
save_path = args.save_path


consumer = KafkaConsumer(
        query_name, 
        bootstrap_servers="localhost:9092",  # Kafka broker address
        auto_offset_reset="earliest",  # Start reading from the beginning
        group_id=None,  # Consumer group ID
        enable_auto_commit=True  # Auto commit offsets
    )

# Start polling messages
idle_timeout = 60
last_message_time = time.time()
received_messages = 0

out = open(f"{save_path}/{query_name}.txt", "w")

while True:
    msg_pack = consumer.poll(timeout_ms=1000)
    if msg_pack:
        for _, messages in msg_pack.items():
            for message in messages:
                #print(message.value.decode('utf-8'))
                out.write(ujson.dumps(ujson.loads(message.value.decode('utf-8'))) + "\n")
                out.flush()

                received_messages += 1
                last_message_time = time.time()
                if received_messages == 1:
                    print(f"Received first message.", end="\r")
    else:
        if time.time() - last_message_time > idle_timeout:
            print(f"Received {received_messages} messages.")
            print(f"Last message received {idle_timeout/60} minutes ago. Exiting.")
            break