from typing import Literal, List, Union
import tempfile
from airflow.utils.decorators import apply_defaults
from airflow.utils.context import Context
from airflow.exceptions import AirflowException, AirflowSkipException
from operators.tile.base_tile import BaseTileOperator
from config import Config
[docs]
class SyncTileOperator(BaseTileOperator):
"""
Operator to sync tileset(s) between tileserver environments.
This operator syncs tiles from the source environment to one or more target
environments; in practice, only work in production mode, to sync tiles between prod and staging.
- prod (kanto, Config.TILESERVER_ALIAS)
- staging (hoenn, Config.TILESERVER_ALIAS_STAGE)
- dev (johto, Config.TILESERVER_ALIAS_DEV)
In non-production mode, is_dry_run is always set to True for logging.
:param src_env: source environment, "prod", "staging", or "dev"
:type src_env: Literal["prod", "staging", "dev"]
:param src_tiles: single path or list of paths to the tilesets (starting at ~/tiles/)
:type src_tiles: Union[str, List[str]]
:param dest_envs: target environments to sync to, any combination of "prod", "staging", and "dev"
:type dest_envs: List[Literal["prod", "staging", "dev"]]
: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 sync operation fails
:raises AirflowSkipException: if not in production mode and not in dry run mode
"""
ui_color = "#7e57c2" # Purple shade
ui_fgcolor = "#FFFFFF"
custom_operator_name = "SyncTileOperator"
template_fields = ["src_tiles", "src_env", "dest_envs", "is_dry_run"]
@apply_defaults
def __init__(
self,
src_tiles: Union[str, List[str]],
src_env: Literal["prod", "staging", "dev"] = "prod",
dest_envs: List[Literal["prod", "staging", "dev"]] = None,
is_dry_run: bool = True,
verify_tile_exists: bool = True,
*args,
**kwargs,
):
# Initialize the parent class with values that will be properly rendered later
super().__init__(
src_env=src_env,
src_tiles=src_tiles, # Will be processed in render_template_fields
is_dry_run=is_dry_run, # Will be processed in render_template_fields
verify_tile_exists=verify_tile_exists,
*args,
**kwargs,
)
self.dest_envs = dest_envs or []
[docs]
def render_template_fields(self, context, jinja_env=None):
"""Override to handle special template rendering cases"""
# Log entry point and context keys
self.log.info(
f"🔍 Entering render_template_fields with context keys: {list(context.keys())}"
)
# Log the parameters and their types before rendering
self.log.info("📊 Parameters before rendering:")
self.log.info(
f" src_tiles: {self.src_tiles} (type: {type(self.src_tiles).__name__})"
)
self.log.info(
f" src_env: {self.src_env} (type: {type(self.src_env).__name__})"
)
self.log.info(
f" dest_envs: {self.dest_envs} (type: {type(self.dest_envs).__name__})"
)
self.log.info(
f" is_dry_run: {self.is_dry_run} (type: {type(self.is_dry_run).__name__})"
)
# First let the parent class (BaseTileOperator/BaseOperator/Airflow's standard templating engine) handle normal rendering
self.log.info("⚙️ Applying standard template rendering via parent class")
super().render_template_fields(context, jinja_env)
# Log the parameters after standard rendering
self.log.info("📊 Parameters after standard rendering:")
self.log.info(
f" src_tiles: {self.src_tiles} (type: {type(self.src_tiles).__name__})"
)
self.log.info(
f" src_env: {self.src_env} (type: {type(self.src_env).__name__})"
)
self.log.info(
f" dest_envs: {self.dest_envs} (type: {type(self.dest_envs).__name__})"
)
self.log.info(
f" is_dry_run: {self.is_dry_run} (type: {type(self.is_dry_run).__name__})"
)
# In non-production mode, always use dry run
if not Config.PRODUCTION:
old_value = self.is_dry_run
self.is_dry_run = True
self.log.info(
f"🔄 Enforcing dry_run=True in non-production mode (was: {old_value})"
)
# Process src_tiles - handle nested lists from XCom
if (
isinstance(self.src_tiles, list)
and self.src_tiles
and isinstance(self.src_tiles[0], list)
):
old_value = self.src_tiles
self.src_tiles = self.src_tiles[0]
self.log.info(
f"🔄 Unnesting src_tiles from {old_value} to {self.src_tiles}"
)
# Process dest_envs - handle nested lists
if (
isinstance(self.dest_envs, list)
and self.dest_envs
and isinstance(self.dest_envs[0], list)
):
old_value = self.dest_envs
self.dest_envs = self.dest_envs[0]
self.log.info(
f"🔄 Unnesting dest_envs from {old_value} to {self.dest_envs}"
)
# Log final parameter states
self.log.info("📊 Final parameter values after processing:")
self.log.info(
f" src_tiles: {self.src_tiles} (type: {type(self.src_tiles).__name__})"
)
self.log.info(
f" dest_envs: {self.dest_envs} (type: {type(self.dest_envs).__name__})"
)
self.log.info(
f" is_dry_run: {self.is_dry_run} (type: {type(self.is_dry_run).__name__})"
)
[docs]
def execute(self, context: Context) -> None:
"""
Sync tileset(s) from source environment to destination environments.
In production mode, performs actual syncing.
In non-production mode, only logs what would happen if in dry run mode.
"""
# Log the context parameters to see what's available
self.log.info(f"Context params: {context.get('params', {})}")
# Validate self.dest_envs
if not self.dest_envs or len(self.dest_envs) == 0:
raise ValueError("At least one destination environment is required")
for env in self.dest_envs:
if env not in ["prod", "staging", "dev"]:
raise ValueError(f"Invalid destination environment: {env}")
if Config.PRODUCTION and env == "dev":
raise ValueError(
f"Sync from staging/prod env to {env} env is non-compliant/not allowed!"
)
# Ensure we're not syncing to the same environment
self.dest_envs = [env for env in self.dest_envs if env != self.src_env]
if not self.dest_envs:
raise ValueError(
"Destination environments must be different from source environment"
)
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 sync operation")
# Get source server alias using BaseTileOperator's tileserver_alias method
source_server = self.tileserver_alias(self.src_env)
self.log.info(
f"Starting SyncTileOperator with {len(self.src_tiles)} {self.src_tiles} tiles from "
f"{self.src_env} ({source_server}) to {', '.join(self.dest_envs)}"
)
if not Config.PRODUCTION:
self.log.info(
f"Note: Currently in development mode - only showing what would happen "
f"if executed in production."
)
# Verify tiles exist before proceeding, but only in production mode
if not self.is_dry_run and Config.PRODUCTION and self.verify_tile_exists:
for src_tile in self.src_tiles:
self.verify_mbtiles_file(src_tile, source_server)
# Get list of tileset names for logging
tileset_names = [
self._get_tileset_name_from_path(path) for path in self.src_tiles
]
if self.is_dry_run:
# Dry run path - just log what would happen
self.log.info(
f"[DRY RUN] Would sync {len(self.src_tiles)} tiles from {self.src_env} ({source_server}) "
f"to {len(self.dest_envs)} destination envs={', '.join(self.dest_envs)}"
)
for idx, src_tile in enumerate(self.src_tiles):
tile_name = tileset_names[idx]
for dest_env in self.dest_envs:
dest_server = self.tileserver_alias(dest_env)
self.log.info(
f"[DRY RUN] Would copy tile '{tile_name}' from {source_server}:{src_tile} "
f"to {dest_server}:{Config.DEFAULT_TS_FLAT_FOLDER}/{tile_name}.mbtiles"
)
self.log.info(
f"[DRY RUN] Would update tileserver config and restart tileserver on {dest_server}"
)
self.log.info(
f"[DRY RUN] Would add tile '{tile_name}' to datahub for environment {dest_env}"
)
else:
# Only proceed with actual execution in production mode
if not Config.PRODUCTION:
self.log.info(
f"Non-production environment detected. Skipping actual file operations."
)
return
# Actual execution path
self.log.info(
f"Syncing {len(self.src_tiles)} tiles from {self.src_env} ({source_server}) "
f"to {len(self.dest_envs)} destination environments"
)
# Process each tile
for idx, src_tile in enumerate(self.src_tiles):
tileset_name = tileset_names[idx]
self.log.info(f"Processing tile: {tileset_name}")
# Create a temporary file to store the tile
with tempfile.NamedTemporaryFile(suffix=".mbtiles") as temp_file:
temp_file_path = temp_file.name
# Download tile from source server to local temp file
self.log.info(
f"Downloading tile from {source_server}:{src_tile} to local temp file"
)
self.scp_from_server(source_server, src_tile, temp_file_path)
# Upload to each destination server
for dest_env in self.dest_envs:
dest_server = self.tileserver_alias(dest_env)
dest_path = f"{Config.DEFAULT_TS_FLAT_FOLDER}/{tileset_name}.mbtiles"
self.log.info(
f"Uploading tile to {dest_server}:{dest_path}"
)
try:
# Create destination directory if needed
self.create_directory_on_server(
dest_server, Config.DEFAULT_TS_FLAT_FOLDER
)
# Upload tile to destination server
self.scp_to_server(
temp_file_path, dest_server, dest_path
)
# Update config and restart tileserver
self.log.info(f"Restarting tileserver on {dest_server}")
self.restart_tileserver(dest_server)
# Add to datahub - temporarily switch to the destination environment
self.log.info(
f"Adding tile {tileset_name} to datahub for environment {dest_env}"
)
# Save current src_env
original_src_env = self.src_env
# Set src_env to destination environment for datahub operations
self.src_env = dest_env
try:
self.add_to_datahub(context, tileset_name)
finally:
# Restore original src_env
self.src_env = original_src_env
except Exception as e:
self.log.error(
f"Error syncing tile '{tileset_name}' to {dest_server}: {e}"
)
try:
# Try to restart the tileserver in case of failure
self.restart_tileserver(dest_server)
except Exception as restart_error:
self.log.error(
f"Failed to restart tileserver after error: {restart_error}"
)
raise AirflowException(
f"Failed to sync tile '{tileset_name}' to {dest_server}: {e}"
)
self.log.info(
f"Successfully synced tile '{tileset_name}' to all destination environments"
)
# Send notification to Slack for success
success_message = f"""
✅ Tile Sync Successful:
From: {self.src_env} ({source_server})
To: {', '.join([f"{env} ({self.tileserver_alias(env)})" for env in self.dest_envs])}
Tiles: {', '.join(tileset_names)}
"""
self.notify_slack(success_message)
# 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"✓ Sync successfully completed for {len(self.src_tiles)} tiles"
)
except AirflowSkipException:
# Re-raise skip exception
raise
except Exception as ex:
self.log.error(f"Error in SyncTileOperator: {ex}")
# Send notification to Slack for failures
error_message = f"""
⚠️ Error in SyncTileOperator: {'in PRODUCTION' if Config.PRODUCTION else 'in Dev'}
From: {self.src_env} ({self.tileserver_alias(self.src_env)})
To: {', '.join([f"{env} ({self.tileserver_alias(env)})" for env in self.dest_envs])}
Tiles Count: {len(self.src_tiles)}: {self.src_tiles[:5]}
{'[DRY RUN] - No actual changes were made' if self.is_dry_run else ''}
Error: {str(ex)}
"""
self.notify_slack(error_message)
raise ex