Integrating Prefect Hooks for Lineage Tracking

Part of: Workflow Hooks in Python Pipelines

Integrating Prefect hooks for lineage tracking requires attaching state-change callbacks to tasks and flows that capture input/output URIs, spatial reference identifiers, transformation parameters, and execution timestamps before persisting them to a structured provenance store. In geospatial pipelines, this means intercepting on_completion, on_failure, and on_running events to emit immutable audit records without modifying core data transformation logic. By decoupling metadata capture from business operations, GIS data stewards and compliance officers can satisfy ISO 19115 metadata requirements, agency data governance mandates, and W3C PROV-O provenance standards while maintaining pipeline performance. For broader architectural context, see how Workflow Hooks in Python Pipelines enable decoupled observability across distributed ETL systems.

How State-Change Callbacks Capture Provenance

Prefect 2.x exposes hook registration at both the task and flow level. Hooks execute synchronously within the worker process immediately after a state transition, granting direct access to the run context, input parameters, and the final State object. This design eliminates the need for intrusive try/except blocks or manual logging calls scattered throughout spatial transformation code. When building Python Automation & Pipeline Integration strategies for raster/vector ETL, hook-based lineage capture ensures that every coordinate transformation, clip operation, or format conversion automatically generates a machine-readable audit trail.

Key advantages of hook-driven capture:

  • Zero-touch instrumentation: Business logic remains pure; metadata extraction happens externally.
  • State-aware execution: Hooks receive the exact State object, enabling conditional routing for success, failure, or cancellation.
  • Context-rich payloads: get_run_context() provides flow/task IDs, parameters, and upstream dependencies without global state pollution.
  • Framework-native reliability: Hooks are managed by the Prefect engine, guaranteeing execution even when tasks raise unhandled exceptions.
Prefect state transitions and the hooks each fires Pending to Running to Completed or Failed or Crashed, with the hook fired at each transition and a note that Crashed bypasses on_failure. Pending Running on_running Completed on_completion Failed on_failure Crashed on_crashed — NOT on_failure infra died, pod evicted Registering only on_failure leaves every infrastructure kill unrecorded — the exact runs an incident review needs.

The Crashed state is the one that catches teams migrating from a simpler try/except mental model. Prefect distinguishes a task that raised — Failed — from one whose infrastructure disappeared beneath it, and the two fire different hooks. A pipeline running on spot instances or a busy Kubernetes cluster will accumulate Crashed runs steadily, and a hook set that registers only on_completion and on_failure records none of them. Those are precisely the runs where an incident reviewer asks what the task had touched before it vanished.

Register on_crashed alongside the other two and emit a record with an explicit status, no output digest, and whatever inputs were known at on_running. That is an honest partial record, and it is far more useful than the silence that otherwise results.

Core Implementation Pattern

The following implementation demonstrates how to capture geospatial lineage metadata using execution context, parameter inspection, and state introspection. Prefect 2.x requires hooks to accept a single State argument and run synchronously within the task/flow lifecycle.

import json
import os
from datetime import datetime, timezone
from typing import Any, Dict, List
from prefect import flow, task, get_run_logger
from prefect.context import get_run_context
from prefect.states import State

# Replace with your agency's lineage API or metadata catalog
LINEAGE_REGISTRY: List[Dict[str, Any]] = []

def capture_geospatial_lineage(state: State) -> None:
    """Prefect hook that extracts execution context and writes lineage records."""
    ctx = get_run_context()
    logger = get_run_logger()

    # Safely extract parameters — available on flow_run context; guard with getattr
    params = getattr(ctx, "parameters", {}) or {}

    lineage_entry = {
        "execution_id": str(getattr(getattr(ctx, "flow_run", None), "id", "unknown")),
        "component_name": getattr(
            getattr(ctx, "task_run", None), "name",
            getattr(getattr(ctx, "flow", None), "name", "unknown")
        ),
        "state": state.type.value,
        "recorded_at": datetime.now(timezone.utc).isoformat(),
        "input_sources": params.get("input_uris", []),
        "output_destination": params.get("output_uri", ""),
        "spatial_reference": params.get("crs", "UNDEFINED"),
        "operation_type": params.get("operation", "unknown"),
        "compliance_framework": os.getenv("AGENCY_LINEAGE_TAG", "ISO_19115_CORE")
    }

    LINEAGE_REGISTRY.append(lineage_entry)
    logger.info(
        "Lineage captured for %s [%s]",
        lineage_entry["component_name"], state.type.value
    )

@task(
    on_completion=[capture_geospatial_lineage],
    on_failure=[capture_geospatial_lineage]
)
def clip_raster(input_uris: List[str], output_uri: str, crs: str, operation: str) -> str:
    logger = get_run_logger()
    logger.info("Clipping %d raster(s) to %s (%s)", len(input_uris), output_uri, crs)
    # Production: integrate rasterio/gdal processing here
    return output_uri

@flow(
    name="agency_spatial_lineage_flow",
    on_completion=[capture_geospatial_lineage],
    on_failure=[capture_geospatial_lineage]
)
def run_geospatial_etl(raw_uris: List[str], processed_uri: str, target_crs: str) -> None:
    clip_raster(
        input_uris=raw_uris,
        output_uri=processed_uri,
        crs=target_crs,
        operation="raster_clip"
    )

if __name__ == "__main__":
    run_geospatial_etl(
        raw_uris=["s3://bucket/raw/aoi.tif"],
        processed_uri="s3://bucket/processed/aoi_clipped.tif",
        target_crs="EPSG:4326"
    )

Configuration & Performance Constraints

When deploying this pattern at scale, consider the following architectural constraints:

  • Hook Execution Order: Prefect runs hooks synchronously in registration order. If you attach multiple callbacks, ensure they are idempotent and avoid blocking network I/O. For high-throughput pipelines, batch lineage writes or push payloads to an async message queue.
  • Context Availability: get_run_context() behaves differently at the flow vs. task level. Flow-level hooks receive flow_run context, while task-level hooks include task_run metadata. Always guard attribute access with getattr() to prevent AttributeError during dry runs or state retries.
  • State Filtering: Not all state transitions warrant lineage records. Filter out CANCELED or RETRYING states if your compliance framework only requires final outcomes. Use state.is_completed() or state.is_failed() for precise control.
  • Parameter Serialization: Prefect automatically serializes parameters, but complex objects (e.g., geopandas.GeoDataFrame) may fail JSON encoding. Pass URIs, CRS strings, and primitive types to hooks, and resolve heavy objects inside the task body. Refer to the official Prefect Task Hooks documentation for lifecycle guarantees.

Getting Run Context Without Threading It Through

Prefect exposes the active run through prefect.context.get_run_context(), which means a hook does not need the flow to pass it deployment identifiers, retry counts or parameters explicitly. That convenience carries one sharp edge: the call raises outside a run context, so a hook unit-tested in isolation fails on the very line that makes it useful in production.

Making a context-reading hook testable A direct-reading hook is untestable outside a run; splitting into a pure function plus a thin context-reading adapter keeps both properties. Hook reads context inline Works in production Raises in every unit test Adapter + pure builder adapter: reads context, calls builder builder: plain args → record dict Test the builder exhaustively; smoke-test the adapter once inside a real flow run. All the logic worth testing lives in the half that needs no orchestrator. The same split works for Airflow listeners, Dagster sensors, and anything else with ambient run state.

This split is worth making early because the builder is where the interesting decisions live — which fields are mandatory, how CRS is read back off the output, what a partial record looks like — and those decisions deserve dense, fast tests. The adapter should be short enough to review by eye: fetch the context, pull four or five values, delegate. Once it is that thin, one integration test that runs a trivial flow and asserts a record appeared is sufficient coverage for it.

Standards Compliance Mapping

Geospatial agencies must align automated lineage capture with established metadata standards. The hook payload above maps directly to:

  • ISO 19115-1:2014: input_sources and output_destination populate the lineage section (LI_Lineage), while recorded_at satisfies process step timestamps.
  • W3C PROV-O: execution_id acts as the prov:Activity identifier, and component_name links to prov:Entity derivatives. For formal validation, export the registry to PROV-N or JSON-LD. Consult the W3C PROV-O specification for exact property mappings and ontology alignment.
  • Agency Governance: Tag records with environment variables (e.g., AGENCY_LINEAGE_TAG) to route metadata to FedRAMP-compliant catalogs or internal data dictionaries.

Flow-Level Versus Task-Level Hooks

What each hook level can answer A flow run containing three tasks, with the flow hook capturing the run envelope and task hooks capturing individual derivations. flow hook — run id, parameters, deployment, total duration task: reproject CRS pair, digests task hook task: clip boundary version task hook task: publish output digest, target task hook Flow hook alone answers: "did the run succeed, and when?" Task hooks answer: "what happened to THIS dataset?"

Register both, and give the flow hook a narrow job. Its value is the envelope — a single run identifier every task record can reference, plus the parameters the run was launched with — and it should not attempt to summarise what the tasks did, because it does not know. Teams that instrument only at flow level end up with lineage at the granularity of “the nightly pipeline ran”, which cannot answer any question about a specific dataset.

The reverse mistake is subtler: instrumenting only at task level loses the parameters the flow was invoked with, which are frequently the thing that distinguishes two otherwise identical runs. A backfill for a different date, or a run with a different boundary version, looks the same at task level unless the flow’s parameters were captured exactly once at the flow level and then referenced by every task record beneath it.

Troubleshooting & Edge Cases

  • Missing Parameters in Hooks: If ctx.parameters returns empty, verify that arguments are explicitly passed as keyword arguments. Prefect’s parameter extraction relies on the function signature matching the invocation.
  • Hook Exceptions Silencing Failures: Unhandled exceptions inside a hook can mask the original task error. Wrap lineage logic in try/except and log failures separately to preserve pipeline observability.
  • Duplicate Records on Retries: Prefect retries tasks by creating new runs. If your lineage store lacks deduplication, you will see multiple entries for the same logical operation. Include state.name and the task run’s run_count (available via ctx.task_run.run_count) in the payload to track retry lineage accurately.
  • Cross-Flow Dependencies: Task-level hooks only see the immediate task context. To capture upstream/downstream flow relationships, attach a flow-level hook that aggregates child run IDs or use Prefect’s artifact system to link execution graphs.

Next Steps

Integrating Prefect hooks for lineage tracking transforms opaque spatial ETL processes into auditable, standards-compliant workflows. By intercepting state transitions at the framework level, teams can enforce ISO 19115 and PROV-O compliance without sacrificing pipeline velocity. Start with synchronous callbacks for validation, then graduate to batched metadata ingestion as execution volumes grow. Validate your output against a PROV-O validator before routing records to production catalogs.