slack.py

import json
import urllib.parse
import requests
import pyarrow.parquet as pq
from io import BytesIO
import re
from datetime import datetime

s3 = boto3.client("s3")

dataset_links = {
    "table1": "nhammad.github.io",
    "table2": "nhammad.github.io"
}

color_codes = {
    "fail": ["#d11919", "rotating_light"],
    "pass": ["#57e334", "white_check_mark"]
}

def send_slack_notification(file_name, status, text):
    slack_url = "https://hooks.slack.com/services/xxxx"
    author_link = dataset_links[file_name]
    color_code = color_codes[status][0]
    emoji = color_codes[status][1]
                 
    payload = {
        "channel": "#channel-name",
        "username": "alerts-bot",

        "attachments": 
        [{
            "mrkdwn_in": ["text"],
            "color": color_code,
            "title": f":{emoji}: {file_name}",
            "title_link": author_link,
            "text": text,
            "fields": 
            [
                {
                    "title": "Timestamp",
                    "value": datetime.utcnow().strftime('%a, %d %b %Y %H:%M:%S UTC'),
                    "short": True
                }
            ],
        }]
    }

    response = requests.post(slack_url, data={"payload": json.dumps(payload)})

    if response.status_code == 200:
        print("Slack message sent successfully")
    else:
        print(
            f"Failed to send Slack message. Status code: {response.status_code}, Response content: {response.text}"  
        )


def read_s3_file(bucket, object_key):
    response = s3.get_object(Bucket=bucket, Key=object_key)
    parquet_data = response["Body"].read()

    # Create a buffer for reading the Parquet data
    buffer = BytesIO(parquet_data)

    # Read the Parquet data into a PyArrow table
    parquet_table = pq.read_table(buffer)

    num_rows = len(parquet_table)

    # obtaining the actual file name from the full object key
    pattern = r'failure_notifications/(.*?)/'
    file_name = (re.search(pattern, object_key)).group(1)

    slack_message = f"Threshold exceeded by: {num_rows} rows."
    
    send_slack_notification(file_name, "fail", slack_message)

    return num_rows


def lambda_handler(event, context):
    s3_event = event["Records"][0]["s3"]
    bucket = s3_event["bucket"]["name"]
    object_key = urllib.parse.unquote_plus(s3_event["object"]["key"])

    print(f"Bucket: {bucket}, Object Key: {object_key}")

    read_s3_file(bucket, object_key)

    return {
        "statusCode": 200,
        "body": json.dumps("Lambda function executed successfully."),
    }