def _get_output_asset_materializations(
asset_key: AssetKey,
asset_partitions: AbstractSet[str],
output: Union[Output, DynamicOutput],
output_def: OutputDefinition,
io_manager_metadata_entries: List[Union[MetadataEntry,
PartitionMetadataEntry]],
) -> Iterator[AssetMaterialization]:
all_metadata = output.metadata_entries + io_manager_metadata_entries
if asset_partitions:
metadata_mapping: Dict[str,
List[Union[MetadataEntry,
PartitionMetadataEntry]]] = {
partition: []
for partition in asset_partitions
}
for entry in all_metadata:
# if you target a given entry at a partition, only apply it to the requested partition
# otherwise, apply it to all partitions
if isinstance(entry, PartitionMetadataEntry):
if entry.partition not in asset_partitions:
raise DagsterInvariantViolationError(
f"Output {output_def.name} associated a metadata entry ({entry}) with the partition "
f"`{entry.partition}`, which is not one of the declared partition mappings ({asset_partitions})."
)
metadata_mapping[entry.partition].append(entry.entry)
else:
for partition in metadata_mapping.keys():
metadata_mapping[partition].append(entry)
for partition in asset_partitions:
with warnings.catch_warnings():
warnings.simplefilter("ignore", category=ExperimentalWarning)
yield AssetMaterialization(
asset_key=asset_key,
partition=partition,
metadata_entries=metadata_mapping[partition],
)
else:
for entry in all_metadata:
if isinstance(entry, PartitionMetadataEntry):
raise DagsterInvariantViolationError(
f"Output {output_def.name} got a PartitionMetadataEntry ({entry}), but "
"is not associated with any specific partitions.")
with warnings.catch_warnings():
warnings.simplefilter("ignore", category=ExperimentalWarning)
yield AssetMaterialization(asset_key=asset_key,
metadata_entries=all_metadata)
def materialization_from_data(data):
if data.get("asset_key"):
return AssetMaterialization(
asset_key=data["asset_key"],
description=data.get("description"),
metadata_entries=list(event_metadata_entries(data.get("metadataEntries")) or []),
)
return Materialization(
label=data["label"],
description=data.get("description"), # enforce?
metadata_entries=list(event_metadata_entries(data.get("metadataEntries")) or []),
)
def materialization_from_data(data):
if data.get('asset_key'):
return AssetMaterialization(
asset_key=data['asset_key'],
description=data.get('description'),
metadata_entries=list(event_metadata_entries(data.get('metadataEntries')) or []),
)
return Materialization(
label=data['label'],
description=data.get('description'), # enforce?
metadata_entries=list(event_metadata_entries(data.get('metadataEntries')) or []),
)
def _stats_records(run_id):
now = time.time()
return [
_event_record(run_id, "A", now - 325, DagsterEventType.STEP_START),
_event_record(
run_id,
"A",
now - 225,
DagsterEventType.STEP_SUCCESS,
StepSuccessData(duration_ms=100000.0),
),
_event_record(run_id, "B", now - 225, DagsterEventType.STEP_START),
_event_record(
run_id,
"B",
now - 175,
DagsterEventType.STEP_FAILURE,
StepFailureData(error=None, user_failure_data=None),
),
_event_record(run_id, "C", now - 175, DagsterEventType.STEP_START),
_event_record(run_id, "C", now - 150, DagsterEventType.STEP_SKIPPED),
_event_record(run_id, "D", now - 150, DagsterEventType.STEP_START),
_event_record(
run_id,
"D",
now - 125,
DagsterEventType.STEP_MATERIALIZATION,
StepMaterializationData(AssetMaterialization(asset_key="mat_1")),
),
_event_record(
run_id,
"D",
now - 100,
DagsterEventType.STEP_EXPECTATION_RESULT,
StepExpectationResultData(
ExpectationResult(success=True, label="exp 1")),
),
_event_record(
run_id,
"D",
now - 75,
DagsterEventType.STEP_MATERIALIZATION,
StepMaterializationData(AssetMaterialization(asset_key="mat_2")),
),
_event_record(
run_id,
"D",
now - 50,
DagsterEventType.STEP_EXPECTATION_RESULT,
StepExpectationResultData(
ExpectationResult(success=False, label="exp 2")),
),
_event_record(
run_id,
"D",
now - 25,
DagsterEventType.STEP_MATERIALIZATION,
StepMaterializationData(AssetMaterialization(asset_key="mat_3")),
),
_event_record(run_id, "D", now, DagsterEventType.STEP_SUCCESS,
StepSuccessData(duration_ms=150000.0)),
]
def test_materialization_assign_label_from_asset_key():
mat = AssetMaterialization(asset_key=AssetKey(['foo', 'bar']))
assert mat.label == 'foo.bar'
def test_materialization():
assert isinstance(AssetMaterialization('foo', 'foo.txt'),
AssetMaterialization)
def test_materialization_assign_label_from_asset_key():
mat = AssetMaterialization(asset_key=AssetKey(["foo", "bar"]))
assert mat.label == "foo.bar"
def test_materialization():
assert isinstance(AssetMaterialization("foo", "foo.txt"), AssetMaterialization)
def _store_output(
step_context: StepExecutionContext,
step_output_handle: StepOutputHandle,
output: Union[Output, DynamicOutput],
input_lineage: List[AssetLineageInfo],
) -> Iterator[DagsterEvent]:
output_def = step_context.solid_def.output_def_named(step_output_handle.output_name)
output_manager = step_context.get_io_manager(step_output_handle)
output_context = step_context.get_output_context(step_output_handle)
manager_materializations = []
manager_metadata_entries: List[Union[PartitionMetadataEntry, MetadataEntry]] = []
# output_manager.handle_output is either a generator function, or a normal function with or
# without a return value. In the case that handle_output is a normal function, we need to
# catch errors should they be raised before a return value. We can do this by wrapping
# handle_output in a generator so that errors will be caught within iterate_with_context.
if not inspect.isgeneratorfunction(output_manager.handle_output):
def _gen_fn():
gen_output = output_manager.handle_output(output_context, output.value)
for event in output_context.consume_events():
yield event
if gen_output:
yield gen_output
handle_output_gen = _gen_fn()
else:
handle_output_gen = output_manager.handle_output(output_context, output.value)
for elt in iterate_with_context(
lambda: solid_execution_error_boundary(
DagsterExecutionHandleOutputError,
msg_fn=lambda: (
f'Error occurred while handling output "{output_context.name}" of '
f'step "{step_context.step.key}":'
),
step_context=step_context,
step_key=step_context.step.key,
output_name=output_context.name,
),
handle_output_gen,
):
for event in output_context.consume_events():
yield event
manager_metadata_entries.extend(output_context.consume_logged_metadata_entries())
if isinstance(elt, DagsterEvent):
yield elt
elif isinstance(elt, AssetMaterialization):
manager_materializations.append(elt)
elif isinstance(elt, (MetadataEntry, PartitionMetadataEntry)):
experimental_functionality_warning(
"Yielding metadata from an IOManager's handle_output() function"
)
manager_metadata_entries.append(elt)
else:
raise DagsterInvariantViolationError(
f"IO manager on output {output_def.name} has returned "
f"value {elt} of type {type(elt).__name__}. The return type can only be "
"one of AssetMaterialization, MetadataEntry, PartitionMetadataEntry."
)
for event in output_context.consume_events():
yield event
manager_metadata_entries.extend(output_context.consume_logged_metadata_entries())
# do not alter explicitly created AssetMaterializations
for materialization in manager_materializations:
if materialization.metadata_entries and manager_metadata_entries:
raise DagsterInvariantViolationError(
f"When handling output '{output_context.name}' of {output_context.solid_def.node_type_str} '{output_context.solid_def.name}', received a materialization with metadata, while context.add_output_metadata was used within the same call to handle_output. Due to potential conflicts, this is not allowed. Please specify metadata in one place within the `handle_output` function."
)
if manager_metadata_entries:
materialization = AssetMaterialization(
asset_key=materialization.asset_key,
description=materialization.description,
metadata_entries=manager_metadata_entries,
partition=materialization.partition,
tags=materialization.tags,
metadata=None,
)
yield DagsterEvent.asset_materialization(step_context, materialization, input_lineage)
asset_key, partitions = _asset_key_and_partitions_for_output(
output_context, output_def, output_manager
)
if asset_key:
for materialization in _get_output_asset_materializations(
asset_key,
partitions,
output,
output_def,
manager_metadata_entries,
):
yield DagsterEvent.asset_materialization(step_context, materialization, input_lineage)
yield DagsterEvent.handled_output(
step_context,
output_name=step_output_handle.output_name,
manager_key=output_def.io_manager_key,
metadata_entries=[
entry for entry in manager_metadata_entries if isinstance(entry, MetadataEntry)
],
)