def test_resource_with_enum_in_schema_configured():
from enum import Enum as PythonEnum
class TestPythonEnum(PythonEnum):
VALUE_ONE = 0
OTHER = 1
DagsterEnumType = Enum(
"TestEnum",
[
EnumValue("VALUE_ONE", TestPythonEnum.VALUE_ONE),
EnumValue("OTHER", TestPythonEnum.OTHER),
],
)
@resource(config_schema={"enum": DagsterEnumType})
def enum_resource(context):
return context.resource_config["enum"]
@configured(enum_resource, {"enum": DagsterEnumType})
def passthrough_to_enum_resource(config):
return {"enum": "VALUE_ONE" if config["enum"] == TestPythonEnum.VALUE_ONE else "OTHER"}
assert_pipeline_runs_with_resource(
passthrough_to_enum_resource, {"enum": "VALUE_ONE"}, TestPythonEnum.VALUE_ONE
)
def shell_solid_config():
return {
'env': Field(
Noneable(Permissive()),
default_value=os.environ.copy(),
is_required=False,
description='An optional dict of environment variables to pass to the subprocess. '
'Defaults to using os.environ.copy().',
),
'output_logging': Field(
Enum(
name='OutputType',
enum_values=[
EnumValue('STREAM', description='Stream script stdout/stderr.'),
EnumValue(
'BUFFER',
description='Buffer shell script stdout/stderr, then log upon completion.',
),
EnumValue('NONE', description='No logging'),
],
),
is_required=False,
default_value='BUFFER',
),
'cwd': Field(
Noneable(str),
default_value=None,
is_required=False,
description='Working directory in which to execute shell script',
),
}
def test_resource_with_enum_in_schema_configured():
from enum import Enum as PythonEnum
class TestPythonEnum(PythonEnum):
VALUE_ONE = 0
OTHER = 1
DagsterEnumType = Enum(
'TestEnum',
[
EnumValue('VALUE_ONE', TestPythonEnum.VALUE_ONE),
EnumValue('OTHER', TestPythonEnum.OTHER),
],
)
@resource(config_schema={'enum': DagsterEnumType})
def enum_resource(context):
return context.resource_config['enum']
@configured(enum_resource, {'enum': DagsterEnumType})
def passthrough_to_enum_resource(config):
return {
'enum':
'VALUE_ONE'
if config['enum'] == TestPythonEnum.VALUE_ONE else 'OTHER'
}
assert_pipeline_runs_with_resource(passthrough_to_enum_resource,
{'enum': 'VALUE_ONE'},
TestPythonEnum.VALUE_ONE)
def test_native_enum_dagster_enum():
dagster_enum = Enum(
'DagsterNativeEnum',
[
EnumValue(config_value='FOO', python_value=NativeEnum.FOO),
EnumValue(config_value='BAR', python_value=NativeEnum.BAR),
],
)
called = {}
@solid(config=dagster_enum)
def dagster_enum_me(context):
assert context.solid_config == NativeEnum.BAR
called['yup'] = True
pipeline_def = PipelineDefinition(name='native_enum_dagster_pipeline',
solid_defs=[dagster_enum_me])
result = execute_pipeline(
pipeline_def, {'solids': {
'dagster_enum_me': {
'config': 'BAR'
}
}})
assert result.success
assert called['yup']
def test_native_enum_dagster_enum():
dagster_enum = Enum(
"DagsterNativeEnum",
[
EnumValue(config_value="FOO", python_value=NativeEnum.FOO),
EnumValue(config_value="BAR", python_value=NativeEnum.BAR),
],
)
called = {}
@solid(config_schema=dagster_enum)
def dagster_enum_me(context):
assert context.solid_config == NativeEnum.BAR
called["yup"] = True
pipeline_def = PipelineDefinition(name="native_enum_dagster_pipeline",
solid_defs=[dagster_enum_me])
result = execute_pipeline(
pipeline_def, {"solids": {
"dagster_enum_me": {
"config": "BAR"
}
}})
assert result.success
assert called["yup"]
def shell_op_config():
return {
"env": Field(
Noneable(Permissive()),
is_required=False,
description="An optional dict of environment variables to pass to the subprocess.",
),
"output_logging": Field(
Enum(
name="OutputType",
enum_values=[
EnumValue("STREAM", description="Stream script stdout/stderr."),
EnumValue(
"BUFFER",
description="Buffer shell script stdout/stderr, then log upon completion.",
),
EnumValue("NONE", description="No logging"),
],
),
is_required=False,
default_value="BUFFER",
),
"cwd": Field(
Noneable(str),
default_value=None,
is_required=False,
description="Working directory in which to execute shell script",
),
}
def test_logger_with_enum_in_schema_using_configured():
from enum import Enum as PythonEnum
class TestPythonEnum(PythonEnum):
VALUE_ONE = 0
OTHER = 1
DagsterEnumType = Enum(
"TestEnum",
[
EnumValue("VALUE_ONE", TestPythonEnum.VALUE_ONE),
EnumValue("OTHER", TestPythonEnum.OTHER),
],
)
it = {}
@logger(config_schema={"enum": DagsterEnumType})
def test_logger(init_context):
assert init_context.logger_config["enum"] == TestPythonEnum.OTHER
it["ran test_logger"] = True
logger_ = logging.Logger("test", level=coerce_valid_log_level("INFO"))
return logger_
@configured(test_logger, {"enum": DagsterEnumType})
def pick_different_enum_value(config):
it["ran pick_different_enum_value"] = True
return {"enum": "OTHER" if config["enum"] == TestPythonEnum.VALUE_ONE else "VALUE_ONE"}
assert_pipeline_runs_with_logger(pick_different_enum_value, {"enum": "VALUE_ONE"})
assert it["ran test_logger"]
assert it["ran pick_different_enum_value"]
def pipeline_with_enum_config():
@solid(
config_schema=Enum(
"TestEnum",
[
EnumValue(config_value="ENUM_VALUE_ONE", description="An enum value."),
EnumValue(config_value="ENUM_VALUE_TWO", description="An enum value."),
EnumValue(config_value="ENUM_VALUE_THREE", description="An enum value."),
],
)
)
def takes_an_enum(_context):
pass
takes_an_enum()
def pipeline_with_enum_config():
@solid(
config_schema=Enum(
'TestEnum',
[
EnumValue(config_value='ENUM_VALUE_ONE', description='An enum value.'),
EnumValue(config_value='ENUM_VALUE_TWO', description='An enum value.'),
EnumValue(config_value='ENUM_VALUE_THREE', description='An enum value.'),
],
)
)
def takes_an_enum(_context):
pass
return takes_an_enum()
def test_enum_in_pipeline_execution():
called = {}
@solid(
config_schema={"int_field": Int, "enum_field": Enum("AnEnum", [EnumValue("ENUM_VALUE")]),}
)
def config_me(context):
assert context.solid_config["int_field"] == 2
assert context.solid_config["enum_field"] == "ENUM_VALUE"
called["yup"] = True
pipeline_def = PipelineDefinition(name="enum_in_pipeline", solid_defs=[config_me])
result = execute_pipeline(
pipeline_def,
{"solids": {"config_me": {"config": {"int_field": 2, "enum_field": "ENUM_VALUE"}}}},
)
assert result.success
assert called["yup"]
with pytest.raises(DagsterInvalidConfigError) as exc_info:
execute_pipeline(
pipeline_def,
{"solids": {"config_me": {"config": {"int_field": 2, "enum_field": "NOPE"}}}},
)
assert (
"Value at path root:solids:config_me:config:enum_field not in enum type AnEnum got NOPE"
in str(exc_info.value)
)
def test_enum_in_pipeline_execution():
called = {}
@solid(
config={'int_field': Int, 'enum_field': Enum('AnEnum', [EnumValue('ENUM_VALUE')]),}
)
def config_me(context):
assert context.solid_config['int_field'] == 2
assert context.solid_config['enum_field'] == 'ENUM_VALUE'
called['yup'] = True
pipeline_def = PipelineDefinition(name='enum_in_pipeline', solid_defs=[config_me])
result = execute_pipeline(
pipeline_def,
{'solids': {'config_me': {'config': {'int_field': 2, 'enum_field': 'ENUM_VALUE'}}}},
)
assert result.success
assert called['yup']
with pytest.raises(DagsterInvalidConfigError) as exc_info:
execute_pipeline(
pipeline_def,
{'solids': {'config_me': {'config': {'int_field': 2, 'enum_field': 'NOPE'}}}},
)
assert 'Error 1: Type failure at path "root:solids:config_me:config:enum_field"' in str(
exc_info.value
)
def define_pipeline_with_enum_config():
@solid(
config_field=Field(
Enum(
'TestEnum',
[
EnumValue(config_value='ENUM_VALUE_ONE', description='An enum value.'),
EnumValue(config_value='ENUM_VALUE_TWO', description='An enum value.'),
EnumValue(config_value='ENUM_VALUE_THREE', description='An enum value.'),
],
)
)
)
def takes_an_enum(_context):
pass
return PipelineDefinition(name='pipeline_with_enum_config', solids=[takes_an_enum])
def test_enum_snap():
enum_snap = snap_from_dagster_type(
Enum(
"CowboyType",
[
EnumValue("good"),
EnumValue("bad"),
EnumValue("ugly"),
],
))
assert enum_snap.given_name == "CowboyType"
assert enum_snap.key == "CowboyType"
assert enum_snap.kind == ConfigTypeKind.ENUM
assert enum_snap.enum_values == [
ConfigEnumValueSnap(value, description=None)
for value in ["good", "bad", "ugly"]
]
assert enum_snap.fields is None
def test_logger_with_enum_in_schema_using_configured():
from enum import Enum as PythonEnum
class TestPythonEnum(PythonEnum):
VALUE_ONE = 0
OTHER = 1
DagsterEnumType = Enum(
'TestEnum',
[
EnumValue('VALUE_ONE', TestPythonEnum.VALUE_ONE),
EnumValue('OTHER', TestPythonEnum.OTHER),
],
)
it = {}
@logger(config_schema={'enum': DagsterEnumType})
def test_logger(init_context):
assert init_context.logger_config['enum'] == TestPythonEnum.OTHER
it['ran test_logger'] = True
logger_ = logging.Logger('test', level=coerce_valid_log_level('INFO'))
return logger_
@configured(test_logger, {'enum': DagsterEnumType})
def pick_different_enum_value(config):
it['ran pick_different_enum_value'] = True
return {
'enum':
'OTHER'
if config['enum'] == TestPythonEnum.VALUE_ONE else 'VALUE_ONE'
}
assert_pipeline_runs_with_logger(pick_different_enum_value,
{'enum': 'VALUE_ONE'})
assert it['ran test_logger']
assert it['ran pick_different_enum_value']
def test_enum_in_pipeline_execution():
called = {}
@solid(
config_field=Field(
Dict(
{
'int_field': Field(Int),
'enum_field': Field(Enum('AnEnum', [EnumValue('ENUM_VALUE')])),
}
)
)
)
def config_me(info):
assert info.config['int_field'] == 2
assert info.config['enum_field'] == 'ENUM_VALUE'
called['yup'] = True
pipeline_def = PipelineDefinition(name='enum_in_pipeline', solids=[config_me])
result = execute_pipeline(
pipeline_def,
{'solids': {'config_me': {'config': {'int_field': 2, 'enum_field': 'ENUM_VALUE'}}}},
)
assert result.success
assert called['yup']
with pytest.raises(PipelineConfigEvaluationError) as exc_info:
execute_pipeline(
pipeline_def,
{'solids': {'config_me': {'config': {'int_field': 2, 'enum_field': 'NOPE'}}}},
)
assert 'Error 1: Type failure at path "root:solids:config_me:config:enum_field"' in str(
exc_info.value
)
def test_config_mapped_enum():
from enum import Enum as PythonEnum
class TestPythonEnum(PythonEnum):
VALUE_ONE = 0
OTHER = 1
DagsterEnumType = Enum(
"TestEnum",
[
EnumValue("VALUE_ONE", TestPythonEnum.VALUE_ONE),
EnumValue("OTHER", TestPythonEnum.OTHER),
],
)
@solid(config_schema={"enum": DagsterEnumType})
def return_enum(context):
return context.solid_config["enum"]
@composite_solid(
config_schema={"num": int},
config_fn=lambda cfg: {
"return_enum": {
"config": {
"enum": "VALUE_ONE" if cfg["num"] == 1 else "OTHER"
}
}
},
)
def wrapping_return_enum():
return return_enum()
@pipeline
def wrapping_return_enum_pipeline():
wrapping_return_enum()
assert (execute_pipeline(
wrapping_return_enum_pipeline,
{
"solids": {
"wrapping_return_enum": {
"config": {
"num": 1
}
}
}
},
).output_for_solid("wrapping_return_enum") == TestPythonEnum.VALUE_ONE)
assert (execute_pipeline(
wrapping_return_enum_pipeline,
{
"solids": {
"wrapping_return_enum": {
"config": {
"num": -11
}
}
}
},
).output_for_solid("wrapping_return_enum") == TestPythonEnum.OTHER)
@solid(config_schema={"num": int})
def return_int(context):
return context.solid_config["num"]
@composite_solid(
config_schema={"enum": DagsterEnumType},
config_fn=lambda cfg: {
"return_int": {
"config": {
"num": 1 if cfg["enum"] == TestPythonEnum.VALUE_ONE else 2
}
}
},
)
def wrap_return_int():
return return_int()
@pipeline
def wrap_return_int_pipeline():
wrap_return_int()
assert (execute_pipeline(
wrap_return_int_pipeline,
{
"solids": {
"wrap_return_int": {
"config": {
"enum": "VALUE_ONE"
}
}
}
},
).output_for_solid("wrap_return_int") == 1)
assert (execute_pipeline(
wrap_return_int_pipeline,
{
"solids": {
"wrap_return_int": {
"config": {
"enum": "OTHER"
}
}
}
},
).output_for_solid("wrap_return_int") == 2)
def test_post_process_config():
scalar_config_type = resolve_to_config_type(String)
assert post_process_config(scalar_config_type, 'foo').value == 'foo'
assert post_process_config(scalar_config_type, 3).value == 3
assert post_process_config(scalar_config_type, {}).value == {}
assert post_process_config(scalar_config_type, None).value is None
enum_config_type = resolve_to_config_type(
Enum('an_enum', [EnumValue('foo'),
EnumValue('bar', python_value=3)]))
assert post_process_config(enum_config_type, 'foo').value == 'foo'
assert post_process_config(enum_config_type, 'bar').value == 3
assert ('config_value should be pre-validated'
in post_process_config(enum_config_type, 'baz').errors[0].message)
assert ('config_value should be pre-validated'
in post_process_config(enum_config_type, None).errors[0].message)
list_config_type = resolve_to_config_type([str])
assert post_process_config(list_config_type, ['foo']).value == ['foo']
assert post_process_config(list_config_type, None).value == []
with pytest.raises(CheckError, match='Null array member not caught'):
assert post_process_config(list_config_type, [None]).value == [None]
nullable_list_config_type = resolve_to_config_type([Noneable(str)])
assert post_process_config(nullable_list_config_type,
['foo']).value == ['foo']
assert post_process_config(nullable_list_config_type,
[None]).value == [None]
assert post_process_config(nullable_list_config_type, None).value == []
composite_config_type = resolve_to_config_type({
'foo':
String,
'bar': {
'baz': [str]
},
'quux':
Field(str, is_required=False, default_value='zip'),
'quiggle':
Field(str, is_required=False),
})
with pytest.raises(CheckError,
match='Missing non-optional composite member'):
post_process_config(composite_config_type, {})
with pytest.raises(CheckError,
match='Missing non-optional composite member'):
post_process_config(composite_config_type, {
'bar': {
'baz': ['giraffe']
},
'quux': 'nimble'
})
with pytest.raises(CheckError,
match='Missing non-optional composite member'):
post_process_config(composite_config_type, {
'foo': 'zowie',
'quux': 'nimble'
})
assert post_process_config(composite_config_type, {
'foo': 'zowie',
'bar': {
'baz': ['giraffe']
},
'quux': 'nimble'
}).value == {
'foo': 'zowie',
'bar': {
'baz': ['giraffe']
},
'quux': 'nimble'
}
assert post_process_config(composite_config_type, {
'foo': 'zowie',
'bar': {
'baz': ['giraffe']
}
}).value == {
'foo': 'zowie',
'bar': {
'baz': ['giraffe']
},
'quux': 'zip'
}
assert post_process_config(composite_config_type, {
'foo': 'zowie',
'bar': {
'baz': ['giraffe']
},
'quiggle': 'squiggle'
}).value == {
'foo': 'zowie',
'bar': {
'baz': ['giraffe']
},
'quux': 'zip',
'quiggle': 'squiggle'
}
nested_composite_config_type = resolve_to_config_type({
'fruts': {
'apple': Field(String),
'banana': Field(String, is_required=False),
'potato': Field(String, is_required=False, default_value='pie'),
}
})
with pytest.raises(CheckError,
match='Missing non-optional composite member'):
post_process_config(nested_composite_config_type, {'fruts': None})
with pytest.raises(CheckError,
match='Missing non-optional composite member'):
post_process_config(nested_composite_config_type,
{'fruts': {
'banana': 'good',
'potato': 'bad'
}})
assert post_process_config(nested_composite_config_type, {
'fruts': {
'apple': 'strawberry'
}
}).value == {
'fruts': {
'apple': 'strawberry',
'potato': 'pie'
}
}
assert post_process_config(nested_composite_config_type, {
'fruts': {
'apple': 'a',
'banana': 'b',
'potato': 'c'
}
}).value == {
'fruts': {
'apple': 'a',
'banana': 'b',
'potato': 'c'
}
}
any_config_type = resolve_to_config_type(Any)
assert post_process_config(any_config_type, {
'foo': 'bar'
}).value == {
'foo': 'bar'
}
assert post_process_config(
ConfigType('gargle', given_name='bargle', kind=ConfigTypeKind.ANY), 3)
selector_config_type = resolve_to_config_type(
Selector({
'one':
Field(String),
'another': {
'foo': Field(String, default_value='bar', is_required=False)
},
'yet_another':
Field(String, default_value='quux', is_required=False),
}))
with pytest.raises(CheckError):
post_process_config(selector_config_type, 'one')
with pytest.raises(ParameterCheckError):
post_process_config(selector_config_type, None)
with pytest.raises(ParameterCheckError,
match='Expected dict with single item'):
post_process_config(selector_config_type, {})
with pytest.raises(CheckError):
post_process_config(selector_config_type, {
'one': 'foo',
'another': 'bar'
})
assert post_process_config(selector_config_type, {
'one': 'foo'
}).value == {
'one': 'foo'
}
assert post_process_config(selector_config_type, {
'one': None
}).value == {
'one': None
}
assert post_process_config(selector_config_type, {
'one': {}
}).value == {
'one': {}
}
assert post_process_config(selector_config_type, {
'another': {}
}).value == {
'another': {
'foo': 'bar'
}
}
singleton_selector_config_type = resolve_to_config_type(
Selector(
{'foo': Field(String, default_value='bar', is_required=False)}))
assert post_process_config(singleton_selector_config_type, None).value == {
'foo': 'bar'
}
permissive_dict_config_type = resolve_to_config_type(
Permissive({
'foo': Field(String),
'bar': Field(String, default_value='baz', is_required=False)
}))
with pytest.raises(CheckError,
match='Missing non-optional composite member'):
post_process_config(permissive_dict_config_type, None)
assert post_process_config(permissive_dict_config_type, {
'foo': 'wow',
'mau': 'mau'
}).value == {
'foo': 'wow',
'bar': 'baz',
'mau': 'mau',
}
'''NOTE: THIS FILE IS AUTO-GENERATED. DO NOT EDIT
@generated
Produced via:
python automation/parse_dataproc_configs.py \
'''
from dagster import Enum, EnumValue
State = Enum(
name='State',
enum_values=[
EnumValue('STATE_UNSPECIFIED', description='Status is unspecified.'),
EnumValue('NEW', description='Status is NEW.'),
EnumValue('NEW_SAVING', description='Status is NEW_SAVING.'),
EnumValue('SUBMITTED', description='Status is SUBMITTED.'),
EnumValue('ACCEPTED', description='Status is ACCEPTED.'),
EnumValue('RUNNING', description='Status is RUNNING.'),
EnumValue('FINISHED', description='Status is FINISHED.'),
EnumValue('FAILED', description='Status is FAILED.'),
EnumValue('KILLED', description='Status is KILLED.'),
],
)
Substate = Enum(
name='Substate',
enum_values=[
EnumValue('UNSPECIFIED', description='The job substate is unknown.'),
EnumValue(
from enum import Enum as PyEnum
from dagster import Enum, EnumValue
EbsVolumeType = Enum(
name='EbsVolumeType', enum_values=[EnumValue('gp2'), EnumValue('io1'), EnumValue('standard')]
)
class EmrClusterState(PyEnum):
Starting = 'STARTING'
Bootstrapping = 'BOOTSTRAPPING'
Running = 'RUNNING'
Waiting = 'WAITING'
Terminating = 'TERMINATING'
Terminated = 'TERMINATED'
TerminatedWithErrors = 'TERMINATED_WITH_ERRORS'
EmrActionOnFailure = Enum(
name='EmrActionOnFailure',
enum_values=[
EnumValue('TERMINATE_JOB_FLOW'),
EnumValue('TERMINATE_CLUSTER'),
EnumValue('CANCEL_AND_WAIT'),
EnumValue('CONTINUE'),
],
)
EmrAdjustmentType = Enum(
name='EmrAdjustmentType',
def test_apply_default_values():
scalar_config_type = resolve_to_config_type(String)
assert apply_default_values(scalar_config_type, 'foo') == 'foo'
assert apply_default_values(scalar_config_type, 3) == 3
assert apply_default_values(scalar_config_type, {}) == {}
assert apply_default_values(scalar_config_type, None) is None
enum_config_type = resolve_to_config_type(
Enum('an_enum', [EnumValue('foo'),
EnumValue('bar', python_value=3)]))
assert apply_default_values(enum_config_type, 'foo') == 'foo'
assert apply_default_values(enum_config_type, 'bar') == 3
with pytest.raises(CheckError,
match='config_value should be pre-validated'):
apply_default_values(enum_config_type, 'baz')
with pytest.raises(CheckError,
match='config_value should be pre-validated'):
apply_default_values(enum_config_type, None)
list_config_type = resolve_to_config_type(List[String])
assert apply_default_values(list_config_type, ['foo']) == ['foo']
assert apply_default_values(list_config_type, None) == []
with pytest.raises(CheckError, match='Null list member not caught'):
assert apply_default_values(list_config_type, [None]) == [None]
nullable_list_config_type = resolve_to_config_type(List[Optional[String]])
assert apply_default_values(nullable_list_config_type, ['foo']) == ['foo']
assert apply_default_values(nullable_list_config_type, [None]) == [None]
assert apply_default_values(nullable_list_config_type, None) == []
composite_config_type = resolve_to_config_type(
Dict({
'foo': Field(String),
'bar': Field(Dict({'baz': Field(List[String])})),
'quux': Field(String, is_optional=True, default_value='zip'),
'quiggle': Field(String, is_optional=True),
}))
with pytest.raises(CheckError,
match='Missing non-optional composite member'):
apply_default_values(composite_config_type, {})
with pytest.raises(CheckError,
match='Missing non-optional composite member'):
apply_default_values(composite_config_type, {
'bar': {
'baz': ['giraffe']
},
'quux': 'nimble'
})
with pytest.raises(CheckError,
match='Missing non-optional composite member'):
apply_default_values(composite_config_type, {
'foo': 'zowie',
'quux': 'nimble'
})
assert apply_default_values(composite_config_type, {
'foo': 'zowie',
'bar': {
'baz': ['giraffe']
},
'quux': 'nimble'
}) == {
'foo': 'zowie',
'bar': {
'baz': ['giraffe']
},
'quux': 'nimble'
}
assert apply_default_values(composite_config_type, {
'foo': 'zowie',
'bar': {
'baz': ['giraffe']
}
}) == {
'foo': 'zowie',
'bar': {
'baz': ['giraffe']
},
'quux': 'zip'
}
assert apply_default_values(composite_config_type, {
'foo': 'zowie',
'bar': {
'baz': ['giraffe']
},
'quiggle': 'squiggle'
}) == {
'foo': 'zowie',
'bar': {
'baz': ['giraffe']
},
'quux': 'zip',
'quiggle': 'squiggle'
}
nested_composite_config_type = resolve_to_config_type(
Dict({
'fruts':
Field(
Dict({
'apple':
Field(String),
'banana':
Field(String, is_optional=True),
'potato':
Field(String, is_optional=True, default_value='pie'),
}))
}))
with pytest.raises(CheckError,
match='Missing non-optional composite member'):
apply_default_values(nested_composite_config_type, {'fruts': None})
with pytest.raises(CheckError,
match='Missing non-optional composite member'):
apply_default_values(nested_composite_config_type,
{'fruts': {
'banana': 'good',
'potato': 'bad'
}})
assert apply_default_values(nested_composite_config_type,
{'fruts': {
'apple': 'strawberry'
}}) == {
'fruts': {
'apple': 'strawberry',
'potato': 'pie'
}
}
assert apply_default_values(
nested_composite_config_type,
{'fruts': {
'apple': 'a',
'banana': 'b',
'potato': 'c'
}}) == {
'fruts': {
'apple': 'a',
'banana': 'b',
'potato': 'c'
}
}
any_config_type = resolve_to_config_type(Any)
assert apply_default_values(any_config_type, {'foo': 'bar'}) == {
'foo': 'bar'
}
with pytest.raises(CheckError, match='Unsupported type'):
assert apply_default_values(
ConfigType('gargle', 'bargle', ConfigTypeKind.REGULAR), 3)
selector_config_type = resolve_to_config_type(
Selector({
'one':
Field(String),
'another':
Field(
Dict({
'foo':
Field(String, default_value='bar', is_optional=True)
})),
'yet_another':
Field(String, default_value='quux', is_optional=True),
}))
with pytest.raises(CheckError):
apply_default_values(selector_config_type, 'one')
with pytest.raises(ParameterCheckError):
apply_default_values(selector_config_type, None)
with pytest.raises(ParameterCheckError,
match='Expected dict with single item'):
apply_default_values(selector_config_type, {})
with pytest.raises(CheckError):
apply_default_values(selector_config_type, {
'one': 'foo',
'another': 'bar'
})
assert apply_default_values(selector_config_type, {'one': 'foo'}) == {
'one': 'foo'
}
assert apply_default_values(selector_config_type, {'one': None}) == {
'one': None
}
assert apply_default_values(selector_config_type, {'one': {}}) == {
'one': {}
}
assert apply_default_values(selector_config_type, {'another': {}}) == {
'another': {
'foo': 'bar'
}
}
singleton_selector_config_type = resolve_to_config_type(
Selector({'foo': Field(String, default_value='bar',
is_optional=True)}))
assert apply_default_values(singleton_selector_config_type, None) == {
'foo': 'bar'
}
permissive_dict_config_type = resolve_to_config_type(
PermissiveDict({
'foo':
Field(String),
'bar':
Field(String, default_value='baz', is_optional=True)
}))
with pytest.raises(CheckError,
match='Missing non-optional composite member'):
apply_default_values(permissive_dict_config_type, None)
assert apply_default_values(permissive_dict_config_type, {
'foo': 'wow',
'mau': 'mau'
}) == {
'foo': 'wow',
'bar': 'baz',
'mau': 'mau',
}
from dagster import Enum, EnumValue
SparkDeployModeCluster = EnumValue('cluster')
SparkDeployModeClient = EnumValue('client')
SparkDeployMode = Enum(
name='SparkDeployMode', enum_values=[SparkDeployModeCluster, SparkDeployModeClient]
)
class SparkSolidError(Exception):
pass
def bash_command_solid(bash_command, name=None, output_encoding=None):
'''Execute a Bash command.
'''
check.str_param(bash_command, 'bash_command')
name = check.opt_str_param(name, 'name', default='bash_solid')
output_encoding = check.opt_str_param(output_encoding,
'output_encoding',
default='utf-8')
@solid(
name=name,
config={
'output_logging':
Field(
Enum(
'OutputType',
[
EnumValue('STREAM',
description='Stream script stdout/stderr.'),
EnumValue(
'BUFFER',
description=
'Buffer bash script stdout/stderr, then log upon completion.',
),
EnumValue('NONE', description='No logging'),
],
),
is_optional=True,
default_value='STREAM',
),
'env':
Field(
Permissive(),
description=
'Environment variables to pass to the child process; if not provided, '
'the current process environment will be passed.',
is_optional=True,
default_value=None,
),
},
)
def _bash_solid(context):
'''This logic is ported from the Airflow BashOperator implementation.
https://github.com/apache/airflow/blob/master/airflow/operators/bash_operator.py
'''
def log_info_msg(log_str):
context.log.info('[bash][{name}] '.format(name=name) + log_str)
tmp_path = seven.get_system_temp_directory()
log_info_msg('using temporary directory: %s' % tmp_path)
env = (context.solid_config['env'] if context.solid_config['env']
is not None else os.environ.copy())
with NamedTemporaryFile(dir=tmp_path, prefix=name) as tmp_file:
tmp_file.write(bytes(bash_command.encode('utf-8')))
tmp_file.flush()
script_location = os.path.abspath(tmp_file.name)
log_info_msg('Temporary script location: {location}'.format(
location=script_location))
def pre_exec():
# Restore default signal disposition and invoke setsid
for sig in ('SIGPIPE', 'SIGXFZ', 'SIGXFSZ'):
if hasattr(signal, sig):
signal.signal(getattr(signal, sig), signal.SIG_DFL)
os.setsid()
log_info_msg(
'Running command: {command}'.format(command=bash_command))
# pylint: disable=subprocess-popen-preexec-fn
sub_process = Popen(
['bash', tmp_file.name],
stdout=PIPE,
stderr=STDOUT,
cwd=tmp_path,
env=env,
preexec_fn=pre_exec,
)
# Stream back logs as they are emitted
if context.solid_config['output_logging'] == 'STREAM':
line = ''
for raw_line in iter(sub_process.stdout.readline, b''):
line = raw_line.decode(output_encoding).rstrip()
log_info_msg(line)
sub_process.wait()
# Collect and buffer all logs, then emit
if context.solid_config['output_logging'] == 'BUFFER':
line = ''
for raw_line in iter(sub_process.stdout.readline, b''):
line += raw_line.decode(output_encoding)
log_info_msg(line)
# no logging in this case
elif context.solid_config['output_logging'] == 'NONE':
pass
log_info_msg('Command exited with return code {retcode}'.format(
retcode=sub_process.returncode))
if sub_process.returncode:
raise Failure(
description='[bash][{name}] Bash command failed'.format(
name=name))
return line
return _bash_solid
SchemaUpdateOption,
SourceFormat,
WriteDisposition,
)
class BigQueryLoadSource(PyEnum):
DataFrame = "DATA_FRAME"
GCS = "GCS"
File = "FILE"
BQCreateDisposition = Enum(
name="BQCreateDisposition",
enum_values=[
EnumValue(CreateDisposition.CREATE_IF_NEEDED),
EnumValue(CreateDisposition.CREATE_NEVER),
],
)
BQPriority = Enum(
name="BQPriority",
enum_values=[
EnumValue(QueryPriority.BATCH),
EnumValue(QueryPriority.INTERACTIVE)
],
)
BQSchemaUpdateOption = Enum(
name="BQSchemaUpdateOption",
enum_values=[
def define_test_enum_type():
return ConfigEnum(name='TestEnum', enum_values=[EnumValue('VALUE_ONE')])
@solid(
config={
'local_filepath':
Field(str,
is_required=True,
description='local file path to get or put.'),
'remote_filepath':
Field(str,
is_required=True,
description='remote file path to get or put.'),
'operation':
Field(
Enum('SFTPOperation',
[EnumValue('GET'), EnumValue('PUT')]),
is_required=False,
default_value='PUT',
description='specify operation \'GET\' or \'PUT\', defaults to PUT',
),
'confirm':
Field(
bool,
is_required=False,
default_value=True,
description=
'specify if the SFTP operation should be confirmed, defaults to True',
),
},
required_resource_keys={'ssh_resource'},
)
def test_config_mapped_enum():
from enum import Enum as PythonEnum
class TestPythonEnum(PythonEnum):
VALUE_ONE = 0
OTHER = 1
DagsterEnumType = Enum(
'TestEnum',
[
EnumValue('VALUE_ONE', TestPythonEnum.VALUE_ONE),
EnumValue('OTHER', TestPythonEnum.OTHER),
],
)
@solid(config={'enum': DagsterEnumType})
def return_enum(context):
return context.solid_config['enum']
@composite_solid(
config={'num': int},
config_fn=lambda _, cfg: {
'return_enum': {'config': {'enum': 'VALUE_ONE' if cfg['num'] == 1 else 'OTHER'}}
},
)
def wrapping_return_enum():
return return_enum()
@pipeline
def wrapping_return_enum_pipeline():
wrapping_return_enum()
assert (
execute_pipeline(
wrapping_return_enum_pipeline,
{'solids': {'wrapping_return_enum': {'config': {'num': 1}}}},
).output_for_solid('wrapping_return_enum')
== TestPythonEnum.VALUE_ONE
)
assert (
execute_pipeline(
wrapping_return_enum_pipeline,
{'solids': {'wrapping_return_enum': {'config': {'num': -11}}}},
).output_for_solid('wrapping_return_enum')
== TestPythonEnum.OTHER
)
@solid(config={'num': int})
def return_int(context):
return context.solid_config['num']
@composite_solid(
config={'enum': DagsterEnumType},
config_fn=lambda _, cfg: {
'return_int': {'config': {'num': 1 if cfg['enum'] == TestPythonEnum.VALUE_ONE else 2}}
},
)
def wrap_return_int():
return return_int()
@pipeline
def wrap_return_int_pipeline():
wrap_return_int()
assert (
execute_pipeline(
wrap_return_int_pipeline,
{'solids': {'wrap_return_int': {'config': {'enum': 'VALUE_ONE'}}}},
).output_for_solid('wrap_return_int')
== 1
)
assert (
execute_pipeline(
wrap_return_int_pipeline,
{'solids': {'wrap_return_int': {'config': {'enum': 'OTHER'}}}},
).output_for_solid('wrap_return_int')
== 2
)
from dagster import Enum, EnumValue
SparkDeployModeCluster = EnumValue("cluster")
SparkDeployModeClient = EnumValue("client")
SparkDeployMode = Enum(
name="SparkDeployMode",
enum_values=[SparkDeployModeCluster, SparkDeployModeClient])
class SparkSolidError(Exception):
pass
class SparkOpError(SparkSolidError):
pass
from enum import Enum as PyEnum
from dagster import Enum, EnumValue
EbsVolumeType = Enum(
name='EbsVolumeType',
enum_values=[EnumValue('gp2'),
EnumValue('io1'),
EnumValue('standard')])
class EmrClusterState(PyEnum):
Starting = 'STARTING'
Bootstrapping = 'BOOTSTRAPPING'
Running = 'RUNNING'
Waiting = 'WAITING'
Terminating = 'TERMINATING'
Terminated = 'TERMINATED'
TerminatedWithErrors = 'TERMINATED_WITH_ERRORS'
EMR_CLUSTER_TERMINATED_STATES = [
EmrClusterState.Terminating,
EmrClusterState.Terminated,
EmrClusterState.TerminatedWithErrors,
]
EMR_CLUSTER_DONE_STATES = EMR_CLUSTER_TERMINATED_STATES + [
EmrClusterState.Waiting
]
PythonObjectDagsterType,
String,
dagster_type_loader,
dagster_type_materializer,
)
from dagster.config.field_utils import Selector
from dagster.core.storage.system_storage import fs_intermediate_storage
from dagster.core.storage.type_storage import TypeStoragePlugin
from dagster.utils import dict_without_keys
from pyspark.sql import DataFrame as NativeSparkDataFrame
WriteModeOptions = Enum(
"WriteMode",
[
EnumValue(
"append", description="Append contents of this :class:`DataFrame` to existing data."
),
EnumValue("overwrite", description="Overwrite existing data."),
EnumValue("ignore", description="Silently ignore this operation if data already exists."),
EnumValue(
"error", description="(default case): Throw an exception if data already exists."
),
EnumValue(
"errorifexists",
description="(default case): Throw an exception if data already exists.",
),
],
)
WriteCompressionTextOptions = Enum(
Permissive,
Selector,
Shape,
String,
TypeCheck,
check,
dagster_type_loader,
dagster_type_materializer,
)
from .utils import DataFrameUtilities, apply_utilities_to_df
WriteCompressionTextOptions = Enum(
"WriteCompressionText",
[
EnumValue("gzip"),
EnumValue("bz2"),
EnumValue("xz"),
],
)
EngineParquetOptions = Enum(
"EngineParquet",
[
EnumValue("auto"),
EnumValue("fastparquet"),
EnumValue("pyarrow"),
],
)
DataFrameReadTypes = {