atlan-collect / submit / consumer.py
consumer.py
Raw
import os, sys
import pika
import json

from pymongo import MongoClient
import pymongo
from send import emailSender

client = MongoClient("host.minikube.internal", 27017)
db = client.submit_form
collection = db.forms

MAX_RETRIES=3

def callback(ch, method, properties, body, retries=1):
    try:
        data = json.loads(body.decode('utf-8'))
        
        
        result = collection.insert_one(data)
        print(f"Form data submitted to MongoDB with ID: {result.inserted_id}")
        ch.basic_ack(delivery_tag=method.delivery_tag)
        
        print(f"Data upload successfultranslation: {result}")
        emailSender.notification(f"Your Data was uploaded successfully with the id: {result.inserted_id}.", subject="Data Submission Successful", reciever_address=data.get('email'))

    except pymongo.errors.WriteError as err:
        print(f"Submission Failed: {err}")
        emailSender.notification(f" failed: {err} (Retry No: {retries})", subject="Data Submission Failed")

        if retries <= MAX_RETRIES:
            callback(ch, method, properties, body, retries=retries+1)
        else:
            print("Max retries exceeded. Message discarded.")
            emailSender.notification(f"Data Submission has failed! Please try again later!", subject="Data Submission Failed")
            ch.basic_ack(delivery_tag=method.delivery_tag)
            return

    except Exception as e:
        # Handle other exceptions
        print(f"Error processing message: {e}")
        # Reject (NACK) the message to requeue it for further processing
        ch.basic_nack(delivery_tag=method.delivery_tag)


def main():

    connection = pika.BlockingConnection(
      pika.ConnectionParameters("rabbitmq")
    )
    channel = connection.channel()

    channel.basic_consume(
        queue=os.environ.get("SUBMIT_QUEUE"), 
        on_message_callback=callback,
    )

    print('Waiting for messages. To exit press CTRL+C')
    try:
        channel.start_consuming()
    except KeyboardInterrupt:
        print("Interrupted")
        channel.stop_consuming()
        connection.close()

if __name__=="__main__":
  try:
    main()
  except KeyboardInterrupt:
    print("Interrupted")
    try:
      sys.exit(0)
    except SystemExit:
      os._exit(0)