Source code for operators.tile.archive_to_s3
import tempfile
from typing import Literal, Optional, 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
[docs]
class ArchiveToS3Operator(BaseTileOperator):
"""
Operator to move tile(s) from the tileserver to S3 and remove them from the server.
:param src_env: source environment, defaults to "main"
:type src_env: Literal["main", "staging"]
:param src_tiles: single path or list of paths to the tilesets (starting at home directory `~/`)
:type src_tiles: Union[str, List[str]]
:param dest_s3: Optional destination S3 path(s). If not provided, standardized paths will be used.
:type dest_s3: Optional[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 before performing the action, defaults to True
:type verify_tile_exists: bool, optional
:raises AirflowException: if the move operation fails
"""
ui_color = "#007acc" # Blue shade
ui_fgcolor = "#FFFFFF"
custom_operator_name = "ArchiveToS3Operator"
template_fields = ["dest_s3", "src_tiles", "src_env", "is_dry_run"]
@apply_defaults
def __init__(
self,
src_tiles: Union[str, List[str]],
src_env: Literal["main", "staging"] = "main",
dest_s3: Optional[str] = None,
is_dry_run: bool = True,
verify_tile_exists: bool = True,
*args,
**kwargs,
):
super().__init__(
src_env=src_env,
src_tiles=src_tiles,
dest_s3=dest_s3,
is_dry_run=is_dry_run,
verify_tile_exists=verify_tile_exists,
*args,
**kwargs,
)
[docs]
def execute(self, context: Context) -> None:
"""
Move tile(s) from the tileserver to S3 and remove them from the server.
"""
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.src_tiles:
raise AirflowException(
"src_tiles is required for archive_to_s3 operation"
)
if len(self.src_tiles) != len(self.dest_s3_list):
raise AirflowException(
f"Number of source tiles ({len(self.src_tiles)}) doesn't match number of destination paths ({len(self.dest_s3_list)})"
)
self.log.info(
f"Starting ArchiveToS3Operator with {len(self.src_tiles)} tiles in environment {self.env_folder}"
)
# Verify tiles exist before proceeding
environment_alias = self.tileserver_alias()
if not self.is_dry_run and self.verify_tile_exists:
for src_tile in self.src_tiles:
self.verify_mbtiles_file(src_tile, environment_alias)
# List of tileset names for logging
tileset_names = [self.get_tileset_name(path) for path in self.dest_s3_list]
self.log.info(f"Moving tiles to S3: {tileset_names}")
if self.is_dry_run:
# Dry run path - just log what would happen
self.log.info(
f"[DRY RUN] Would move {len(self.src_tiles)} tiles from {environment_alias} to S3"
)
for idx, src_tile in enumerate(self.src_tiles):
self.log.info(
f"[DRY RUN] Would copy {src_tile} to {self.dest_s3_list[idx]}"
)
self.log.info(
f"[DRY RUN] Would remove {src_tile} from {environment_alias}"
)
self.log.info(f"[DRY RUN] Would stop tileserver on {environment_alias}")
self.log.info(
f"[DRY RUN] Would update tileserver config, then restart tileserver {environment_alias}"
)
self.log.info(
f"[DRY RUN] Would move tiles to s3: {', '.join(tileset_names)}"
)
else:
# Actual execution path
self.log.info(
f"Moving {len(self.src_tiles)} tiles from {environment_alias} to S3"
)
try:
# Process each file
for idx, src_tile in enumerate(self.src_tiles):
dest_s3 = self.dest_s3_list[idx]
tileset_name = self.get_tileset_name(src_tile)
self.log.info(
f"Moving tile {tileset_name} from {environment_alias}:{src_tile} to {dest_s3}"
)
# Copy to S3 using temp file
with tempfile.NamedTemporaryFile(
suffix=".mbtiles"
) as temp_file:
temp_file_path = temp_file.name
# Copy file from server to local temp
self.scp_from_server(
environment_alias, src_tile, temp_file_path
)
# Upload to S3
self.s3_upload(temp_file_path, dest_s3)
# Remove from server
self.execute_ssh_command(
environment_alias,
f"rm -f {src_tile}",
f"remove file from server: {src_tile}",
)
# Restart tileserver ONCE after all files are processed
self.restart_tileserver(environment_alias)
self.log.info(
f"Successfully moved {len(self.src_tiles)} tiles to S3"
)
except Exception as e:
# Attempt to restart the tileserver in case of failure
self.log.error(f"Error during move operation: {e}")
try:
self.restart_tileserver(environment_alias)
except Exception as restart_error:
self.log.error(
f"Failed to restart tileserver after error: {restart_error}"
)
raise AirflowException(f"Failed to move tiles to S3: {e}")
# Remove from datahub after server operations
for tileset_name in tileset_names:
try:
self.remove_from_datahub(tileset_name)
except Exception as e:
self.log.error(
f"Error removing tileset {tileset_name} from datahub: {e}"
)
# 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"✓ Move to S3 successfully completed for {len(self.src_tiles)} tiles"
)
except Exception as ex:
self.log.error(f"Error in ArchiveToS3Operator: {ex}")
# Send notification to Slack for failures
error_message = f"""
⚠️ Error in ArchiveToS3Operator:
Environment: {self.env_folder}
Tiles Count: {len(self.src_tiles)}
{'[DRY RUN] - No actual changes were made' if self.is_dry_run else ''}
Error: {str(ex)}
"""
self.notify_slack(error_message)
raise ex