from typing import Literal
from config import Config
from airflow.providers.postgres.hooks.postgres import PostgresHook
from airflow.operators.bash import BashOperator
from airflow.exceptions import AirflowException
from traceback import print_exc
from subprocess import check_output
import time, os
import json
[docs]
class ECSUtil:
"""
ECS utilities for running, monitoring, and managing tasks
"""
[docs]
@staticmethod
def wait_for_task(
task_arn: str,
cluster: str = Config.ECS_CLUSTER,
timeout: int = Config.MAX_ITERATIONS,
sleep_time: int = Config.SLEEP_TIME,
):
"""
Wait for a task to finish
:param task_arn: the task arn
:param cluster: the cluster
:param timeout: the timeout
:param sleep_time: the sleep time
:raises Exception: if the task times out
"""
import boto3
client = boto3.client("ecs")
task_response = client.list_tasks(cluster=cluster)
is_running = task_arn in task_response["taskArns"]
count = 0
while is_running:
print(f"Task {task_arn} is running: {count}...")
count += 1
time.sleep(sleep_time)
task_response = client.list_tasks(cluster=cluster)
is_running = task_arn in task_response["taskArns"]
if count > timeout:
print('Killing the task')
client.stop_task(
cluster=cluster,
task=task_arn,
reason='Stopping task due to Airflow timeout'
)
raise Exception("Task timed out")
[docs]
@staticmethod
def run_task(
container_name: str,
task_definition: str,
env_vars: list[dict[str, str]], # parametrize new containers with env variables
cluster: str = Config.ECS_CLUSTER,
subnets: list[
str
] = Config.SUBNETS, # spreading across subnets allows us to allocate resources more efficiently, reduces start failures
security_groups: list[
str
] = Config.ECS_SECURITY_GROUPS, # security groups for the task
timeout: int = Config.MAX_ITERATIONS, # sometimes the task will fail to start and we need to retry though
sleep_time: int = Config.SLEEP_TIME,
task_suffix: str = "", # add addl options for differently sized default containers
wait: bool = True,
check_status: bool = True,
wait_timeout: int = Config.MAX_ITERATIONS,
):
"""
Run a task on ECS
:param container_name: the container name
:param task_definition: the task definition
:param env_vars: the env vars to pass to the task, parametrize new containers with env variables
:param cluster: the cluster to run the task on
:param subnets: the subnets that one may run the task on, spreading across subnets allows us to allocate resources more efficiently, reduces start failures
:param timeout: the timeout, sometimes the task will fail to start and we need to retry
:param sleep_time: the sleep time between timeout retries
:param task_suffix: add addl options for differently sized default containers
:param wait: whether or not to wait for the task to finish
:param check_status: whether or not to check the status of the task
:param wait_timeout: the timeout that airflow will wait to announce task as timed-out
:raises Exception: if the task fails
:return: the task arn
:rtype: _type_
"""
import boto3
client = boto3.client("ecs")
target_task_definition = (
task_definition
if task_suffix in ["small", ""]
else f"{task_definition}_{task_suffix}"
)
for iteration in range(timeout):
response = client.run_task(
cluster=cluster,
count=1,
launchType="FARGATE",
networkConfiguration={
"awsvpcConfiguration": {
"subnets": subnets,
"securityGroups": security_groups,
"assignPublicIp": "ENABLED",
}
},
overrides={
"containerOverrides": [
{
"name": container_name,
"environment": env_vars,
},
],
},
taskDefinition=target_task_definition,
)
try:
task_arn = response["tasks"][0]["taskArn"]
break
except Exception as e:
task_arn = "NONE"
print_exc()
print(f"Error running task: {e}")
print(response)
print(f"Iteration: {iteration}...")
time.sleep(sleep_time)
print(f"Running task: {task_arn}")
task_parts = task_arn.split("/")
if len(task_parts) == 3:
print(
f"Task log Url: https://us-east-1.console.aws.amazon.com/ecs/v2/clusters/{task_parts[1]}/tasks/{task_parts[2]}/logs?region=us-east-1"
)
print(
f"CloudWatch Url: https://us-east-1.console.aws.amazon.com/cloudwatch/home?region=us-east-1#logEventViewer:group=/ecs/{target_task_definition};stream=ecs/{container_name}/{task_parts[2]}"
)
if wait:
print(f"Waiting for task {task_arn} to finish...")
ECSUtil.wait_for_task(task_arn=task_arn, timeout=wait_timeout)
print(f"Finished waiting for Task {task_arn}")
if check_status:
time.sleep(5)
logs = ECSUtil.check_ecs_status(task_arn, target_task_definition)
if isinstance(logs, str):
print(logs)
return task_arn
[docs]
@staticmethod
def check_ecs_status(
task_arn: str,
task_definition: str,
cluster: str = Config.ECS_CLUSTER,
log_success: bool = False,
log_failure: bool = True,
return_logs: bool = False,
strict_error: bool = True,
):
"""
Check the status of a task on ECS
:param task_arn: the task arn
:param task_definition: the task definition
:param cluster: the cluster
:param log_success: whether or not to log success
:param log_failure: whether or not to log failure
:param return_logs: whether or not to return logs
:param strict_error: whether or not to raise an error
:raises AirflowException: if the task fails
:return: the logs
"""
logs = ECSUtil.get_ecs_logs(task_arn, task_definition, cluster)
if "ReadyECSWorkerTaskSuccess" in logs:
if log_success:
print(logs)
return True if not return_logs else logs
elif "ReadyECSWorkerTaskFailure" in logs:
if log_failure:
print(logs)
raise AirflowException(f"Task failed: {task_arn}")
else:
if strict_error:
raise AirflowException(f"Task failed: {task_arn}")
return logs
[docs]
@staticmethod
def get_ecs_logs(
task_arn: str,
task_definition: str,
cluster: str = Config.ECS_CLUSTER,
sleep_time: int = 5,
iterations: int = 6,
) -> str:
"""
Get the logs for a task from ECS using the cli
:param task_arn: the task arn
:param task_definition: the task definition
:param cluster: the cluster that the task is on
:param sleep_time: the sleep time between retries
:param iterations: the number of iterations to try
:raises Exception: if the logs cannot be retrieved
:return: the logs
"""
for i in range(iterations):
try:
return check_output(
f"ecs-cli logs --task-id {task_arn.rsplit('/', 1)[1]} -c {cluster} --task-def {task_definition}".split()
).decode()
except Exception as e:
print(f"Error {i} getting logs: {e}")
time.sleep(sleep_time)
return ""
[docs]
@staticmethod
def run_ecs_h3(
output_tablename: str,
aws_query: str,
agg_dict: dict[str, str],
export_path: str,
partition_cols: list[str] = [],
overwrite: bool = False,
resolutions: list[int] = [5, 6, 7, 8, 9],
geometry_column: str = "lat,lon",
s3_bucket: str = Config.DATA_LAKE_BUCKET,
database: str = Config.DATA_LAKE_DB,
workgroup: str = Config.DEFAULT_WORKGROUP,
wait: bool = True,
task_suffix: Literal["", "large", "xlarge"] = "",
):
"""Run a athena query on ECS h3 aggregation machine, and save the results as table in s3
:param output_tablename: the destination table name in athena, it will have format: c_h3_{output_tablename}, with extra cols: cell_id (h3 index id), cell_geom (geometry), resolution, ...
:param aws_query: the query string to run on athena. Use `SELECT col_x AS col_y` to rename columns as needed
:param agg_dict: the mapping of column-aggregation method; acceptable aggregation methods includes:
`count`, `median` (AND LATER: `sum`, `mean`, `std`, `var`, `min`, `max`)
:param partition_cols: the additional columns to partition the result table,
in format: h3_res/additional-partition-cols.
NOTE: these columns must be included in the query; they will be used for partitioning the result table output;
they will not be used for aggregation!
:param overwrite: To overwrite existing data (table or partition) in the destination table if already exists
:param resolutions: h3 resolutions to use as aggregation unit
:param geometry_column: the column that includes `wkt` format of the data.
if it is a comma seperated text we split it to lat and lon and read it as point
:param export_path: the S3 path to save result table, default is `curated/h3agg/c_h3_{output_tablename}/resolution=<>/[...]/.parquet`
:param s3_bucket: location to store temp UNLOAD parquet and result parquets
:param database: the datatable in athena
:param wr_workgroup: the workgroup in athena
:param wait: whether or not to wait for the task to finish
:param task_suffix: the size of ECS machine to use for processing
:return: the task arn
"""
from json import dumps
table_name = f"c_h3_{output_tablename}"
print(
f""" ########################### LOGGING:
Full Athena Table name: {table_name}
S3 location: {export_path=}"""
)
task_arn = ECSUtil.run_task(
container_name=Config.ECS_H3_CONTAINER_NAME,
task_definition=Config.ECS_H3_TASK_DEFINITION,
env_vars=[
{"name": "OUTPUT_TABLE_NAME", "value": table_name},
{"name": "AWS_QUERY", "value": aws_query},
{"name": "OVERWRITE", "value": str(overwrite)},
{"name": "AGGREGATE_DICT", "value": dumps(agg_dict)},
{"name": "RESOLUTIONS", "value": ",".join(map(str, resolutions))},
# NOTE: passing partition_cols as a string (it could be empty string ''), so that we can pass it as a single env var
{"name": "PARTITION_COLS", "value": ",".join(partition_cols)},
{"name": "GEOMETRY_COLUMN", "value": geometry_column},
{"name": "S3_PATH", "value": export_path},
{"name": "S3_BUCKET", "value": s3_bucket},
{"name": "WR_DATABASE", "value": database},
{"name": "WR_WORKGROUP", "value": workgroup},
],
task_suffix=task_suffix,
wait=wait,
)
return task_arn
[docs]
@staticmethod
def run_hld_model(
statefp: str,
storage_prefix: str,
sql: str,
bucket: str = Config.DATA_LAKE_BUCKET,
datalake_db: int = Config.DATA_LAKE_DB,
workgroup: str = Config.DEFAULT_WORKGROUP,
wait: bool = True,
task_suffix: str = "",
container_name: str = Config.ECS_HLD_CONTAINER_NAME,
task_definition: str = Config.ECS_HLD_TASK_DEFINITION,
):
"""Run HLD setup process
:param statefp: state fip to process
:type statefp: str
:param storage_prefix: the storage prefix
:type storage_prefix: str
:param sql: the sql query to run
:type sql: str
:param bucket: the s3 bucket to store setup network, defaults to Config.DATA_LAKE_BUCKET
:type bucket: str, optional
:param datalake_db: the datalake db, defaults to Config.DATA_LAKE_DB
:type datalake_db: int, optional
:param workgroup: the workgroup, defaults to Config.DEFAULT_WORKGROUP
:type workgroup: str, optional
:param wait: whether or not to wait for the task to finish, defaults to True
:type wait: bool, optional
:param task_suffix: the task suffix, e.g. "large", defaults to ""
:type task_suffix: str, optional
:param container_name: the container name, defaults to Config.ECS_HLD_CONTAINER_NAME
:type container_name: str, optional
:param task_definition: the task definition, defaults to Config.ECS_HLD_TASK_DEFINITION
:type task_definition: str, optional
:return: task ARN
:rtype: _type_
"""
task_arn = ECSUtil.run_task(
container_name=container_name,
task_definition=task_definition,
env_vars=[
{"name": "MODEL", "value": "airflow"},
{"name": "ECS_AGENT", "value": "airflow"},
{"name": "STATEFP", "value": str(statefp)},
{"name": "ATHENA_DATABASE", "value": datalake_db},
{"name": "ATHENA_WORKGROUP", "value": workgroup},
{"name": "BUCKET", "value": bucket},
{"name": "PREFIX", "value": storage_prefix},
{"name": "SQL", "value": sql},
],
task_suffix=task_suffix,
wait=wait,
)
return task_arn
[docs]
@staticmethod
def run_hld_setup(
statefp: str,
bucket: str = Config.DATA_LAKE_BUCKET,
datalake_db: int = Config.DATA_LAKE_DB,
workgroup: str = Config.DEFAULT_WORKGROUP,
wait: bool = True,
task_suffix: str = "",
container_name: str = Config.ECS_HLD_CONTAINER_NAME,
task_definition: str = Config.ECS_HLD_TASK_DEFINITION,
):
"""
Run HLD setup process
:param statefp: state fip to process
:type statefp: str
:param bucket: the s3 bucket to store setup network, defaults to Config.DATA_LAKE_BUCKET
:type bucket: str, optional
:param datalake_db: the datalake db, defaults to Config.DATA_LAKE_DB
:type datalake_db: int, optional
:param workgroup: the workgroup, defaults to Config.DEFAULT_WORKGROUP
:type workgroup: str, optional
:param wait: whether or not to wait for the task to finish, defaults to True
:type wait: bool, optional
:param task_suffix: the task suffix, e.g. "large", defaults to ""
:type task_suffix: str, optional
:param container_name: the container name, defaults to Config.ECS_HLD_CONTAINER_NAME
:type container_name: str, optional
:param task_definition: the task definition, defaults to Config.ECS_HLD_TASK_DEFINITION
:type task_definition: str, optional
:return: task arn
:rtype: _type_
"""
task_arn = ECSUtil.run_task(
container_name=container_name,
task_definition=task_definition,
env_vars=[
{"name": "MODEL", "value": "setup"},
{"name": "STATEFP", "value": str(statefp)},
{"name": "ATHENA_DATABASE", "value": datalake_db},
{"name": "ATHENA_WORKGROUP", "value": workgroup},
{"name": "BUCKET", "value": bucket},
],
task_suffix=task_suffix,
wait=wait,
)
return task_arn
[docs]
@staticmethod
def run_s3_parquet_to_pg(
table: str,
key: str,
drop_table: bool = False,
suffix: str = "",
bucket: str = Config.DATA_LAKE_BUCKET,
batch_size: int = Config.BATCH_SIZE,
conn_id: str = Config.READY_DATA_DB,
wait: bool = True,
task_suffix: str = "",
container_name: str = Config.ECS_S3_PARQUET_PG_CONTAINER_NAME,
task_definition: str = Config.ECS_S3_PARQUET_PG_TASK_DEFINITION,
local: bool = False,
use_pg_copy=False,
on_conflict_sql: str = None,
upsert_origin_table: str = None,
):
"""
Run an ECS task to copy a parquet file from s3 to postgres in a set of chunks
:param table: the table name for the end table
:type table: str
:param key: the key of the parquet file
:type key: str
:param drop_table: whether or not to drop the table before copying
:type drop_table: bool
:param suffix: the suffix for the CREATE TABLE query that runs in ECS
:type suffix: str
:param bucket: the bucket of the parquet file
:type bucket: str
:param batch_size: the batch size for the transfer in pandas
:type batch_size: int
:param conn_id: the connection id
:type conn_id: str
:param wait: whether or not to wait for the task to finish
:type wait: bool
:param task_suffix: the task suffix, e.g. "large"
:type task_suffix: str
:param container_name: the container name
:type container_name: str
:param task_definition: the task definition
:type task_definition: str
:param local: whether or not to run the task locally (deprecated)
:type local: bool
:param use_pg_copy: experimental use of pg_copy
:type use_pg_copy: bool
:return: the task arn
:rtype: _type_
"""
pg_hook = PostgresHook(postgres_conn_id=conn_id)
conn = pg_hook.get_uri()
if drop_table and not suffix:
raise ValueError("You must provide a suffix to drop the table")
if local:
pass
else:
env_vars = [
{"name": "TGT_TABLE", "value": f"{table}"},
{"name": "CONN", "value": f"{conn}"},
{"name": "TGT_BUCKET", "value": f"{bucket}"},
{"name": "TGT_KEY", "value": f"{key}"},
{"name": "CREATE_SUFFIX", "value": f"{suffix}"},
{"name": "BATCH_SIZE", "value": f"{batch_size}"},
{"name": "DROP_TABLE", "value": f"{drop_table}"},
{"name": "USE_PG_COPY", "value": f"{use_pg_copy}"},
]
if upsert_origin_table:
env_vars.append(
{"name": "UPSERT_FROM", "value": f"{upsert_origin_table}"}
)
if on_conflict_sql:
env_vars.append({"name": "ON_CONFLICT", "value": on_conflict_sql})
task_arn = ECSUtil.run_task(
container_name=container_name,
task_definition=task_definition,
env_vars=env_vars,
task_suffix=task_suffix,
wait=wait,
wait_timeout=3 * Config.MAX_ITERATIONS,
)
return task_arn
[docs]
@staticmethod
def run_mlflow_model(
model_uri: str,
params: dict[str, str],
input: dict[str, str],
s3_root: str = None,
output_athena_dataset_name: str = None,
wr_database: str = Config.DATA_LAKE_DB,
wait: bool = True,
task_suffix: str = "",
container_name: str = Config.ECS_MLFLOW_CONTAINER_NAME,
task_definition: str = Config.ECS_MLFLOW_TASK_DEFINITION,
timeout: int = Config.MAX_ITERATIONS,
):
params_string = json.dumps(params)
input_string = json.dumps(input)
env_vars = [
{"name": "MODEL_URI", "value": f"{model_uri}"},
{"name": "PARAMS", "value": f"{params_string}"},
{"name": "INPUT", "value": f"{input_string}"},
{"name": "DATALAKE_DB", "value": f"{wr_database}"},
]
if output_athena_dataset_name:
env_vars.append(
{
"name": "OUTPUT_ATHENA_DATASET_NAME",
"value": f"{output_athena_dataset_name}",
}
)
if s3_root:
env_vars.append({"name": "S3_ROOT", "value": f"{s3_root}"})
task_arn = ECSUtil.run_task(
container_name=container_name,
task_definition=task_definition,
env_vars=env_vars,
task_suffix=task_suffix,
wait=wait,
timeout=timeout,
)
return task_arn
[docs]
@staticmethod
def run_graphify_model(
project_name: str,
github_path: str,
wait: bool = True,
task_suffix: str = "",
container_name: str = Config.ECS_GRAPHIFY_CONTAINER_NAME,
task_definition: str = Config.ECS_GRAPHIFY_TASK_DEFINITION,
timeout: int = Config.MAX_ITERATIONS,
):
env_vars = [
{"name": "PROJECT_NAME", "value": f"{project_name}"},
{"name": "GITHUB_PATH", "value": f"{github_path}"}
]
task_arn = ECSUtil.run_task(
container_name=container_name,
task_definition=task_definition,
env_vars=env_vars,
task_suffix=task_suffix,
wait=wait,
timeout=timeout,
)
return task_arn
[docs]
@staticmethod
def run_athena_to_geojson(
query: str,
title: str,
geometry_column: str = "geometry",
wr_database: str = Config.DATA_LAKE_DB,
wr_workgroup: str = Config.DEFAULT_WORKGROUP,
source_epsg: str = "EPSG:4326",
export_path_template: str = "{title}.geojson",
wait: bool = True,
task_suffix: str = "",
column_mapping: dict[str, str] = {},
label_columns: list[str] = [],
container_name: str = Config.ECS_ATHENA_TO_GEOJSON_CONTAINER_NAME,
task_definition: str = Config.ECS_ATHENA_TO_GEOJSON_TASK_DEFINITION,
):
"""
Run a athena query and save the results as a geojson via ecs
:param query: the query string to run on athena. use `AS` to rename columns
:type query: str
:param title: the title of the job. we use this as output file name too
:type title: str
:param wr_database: the datatable in athena
:type wr_database: str
:param wr_workgroup: the workgroup in athena
:type wr_workgroup: str
:param geometry_column: the column that includes `wkt` format of the data. if it is a comma seperated text we split it to lat and lon and read it as point
:type geometry_column: str
:param source_epsg: the epsg of the source data
:type source_epsg: str
:param export_path_template: the template for the export path
:type export_path_template: str
:param wait: whether or not to wait for the task to finish
:type wait: bool
:param task_suffix: the task suffix, e.g. "large"
:type task_suffix: str
:param column_mapping: the column mapping
:type column_mapping: dict[str, str]
:param container_name: the container name
:type container_name: str
:param task_definition: the task definition
:type task_definition: str
:return: the task arn
:rtype: _type_
"""
# just set the export path
# dev and prod path are the same, this is guaranteed to be in the config now
export_path = (
f"{Config.ECS_MOUNT_PATH}/{export_path_template.format(title=title)}"
)
column_map_string = json.dumps(column_mapping)
task_arn = ECSUtil.run_task(
container_name=container_name,
task_definition=task_definition,
env_vars=[
{"name": "TILE", "value": f"{title}"},
{"name": "TILE_SQL_QUERY", "value": f"{query}"},
{"name": "WR_DATABASE", "value": f"{wr_database}"},
{"name": "WR_WORKGROUP", "value": f"{wr_workgroup}"},
{"name": "GEOMETRY_COLUMN", "value": f"{geometry_column}"},
{"name": "SOURCE_EPSG", "value": f"{source_epsg}"},
{"name": "EXPORT_PATH", "value": f"{export_path}"},
{"name": "COLUMN_MAPPING", "value": f"{column_map_string}"},
{"name": "LABEL_COLUMNS", "value": f"{','.join(label_columns)}"},
],
task_suffix=task_suffix,
wait=wait,
)
return task_arn
[docs]
@staticmethod
def run_geopandas(
query,
title: str,
conn_id: str = Config.READY_DATA_DB,
wait: bool = True,
task_suffix: str = "",
container_name: str = Config.ECS_GEOPANDAS_CONTAINER_NAME,
task_definition: str = Config.ECS_GEOPANDAS_TASK_DEFINITION,
local: bool = False,
):
"""
(More complete and accurate name: run_postgres_to_geojson) Run a Postgres query, and save response as geojson file via ecs
In contrast to run_athena_to_geojson; Neither methods read the web response into a geopandas dataframe
Under the hood: https://github.com/ready/terraform/blob/main/ecs_geopandas.tf and ECSUtil.run_task
:param query: the query to run
:param title: the title of the end file
:param conn_id: the connection id
:param wait: whether or not to wait for the task to finish
:param task_suffix: the task suffix, e.g. "large"
:param container_name: the container name
:param task_definition: the task definition
:param local: whether or not to run the task locally (deprecated)
:return: the task arn
:rtype: _type_
"""
import geopandas as gpd
pg_hook = PostgresHook(postgres_conn_id=conn_id)
conn = pg_hook.get_uri()
eng = pg_hook.get_sqlalchemy_engine()
if local:
geojson_path = f"/tmp/{title}.geojson"
df = gpd.read_postgis(
query,
con=eng,
)
gdf = gpd.GeoDataFrame(df)
gdf.to_file(geojson_path, driver="GeoJSON")
else:
task_arn = ECSUtil.run_task(
container_name=container_name,
task_definition=task_definition,
env_vars=[
{"name": "TILE_TITLE", "value": f"{title}"},
{"name": "CONN", "value": f"{conn}"},
{"name": "TILE_QUERY", "value": f"{query}"},
],
task_suffix=task_suffix,
wait=wait,
)
return task_arn
[docs]
@staticmethod
def run_tippecanoe(
title: str,
description: str = "",
cmd: str = "tippecanoe --temporary-directory={mount_path} -zg -Z0 --extend-zooms-if-still-dropping --drop-densest-as-needed --force -o {mbtiles_path} -l {title} -n {title_description} -N {description} {geojson_path}",
geojson_template: str = "{mount_path}/{title}.geojson",
mbtiles_template: str = "{mount_path}/{title}.mbtiles",
wait: bool = True,
container_name: str = Config.ECS_TIPPECANOE_CONTAINER_NAME,
task_definition: str = Config.ECS_TIPPECANOE_TASK_DEFINITION,
task_suffix: str = "",
local: bool = False,
):
"""
Run a bash command on a geojson file and save the results as an mbtiles file
via a container with tippecanoe installed in ecs
:param title: the title of the file
:param description: the description of the tileset. We will automatically add latest_update to the tileset description for the records
:param cmd: the command to run in ecs, runs in a container with tippecanoe dependencies installed
:param geojson_template: the template for the geojson file path
:param mbtiles_template: the template for the mbtiles file path
:param wait: whether or not to wait for the task to finish
:param container_name: the container name
:param task_definition: the task definition
:param task_suffix: the task suffix, e.g. "large" or "xlarge"
:param local: whether or not to run the task locally (deprecated)
:return: the task arn
:rtype: _type_
"""
geojson_path = geojson_template.format(
title=title, mount_path=Config.ECS_MOUNT_PATH
)
mbtiles_path = mbtiles_template.format(
title=title, mount_path=Config.ECS_MOUNT_PATH
)
current_date_time = Config.pendulum.now("US/Eastern").format(
"YYYY-MM-DD HH:mm:ss zz"
)
current_date_time_short = Config.pendulum.now("US/Eastern").format(
"YYYY-MM-DD@HH:mmzz"
)
title_description = (
f"{title}: {description} (Last_updated: {current_date_time_short})"
if description
else f"{title} (Last_updated: {current_date_time_short})"
)
# single <> double quotes https://stackoverflow.com/questions/6697753/difference-between-single-and-double-quotes-in-bash
description_last_updated = (
f"{description} (Last_updated: {current_date_time})"
if description
else f"Last_updated: {current_date_time}"
)
command = cmd.format(
mount_path=Config.ECS_MOUNT_PATH,
geojson_path=geojson_path,
mbtiles_path=mbtiles_path,
title=f"'{title}'",
description=f"'{description_last_updated}'",
title_description=f"'{title_description}'",
)
if local:
command = command.replace(Config.ECS_MOUNT_PATH, "/tmp")
print(f'Local command: "{command}"')
BashOperator(
task_id=f"local_geopandas_{title}",
bash_command=command,
).execute(),
else:
print(f'Tippecanoe command: "{command}"')
task_arn = ECSUtil.run_task(
container_name=container_name,
task_definition=task_definition,
env_vars=[
{"name": "TGT_COMMAND", "value": f"''{command}''"},
],
task_suffix=task_suffix,
wait=wait,
)
return task_arn
[docs]
@staticmethod
def run_pgdump_command(
cmd: str,
wait: bool = True,
container_name: str = Config.ECS_PGDUMP_CONTAINER_NAME,
task_definition: str = Config.ECS_PGDUMP_TASK_DEFINITION,
):
"""
Run a pgdump command via bash
:param cmd: fully customized pgdump command to run in ecs
:param src_conn_id: the connection id (NOT full conn str!) of the source db
:return: the task arn
:rtype: _type_
"""
print(f"""cur_tgt_cmd: {cmd} """)
task_arn = ECSUtil.run_task(
container_name=container_name,
task_definition=task_definition,
env_vars=[
{"name": "TGT_COMMAND", "value": f"''{cmd}''"},
],
wait=wait,
)
return task_arn
[docs]
@staticmethod
def run_tile_join(
name: str,
description: str,
input_mbtiles_stub: str,
wait: bool = True,
tippecanoe_cmd: 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}'",
container_name: str = Config.ECS_TIPPECANOE_CONTAINER_NAME,
task_definition: str = Config.ECS_TIPPECANOE_TASK_DEFINITION,
task_suffix: str = "",
):
"""
Run a bash command on a set of mbtiles to join them together [under dev]
via a container with tippecanoe installed in ecs
:param name: the name of the output tileset
:param description: the description of the output tileset. We will automatically add latest_update to the tileset description for the records
:param input_mbtiles_stub: basename for input mbtiles files. Will be expanded to find all <basename>_*.mbtiles to join
:param wait: whether or not to wait for the task to finish
:param container_name: the container name
:param task_definition: the task definition
:param task_suffix: the task suffix, e.g. "large", defaults to small
:return: the task arn
:rtype: str
"""
current_date_time = Config.pendulum.now("US/Eastern").format(
"YYYY-MM-DD HH:mm:ss"
)
last_updated = f"(Last_updated: {current_date_time})"
# globbed_command = "find {mount_path} -name '{input_mbtiles_stub}_*.mbtiles' | xargs tile-join -pk -f --overzoom -o {mount_path}/{input_mbtiles_stub}.mbtiles -n '{name}' -N '{description}'"
# globbed_command = "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}'"
description_with_last_updated = f"{description} {last_updated}"
command = tippecanoe_cmd.format(
mount_path=Config.ECS_MOUNT_PATH,
input_mbtiles_stub=input_mbtiles_stub,
name=name,
description=description_with_last_updated,
)
print(f'Tile-join command: "{command}"')
task_arn = ECSUtil.run_task(
container_name=container_name,
task_definition=task_definition,
env_vars=[
{"name": "TGT_COMMAND", "value": f"''{command}''"},
],
task_suffix=task_suffix,
wait=wait,
)
return task_arn
[docs]
def sql_to_tile(
query: str,
title: str,
conn_id: str = Config.READY_DATA_DB,
task_suffix: str = "",
) -> tuple[str, str]:
"""Run a SQL query and generate a tileset from the results. this should probably belong to another interface and have multiple methods of creating tiles e.g. create locally with TilerUtil.create_tile
:param query: the query to run
:type query: str
:param title: the title of the end file
:type title: str
:param conn_id: the connection id, defaults to Config.READY_DATA_DB
:type conn_id: str, optional
:param task_suffix: the task suffix, e.g. "large", defaults to ""
:type task_suffix: str, optional
:return: the arns of the geojson and tippecanoe tasks
:rtype: tuple[str, str]
"""
geo_arn = ECSUtil.run_geopandas(
query=query,
title=title,
conn_id=conn_id,
task_suffix=task_suffix,
)
tippecanoe_arn = ECSUtil.run_tippecanoe(
title=title,
task_suffix=task_suffix,
)
return (geo_arn, tippecanoe_arn)
@staticmethod
def _swap():
"""
restart other service
wait for it to start
lookup IP of new task
list targets
drain existing tgt
attach IP to corresponding tgt group
wait until it's healthy
modify listener to add new tgt group
modify listener to remove old target group
basically deprecated
"""
TASK_GROUP_COMMON_STR = "ts-gl" # assumes this string is in your tasks, task definitions, tgt groups, and load balancers
ALT_STR = "alt" # assumes this string is in alt tasks, task definitions, tgt groups, and load balancers
CLUSTER = "dev" if not Config.PRODUCTION else ""
LOCK_PATH = "/tmp/tsgl-lock"
def check_lock() -> bool:
"""
helper for file path boolean
"""
return os.path.isfile(LOCK_PATH)
def get_tgt_task(tasks, use_alt):
"""
this is kind of a crazy list comprehension sorry
"""
return [
ts_task
for ts_task in tasks
if TASK_GROUP_COMMON_STR in ts_task["group"]
and (
(use_alt and ALT_STR in ts_task["group"])
or (not use_alt and ALT_STR not in ts_task["group"])
)
][0]
use_alt = not check_lock()
import boto3
ecs = boto3.client("ecs")
old_task = get_tgt_task(
ecs.describe_tasks(
cluster=CLUSTER, tasks=ecs.list_tasks(cluster=CLUSTER)["taskArns"]
)["tasks"],
use_alt,
)
old_task_arn = old_task["taskArn"]
old_ip = old_task["containers"][0]["networkInterfaces"][0]["privateIpv4Address"]
print(old_task_arn)
ecs.stop_task(cluster=CLUSTER, task=old_task_arn)
# counter = 0
# while [ # wait for service to bring new task up by itself
# task
# for task in ecs.describe_tasks(cluster=CLUSTER, tasks=ecs.list_tasks(cluster=CLUSTER)['taskArns'])['tasks']
# if TASK_GROUP_COMMON_STR in task['group']
# ]< 2:
# time.sleep(1)
# counter += 1
# if counter >= RETRY_THRESHOLD:
# break
# active_ip = get_tgt_task(
# ecs.describe_tasks(cluster=CLUSTER, tasks=ecs.list_tasks(cluster=CLUSTER)['taskArns'])['tasks'],
# use_alt
# )['containers'][0]['networkInterfaces'][0]['privateIpv4Address']
# elb = boto3.client('elbv2')
# active_target_group_arn = [
# group
# for group in elb.describe_target_groups()['TargetGroups']
# if TASK_GROUP_COMMON_STR in group['TargetGroupName']
# and not group['LoadBalancerArns']
# ][0]['TargetGroupArn']
# elb_arn = [
# each_elb
# for each_elb in elb.describe_load_balancers()['LoadBalancers']
# if TASK_GROUP_COMMON_STR in each_elb['LoadBalancerName'] and (
# Config.PRODUCTION or 'dev' in each_elb['LoadBalancerName']
# )
# ][0]['LoadBalancerArn']
# listener_arn = elb.describe_listeners(LoadBalancerArn=elb_arn)['Listeners'][0]['ListenerArn']
# elb.deregister_targets(
# TargetGroupArn=active_target_group_arn,
# Targets=
# [
# {
# 'Id': old_ip,
# }
# ]
# )
# elb.register_targets(
# TargetGroupArn=active_target_group_arn,
# Targets=
# [
# {
# 'Id': active_ip,
# 'Port': 8080
# }
# ]
# )
# do i need to wait here for a bit?
# elb.modify_listener(
# ListenerArn=listener_arn,
# DefaultActions=[
# {
# 'ForwardConfig': {
# 'TargetGroupStickinessConfig': {
# 'DurationSeconds': 3600,
# 'Enabled': False
# },
# 'TargetGroups': [
# {
# 'TargetGroupArn': active_target_group_arn,
# 'Weight': 1
# }
# ]
# },
# 'Order': 1,
# 'TargetGroupArn': active_target_group_arn,
# 'Type': 'forward'
# }
# ]