Source code for operators.tile.remove_in_s3

from typing import List, Union
from airflow.utils.decorators import apply_defaults
from airflow.utils.context import Context
from airflow.exceptions import AirflowException
from operators.tile.base_tile import BaseTileOperator




class RemoveInS3Operator(BaseTileOperator):
    """
    Operator to delete tile(s) from S3 without affecting the tileserver.

    :param dest_s3: path(s) to the S3 files to delete
    :type dest_s3: Union[str, List[str]]
    :param is_dry_run: whether to perform a dry run, defaults to True
    :type is_dry_run: bool, optional
    :param verify_tile_exists: whether to verify if the tile exists in S3 before deleting, defaults to True
    :type verify_tile_exists: bool, optional

    :raises AirflowException: if the deletion fails
    """

    ui_color = "#FFA500"  # Orange shade
    ui_fgcolor = "#000000"
    custom_operator_name = "RemoveInS3Operator"
    template_fields = ["dest_s3", "is_dry_run"]

    @apply_defaults
    def __init__(
        self,
        dest_s3: Union[str, List[str]],
        is_dry_run: bool = True,
        verify_tile_exists: bool = True,
        *args,
        **kwargs,
    ):
        super().__init__(
            src_env=None,  # Not needed for S3-only operations
            src_tiles=None,  # Not needed for S3-only operations
            dest_s3=dest_s3 if isinstance(dest_s3, str) else None,
            is_dry_run=is_dry_run,
            verify_tile_exists=verify_tile_exists,
            *args,
            **kwargs,
        )

        # Handle list of S3 paths directly
        if isinstance(dest_s3, list):
            self.dest_s3_list = [self._ensure_s3_path(path) for path in dest_s3]



    def execute(self, context: Context) -> None:
        """
        Delete tile(s) from S3.
        """
        try:
            # Add clear indication of dry run mode at the beginning
            if self.is_dry_run:
                self.log.info(
                    f"╔═════════════════════════════════════════════════════════════════╗"
                )
                self.log.info(
                    f"║ 🔍 DRY RUN MODE: No actual changes will be made to files or systems ║"
                )
                self.log.info(
                    f"╚═════════════════════════════════════════════════════════════════╝"
                )

            if not self.dest_s3_list:
                raise AirflowException("dest_s3 is required for remove_in_s3 operation")

            self.log.info(
                f"Starting RemoveInS3Operator with {len(self.dest_s3_list)} S3 files"
            )

            # Verify S3 files exist before proceeding
            if not self.is_dry_run and self.verify_tile_exists:
                for s3_path in self.dest_s3_list:
                    self.verify_s3_file(s3_path)

            # List of tileset names for logging
            tileset_names = [self.get_tileset_name(path) for path in self.dest_s3_list]
            self.log.info(f"Removing tiles from S3: {tileset_names}")

            if self.is_dry_run:
                # Dry run path - just log what would happen
                for s3_path in self.dest_s3_list:
                    tileset_name = self.get_tileset_name(s3_path)
                    self.log.info(
                        f"[DRY RUN] Would delete tile {tileset_name} from S3: {s3_path}"
                    )
                self.log.info(
                    f"[DRY RUN] Would delete {len(self.dest_s3_list)} tiles from S3"
                )
            else:
                # Actual execution path
                self.log.info(f"Deleting {len(self.dest_s3_list)} tiles from S3")

                errors = []
                for s3_path in self.dest_s3_list:
                    tileset_name = self.get_tileset_name(s3_path)
                    try:
                        self.log.info(
                            f"Deleting tile {tileset_name} from S3: {s3_path}"
                        )
                        self.s3_delete(s3_path)
                        self.log.info(
                            f"Successfully deleted tile {tileset_name} from S3"
                        )
                    except Exception as e:
                        errors.append(f"Failed to delete {s3_path}: {str(e)}")
                        self.log.error(f"Error deleting {s3_path}: {e}")

                if errors:
                    raise AirflowException(
                        f"Failed to delete some files from S3: {'; '.join(errors)}"
                    )

                self.log.info(
                    f"Successfully deleted {len(self.dest_s3_list)} tiles from S3"
                )

            # Add clear indication of completion
            if self.is_dry_run:
                self.log.info(
                    f"╔════════════════════════════════════════════════════════════════╗"
                )
                self.log.info(
                    f"║ ✓ DRY RUN COMPLETED: No actual changes were made               ║"
                )
                self.log.info(
                    f"╚════════════════════════════════════════════════════════════════╝"
                )
            else:
                self.log.info(
                    f"✓ S3 deletion successfully completed for {len(self.dest_s3_list)} files"
                )

        except Exception as ex:
            self.log.error(f"Error in RemoveInS3Operator: {ex}")

            # Send notification to Slack for failures
            error_message = f"""
            ⚠️ Error in RemoveInS3Operator:
            Files Count: {len(self.dest_s3_list)}
            {'[DRY RUN] - No actual changes were made' if self.is_dry_run else ''}
            Error: {str(ex)}
            """
            self.notify_slack(error_message)

            raise ex