Source code for operators.tile.list_tile
import re
from typing import List, Optional, Literal
from airflow.utils.decorators import apply_defaults
from airflow.utils.context import Context
from airflow.exceptions import AirflowException
from config import Config
from util import Util
from operators.tile.base_tile import BaseTileOperator
[docs]
class ListTileOperator(BaseTileOperator):
"""
Operator for listing tiles based on input selection method (titles or regex).
This operator can scan both the tileserver and S3 for available tiles.
Usage tip: Enter regex `.*` to list all tiles.
:param input_select: Method to select tiles - either "titles" or "regex"
:param regex_pattern: Regex pattern to match tile names (used when input_select is "regex")
:param search_s3: Whether to search in S3 instead of tileserver
"""
ui_color = "#5DADE2" # Blue shade
ui_fgcolor = "#FFFFFF"
custom_operator_name = "ListTileOperator"
template_fields = [
"src_tiles",
"src_env",
"input_select",
"regex_pattern",
"search_s3",
]
@apply_defaults
def __init__(
self,
input_select: Literal["titles", "regex"] = "titles",
regex_pattern: Optional[str] = None,
search_s3: bool = False,
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.input_select = input_select
self.regex_pattern = regex_pattern
self.search_s3 = search_s3
[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" input_select: {self.input_select} (type: {type(self.input_select).__name__})"
)
self.log.info(
f" regex_pattern: {self.regex_pattern} (type: {type(self.regex_pattern).__name__})"
)
self.log.info(
f" search_s3: {self.search_s3} (type: {type(self.search_s3).__name__})"
)
# Let parent class 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" input_select: {self.input_select} (type: {type(self.input_select).__name__})"
)
self.log.info(
f" regex_pattern: {self.regex_pattern} (type: {type(self.regex_pattern).__name__})"
)
self.log.info(
f" search_s3: {self.search_s3} (type: {type(self.search_s3).__name__})"
)
# Process src_tiles
if (
isinstance(self.src_tiles, list)
and len(self.src_tiles) == 1
and not self.src_tiles[0]
):
# If src_tiles is a list of one empty element, set it to an empty list
self.src_tiles = []
self.log.info("🔄 src_tiles was an empty list, now set to []")
elif (
isinstance(self.src_tiles, list)
and self.src_tiles
and isinstance(self.src_tiles[0], list)
):
# handle nested lists from XCom
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}"
)
# 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" src_env: {self.src_env} (type: {type(self.src_env).__name__})"
)
self.log.info(
f" input_select: {self.input_select} (type: {type(self.input_select).__name__})"
)
self.log.info(
f" regex_pattern: {self.regex_pattern} (type: {type(self.regex_pattern).__name__})"
)
self.log.info(
f" search_s3: {self.search_s3} (type: {type(self.search_s3).__name__})"
)
# Validate input parameters after template rendering
if self.input_select == "titles" and not self.src_tiles:
raise AirflowException(
"No tile titles provided when using 'titles' selection method"
)
elif self.input_select == "regex" and not self.regex_pattern:
raise AirflowException(
"No regex pattern provided when using 'regex' selection method"
)
[docs]
def execute(self, context: Context) -> List[str]:
"""
Execute the tile listing operation.
:param context: The Airflow task context
:return: A list of matched tile paths
"""
self.log.info(f"Listing tiles with method: {self.input_select}")
self.log.info(f"Titles to search: {self.src_tiles}")
self.log.info(f"Source environment: {self.src_env}")
self.log.info(f"Search location: {'S3' if self.search_s3 else 'Tileserver'}")
# Get all available tiles
all_tiles = self._get_available_tiles()
self.log.info(f"Found {len(all_tiles)} tiles in total")
# Filter tiles based on selection method
matched_tiles = self._filter_tiles(all_tiles)
self.log.info(
f"Matched {len(matched_tiles)} tiles using {self.input_select} method"
)
# Format tile paths according to the expected format
formatted_tiles = self._format_tile_paths(matched_tiles)
# Display the list of matched tiles
self.log.info(f"Matched {len(formatted_tiles)} tiles:")
for i, path in enumerate(formatted_tiles, 1):
self.log.info(f" {i}. {path}")
return formatted_tiles
def _get_available_tiles(self) -> List[str]:
"""
Get all available tiles from either S3 or the tileserver.
:return: A list of available tile names
"""
if self.search_s3:
return self._list_s3_tiles()
else:
return self._list_tileserver_tiles()
def _list_tileserver_tiles(self) -> List[str]:
"""
List tiles available on the tileserver.
:return: List of tile names on tileserver
"""
# Get environment name from config mappings
for env, config in Config.TILE_ENVS_TS_AND_S3.items():
if config["ts"] == self.src_env:
env_name = env
break
else:
env_name = self.src_env # Fallback to src_env if not found
# Use Util to get tileset info for the environment
try:
tiles_info = Util.get_tileset_info(env_name)
return list(tiles_info.keys())
except Exception as e:
self.log.error(f"Error getting tileset info: {e}")
raise AirflowException(f"Failed to list tiles on tileserver: {e}")
def _list_s3_tiles(self) -> List[str]:
"""
List tiles available in S3 using a simple AWS CLI command.
:return: List of tile names in S3
"""
# Get S3 parent folder for the environment based on TILE_ENVS_TS_AND_S3
s3_parent_folder = None
for _, config in Config.TILE_ENVS_TS_AND_S3.items():
if config["s3"] == self.src_env:
s3_parent_folder = config["s3"]
break
if not s3_parent_folder:
self.log.error(
f"Cannot find S3 folder mapping for environment {self.src_env}"
)
self.log.info(f"Available mappings: {Config.TILE_ENVS_TS_AND_S3}")
raise AirflowException(
f"Cannot find S3 folder for environment {self.src_env}"
)
s3_prefix = f"s3://{Config.ARCHIVE_TILES_S3_BUCKET}/{s3_parent_folder}/"
self.log.info(f"Listing tiles from S3 bucket at: {s3_prefix}")
try:
# Use run_command to execute AWS CLI
cmd = f"aws s3 ls {s3_prefix} | grep \\.mbtiles || true"
output = self.run_command(
cmd, f"list mbtiles files in S3 bucket", check_output=True
)
all_tiles = []
if output:
for line in output.splitlines():
# AWS S3 ls output format: 2025-04-28 12:34:56 12345 filename.mbtiles
parts = line.strip().split()
if len(parts) >= 4: # Format: date time size filename
filename = " ".join(parts[3:]).replace(
" ", "_"
) # Handle filenames with spaces
if filename.endswith(".mbtiles"):
all_tiles.append(filename)
self.log.info(f"Found {len(all_tiles)} .mbtiles files in S3")
return all_tiles
except Exception as e:
self.log.error(f"Error listing S3 tiles: {e}")
raise AirflowException(f"Failed to list S3 objects: {e}")
def _filter_tiles(self, all_tiles: List[str]) -> List[str]:
"""
Filter tiles based on the selected input method.
:param all_tiles: List of all available tiles
:return: List of matched tiles
"""
if self.input_select == "titles":
# Filter by specific titles
src_tile_names = [
tile.split("/")[-1].replace(".mbtiles", "")
for tile in (self.src_tiles or [])
]
# Match either exact name or name with .mbtiles extension
matched_tiles = [
tile_id
for tile_id in all_tiles
if tile_id in src_tile_names
or tile_id.replace(".mbtiles", "") in src_tile_names
]
self.log.info(
f"Matched {len(matched_tiles)} tiles from {len(src_tile_names)} provided titles"
)
return matched_tiles
elif self.input_select == "regex":
# Filter by regex pattern
try:
pattern = re.compile(self.regex_pattern)
matched_tiles = [
tile_id for tile_id in all_tiles if pattern.search(tile_id)
]
self.log.info(
f"Pattern '{self.regex_pattern}' matched {len(matched_tiles)} tiles"
)
return matched_tiles
except re.error as e:
raise AirflowException(
f"Invalid regex pattern '{self.regex_pattern}': {e}"
)
return []
def _format_tile_paths(self, matched_tiles: List[str]) -> List[str]:
"""
Format matched tile names into full paths based on search location.
:param matched_tiles: List of matched tile names
:return: List of formatted tile paths
"""
formatted_paths = []
for tile_id in matched_tiles:
# Make sure tile has .mbtiles extension
if not tile_id.endswith(".mbtiles"):
tile_id = f"{tile_id}.mbtiles"
if self.search_s3:
# Format S3 path
for env, config in Config.TILE_ENVS_TS_AND_S3.items():
if config["ts"] == self.src_env:
s3_parent_folder = config["s3"]
s3_path = f"s3://{Config.ARCHIVE_TILES_S3_BUCKET}/{s3_parent_folder}/{tile_id}"
formatted_paths.append(s3_path)
break
else:
# Format tileserver path
formatted_paths.append(f"{Config.DEFAULT_TS_FLAT_FOLDER}/{tile_id}")
return formatted_paths