Source code for operators.tile.recover_from_s3
import os
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 RecoverFromS3Operator(BaseTileOperator):
"""
Operator to recover tile(s) from S3 to the tileserver (keeping the files in S3).
:param src_env: destination environment, defaults to "main"
:type src_env: Literal["main", "staging"]
:param src_tiles: single path or list of paths where the tiles should be placed on the server
:type src_tiles: Union[str, List[str]]
:param dest_s3: path(s) to the S3 source files
: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 performing the action, defaults to True
:type verify_tile_exists: bool, optional
:raises AirflowException: if the recovery fails
"""
ui_color = "#4CAF50" # Green shade
ui_fgcolor = "#FFFFFF"
custom_operator_name = "RecoverFromS3Operator"
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:
"""
Recover tile(s) from S3 to the tileserver.
"""
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 or not self.dest_s3_list:
raise AirflowException(
"src_tiles and dest_s3 paths are required for recover_from_s3 operation"
)
if len(self.src_tiles) != len(self.dest_s3_list):
raise AirflowException(
f"Number of destination tiles ({len(self.src_tiles)}) doesn't match number of source S3 paths ({len(self.dest_s3_list)})"
)
self.log.info(
f"Starting RecoverFromS3Operator with {len(self.src_tiles)} tiles in environment {self.env_folder}"
)
# 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)
environment_alias = self.tileserver_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"Recovering tiles from S3: {tileset_names}")
if self.is_dry_run:
# Dry run path - just log what would happen
self.log.info(
f"[DRY RUN] Would recover {len(self.dest_s3_list)} tiles from S3 to {environment_alias}"
)
for idx, dest_file in enumerate(self.src_tiles):
s3_path = self.dest_s3_list[idx]
dest_dir = os.path.dirname(dest_file)
self.log.info(
f"[DRY RUN] Would download {s3_path} to local for recovery"
)
self.log.info(
f"[DRY RUN] Would create directory {dest_dir} if needed"
)
self.log.info(
f"[DRY RUN] Would copy file to {environment_alias}:{dest_file}"
)
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 recover tiles from s3: {', '.join(tileset_names)}"
)
else:
# Actual execution path
self.log.info(
f"Recovering {len(self.dest_s3_list)} tiles from S3 to {environment_alias}"
)
try:
# Process each file
for idx, dest_file in enumerate(self.src_tiles):
s3_path = self.dest_s3_list[idx]
tileset_name = self.get_tileset_name(dest_file)
dest_dir = os.path.dirname(dest_file)
self.log.info(
f"Recovering tile {tileset_name} from {s3_path} to {environment_alias}:{dest_file}"
)
# Create destination directory if needed
self.create_directory_on_server(environment_alias, dest_dir)
# Copy from S3 to server using temp file
with tempfile.NamedTemporaryFile(
suffix=".mbtiles"
) as temp_file:
temp_file_path = temp_file.name
# Download from S3 to temp file
self.s3_download(s3_path, temp_file_path)
# Copy to server
self.scp_to_server(
temp_file_path, environment_alias, dest_file
)
# Restart tileserver ONCE after all files are processed
self.restart_tileserver(environment_alias)
self.log.info(
f"Successfully recovered {len(self.dest_s3_list)} tiles from S3"
)
except Exception as e:
# Attempt to restart the tileserver in case of failure
self.log.error(f"Error during recover 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 recover tiles from S3: {e}")
# Add to datahub after server operations
for tileset_name in tileset_names:
try:
self.add_to_datahub(context, tileset_name)
except Exception as e:
self.log.error(
f"Error adding tileset {tileset_name} to 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"✓ Recovery from S3 successfully completed for {len(self.src_tiles)} tiles"
)
except Exception as ex:
self.log.error(f"Error in RecoverFromS3Operator: {ex}")
# Send notification to Slack for failures
error_message = f"""
⚠️ Error in RecoverFromS3Operator:
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