import shutil
import os
import uuid
from config import Config
from util import Util
from typing import Literal
from pathlib import Path
from functools import cached_property
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.utils.context import Context
from airflow.exceptions import AirflowException, AirflowSkipException
from airflow.hooks.subprocess import SubprocessHook
from datahub_airflow_plugin.entities import Urn
from airflow.providers.slack.hooks.slack import SlackHook
[docs]
class BaseTileDeployOperator(BaseOperator):
"""
A Base operator to deploy a tile to tileserver
"""
ui_color = "#9933FF"
ui_fgcolor = "#FFFFFF"
custom_operator_name = "BaseTileDeployOperator"
template_fields_renderers = {
"deploy_tiles": "bool",
"copy_main_to_stage": "bool",
}
@apply_defaults
def __init__(
self,
title: str,
copy_main_to_stage: bool = False,
description: str = "",
deploy_tiles: bool = False,
deploy_environment: Literal["main", "staging"] = "main",
inlets: list[Urn] = [],
data_product_name: str | None = None,
keep_tiles_on_mount_folder: bool = False,
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.title = title
self.copy_main_to_stage = copy_main_to_stage
self.tile_server_folder = "" # tile_server_folder.lstrip('/') Disabled for now
self.deploy_tiles = deploy_tiles
self.description = description
self.tileserver_base_path = Config.DEFAULT_TS_FLAT_FOLDER
self.dag_root = Config.DAG_ROOT
self.deploy_environment = deploy_environment
self.inlets = inlets
self.data_product_name = data_product_name
self.keep_tiles_on_mount_folder = keep_tiles_on_mount_folder
@cached_property
def subprocess_hook(self):
"""Returns hook for running the bash command."""
return SubprocessHook()
[docs]
def copy_mbtiles(
self,
environment_alias: str,
) -> None:
"""
copy mbtile files to the tileserver
:param environment_alias: the tile server's ailas
:type environment_alias: str
:return: None
:raises AirflowException: if the command fails
"""
dest_path = f"{self.tileserver_base_path}/{self.tile_server_folder}".rstrip(
"/"
) # remove trail if exist
local_path_prefix = f"{Config.MOUNT_PATH}/{self.title}"
# make directory if it does not exist
mkdir_ssh_command = f"ssh {environment_alias} 'mkdir -p {dest_path}'"
self.run_bash_command(mkdir_ssh_command)
# copy file
scp_tile_copy_command = (
f"scp {local_path_prefix}.mbtiles {environment_alias}:{dest_path}"
)
self.run_bash_command(scp_tile_copy_command)
# update config
ssh_update_config_command = f"ssh {environment_alias} '~/update_config.py'"
self.run_bash_command(ssh_update_config_command)
# restart server
ssh_update_config_command = (
f"ssh {environment_alias} 'sudo systemctl restart tileserver'"
)
self.run_bash_command(ssh_update_config_command)
[docs]
def run_bash_command(
self,
command: str,
) -> None:
"""
runs a command with subprocess hook
:param command: the command
:type command: str
:raises AirflowException: if the command fails
"""
system_env = os.environ.copy()
bash_path = shutil.which("bash") or "bash"
# self.log.info(f"deploy mbtile file using command:{command}")
bash_result = self.subprocess_hook.run_command(
command=[bash_path, "-c", command], env=system_env
)
if bash_result.exit_code:
raise AirflowException(
f"Bash command failed. The command returned a non-zero exit code {bash_result.exit_code}."
)
[docs]
def remove_source_mbtiles_and_geojson(self) -> None:
"""
removes files from efs-mount drive
:return: None
:raises AirflowException: if the command fails
"""
local_path_prefix = f"{Config.MOUNT_PATH}/{self.title}"
rm_source_mbtiles_command = f"rm -f {local_path_prefix}.mbtiles && rm -f {local_path_prefix}.geojson && rm -f {local_path_prefix}_labels.geojson"
self.run_bash_command(rm_source_mbtiles_command)
[docs]
def copy_to_stage(self):
"""Handles copying the mbtiles to the staging tileserver. only works on production"""
if Config.PRODUCTION and self.deploy_environment == "main":
if not Util.tileset_exist_on_server(
self.title, Config.TILESERVER_ALIAS_STAGE
):
self.log.info(
f"ALERT: Tileset {self.title} does not exist on staging => copying (as new) from production to staging!"
)
self.copy_mbtiles(environment_alias=Config.TILESERVER_ALIAS_STAGE)
elif self.copy_main_to_stage:
self.log.info(
f"ALERT: Tileset {self.title} already exists on staging => now overwriting..."
)
self.copy_mbtiles(environment_alias=Config.TILESERVER_ALIAS_STAGE)
[docs]
def set_map_name(self, context: Context) -> None:
title = self.title
if not Config.PRODUCTION:
tile_envs = "dev_only"
elif self.copy_main_to_stage:
tile_envs = "prod+stage"
else:
tile_envs = "prod_only"
if self.map_name == "":
self.map_name = f"{tile_envs}.{title}"
# add map_index_template
context["params"]["map_name"] = self.map_name
context["map_index_template"] = "{{ params.map_name }}"
[docs]
def process_tile(self, context: Context) -> None:
raise NotImplementedError(
"process_tile method should be implemented in the child class"
)
[docs]
def deploy_tilset(self) -> None:
"""
deploys the tileset to the tileserver
:return: None
:raises AirflowException: if the command fails
"""
self.copy_mbtiles(environment_alias=self.tileserver_alias())
self.copy_to_stage()
[docs]
def validate_deployment(self, context: Context) -> None:
message_full = (
f"""tile-sever: {self.tileserver_alias()} \n tileset: {self.title}\n"""
)
if not Util.tileserver_is_healthy(self.tileserver_alias()):
message_full = f"""{message_full}\n notify: <@majid> \n status: failed to restart after 2 minutes"""
self.notify_slack(message_full)
raise AirflowException(f"Tileserver failed to start")
else:
message_full = f"{message_full}\n status: deployed successfully"
self.notify_slack(message_full)
if not Util.tileset_exist_on_server(self.title, self.tileserver_alias()):
raise AirflowException(f"cannot find tileset on server")
self.log.info(f"Tileset exist on server")
[docs]
def post_deployment_actions(self, context: Context) -> None:
pass
[docs]
def execute(
self,
context: Context,
) -> None:
"""
:param context: the context
:type context: Context
"""
import time
self.set_map_name(context)
try:
self.process_tile(context)
print(
f"DEBUG: TILING environment {self.deploy_environment=}; {Config.PRODUCTION=}; {self.title=}; {self.copy_main_to_stage=}; "
)
if not Util.is_mbtile_valid(f"{Config.MOUNT_PATH}/{self.title}.mbtiles"):
raise AirflowException(
f"Error generating tile, {self.title}.mbtiles is not valid"
)
if self.deploy_tiles:
self.deploy_tilset()
# to make sure server is up and running
time.sleep(120)
self.validate_deployment(context)
self.add_to_datahub(context=context)
if not self.keep_tiles_on_mount_folder:
pass
# self.remove_source_mbtiles_and_geojson()
self.log.info(f"Tile generation finished. Tileset name: {self.title}")
except Exception as ex:
self.log.info(f"Error generating tile, {ex}")
raise ex
finally:
self.post_deployment_actions(context=context)
[docs]
def on_kill(self) -> None:
self.subprocess_hook.send_sigterm()
[docs]
def add_to_datahub(
self,
context: Context,
) -> None:
tileset_urn = self.tileset_urn().urn
print(f"DEBUG: tileset_urn:{tileset_urn}")
env = self.datahub_environment()
assertion_urn = Util.get_tileset_assertion_urn(self.title, env)
# Add dataset
try:
Util.add_tile_set_to_datahub(
f"https://{self.tileserver_alias()}.ready.net",
self.title,
tileset_urn,
)
# 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 datahubhub lineage
inlets, outlets = self.get_lineage()
context["ti"].task.inlets = inlets
context["ti"].task.outlets = outlets
# 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")
if self.data_product_name:
domain_urn = Util.domain_urn()
Util.get_data_product(self.data_product_name, domain_urn)
Util.add_assets_to_data_product(
data_product_id=self.data_product_name,
assets=[tileset_urn],
)
except Exception as ex:
self.log.info(f"Error adding tileset to datahub, {ex}")
# Util.add_tileset_assertions(
# tileset_urn, assertion_urn, "FAILURE"
# )
[docs]
def tileserver_alias(self) -> str:
return (
Config.TILESERVER_ALIAS_STAGE
if self.deploy_environment == "staging"
else Config.TILESERVER_ALIAS
)
[docs]
def notify_slack(self, message) -> None:
slack_hook = SlackHook(slack_conn_id=Config.SLACK_DEFAULT_API)
slack_hook.call(
"chat.postMessage",
json={
"channel": "#tileserver",
"text": message,
},
)
[docs]
def datahub_environment(self) -> str:
if self.deploy_environment == "staging":
return "STG"
return (
"PROD" if Config.PRODUCTION else "DEV"
) # if not production, it will be dev
[docs]
def tileset_urn(self) -> Urn:
from datahub_airflow_plugin.entities import Urn
import datahub.emitter.mce_builder as builder
dataset_urn = builder.make_dataset_urn(
"tileservergl", self.title, self.datahub_environment()
)
print(f"DEBUG: dataset_urn:{dataset_urn}")
return Urn(dataset_urn)
[docs]
def get_lineage(self) -> tuple[list[Urn], list[Urn]]:
outlets = [self.tileset_urn()]
inlets = []
if self.inlets:
for urn in inlets:
if isinstance(urn, str):
inlets.append(Urn(urn))
else:
inlets.append(urn)
return inlets, outlets
[docs]
class JoinTileOperator(BaseTileDeployOperator):
"""
An operator to use join_tile tippecaneo command to join multiple tiles
"""
ui_color = "#9933FF"
ui_fgcolor = "#FFFFFF"
custom_operator_name = "JoinTileOperator"
@apply_defaults
def __init__(
self,
title: str,
input_mbtiles_stub: str,
copy_main_to_stage: bool = False,
description: str = "",
deploy_tiles: bool = False,
task_suffix: Literal["large", "small", "xlarge"] = "large",
deploy_environment: Literal["main", "staging"] = "main",
inlets: list[Urn] = [],
tippecanoe_cmd_template: str = "find {mount_path} -name '{input_mbtiles_stub}_*.mbtiles' | xargs tile-join -z 15 -pk -f -o {mount_path}/{input_mbtiles_stub}.mbtiles -n '{name}' -N '{description}'",
map_name: str = "",
data_product_name: str | None = None,
keep_tiles_on_mount_folder: bool = False,
*args,
**kwargs,
):
self.title = title
self.copy_main_to_stage = copy_main_to_stage
self.tile_server_folder = "" # tile_server_folder.lstrip('/') Disabled for now
self.deploy_tiles = deploy_tiles
self.description = description
self.tileserver_base_path = Config.DEFAULT_TS_FLAT_FOLDER
self.dag_root = Config.DAG_ROOT
self.deploy_environment = deploy_environment
self.inlets = inlets
self.data_product_name = data_product_name
self.keep_tiles_on_mount_folder = keep_tiles_on_mount_folder
self.input_mbtiles_stub = input_mbtiles_stub
self.task_suffix = task_suffix
self.tippecanoe_cmd_template = tippecanoe_cmd_template
self.map_name = map_name
super().__init__(
title=title,
copy_main_to_stage=copy_main_to_stage,
description=description,
deploy_tiles=deploy_tiles,
deploy_environment=deploy_environment,
inlets=inlets,
data_product_name=data_product_name,
keep_tiles_on_mount_folder=keep_tiles_on_mount_folder,
*args,
**kwargs,
)
[docs]
def process_tile(self, context: Context) -> None:
"""
process the tile using tippecanoe command
:param context: the context
:type context: Context
"""
Util.run_tile_join(
name=self.title,
input_mbtiles_stub=self.input_mbtiles_stub,
description=self.description,
task_suffix=self.task_suffix,
tippecanoe_cmd=self.tippecanoe_cmd_template,
)