data/dataAccessor.py

import requests
from data.task import Task
from util.slack import Slack

comic_url = "http://internal-k8s-gamma-internal-ea8e32da94-1997933257.ap-south-1.elb.amazonaws.com:80"


def updateSource(sourceId, userId, state):
    print("update source is called")
    url = comic_url + f"/comic-crecoai/source/{sourceId}"
    headers = {"Content-Type": "application/json", "user-id": str(userId)}

    data = {"state": state}

    try:
        response = requests.patch(url, headers=headers, json=data, timeout=10)
        print("update source response", response)
    except requests.exceptions.Timeout:
        print("Request timed out while updating source")
    except requests.exceptions.RequestException as e:
        print(f"Error while updating source: {e}")

    return


def saveGeneratedImages(sourceId, userId):
    print("save generation called")
    url = comic_url + "/comic-crecoai/source/" + str(sourceId) + "/generatedImages"
    headers = {"Content-Type": "application/json", "user-id": str(userId)}
    data = {"state": "ACTIVE"}

    try:
        requests.patch(url, headers=headers, json=data)
        # print("save generation response", response)
    except requests.exceptions.Timeout:
        print("Request timed out while saving image")
    except requests.exceptions.RequestException as e:
        print("Failed to mark source as active: ", e)
        return
    return


def update_db(func):
    def caller(*args, **kwargs):
        if type(args[0]) is not Task:
            raise Exception("First argument must be a Task object")
        task = args[0]
        try:
            updateSource(task.get_sourceId(), task.get_userId(), "INPROGRESS")
            rargs = func(*args, **kwargs)
            updateSource(task.get_sourceId(), task.get_userId(), "COMPLETED")
            saveGeneratedImages(task.get_sourceId(), task.get_userId())
            return rargs
        except Exception as e:
            print("Error processing image: {}".format(str(e)))
            slack = Slack()
            slack.error_alert(task, e)
            updateSource(task.get_sourceId(), task.get_userId(), "FAILED")

    return caller