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) ], )