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

from data_uploader import data_uploader
from send import emailSender

gc = pygsheets.authorize(service_file='credentials.json')
sh = gc.open_by_url("https://docs.google.com/spreadsheets/d/16S5OUFYwRu7hkYDDMadkRvKCuIQ1-42Ir08VfrlAIKo/edit?usp=sharing") 

print(sh.id)
print(sh.worksheets())

worksheet = sh.sheet1

MAX_RETRIES = 3

def process_message(ch, method, properties, body, retries=1):
    if retries > MAX_RETRIES:
        print("Max retries exceeded. Message discarded.")
        emailSender.notification(f"Data Submission has failed", reciever=json.loads(body)["email"])
        ch.basic_ack(delivery_tag=method.delivery_tag)
        return
    
    result, err = data_uploader.insert_data(worksheet, body)
    if err:
        print(f"Data upload Failed: {err}")
        emailSender.notification(f"Spreadsheet Data Upload failed: {err} (Retry: {retries})", reciever=json.loads(body)["email"])
        process_message(ch, method, properties, body, retries=retries+1)
    else:
      print(f"Data upload Failed: {err}")
      emailSender.notification(f"{result}", reciever=json.loads(body)["email"])
      ch.basic_ack(delivery_tag=method.delivery_tag)

def main():
  connection = pika.BlockingConnection(
    pika.ConnectionParameters(host="rabbitmq")
  )

  channel = connection.channel()
  
  channel.basic_consume(
    queue=os.environ.get("SPREADSHEET_QUEUE"),
    on_message_callback = process_message,
  )

  print("Waiting for messages. To exit press CTRL+C")

  channel.start_consuming()

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