result_lineage.py

# /// script
# dependencies = ["prefect"]
# ///

import webbrowser
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any

from prefect import flow, task
from prefect.filesystems import LocalFileSystem
from prefect.runtime import flow_run as flow_run_context
from prefect.settings import PREFECT_UI_URL

if not (STORAGE_PATH := Path(__file__).parent / "testing").exists():
    STORAGE_PATH.mkdir(parents=True, exist_ok=True)

STORAGE = LocalFileSystem(basepath=str(STORAGE_PATH))
STORAGE.save("test-local-storage", overwrite=True)


@dataclass
class Asset:
    id: str
    type: str
    discovery_context: dict[str, Any] = field(default_factory=dict)


@task(result_storage_key="{parameters[domain]}-ip-{parameters[run_id]}")
def discover_ip(domain: str, run_id: str) -> Asset:
    """Simulate discovering an IP address through DNS resolution of a domain."""
    return Asset(
        id="93.184.216.34",
        type="ip",
        discovery_context={
            "run_id": run_id,
            "method": "dns_resolution",
            "discovered_at": datetime.now().isoformat(),
            "discovered_by": "a_record",
            "source_domain": domain,
        },
    )


@task(result_storage_key="{parameters[domain]}-subdomain-{parameters[run_id]}")
def discover_subdomain(domain: str, run_id: str) -> Asset:
    """Simulate discovering a subdomain through DNS resolution of a domain."""
    return Asset(
        id="api.example.com",
        type="subdomain",
        discovery_context={
            "run_id": run_id,
            "method": "dns_resolution",
            "discovered_at": datetime.now().isoformat(),
            "discovered_by": "cname_record",
            "source_domain": domain,
        },
    )


@flow(persist_result=True, result_storage=STORAGE)
def simulate_asset_discovery(domain: str = "example.com"):
    """Simulate the discovery of assets through DNS resolution."""
    run_id = str(flow_run_context.id)[:3]  # some arbitrary id for bookkeeping
    ip = discover_ip(domain, run_id)
    subdomain = discover_subdomain(domain, run_id)
    print(ip)
    print(subdomain)


if __name__ == "__main__":
    simulate_asset_discovery()
    webbrowser.open_new_tab(f"{PREFECT_UI_URL}/explore")