slack_connection.py

import os
import requests
from datetime import datetime
from dotenv import load_dotenv
from pathlib import Path

load_dotenv(
    Path.home() / ".environment_variables" / Path(__file__).parent.name / ".env"
)


slack_channel_ids = {
    'audience-requests': 'C06EUDJMEQ7',  # persona engine bot
    'media-optimization-bot': 'C07M53QQDLZ',
    'matthew-test': 'C07MG3KM848'  # testing, to be removed later
}

slack_developer_ids = {
    'channel': '<!channel>',

    'vlad': '<@D066VRQLVR9>',
    'hassan': '<@D067F9N9E11>',

    'matthew': '<@U0511RUAK9N>',
    'preet': '<@U03EFJE8GL9>',
    'elyes': '<@U068HJ3SMU3>',

    'udeet': '<@U03E02MJ3FF>',
    'yuxiang': '<@U068HQZP278>'
}


class SlackNotify:
    def __init__(self, channel_id, developer_id, endpoint_name, input_object):
        self.slack_access_token = os.getenv('slack_token')
        self.slack_test_mode = os.getenv('slack_test_mode')
        self.slack_url = 'https://slack.com/api/chat.postMessage'

        self.channel_id = channel_id
        self.developer_id = developer_id
        self.endpoint_name = endpoint_name
        self.input_object = input_object

        # set up the headers
        self.headers = {
            'Content-type': 'application/json',
            'Authorization': f'Bearer {self.slack_access_token}'
        }

    def slack_server_status_bot(self, notify_slack, status):
        # in case notify is turned off (slack_test_mode=1), please turn back on after testing, variable to be updated in .env file
        if self.slack_test_mode == '1':
            return

        # line 1 of the message
        current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        message = f"{self.developer_id} \n"

        if status == 'error':
            message += f"\n :red_circle: *`ERROR - {current_time}`* \n\n"
        elif status == 'success':
            message += f"\n :white_check_mark: *`SUCCESS - {current_time}`* \n\n"
        elif status == 'danger':
            message += f"\n :zap: *`WARNING - {current_time}`* \n"
        else:
            message += f"\n :information_source: *`INFO - {current_time}`* \n\n"

        # prepare the payload
        payload = {
            "channel": self.channel_id,
            "link_names": 1,
            "mrkdwn": 1,
            "text": f"{message}```{notify_slack}```"
        }

        try:
            # send the POST request
            response = requests.post(self.slack_url, json=payload, headers=self.headers)

            # check for errors in the response
            response.raise_for_status()
            print(response.json())

        except requests.exceptions.RequestException as e:
            # handle any request exceptions/errors
            print(f'Error occurred: {e}')

        return

    def run_success_function(self, info_string="Test passed successfully"):
        notify_slack = f"Endpoint: {self.endpoint_name} | Input: {self.input_object} | {info_string}"
        status = 'success'
        self.slack_server_status_bot(notify_slack, status)
        return

    def run_info_function(self, info_string):
        notify_slack = f"Endpoint: {self.endpoint_name} | Input: {self.input_object} | {info_string}"
        status = 'info'
        self.slack_server_status_bot(notify_slack, status)
        return

    # errors - any broken code, any missing data (required by frontend)
    def run_error_function(self, full_stacktrace):
        notify_slack = f"Endpoint: {self.endpoint_name} | Input: {self.input_object} \n{full_stacktrace}"
        status = 'error'
        self.slack_server_status_bot(notify_slack, status)
        return

    # warnings - any try catch block that is not an "error" as described above
    def run_warning_function(self, full_stacktrace):
        notify_slack = f"Endpoint: {self.endpoint_name} | Input: {self.input_object} \n{full_stacktrace}"
        status = 'danger'
        self.slack_server_status_bot(notify_slack, status)
        return


if __name__ == "__main__":
    # 1. import level - libraries and models
    ob1 = SlackNotify(channel_id=slack_channel_ids['matthew-test'], developer_id=slack_developer_ids['matthew'],
                      endpoint_name="ML: test_model", input_object="[Import Section]")
    #ob1.run_error_function(full_stacktrace="error data will be displayed here")

    # 2. main function level - the main function that controls the endpoint
    ob1 = SlackNotify(channel_id=slack_channel_ids['matthew-test'], developer_id=slack_developer_ids['matthew'],
                      endpoint_name="ML: test_model", input_object="{'inputId': 1}")
    #ob1.run_error_function(full_stacktrace="error data will be displayed here")

    # 3. primary function level - internal modules called within the main function
    #ob1.run_error_function(full_stacktrace="error data will be displayed here")

    # 4. secondary function level - internal modules called within primary functions
    #ob1.run_warning_function(full_stacktrace="warning data will be displayed here")
    #or
    #ob1.run_error_function(full_stacktrace="error data will be displayed here")