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