from typing import Any

from airflow.models import BaseOperator
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
from airflow.utils.context import Context

from common.config.app import S3_ARTIFACTS_BUCKET
from common.macros.generic import get_run_key
from common.utils.s3 import S3


class CleanupOperator(BaseOperator):

    def __init__(self, aws_conn_id: str = S3Hook.default_conn_name, **kwargs):
        super().__init__(**kwargs)
        self.aws_conn_id = aws_conn_id

    def execute(self, context: Context) -> Any:
        s3_hook = S3Hook(aws_conn_id=self.aws_conn_id).get_conn()
        s3 = S3(s3_hook, run_id=get_run_key(context["dag_run"]), bucket_name=S3_ARTIFACTS_BUCKET)

        s3.wipe_by_prefix("/")