Using Asset tag detection for conditional behavior

[clayheaton]

I have a situation where I have an asset that needs to materialize every 15 minutes. It's OK if it skips one (or even two) of these slots. The main restriction is that it cannot materialize if another run is materializing assets that rely on the table that it generates.

My goal was to enable this frequently-updating asset to skip it's materialization if any active runs use the table it creates. Hence, I wanted a way to tag assets so that the frequent-updater could detect if they're running.

Note that I'm self-hosting Dagster and using a docker compose setup with multiple code locations. This works across code locations. I'm using the Snowflake IO Manager, but it should work with any IO Manager.

Here's what I did... I wrote this function (that might need some error-checking... not sure I'm catching all cases) that lives in each code location that needs to run this type of check. I have a shared_functions directory for stuff like this, though I should turn it into a module I can install. That's a story for another day.

import requests
from typing import List, Dict
from dagster.core.storage.pipeline_run import DagsterRunStatus, RunsFilter

DAGSTER_URL = "http://your-dagster-url/graphql"

def check_tags_are_active(tags_to_search: List[Dict[str, str]], context) -> bool:
    """
    Checks if any of the assets in the active Dagster runs have tags that match
    any tag in `tags_to_search`. Returns True if a match is found, False otherwise.
    """
    # Get active Dagster runs
    active_runs = context.instance.get_runs(
        filters=RunsFilter(
            statuses=[
                DagsterRunStatus.STARTED,
                DagsterRunStatus.STARTING,
            ]
        )
    )

    for run in active_runs:
        # Get asset keys either from step_keys_to_execute (for __ASSET_JOB) or from job assets
        if run.job_name == "__ASSET_JOB" and run.step_keys_to_execute:
            asset_keys = run.step_keys_to_execute  # Directly from the run
        else:
            # Construct GraphQL query to fetch the job's asset keys
            query = """
            query GetJobAssets($jobName: String!) {
              jobOrError(jobName: $jobName) {
                ... on Job {
                  assetNodes {
                    assetKey {
                      path
                    }
                  }
                }
              }
            }
            """
            variables = {"jobName": run.job_name}
            response = requests.post(DAGSTER_URL, json={"query": query, "variables": variables})
            if response.status_code != 200:
                print(f"Error fetching job assets for {run.job_name}: {response.status_code}")
                continue

            data = response.json()
            try:
                asset_nodes = data["data"]["jobOrError"]["assetNodes"]
                asset_keys = ["/".join(asset["assetKey"]["path"]) for asset in asset_nodes]
            except KeyError:
                asset_keys = []

        # Fetch the tags for the asset keys
        if asset_keys:
            asset_tags_dict = {}
            for asset_key in asset_keys:
                # GraphQL query to get tags for each asset
                query = f"""
                query {{
                  assetNodeOrError(assetKey: {{ path: ["{asset_key}"] }}) {{
                    ... on AssetNode {{
                      tags {{
                        key
                        value
                      }}
                    }}
                  }}
                }}
                """
                response = requests.post(DAGSTER_URL, json={"query": query})
                if response.status_code != 200:
                    print(f"Error fetching tags for {asset_key}: {response.status_code}")
                    continue

                data = response.json()
                try:
                    tags = data["data"]["assetNodeOrError"]["tags"]
                    formatted_tags = [{tag["key"]: tag["value"]} for tag in tags]
                    asset_tags_dict[asset_key] = formatted_tags
                except KeyError:
                    continue

            # Check if any of the asset tags match the tags_to_search
            for asset_key, tags in asset_tags_dict.items():
                for tag in tags:
                    if any(tag_to_search.items() <= tag.items() for tag_to_search in tags_to_search):
                        return True

    return False

Then what I do is use the tags kwarg in the @asset decorator, like this:

@asset(
    tags={"my_key": "my_value"},
)
def my_asset(context: AssetExecutionContext) -> None:
    """
    If this asset is materializing or part of an active run, don't materialize the frequent-updater.
    """
    # business logic
    return None

In my frequent updater, I can do this:

from shared_functions import check_tags_are_active

@asset(
    output_required=False,
)
def frequent_updater(context: AssetExecutionContext) -> pd.DataFrame | None:
    """
    Creates an important database table.
    Skips materialization if an asset with tags={"my_key": "my_value"} is running.
    """

    # Check if there are any active tags that block this asset's execution
    tags_to_check = [
        {"my_key": "my_value"},
    ]
    active = check_tags_are_active(tags_to_check, context)

    if active:
        context.log.warn("Active tags found. Skipping materialization.")

    else:
        context.log.info("No active tags found. Proceeding with materialization.")

        # Create the DataFrame with the data to store as a table
        df = get_a_dataframe_somewhere()

        # Yield the DataFrame as an Output
        yield Output(df, output_name="result")

It's quite convenient and I can see using it for other purposes. I'd prefer not to have to use the GraphQL API, but I couldn't figure out another way to do it.

[geoHeil]

Why did you not set concurrency limitations for this particular asset?

[clayheaton]

I have dozens of other flows running around the clock, in 8 different code locations, and only a handful of assets in a few of them should block this one. It is very fast to materialize, so resources aren’t much of an issue. This asset materializes as the last step of a longer flow, the rest of which is fetching data from a 3rd party into an important internal tool and which cannot be blocked. The behavior I desired was for this one to conditionally skip only if one of the blocking ones is running. The blocking ones materialize based on an observable source asset with an erratic refresh schedule.