def test_s3_intermediate_storage_with_composite_type_storage_plugin(s3_bucket):
run_id = make_new_run_id()
intermediate_storage = S3IntermediateStorage(
run_id=run_id,
s3_bucket=s3_bucket,
type_storage_plugin_registry=TypeStoragePluginRegistry([
(RuntimeString, FancyStringS3TypeStoragePlugin)
]),
)
with yield_empty_pipeline_context(run_id=run_id) as context:
with pytest.raises(check.NotImplementedCheckError):
intermediate_storage.set_intermediate(
context,
resolve_dagster_type(List[String]),
StepOutputHandle("obj_name"),
[
"hello",
],
)
def test_gcs_intermediate_store_with_type_storage_plugin(gcs_bucket):
run_id = str(uuid.uuid4())
intermediate_store = GCSIntermediateStore(
run_id=run_id,
gcs_bucket=gcs_bucket,
type_storage_plugin_registry=TypeStoragePluginRegistry(
{RuntimeString.inst(): FancyStringGCSTypeStoragePlugin}),
)
with yield_empty_pipeline_context(run_id=run_id) as context:
try:
intermediate_store.set_value('hello', context,
RuntimeString.inst(), ['obj_name'])
assert intermediate_store.has_object(context, ['obj_name'])
assert (intermediate_store.get_value(context, RuntimeString.inst(),
['obj_name']) == 'hello')
finally:
intermediate_store.rm_object(context, ['obj_name'])
def test_file_system_intermediate_store_with_type_storage_plugin():
run_id = str(uuid.uuid4())
# FIXME need a dedicated test bucket
intermediate_store = FileSystemIntermediateStore(
run_id=run_id,
type_storage_plugin_registry=TypeStoragePluginRegistry(
{RuntimeString.inst(): FancyStringFilesystemTypeStoragePlugin}),
)
with yield_empty_pipeline_context(run_id=run_id) as context:
try:
intermediate_store.set_value('hello', context,
RuntimeString.inst(), ['obj_name'])
assert intermediate_store.has_object(context, ['obj_name'])
assert (intermediate_store.get_value(context, RuntimeString.inst(),
['obj_name']) == 'hello')
finally:
intermediate_store.rm_object(context, ['obj_name'])
def test_s3_intermediate_store_with_type_storage_plugin(s3_bucket):
run_id = make_new_run_id()
intermediate_store = S3IntermediateStore(
run_id=run_id,
s3_bucket=s3_bucket,
type_storage_plugin_registry=TypeStoragePluginRegistry([
(RuntimeString, FancyStringS3TypeStoragePlugin)
]),
)
with yield_empty_pipeline_context(run_id=run_id) as context:
try:
intermediate_store.set_value('hello', context, RuntimeString,
['obj_name'])
assert intermediate_store.has_object(context, ['obj_name'])
assert intermediate_store.get_value(context, RuntimeString,
['obj_name']) == 'hello'
finally:
intermediate_store.rm_object(context, ['obj_name'])
def test_adls2_intermediate_store_with_type_storage_plugin(storage_account, file_system):
run_id = make_new_run_id()
intermediate_store = ADLS2IntermediateStore(
adls2_client=get_adls2_client(storage_account),
blob_client=get_blob_client(storage_account),
run_id=run_id,
file_system=file_system,
type_storage_plugin_registry=TypeStoragePluginRegistry(
[(RuntimeString, FancyStringS3TypeStoragePlugin)]
),
)
with yield_empty_pipeline_context(run_id=run_id) as context:
try:
intermediate_store.set_value('hello', context, RuntimeString, ['obj_name'])
assert intermediate_store.has_object(context, ['obj_name'])
assert intermediate_store.get_value(context, RuntimeString, ['obj_name']) == 'hello'
finally:
intermediate_store.rm_object(context, ['obj_name'])
def test_file_system_intermediate_storage_with_type_storage_plugin():
run_id, instance, intermediate_storage = define_intermediate_storage(
type_storage_plugin_registry=TypeStoragePluginRegistry([
(RuntimeString, FancyStringFilesystemTypeStoragePlugin)
]), )
with yield_empty_pipeline_context(run_id=run_id,
instance=instance) as context:
try:
intermediate_storage.set_intermediate(context, RuntimeString,
StepOutputHandle('obj_name'),
'hello')
assert intermediate_storage.has_intermediate(
context, StepOutputHandle('obj_name'))
assert (intermediate_storage.get_intermediate(
context, RuntimeString,
StepOutputHandle('obj_name')) == 'hello')
finally:
intermediate_storage.rm_intermediate(context,
StepOutputHandle('obj_name'))
def test_file_system_intermediate_storage_with_composite_type_storage_plugin():
run_id, _, intermediate_storage = define_intermediate_storage(
type_storage_plugin_registry=TypeStoragePluginRegistry([
(RuntimeString, FancyStringFilesystemTypeStoragePlugin)
]), )
with yield_empty_pipeline_context(run_id=run_id) as context:
with pytest.raises(check.NotImplementedCheckError):
intermediate_storage.set_intermediate(
context, resolve_dagster_type(List[String]),
StepOutputHandle('obj_name'), ['hello'])
with yield_empty_pipeline_context(run_id=run_id) as context:
with pytest.raises(check.NotImplementedCheckError):
intermediate_storage.set_intermediate(
context,
resolve_dagster_type(Optional[String]),
StepOutputHandle('obj_name'),
['hello'],
)
with yield_empty_pipeline_context(run_id=run_id) as context:
with pytest.raises(check.NotImplementedCheckError):
intermediate_storage.set_intermediate(
context,
resolve_dagster_type(List[Optional[String]]),
StepOutputHandle('obj_name'),
['hello'],
)
with yield_empty_pipeline_context(run_id=run_id) as context:
with pytest.raises(check.NotImplementedCheckError):
intermediate_storage.set_intermediate(
context,
resolve_dagster_type(Optional[List[String]]),
StepOutputHandle('obj_name'),
['hello'],
)
def test_file_system_intermediate_store_with_type_storage_plugin():
run_id = str(uuid.uuid4())
instance = DagsterInstance.ephemeral()
intermediate_store = build_fs_intermediate_store(
instance.intermediates_directory,
run_id=run_id,
type_storage_plugin_registry=TypeStoragePluginRegistry(
{RuntimeString: FancyStringFilesystemTypeStoragePlugin}),
)
with yield_empty_pipeline_context(run_id=run_id,
instance=instance) as context:
try:
intermediate_store.set_value('hello', context, RuntimeString,
['obj_name'])
assert intermediate_store.has_object(context, ['obj_name'])
assert intermediate_store.get_value(context, RuntimeString,
['obj_name']) == 'hello'
finally:
intermediate_store.rm_object(context, ['obj_name'])
def __init__(self,
s3_bucket,
run_id,
s3_session=None,
type_storage_plugin_registry=None):
check.str_param(s3_bucket, 's3_bucket')
check.str_param(run_id, 'run_id')
object_store = S3ObjectStore(s3_bucket, s3_session=s3_session)
def root_for_run_id(r_id):
return object_store.key_for_paths(['dagster', 'storage', r_id])
super(S3IntermediateStore, self).__init__(
object_store,
root_for_run_id=root_for_run_id,
run_id=run_id,
type_storage_plugin_registry=check.inst_param(
type_storage_plugin_registry if type_storage_plugin_registry
else TypeStoragePluginRegistry(types_to_register={}),
'type_storage_plugin_registry',
TypeStoragePluginRegistry,
),
)
def test_file_system_intermediate_store_with_composite_type_storage_plugin():
run_id = str(uuid.uuid4())
intermediate_store = build_fs_intermediate_store(
DagsterInstance.ephemeral().intermediates_directory,
run_id=run_id,
type_storage_plugin_registry=TypeStoragePluginRegistry([
(RuntimeString, FancyStringFilesystemTypeStoragePlugin)
]),
)
with yield_empty_pipeline_context(run_id=run_id) as context:
with pytest.raises(check.NotImplementedCheckError):
intermediate_store.set_value(['hello'], context,
resolve_dagster_type(List[String]),
['obj_name'])
with yield_empty_pipeline_context(run_id=run_id) as context:
with pytest.raises(check.NotImplementedCheckError):
intermediate_store.set_value(['hello'], context,
resolve_dagster_type(
Optional[String]), ['obj_name'])
with yield_empty_pipeline_context(run_id=run_id) as context:
with pytest.raises(check.NotImplementedCheckError):
intermediate_store.set_value(['hello'], context,
resolve_dagster_type(
List[Optional[String]]),
['obj_name'])
with yield_empty_pipeline_context(run_id=run_id) as context:
with pytest.raises(check.NotImplementedCheckError):
intermediate_store.set_value(['hello'], context,
resolve_dagster_type(
Optional[List[String]]),
['obj_name'])
def __init__(self,
gcs_bucket,
run_id,
client=None,
type_storage_plugin_registry=None):
check.str_param(gcs_bucket, 'gcs_bucket')
check.str_param(run_id, 'run_id')
object_store = GCSObjectStore(gcs_bucket, client=client)
def root_for_run_id(r_id):
return object_store.key_for_paths(['dagster', 'storage', r_id])
super(GCSIntermediateStore, self).__init__(
object_store,
root_for_run_id=root_for_run_id,
run_id=run_id,
type_storage_plugin_registry=check.inst_param(
type_storage_plugin_registry if type_storage_plugin_registry
else TypeStoragePluginRegistry(types_to_register={}),
'type_storage_plugin_registry',
TypeStoragePluginRegistry,
),
)