import os
import shutil
from typing import Literal, Optional, List, Union
from config import Config
from util import Util
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.utils.context import Context
from airflow.exceptions import AirflowException
from airflow.hooks.subprocess import SubprocessHook
from datahub_airflow_plugin.entities import Urn
from airflow.providers.slack.hooks.slack import SlackHook
[docs]
class BaseTileOperator(BaseOperator):
"""
Base operator for tile operations that provides common utilities for
working with tilesets on servers and S3.
"""
@apply_defaults
def __init__(
self,
src_tiles: Optional[Union[str, List[str]]] = None,
src_env: Optional[Literal["main", "staging", "prod", "dev"]] = "main",
dest_s3: Optional[str] = None,
is_dry_run: bool = True,
verify_tile_exists: bool = True,
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.src_env = src_env
self.is_dry_run = is_dry_run
self.verify_tile_exists = verify_tile_exists
self.subprocess_hook = SubprocessHook()
# Convert single src_tile to list if provided
if isinstance(src_tiles, str):
self.src_tiles = [src_tiles]
else:
self.src_tiles = src_tiles
self.env_folder = self._get_env_folder()
# Generate standardized destination S3 paths
if dest_s3 and len(self.src_tiles or []) > 1:
self.log.warning(
"dest_s3 provided with multiple src_tiles will be ignored. Using standardized paths."
)
dest_s3 = None
self.dest_s3_list = []
if dest_s3:
self.dest_s3 = self._ensure_s3_path(dest_s3)
self.dest_s3_list = [self.dest_s3]
else:
self.dest_s3 = None
if self.src_tiles:
for src_tile in self.src_tiles:
tileset_name = self._get_tileset_name_from_path(src_tile)
s3_path = f"s3://{Config.ARCHIVE_TILES_S3_BUCKET}/{self.env_folder}/{tileset_name}.mbtiles"
self.dest_s3_list.append(s3_path)
# =========================================================================
# UTILITY METHODS
# =========================================================================
def _get_env_folder(self) -> str:
"""
Get the S3 folder name (tileserver env) based on the src_env and Config.PRODUCTION
:return: S3 folder name (dev, stage, prod)
"""
if self.src_env == "staging":
return "staging"
elif Config.PRODUCTION:
return "prod"
else:
return "dev"
def _ensure_s3_path(self, path: str) -> str:
"""
Ensure the path has the s3:// prefix.
:param path: S3 path
:return: Path with s3:// prefix
"""
if not path.startswith("s3://"):
return f"s3://{path}"
return path
def _get_tileset_name_from_path(self, path: str) -> str:
"""
Extract the tileset name from a file path.
:param path: File path (can be local path or S3 path)
:return: Tileset name without extension
"""
if path.startswith("s3://"):
# Extract from S3 path
s3_path = path.replace("s3://", "")
filename = s3_path.split("/")[-1]
else:
# Extract from local path
filename = os.path.basename(path)
# Remove extension if it exists
return os.path.splitext(filename)[0]
[docs]
def get_tileset_name(self, path: Optional[str] = None) -> str:
"""
Extract the tileset name from the path or dest_s3.
:param path: Optional path to extract name from, defaults to first src_tile/dest_s3
:return: Tileset name without extension
"""
if path:
return self._get_tileset_name_from_path(path)
elif self.src_tiles and len(self.src_tiles) > 0:
return self._get_tileset_name_from_path(self.src_tiles[0])
else:
return self._get_tileset_name_from_path(self.dest_s3_list[0])
[docs]
def tileserver_alias(self, env: Optional[str] = None) -> str:
"""
Get the appropriate tileserver alias based on the environment.
In production:
- "prod"/"main" -> Config.TILESERVER_ALIAS (kanto)
- "staging" -> Config.TILESERVER_ALIAS_STAGE (hoenn)
- "dev" -> Config.TILESERVER_ALIAS_DEV (johto)
In development:
- All environments default to Config.TILESERVER_ALIAS (johto)
:param env: Optional environment override ("prod", "staging", "dev", "main")
If not provided, uses self.src_env
:return: Server alias from Config
"""
# Use provided env or fall back to self.src_env
env_to_use = env or self.src_env
if env_to_use is None:
raise ValueError("Environment is required for this operation")
# Normalize environment names
normalized_env = env_to_use.lower()
# In dev environment, everything maps to johto
if not Config.PRODUCTION:
return Config.TILESERVER_ALIAS # johto
# In production, map to appropriate server aliases
if normalized_env in ["staging"]:
return Config.TILESERVER_ALIAS_STAGE # hoenn
elif normalized_env in ["dev"]:
return Config.TILESERVER_ALIAS_DEV # johto
elif normalized_env in ["prod", "main"]:
return Config.TILESERVER_ALIAS # kanto
else:
raise ValueError(f"Unknown environment: {env_to_use}")
def _get_datahub_env(self) -> str:
"""
Get the appropriate datahub environment based on the source environment.
"""
if self.src_env == "staging":
return "STG"
return "PROD" if Config.PRODUCTION else "DEV"
[docs]
def tileset_urn(self, tileset_name: str) -> Urn:
"""Get the URN for the tileset in datahub."""
import datahub.emitter.mce_builder as builder
dataset_urn = builder.make_dataset_urn(
"tileservergl", tileset_name, self._get_datahub_env()
)
self.log.info(f"Dataset URN: {dataset_urn}")
return Urn(dataset_urn)
# =========================================================================
# COMMAND EXECUTION FUNCTIONS
# =========================================================================
[docs]
def run_command(self, command: str, description: str, check_output: bool = False):
"""Run a command with handling for dry run mode."""
if self.is_dry_run:
self.log.info(f"[DRY RUN] Would {description}: {command}")
return "" if check_output else None
self.log.info(f"EXECUTING: {command}")
system_env = os.environ.copy()
bash_path = shutil.which("bash") or "bash"
bash_result = self.subprocess_hook.run_command(
command=[bash_path, "-c", command], env=system_env
)
if bash_result.exit_code:
raise AirflowException(
f"Command failed with exit code {bash_result.exit_code}. "
f"Output: {bash_result.output}"
)
if check_output:
return bash_result.output
[docs]
def execute_ssh_command(
self, server: str, command: str, description: str, check_output: bool = False
):
"""Execute an SSH command on a remote server."""
ssh_command = f"ssh {server} '{command}'"
return self.run_command(ssh_command, description, check_output)
[docs]
def restart_tileserver(self, server: str):
"""Stop, update inventory config (update_config.py), and restart the tileserver on a server."""
# Stop the tileserver on a server
self.execute_ssh_command(
server, "sudo systemctl stop tileserver", f"stop tileserver on {server}"
)
# Update config
self.execute_ssh_command(
server, "~/update_config.py", f"update tileserver config on {server}"
)
# Restart tileserver
self.execute_ssh_command(
server,
"sudo systemctl restart tileserver",
f"restart tileserver on {server}",
)
[docs]
def scp_to_server(self, local_path: str, server: str, remote_path: str):
"""Copy a file to a server using SCP."""
scp_command = f"scp {local_path} {server}:{remote_path}"
self.run_command(scp_command, f"copy file to {server}:{remote_path}")
[docs]
def scp_from_server(self, server: str, remote_path: str, local_path: str):
"""Copy a file from a server using SCP."""
scp_command = f"scp {server}:{remote_path} {local_path}"
self.run_command(scp_command, f"copy file from {server}:{remote_path}")
[docs]
def s3_upload(self, local_path: str, s3_path: str):
"""Upload a file to S3."""
s3_command = f"aws s3 cp {local_path} {s3_path}"
self.run_command(s3_command, f"upload to S3: {s3_path}")
[docs]
def s3_download(self, s3_path: str, local_path: str):
"""Download a file from S3."""
s3_command = f"aws s3 cp {s3_path} {local_path}"
self.run_command(s3_command, f"download from S3: {s3_path}")
[docs]
def s3_delete(self, s3_path: str):
"""Delete a file from S3."""
s3_command = f"aws s3 rm {s3_path}"
self.run_command(s3_command, f"delete from S3: {s3_path}")
[docs]
def s3_check_exists(self, s3_path: str) -> bool:
"""Check if a file exists in S3."""
if self.is_dry_run:
self.log.info(f"[DRY RUN] Would check if S3 file exists: {s3_path}")
return True
try:
self.run_command(
f"aws s3 ls {s3_path}",
f"check if S3 file exists: {s3_path}",
check_output=True,
)
return True
except AirflowException:
return False
[docs]
def check_file_exists_on_server(self, server: str, file_path: str) -> bool:
"""Check if a file exists on a server."""
if self.is_dry_run:
self.log.info(f"[DRY RUN] Would check if file exists: {server}:{file_path}")
return True
result = self.execute_ssh_command(
server,
f'if [ -f {file_path} ]; then echo "exists"; else echo "not found"; fi',
f"check if file exists: {file_path}",
check_output=True,
)
return result and "exists" in result
[docs]
def create_directory_on_server(self, server: str, directory: str):
"""Create a directory on a server if it doesn't exist."""
self.execute_ssh_command(
server, f"mkdir -p {directory}", f"create directory if needed: {directory}"
)
# =========================================================================
# VERIFICATION METHODS
# =========================================================================
[docs]
def verify_mbtiles_file(self, file_path: str, server_alias: str) -> bool:
"""
Verify that an mbtiles file exists on a server.
"""
if self.is_dry_run:
self.log.info(
f"[DRY RUN] Would verify mbtiles file: {server_alias}:{file_path}"
)
return True
self.log.info(f"Verifying mbtiles file: {file_path} on {server_alias}")
# Check if file has .mbtiles extension
if ".mbtiles" not in file_path:
error_msg = f"File does not have .mbtiles extension: {file_path}"
self.log.error(error_msg)
raise AirflowException(error_msg)
# Check if file exists on server
if not self.check_file_exists_on_server(server_alias, file_path):
error_msg = f"File not found on server: {server_alias}:{file_path}"
self.log.error(error_msg)
raise AirflowException(error_msg)
# Check if the mbtile is valid
tileset_name = self.get_tileset_name(file_path)
if not Util.is_mbtile_valid(f"{Config.MOUNT_PATH}/{tileset_name}.mbtiles"):
error_msg = f"File is not a valid mbtile: {file_path}"
self.log.error(error_msg)
raise AirflowException(error_msg)
return True
[docs]
def verify_s3_file(self, s3_path: str) -> bool:
"""
Verify that a file exists in S3.
"""
if self.is_dry_run:
self.log.info(f"[DRY RUN] Would verify S3 file: {s3_path}")
return True
self.log.info(f"Verifying S3 file: {s3_path}")
# Check if file has .mbtiles extension
if ".mbtiles" not in s3_path:
error_msg = f"S3 path does not have .mbtiles extension: {s3_path}"
self.log.error(error_msg)
raise AirflowException(error_msg)
# Check if file exists in S3
if not self.s3_check_exists(s3_path):
error_msg = f"File not found in S3: {s3_path}"
self.log.error(error_msg)
raise AirflowException(error_msg)
return True
# =========================================================================
# DATAHUB INTEGRATION METHODS
# =========================================================================
[docs]
def remove_from_datahub(self, tileset_name: str) -> None:
"""Remove the tileset from datahub."""
if self.is_dry_run:
self.log.info(f"[DRY RUN] Would remove tileset {tileset_name} from datahub")
return
self.log.info(f"Removing tileset {tileset_name} from datahub")
tileset_urn = self.tileset_urn(tileset_name).urn
try:
Util.remove_dataset(tileset_urn)
self.log.info(f"Successfully removed tileset {tileset_name} from datahub")
except Exception as ex:
self.log.error(f"Error removing tileset from datahub: {ex}")
[docs]
def add_to_datahub(self, context: Context, tileset_name: str) -> None:
"""
Add the tileset to datahub with tags, and ownership information
"""
import time
if self.is_dry_run:
self.log.info(f"[DRY RUN] Would add tileset {tileset_name} to datahub")
self.log.info(
f"[DRY RUN] Would add dataset info, owner, tags, and assertions"
)
return
self.log.info(f"Adding tileset {tileset_name} to datahub")
try:
tileset_urn = self.tileset_urn(tileset_name).urn
env = self._get_datahub_env()
assertion_urn = Util.get_tileset_assertion_urn(tileset_name, env)
# Add dataset with retries for JSON error
delays = [2, 5, 10] # Retry delays in seconds
last_error = None
for delay in delays:
try:
self.log.info(
f"Waiting {delay}s for tileserver to process the tile..."
)
time.sleep(delay)
# Add dataset
Util.add_tile_set_to_datahub(
f"https://{self.tileserver_alias()}.ready.net",
tileset_name,
tileset_urn,
)
last_error = None
break
except Exception as e:
last_error = e
if "Expecting value" not in str(e):
break # Only retry for JSON parsing errors
self.log.warning(f"JSON error while adding to datahub: {e}")
if last_error:
self.log.error(
f"Failed to add tileset metadata after retries: {last_error}"
)
self.log.info(
"Continuing with other datahub operations despite metadata error"
)
# Add owner
if len(context["dag"].default_args["email"]) > 0:
owner = context["dag"].default_args["email"][0]
user_urn = Util.user_urn(owner)
Util.add_owner_to_dataset(dataset_urn=tileset_urn, owner_urn=user_urn)
# Add tags
tileset_tag_urn = Util.create_tag(
"tileset", "A tileset to render on the map"
)
tileservergl_tag_urn = Util.create_tag(
"tileservergl", "A tile rendering engine to serve the tilesets"
)
Util.add_tag_to_dataset(dataset_urn=tileset_urn, tag_urn=tileset_tag_urn)
Util.add_tag_to_dataset(
dataset_urn=tileset_urn, tag_urn=tileservergl_tag_urn
)
# Add assertions
Util.add_tileset_assertions(tileset_urn, assertion_urn, "SUCCESS")
self.log.info(f"Successfully added tileset {tileset_name} to datahub")
except Exception as ex:
self.log.error(f"Error adding tileset to datahub: {ex}")
self.log.info(
f"Note: Tile was successfully recovered regardless of datahub registration status"
)
[docs]
def notify_slack(self, message: str) -> None:
"""Send a notification to Slack"""
if self.is_dry_run:
self.log.info(f"[DRY RUN] Would send Slack notification: {message}")
return
try:
slack_hook = SlackHook(slack_conn_id=Config.SLACK_DEFAULT_API)
slack_hook.call(
"chat.postMessage",
json={
"channel": "#tileserver",
"text": message,
},
)
except Exception as ex:
self.log.error(f"Error sending Slack notification: {ex}")
[docs]
def on_kill(self) -> None:
"""Handle task kill event when Airflow kills the task instance"""
self.log.info("Task is being killed, performing cleanup...")
try:
# Only attempt to restart tileserver if:
# 1. We're not in dry run mode
# 2. We have a source environment
if not self.is_dry_run and self.src_env is not None:
environment_alias = self.tileserver_alias()
self.log.info(
f"Attempting to restart tileserver on {environment_alias}"
)
try:
# Try to restart tileserver without failing if it doesn't work
self.restart_tileserver(environment_alias)
self.log.info(
f"Successfully restarted tileserver on {environment_alias}"
)
except Exception as e:
self.log.error(f"Failed to restart tileserver on kill: {e}")
except Exception as ex:
self.log.error(f"Error during on_kill cleanup: {ex}")
# Always terminate any running subprocess as the final step
self.subprocess_hook.send_sigterm()
self.log.info("Task kill cleanup completed")