Source code for operators.slack_to_s3

from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.providers.slack.operators.slack import SlackHook
from config import Config
import json

DEFAULT_S3_BUCKET = Config.DATA_LAKE_BUCKET




class SlackToS3Operator(BaseOperator):
    """
    pull targetted slack history to s3
    for channels that ReadyAirflow is added to
    """

    ui_color = "#eb4034"

    @apply_defaults
    def __init__(
        self,
        slack_channel_id: str,
        slack_users: list[str],
        s3_key: str,
        slack_conn: str = Config.SLACK_DEFAULT_API,
        s3_bucket_name: str = DEFAULT_S3_BUCKET,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.slack_channel_id = slack_channel_id
        self.slack_users = slack_users
        self.s3_bucket_name = s3_bucket_name
        self.s3_key = s3_key
        self.slack_conn = slack_conn



    def execute(self, context):
        """
        Gets all threads where specified users are mentioned,
        saves each thread as a separate JSON file to S3.
        """
        import boto3

        slack_hook = SlackHook(slack_conn_id=Config.SLACK_DEFAULT_API)
        threads = {}
        cursor = None

        try:
            while True:
                params = {"channel": self.slack_channel_id, "limit": 1000}

                if cursor:
                    params["cursor"] = cursor

                response = slack_hook.call("conversations.history", json=params)

                if not response.get("ok"):
                    raise Exception(f"Slack API error: {response.get('error')}")

                # Look for messages that mention our users
                for msg in response.get("messages", []):
                    mentions = [
                        user
                        for user in self.slack_users
                        if f"<@{user}>" in msg.get("text", "")
                    ]

                    if mentions:
                        thread_ts = msg.get("thread_ts") or msg.get("ts")
                        if thread_ts and thread_ts not in threads:
                            # Print debug info
                            print(
                                f"Fetching thread {thread_ts} for channel {self.slack_channel_id}"
                            )

                            thread_params = {
                                "channel": self.slack_channel_id,
                                "ts": thread_ts,
                            }

                            thread_response = slack_hook.call(
                                "conversations.replies", params=thread_params
                            )

                            if thread_response.get("ok"):
                                thread_messages = thread_response.get("messages", [])
                                threads[thread_ts] = {
                                    "parent_message": (
                                        thread_messages[0] if thread_messages else msg
                                    ),
                                    "replies": (
                                        thread_messages[1:] if thread_messages else []
                                    ),
                                    "mentioned_users": mentions,
                                }
                            else:
                                print(
                                    f"Error fetching thread: {thread_response.get('error')}"
                                )

                cursor = response.get("response_metadata", {}).get("next_cursor")

                if not cursor:
                    break

            if not threads:
                print(
                    f"No threads found mentioning users {self.slack_users} in channel {self.slack_channel_id}"
                )
                return

            # Initialize S3 client
            s3_client = boto3.client("s3")

            # Format all threads into a list
            thread_list = [
                {
                    "thread_ts": thread_ts,
                    "parent": {
                        "text": thread_data["parent_message"].get("text", ""),
                        "sender": thread_data["parent_message"].get("user", ""),
                        "ts": thread_data["parent_message"].get("ts", ""),
                    },
                    "replies": [
                        {
                            "text": reply.get("text", ""),
                            "sender": reply.get("user", ""),
                            "ts": reply.get("ts", ""),
                        }
                        for reply in thread_data["replies"]
                    ],
                }
                for thread_ts, thread_data in threads.items()
            ]

            # Upload single JSON file with all threads
            s3_client.put_object(
                Bucket=self.s3_bucket_name,
                Key=f"{self.s3_key.rstrip('/')}/threads.json",
                Body=json.dumps(thread_list, indent=2),
                ContentType="application/json",
            )

            print(f"Successfully uploaded {len(thread_list)} threads to S3")

        except Exception as e:
            raise Exception(f"Error transferring threads to S3: {str(e)}")