def test_basic_yield_multiple_materializations():
SomeDagsterType = create_any_type(
name='SomeType',
output_materialization_config=yield_two_materializations)
@lambda_solid(output_def=OutputDefinition(SomeDagsterType))
def return_one():
return 1
pipeline_def = PipelineDefinition(name='single_int_output_pipeline',
solid_defs=[return_one])
result = execute_pipeline(pipeline_def,
environment_dict={
'solids': {
'return_one': {
'outputs': [{
'result': 2
}]
}
}
})
assert result.success
event_types = [event.event_type_value for event in result.event_list]
assert 2 == (sum([
True for event_type in event_types
if event_type == DagsterEventType.STEP_MATERIALIZATION.value
]))
def test_user_error_boundary_storage_plugin():
class CustomStoragePlugin(TypeStoragePlugin): # pylint: disable=no-init
@classmethod
def compatible_with_storage_def(cls, system_storage_def):
return True
@classmethod
def set_object(cls, intermediate_store, obj, context, runtime_type,
paths):
raise UserError()
@classmethod
def get_object(cls, intermediate_store, context, runtime_type, paths):
raise UserError()
CustomDagsterType = create_any_type(name='CustomType',
auto_plugins=[CustomStoragePlugin])
@solid(output_defs=[OutputDefinition(CustomDagsterType)])
def output_solid(_context):
return 'hello'
@pipeline
def plugin_pipeline():
output_solid()
pipeline_result = execute_pipeline(plugin_pipeline,
{'storage': {
'filesystem': {}
}},
raise_on_error=False)
assert not pipeline_result.success
def test_basic_yield_multiple_materializations():
SomeDagsterType = create_any_type(name="SomeType",
materializer=yield_two_materializations)
@lambda_solid(output_def=OutputDefinition(SomeDagsterType))
def return_one():
return 1
pipeline_def = PipelineDefinition(name="single_int_output_pipeline",
solid_defs=[return_one])
result = execute_pipeline(
pipeline_def,
run_config={"solids": {
"return_one": {
"outputs": [{
"result": 2
}]
}
}})
assert result.success
event_types = [event.event_type_value for event in result.event_list]
assert 2 == (sum([
True for event_type in event_types
if event_type == DagsterEventType.ASSET_MATERIALIZATION.value
]))
def test_user_error_boundary_storage_plugin():
class CustomStoragePlugin(TypeStoragePlugin): # pylint: disable=no-init
@classmethod
def compatible_with_storage_def(cls, intermediate_storage_def):
return True
@classmethod
def set_intermediate_object(cls, intermediate_storage, context,
dagster_type, step_output_handle, value):
raise UserError()
@classmethod
def get_intermediate_object(cls, intermediate_storage, context,
dagster_type, step_output_handle):
raise UserError()
CustomDagsterType = create_any_type(name="CustomType",
auto_plugins=[CustomStoragePlugin])
@solid(output_defs=[OutputDefinition(CustomDagsterType)])
def output_solid(_context):
return "hello"
@pipeline
def plugin_pipeline():
output_solid()
pipeline_result = execute_pipeline(plugin_pipeline,
{"storage": {
"filesystem": {}
}},
raise_on_error=False)
assert not pipeline_result.success
def test_user_error_boundary_output_materialization():
@output_materialization_config(String)
def materialize(context, *_args, **_kwargs):
raise UserError()
CustomDagsterType = create_any_type(
name='CustomType', output_materialization_config=materialize)
@solid(output_defs=[OutputDefinition(CustomDagsterType)])
def output_solid(_context):
return 'hello'
@pipeline
def output_materialization_pipeline():
output_solid()
pipeline_result = execute_pipeline(
output_materialization_pipeline,
{'solids': {
'output_solid': {
'outputs': [{
'result': 'hello'
}]
}
}},
raise_on_error=False,
)
assert not pipeline_result.success
def define_materialization_pipeline(should_require_resources=True,
resources_initted=None):
if resources_initted is None:
resources_initted = {}
@resource
def resource_a(_):
resources_initted['a'] = True
yield 'A'
@dagster_type_materializer(
String,
required_resource_keys={'a'} if should_require_resources else set())
def materialize(context, *_args, **_kwargs):
assert context.resources.a == 'A'
return AssetMaterialization('hello')
CustomDagsterType = create_any_type(name='CustomType',
materializer=materialize)
@solid(output_defs=[OutputDefinition(CustomDagsterType)])
def output_solid(_context):
return 'hello'
@pipeline(mode_defs=[ModeDefinition(resource_defs={'a': resource_a})])
def output_pipeline():
output_solid()
return output_pipeline
def test_user_error_boundary_output_materialization():
@dagster_type_materializer(String)
def materialize(context, *_args, **_kwargs):
raise UserError()
CustomDagsterType = create_any_type(name="CustomType",
materializer=materialize)
@solid(output_defs=[OutputDefinition(CustomDagsterType)])
def output_solid(_context):
return "hello"
@pipeline
def output_materialization_pipeline():
output_solid()
pipeline_result = execute_pipeline(
output_materialization_pipeline,
{"solids": {
"output_solid": {
"outputs": [{
"result": "hello"
}]
}
}},
raise_on_error=False,
)
assert not pipeline_result.success
def define_plugin_pipeline(
should_require_resources=True,
resources_initted=None,
compatible_storage=True,
mode_defines_resource=True,
):
if resources_initted is None:
resources_initted = {}
class CustomStoragePlugin(TypeStoragePlugin): # pylint: disable=no-init
@classmethod
def compatible_with_storage_def(cls, _):
return compatible_storage
@classmethod
def set_intermediate_object(cls, intermediate_storage, context,
dagster_type, step_output_handle, value):
assert context.resources.a == "A"
return intermediate_storage.set_intermediate_object(
dagster_type, step_output_handle, value)
@classmethod
def get_intermediate_object(cls, intermediate_storage, context,
dagster_type, step_output_handle):
assert context.resources.a == "A"
return intermediate_storage.get_intermediate_object(
dagster_type, step_output_handle)
@classmethod
def required_resource_keys(cls):
return {"a"} if should_require_resources else set()
@resource
def resource_a(_):
resources_initted["a"] = True
yield "A"
CustomDagsterType = create_any_type(name="CustomType",
auto_plugins=[CustomStoragePlugin])
@solid(output_defs=[OutputDefinition(CustomDagsterType)])
def output_solid(_context):
return "hello"
if mode_defines_resource:
mode_defs = [ModeDefinition(resource_defs={"a": resource_a})]
else:
mode_defs = [ModeDefinition()]
@pipeline(mode_defs=mode_defs)
def plugin_pipeline():
output_solid()
return plugin_pipeline
def test_basic_bad_output_materialization():
SomeDagsterType = create_any_type(name="SomeType", materializer=return_int)
@lambda_solid(output_def=OutputDefinition(SomeDagsterType))
def return_one():
return 1
pipeline_def = PipelineDefinition(name="single_int_output_pipeline", solid_defs=[return_one])
with pytest.raises(
DagsterInvariantViolationError, match="You must return an AssetMaterialization"
):
execute_pipeline(
pipeline_def, run_config={"solids": {"return_one": {"outputs": [{"result": 2}]}}}
)
def define_plugin_pipeline(should_require_resources=True,
resources_initted=None,
compatible_storage=True):
if resources_initted is None:
resources_initted = {}
class CustomStoragePlugin(TypeStoragePlugin): # pylint: disable=no-init
@classmethod
def compatible_with_storage_def(cls, system_storage_def):
return compatible_storage
@classmethod
def set_object(cls, intermediate_store, obj, context, dagster_type,
paths):
assert context.resources.a == 'A'
return intermediate_store.set_object(obj, context, dagster_type,
paths)
@classmethod
def get_object(cls, intermediate_store, context, dagster_type, paths):
assert context.resources.a == 'A'
return intermediate_store.get_object(context, dagster_type, paths)
@classmethod
def required_resource_keys(cls):
return {'a'} if should_require_resources else set()
@resource
def resource_a(_):
resources_initted['a'] = True
yield 'A'
CustomDagsterType = create_any_type(name='CustomType',
auto_plugins=[CustomStoragePlugin])
@solid(output_defs=[OutputDefinition(CustomDagsterType)])
def output_solid(_context):
return 'hello'
@pipeline(mode_defs=[ModeDefinition(resource_defs={'a': resource_a})])
def plugin_pipeline():
output_solid()
return plugin_pipeline
def test_materialize_missing_resource_fails():
@dagster_type_materializer(String, required_resource_keys={"a"})
def materialize(context, *_args, **_kwargs):
assert context.resources.a == "A"
return AssetMaterialization("hello")
CustomType = create_any_type(name="CustomType", materializer=materialize)
@solid(output_defs=[OutputDefinition(CustomType, "custom_type")])
def custom_type_solid(_):
return "A"
with pytest.raises(
DagsterInvalidDefinitionError,
match='required by the materializer on type "CustomType"'):
@pipeline
def _type_check_pipeline():
custom_type_solid()
def test_basic_bad_output_materialization():
SomeDagsterType = create_any_type(name='SomeType', materializer=return_int)
@lambda_solid(output_def=OutputDefinition(SomeDagsterType))
def return_one():
return 1
pipeline_def = PipelineDefinition(name='single_int_output_pipeline',
solid_defs=[return_one])
with pytest.raises(DagsterInvariantViolationError,
match='You must return an AssetMaterialization'):
execute_pipeline(pipeline_def,
run_config={
'solids': {
'return_one': {
'outputs': [{
'result': 2
}]
}
}
})
def define_composite_materialization_pipeline(
should_require_resources=True, resources_initted=None
):
if resources_initted is None:
resources_initted = {}
@resource
def resource_a(_):
resources_initted['a'] = True
yield 'A'
@output_materialization_config(
String, required_resource_keys={'a'} if should_require_resources else set()
)
def materialize(context, *_args, **_kwargs):
assert context.resources.a == 'A'
return Materialization('hello')
CustomDagsterType = create_any_type(
name='CustomType', output_materialization_config=materialize
)
@solid(output_defs=[OutputDefinition(CustomDagsterType)])
def output_solid(_context):
return 'hello'
wrap_solid = CompositeSolidDefinition(
name="wrap_solid",
solid_defs=[output_solid],
output_mappings=[OutputDefinition(CustomDagsterType).mapping_from('output_solid')],
)
@pipeline(mode_defs=[ModeDefinition(resource_defs={'a': resource_a})])
def output_pipeline():
wrap_solid()
return output_pipeline
def define_composite_materialization_pipeline(should_require_resources=True,
resources_initted=None):
if resources_initted is None:
resources_initted = {}
@resource
def resource_a(_):
resources_initted["a"] = True
yield "A"
@dagster_type_materializer(
String,
required_resource_keys={"a"} if should_require_resources else set())
def materialize(context, *_args, **_kwargs):
assert context.resources.a == "A"
return AssetMaterialization("hello")
CustomDagsterType = create_any_type(name="CustomType",
materializer=materialize)
@solid(output_defs=[OutputDefinition(CustomDagsterType)])
def output_solid(_context):
return "hello"
wrap_solid = CompositeSolidDefinition(
name="wrap_solid",
solid_defs=[output_solid],
output_mappings=[
OutputDefinition(CustomDagsterType).mapping_from("output_solid")
],
)
@pipeline(mode_defs=[ModeDefinition(resource_defs={"a": resource_a})])
def output_pipeline():
wrap_solid()
return output_pipeline
assert result.success
with open(filename_one, 'r') as ff:
value = json.loads(ff.read())
assert value == {'value': 1}
with open(filename_two, 'r') as ff:
value = json.loads(ff.read())
assert value == {'value': 1}
@output_materialization_config(Int)
def return_int(*_args, **_kwargs):
return 1
SomeDagsterType = create_any_type(name='SomeType', output_materialization_config=return_int)
def test_basic_bad_output_materialization():
@lambda_solid(output_def=OutputDefinition(SomeDagsterType))
def return_one():
return 1
pipeline_def = PipelineDefinition(name='single_int_output_pipeline', solid_defs=[return_one])
with pytest.raises(DagsterInvariantViolationError, match='You must return a Materialization'):
execute_pipeline(
pipeline_def, environment_dict={'solids': {'return_one': {'outputs': [{'result': 2}]}}}
)
adls2_resource,
create_adls2_client,
)
from dagster_azure.blob import create_blob_client
class UppercaseSerializationStrategy(SerializationStrategy): # pylint: disable=no-init
def serialize(self, value, write_file_obj):
return write_file_obj.write(bytes(value.upper().encode("utf-8")))
def deserialize(self, read_file_obj):
return read_file_obj.read().decode("utf-8").lower()
LowercaseString = create_any_type(
"LowercaseString",
serialization_strategy=UppercaseSerializationStrategy("uppercase"),
)
nettest = pytest.mark.nettest
def define_inty_pipeline(should_throw=True):
@lambda_solid
def return_one():
return 1
@lambda_solid(input_defs=[InputDefinition("num", Int)],
output_def=OutputDefinition(Int))
def add_one(num):
return num + 1
from dagster.core.types.dagster_type import String as RuntimeString
from dagster.core.types.dagster_type import create_any_type, resolve_dagster_type
from dagster.core.utils import make_new_run_id
from dagster.utils.test import yield_empty_pipeline_context
class UppercaseSerializationStrategy(SerializationStrategy): # pylint: disable=no-init
def serialize(self, value, write_file_obj):
return write_file_obj.write(bytes(value.upper().encode('utf-8')))
def deserialize(self, read_file_obj):
return read_file_obj.read().decode('utf-8').lower()
LowercaseString = create_any_type(
'LowercaseString',
serialization_strategy=UppercaseSerializationStrategy('uppercase'),
)
def aws_credentials_present():
return os.getenv('AWS_ACCESS_KEY_ID') and os.getenv(
'AWS_SECRET_ACCESS_KEY')
nettest = pytest.mark.nettest
def define_inty_pipeline(should_throw=True):
@lambda_solid
def return_one():
return 1
