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 CreateTileOperator(BaseOperator):
"""
An operator to generate tiles
:param source_sql: The sql query to run against database, if it ends with `.sql` we read the file relative to the `template` path
:type source_sql: str
:param title: The title of the file,
:type title: str
:param copy_main_to_stage: Copy the generated mbtiles to the staging tileserver. Only works on production
:type copy_main_to_stage: bool = False
:param description: The description of the file,
:type description: str
:param tile_server_folder: Path of on the tileserver to store mbtiles. Disabled on purpose for future
:type tile_server_folder: str
:param geometry_column: add your geometry column name here. It can be comma separated if you have two columns like lat, lon (with the same order).,
:type geometry_column: str = "geometry"
:param deploy_tiles: Deploy tiles on the tileserver
:type deploy_tiles: bool = False
:param workgroup: workgroup information to read sql query from it
:type workgroup: str = Config.DEFAULT_WORKGROUP
:param tgt_db: tgt_db information to read sql query from it
:type tgt_db: str = Config.DATA_LAKE_DB
:param template: the path for the sql query template.
:type template: str
:param task_suffix: the scale of the ecs machine. Xlarge is only supported for the tippecaneo ecs and it will fall back to large for ecs_athena.
:type task_suffix: Literal["large", "small", "xlarge"] = "small"
:param deploy_environment: This will control the environment to deploy. on airflow dev `main=dev` and on prod `main=prod` and staging is always staging.
:type deploy_environment: Literal["main", "staging"] = "main"
:param tippecanoe_cmd_template: the command template to send to tippecaneo.Don't change if you are not sure what are the params
:param map_name: the name of the map. if not set it will be set to {tile_envs}.{title}
:type tippecanoe_cmd_template: str = "tippecanoe -zg -Z0 --extend-zooms-if-still-dropping --drop-densest-as-needed --force -o {mbtiles_path} -l {title} -n {title_description} -N {description} {geojson_path}"
:param inlets: An array of inlets to this task. use `Util.lineage_urn` to create the urns they should be type of `Urn`
:type inlets: list[Urn]
:param label_columns: the columns to be used as labels in the map. it can be an array. mainly use this if your data are polygons and you want to show the name of the polygon on the map. NOTE: If you apply column mapping the new column name should be used here.
:type label_columns: list[str]
:param data_product_name: the name of the data product to add the tileset to. if not set it will not be added
:type data_product_name: str | None = None
:param keep_tiles_on_mount_folder: if set to true, the mbtiles will not be removed from the mount folder. mainly used for debugging
:type keep_tiles_on_mount_folder: bool = False
:raises ex: if is_athena is set to false. pg is not yet supported
"""
ui_color = "#9933FF"
ui_fgcolor = "#FFFFFF"
custom_operator_name = "CreateTileOperator"
template_fields_renderers = {
"source_sql": "sql",
"deploy_tiles": "bool",
"copy_main_to_stage": "bool",
}
template_ext = (".sql",)
# templating is meant for minimal variables
template_fields = (
"source_sql",
"title",
"task_suffix",
"deploy_tiles",
"description",
"deploy_environment",
)
@apply_defaults
def __init__(
self,
source_sql: str,
title: str,
copy_main_to_stage: bool = False,
description: str = "",
tile_server_folder: str = "",
geometry_column: str = "geometry",
deploy_tiles: bool = False,
task_suffix: Literal["large", "small", "xlarge"] = "large",
deploy_environment: Literal["main", "staging"] = "main",
column_mapping: dict[str, str] = {},
workgroup: str = Config.DEFAULT_WORKGROUP,
tgt_db: str = Config.DATA_LAKE_DB,
template: str = "",
label_columns: list[str] = [],
inlets: list[Urn] = [],
tippecanoe_cmd_template: str = "tippecanoe -zg -Z0 --extend-zooms-if-still-dropping --drop-densest-as-needed --force -o {mbtiles_path} -l {title} -n {title_description} -N {description} {geojson_path}",
map_name: str = "",
data_product_name: str | None = None,
keep_tiles_on_mount_folder: bool = False,
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.geometry_column = geometry_column
self.title = title
self.copy_main_to_stage = copy_main_to_stage
self.task_suffix = task_suffix
self.workgroup = workgroup
self.tgt_db = tgt_db
self.tile_server_folder = "" # tile_server_folder.lstrip('/') Disabled for now
self.deploy_tiles = deploy_tiles
self.column_mapping = column_mapping
self.description = description
self.template = template
self.source_sql = source_sql
self.tileserver_base_path = Config.DEFAULT_TS_FLAT_FOLDER
self.dag_root = Config.DAG_ROOT
filename = str(uuid.uuid4().hex)
self.file_path = f"{Config.MOUNT_PATH}/{filename}.sql"
self.deploy_environment = deploy_environment
self.tippecanoe_cmd_template = tippecanoe_cmd_template
self.inlets = inlets
self.label_columns = label_columns
self.map_name = map_name
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 execute(
self,
context: Context,
) -> None:
"""
main call to generate geojson -> tile
:param context: the context
:type context: Context
"""
import awswrangler as wr
import sqlparse
import time
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 }}"
column_mapping = self.column_mapping
task_suffix = "" if self.task_suffix == "small" else self.task_suffix
formatted_sql = sqlparse.format(
self.source_sql, reindent=True, keyword_case="upper"
)
self.log.info(f"self.source_sql:\n{formatted_sql}")
self.log.info(f"self.task_suffix:{self.task_suffix}")
if self.source_sql.lower().endswith(".sql"):
path_to_sql = f"{self.dag_root}/{self.template}/{self.source_sql}"
self.log.info(f"opening sql file {path_to_sql}")
sql_file_path = Path(path_to_sql)
self.source_sql = sql_file_path.read_text()
# write sql file to the efs-mount
df_size = wr.athena.read_sql_query(
f"with temp_table as ({self.source_sql}) select count(*) as count from temp_table limit 10",
workgroup=self.workgroup,
database=self.tgt_db,
)
if not df_size["count"][0]:
raise AirflowSkipException(
"Skipping the task since query returns empty data frame"
)
with open(self.file_path, "w") as f:
f.write(self.source_sql)
try:
Util.run_athena_to_geojson(
query=self.file_path,
title=title,
geometry_column=self.geometry_column,
wr_workgroup=self.workgroup,
wr_database=self.tgt_db,
task_suffix=task_suffix,
column_mapping=column_mapping,
label_columns=self.label_columns,
)
self.log.info(f"Start tippecanoe")
Util.run_tippecanoe(
title=title,
task_suffix=task_suffix,
description=self.description,
cmd=self.tippecanoe_cmd_template,
)
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.copy_mbtiles(environment_alias=self.tileserver_alias())
self.copy_to_stage()
# to make sure server is up and running
time.sleep(120)
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)
self.add_to_datahub(context=context)
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")
if not self.keep_tiles_on_mount_folder:
self.remove_source_mbtiles_and_geojson()
self.log.info(f"Tile generation finished. Tileset name: {title}")
except Exception as ex:
self.log.info(f"Error generating tile, {ex}")
raise ex
finally:
Path(self.file_path).unlink(missing_ok=True)
[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