Source code for bhf_smarthealth.utils.workflow_utils
"""This module contains utility functions for common Snakemake workflow actions"""
import re
import logging
import boto3
from pathlib import Path
from dataclasses import dataclass
from enum import StrEnum, auto
from datetime import date, timedelta, datetime
from snakemake.io import expand
from typing import Any
from uuid import uuid4
from bhf_smarthealth.models import BucketNames, PartitionFreqs
from bhf_smarthealth.config import CLEAN_TABLE_REGISTRY
logger = logging.getLogger(__name__)
[docs]
@dataclass
class ValidTargets:
"""Class to hold valid date and type target combinations"""
date: str
type: str
[docs]
class SENTINEL_FLAGS(StrEnum):
"""Configure naming of Snakemake sentinel flags"""
PROCESSING_COMPLETE = auto()
UPLOAD_RAW_COMPLETE = auto()
UPLOAD_INTERMEDIATE_COMPLETE = auto()
DATABASE_LOAD_COMPLETE = auto()
DATASTREAM_SUCCESS = auto()
PIPELINE_ERROR = auto()
PIPELINE_SUCCESS = auto()
[docs]
def read_or_generate_run_id(filepath: Path | str) -> str:
"""Check whether a run_id file has been generated by a snakemake pipeline run, else generate a new run_id file. This is required to ensure persistent run_id state throughout a pipeline run.
Parameters
----------
filepath : Path | str
Path, or string of path, to the lockfile that contains, or should contain, the run_id
Returns
-------
str
String of the generated run_id
Warning
-------
The output file must be deleted by the pipeline onsuccess or onerror to ensure a new run_id is generated.
"""
if isinstance(filepath, str):
filepath = Path(filepath)
# read existing run_id
if filepath.exists():
with open(filepath, "r") as f:
run_id = f.read().strip()
# generate new run_id
else:
filepath.parent.mkdir(parents=True, exist_ok=True) # ensure dir exists
timestamp = datetime.now().strftime("%Y%m%d%H%M")
short_id = uuid4().hex[
:6
] # 0.03% collision chance for 100 runs within 1 minute
run_id = f"{timestamp}_{short_id}"
with open(filepath, "w") as f:
f.write(run_id)
return run_id
[docs]
def get_bucket_credentials(profile: str) -> dict[str, str]:
"""Retrieve credentials from the s3 object storage associated with the specified 'profile' as configured in the .aws/config file
Parameters
----------
profile : str
String of profile name as configured through AWS CLI
Returns
-------
dict[str, str]
Dict containing bucket credentials
Raises
------
RuntimeError
Raised if error retrieveing credentials for the specified profile
"""
try:
creds = (
boto3.Session(profile_name=profile)
.get_credentials()
.get_frozen_credentials()
)
return {"access_key": creds.access_key, "secret_key": creds.secret_key}
except Exception as e:
raise RuntimeError(f"Could not load credentials for profile {profile}: {e}")
[docs]
def validate_date_to_process(date_to_process: str | date) -> date:
"""Validates that the date specified in the Snakemake config.yaml file is a valid ISO date or the string "today".
Parameters
----------
date_to_process : str | date
The DATE_TO_PROCESS config specified in the Snakemake config file
Returns
------
validated_date : date
A date object of the validated date to process
Raises
------
ValueError
Raised if the specified date is not a valid ISO 8601 date or 'today'
"""
if date_to_process == "today":
validated_date = date.today()
elif isinstance(date_to_process, date):
validated_date = date_to_process
else:
raise ValueError(
f"Specified DATE_TO_PROCESS {date_to_process} in config is not valid. Either enter 'today' or a date in the form 'YYYY-MM-DD'"
)
return validated_date
[docs]
def generate_dates(start_date: date, deltas: list[int]) -> list[str]:
"""Generates dates for days before the given start date, as specified by the values in 'deltas'
'deltas' must contain values that are greater than 0, values less than or equal to 0 will be ignored
Duplicate values in deltas will also be ignored.
Parameters
----------
start_date : date
Date to generate other dates relative to
deltas : list[int]
List of deltas representing the interval (in days) for dates prior to the start date
Returns
-------
list[str]
List of ISO8601 formatted date strings
"""
if len(deltas) > 1 and not all(delta > 0 for delta in deltas):
logger.warning(
"values less than 1 in deltas are invalid and will be ignored when generating dates"
)
if len(deltas) == 1 and deltas[0] == 0:
return [start_date.strftime("%Y%m%d")]
dates = (start_date - timedelta(delta) for delta in deltas if delta > 0)
# Use set comprehension to remove duplicate values
dates_str = {dt.strftime("%Y%m%d") for dt in dates}
logger.debug(f"Generated {len(list(dates_str))} dates from {len(deltas)} deltas")
return list(dates_str)
[docs]
def get_valid_day_deltas(
data_type, all_dates, targets, wildcards, checkpoints, storage_provider
):
"""
1. Runs the checkpoint to list archive contents.
2. Parses the text file to find available dates.
3. Returns targets only for valid dates.
"""
logger.info("Inspecting archive contents to determine valid deltas to process")
# Force run checkpoint
checkpoint_output = checkpoints.inspect_archive_contents.get(**wildcards).output[0]
logger.debug(f"Checking for {all_dates} in {checkpoint_output}")
# Read the created 7zip listing file
with open(checkpoint_output, "r") as f:
archive_content = f.read()
# Check which dates exist in the archive listing
available_dates = {
d for d in all_dates if re.search(rf"{d}.*/{data_type}", archive_content)
}
# Compare with requested dates
requested_set = set(all_dates)
valid_dates = requested_set.intersection(available_dates)
missing_dates = requested_set.difference(available_dates)
if missing_dates:
logger.warning(
f"The following requested dates were not found in the archive and will be skipped: {sorted(missing_dates)}"
)
logger.info(
f"The following requested dates will be processed: {sorted(valid_dates)}"
)
# Return targets
target_files = expand(targets, date=valid_dates)
return storage_provider(target_files)
[docs]
def get_valid_clean_targets(
config: dict[str, Any],
) -> list[ValidTargets]:
"""Return a list of ValidTargets objects that comprise the valid table,date combinations determined from the clean table registry and configured day deltas for each stream.
Parameters
----------
config : dict[str, Any]
Snakemake config dict (should contain the streams and their associated day deltas)
Returns
-------
list[ValidTargets]
List of ValidTargets
"""
active_streams = [
stream
for stream in config["STREAMS"]
if config["STREAMS"][stream]["ACTIVE"] and stream != "CLEAN"
]
logger.info(f"The following streams are active: {sorted(active_streams)}")
valid_targets = []
for stream in active_streams:
target_tables = [
table
for table, properties in CLEAN_TABLE_REGISTRY.items()
if stream.lower() == properties.source_stream
and properties.bucket == BucketNames.CLEAN
]
# currently assumes all day deltas are valid
# TODO: enable state-aware dates i.e. know if day deltas were missing from Daiser archive
dates = generate_dates(
config["DATE_TO_PROCESS"], config["STREAMS"][stream]["DAY_DELTAS"]
)
valid_targets.extend(
[
ValidTargets(date=date, type=table)
for date in dates
for table in target_tables
]
)
return valid_targets
[docs]
def generate_s3_key(wildcards):
"""
Parses snakemake wildcards to generate file key depending on table partitioning
"""
table_name = wildcards.type
if CLEAN_TABLE_REGISTRY[table_name].partition_freq == PartitionFreqs.DAILY:
date_str = wildcards.date
date_obj = datetime.strptime(date_str, "%Y%m%d")
year = date_obj.year
month = date_obj.month
day = date_obj.day
return f"clean/{table_name}/{year}/{month}/{day}/clean_{table_name}_{date_str}.parquet"
if not CLEAN_TABLE_REGISTRY[table_name].partition_freq:
return f"clean/{table_name}/clean_{table_name}.parquet"
