def test_get_data_store_instance_from_pool_without_pool(self):
with self.assertRaises(ValueError) as cm:
get_data_store_instance('@dir')
self.assertEqual('store_pool must be given,'
' with store_id ("@dir") referring'
' to a configured store',
f'{cm.exception}')
def test_get_data_store_instance_new_inst(self):
instance = get_data_store_instance('memory')
self.assertIsInstance(instance, DataStoreInstance)
self.assertIsInstance(instance.store, DataStore)
instance2 = get_data_store_instance('memory')
self.assertIsNot(instance, instance2)
self.assertIsNot(instance.store, instance2.store)
def test_get_data_store_instance_from_pool(self):
pool = DataStorePool({
'dir':
DataStoreConfig('directory', store_params=dict(base_dir='.'))
})
instance = get_data_store_instance('@dir', store_pool=pool)
self.assertIsInstance(instance.store, DirectoryDataStore)
instance2 = get_data_store_instance('@dir', store_pool=pool)
self.assertIs(instance, instance2)
def test_get_data_store_instance_new_inst(self):
instance = get_data_store_instance('directory',
store_params=dict(base_dir='.'))
self.assertIsInstance(instance, DataStoreInstance)
self.assertIsInstance(instance.store, DirectoryDataStore)
instance2 = get_data_store_instance('directory',
store_params=dict(base_dir='.'))
self.assertIsNot(instance, instance2)
self.assertIsNot(instance.store, instance2.store)
def test_get_data_store_instance_from_pool(self):
pool = DataStorePool({
'dir': DataStoreConfig('file',
store_params=dict(root='.'))
})
instance = get_data_store_instance('@dir', store_pool=pool)
self.assertTrue(hasattr(instance.store, 'root'))
# noinspection PyUnresolvedReferences
self.assertTrue(os.path.isabs(instance.store.root))
self.assertTrue(os.path.isdir(instance.store.root))
instance2 = get_data_store_instance('@dir', store_pool=pool)
self.assertIs(instance, instance2)
def test_get_data_store_instance_from_pool_with_params(self):
pool = DataStorePool({
'@dir':
DataStoreConfig('directory', store_params=dict(base_dir='.'))
})
with self.assertRaises(ValueError) as cm:
get_data_store_instance('@dir',
store_pool=pool,
store_params={'thres': 5})
self.assertEqual(
'store_params cannot be given, with store_id ("@dir") referring to a configured store',
f'{cm.exception}')
def test_get_data_store_instance_from_pool_with_params(self):
pool = DataStorePool({
'@dir': DataStoreConfig('file',
store_params=dict(root='.'))
})
with self.assertRaises(ValueError) as cm:
get_data_store_instance(
'@dir', store_pool=pool, store_params={'auto_mkdir': True}
)
self.assertEqual('store_params cannot be given,'
' with store_id ("@dir") referring'
' to a configured store',
f'{cm.exception}')
def write_cube(self, cube: xr.Dataset,
gm: GridMapping) -> Tuple[str, xr.Dataset]:
output_config = self._output_config
dataset = encode_cube(cube, grid_mapping=gm)
with observe_dask_progress('writing cube', 100):
write_params = output_config.write_params or {}
store_params = output_config.store_params or {}
if output_config.store_id:
store_instance = get_data_store_instance(
output_config.store_id,
store_params=store_params,
store_pool=self._store_pool)
writer = store_instance.store
write_params.update(writer_id=output_config.writer_id,
**write_params)
else:
writer = new_data_writer(output_config.writer_id)
write_params.update(**store_params, **write_params)
if not dataset.attrs.get('title'):
# Set fallback title, so we can distinguish
# datasets from stores in xcube-viewer
dataset = dataset.assign_attrs(title=output_config.data_id)
data_id = writer.write_data(dataset,
data_id=output_config.data_id,
replace=output_config.replace or False,
**write_params)
return data_id, dataset
def _describe_dataset(self, input_config: InputConfig) \
-> DatasetDescriptor:
opener_id = input_config.opener_id
store_params = input_config.store_params or {}
if input_config.store_id:
store_instance = get_data_store_instance(
input_config.store_id,
store_params=store_params,
store_pool=self._store_pool
)
opener = store_instance.store
else:
opener = new_data_opener(opener_id)
try:
descriptor = opener.describe_data(input_config.data_id,
data_type=DATASET_TYPE)
except DataStoreError as dse:
raise CubeGeneratorError(f'{dse}',
status_code=400) from dse
if not isinstance(descriptor, DatasetDescriptor):
raise RuntimeError(f'internal error: data store '
f'"{input_config.store_id}": '
f'expected DatasetDescriptor but got '
f'a {type(descriptor)}')
return descriptor
def open_cube(self, input_config: InputConfig) -> TransformedCube:
cube_config = self._cube_config
cube_params = cube_config.to_dict()
opener_id = input_config.opener_id
store_params = input_config.store_params or {}
open_params = input_config.open_params or {}
with observe_progress('reading cube', 3) as observer:
try:
if input_config.store_id:
store_instance = get_data_store_instance(
input_config.store_id,
store_params=store_params,
store_pool=self._store_pool)
store = store_instance.store
if opener_id is None:
opener_id = self._get_opener_id(input_config, store)
opener = store
open_params = dict(open_params)
open_params['opener_id'] = opener_id
else:
opener = new_data_opener(opener_id)
open_params = dict(open_params)
open_params.update(store_params)
open_params_schema = opener.get_open_data_params_schema(
input_config.data_id)
dataset_open_params = {
k: v
for k, v in cube_params.items()
if k in open_params_schema.properties
}
observer.worked(1)
dataset = opener.open_data(input_config.data_id, **open_params,
**dataset_open_params)
observer.worked(1)
except DataStoreError as dse:
raise CubeGeneratorError(f'{dse}', status_code=400) from dse
# Turn dataset into cube and grid_mapping
try:
cube, gm, _ = decode_cube(dataset, normalize=True)
except DatasetIsNotACubeError as e:
raise CubeGeneratorError(f'{e}') from e
observer.worked(1)
if dataset_open_params:
drop_names = [
k for k in dataset_open_params.keys()
if k not in _STEADY_CUBE_CONFIG_NAMES
]
cube_config = cube_config.drop_props(drop_names)
return cube, gm, cube_config
def open_cubes(input_configs: Sequence[InputConfig],
cube_config: CubeConfig,
store_pool: DataStorePool = None):
cubes = []
all_cube_params = cube_config.to_dict()
with observe_progress('Opening input(s)', len(input_configs)) as progress:
for input_config in input_configs:
open_params = {}
opener_id = input_config.opener_id
if input_config.store_id:
store_instance = get_data_store_instance(
input_config.store_id,
store_params=input_config.store_params,
store_pool=store_pool)
store = store_instance.store
if opener_id is None:
opener_ids = store.get_data_opener_ids(
data_id=input_config.data_id,
type_specifier=TYPE_SPECIFIER_CUBE)
if not opener_ids:
raise DataStoreError(
f'Data store "{input_config.store_id}" does not support data cubes'
)
opener_id = opener_ids[0]
opener = store
open_params.update(opener_id=opener_id,
**input_config.open_params)
else:
opener = new_data_opener(opener_id)
open_params.update(**input_config.store_params,
**input_config.open_params)
open_params_schema = opener.get_open_data_params_schema(
input_config.data_id)
cube_params = {
k: v
for k, v in all_cube_params.items()
if k in open_params_schema.properties
}
cube = opener.open_data(input_config.data_id, **open_params,
**cube_params)
cubes.append(cube)
progress.worked(1)
return cubes
def write_cube(cube: xr.Dataset,
output_config: OutputConfig,
store_pool: DataStorePool = None) -> str:
with observe_progress('Writing output', 1) as progress:
write_params = dict()
if output_config.store_id:
store_instance = get_data_store_instance(output_config.store_id,
store_params=output_config.store_params,
store_pool=store_pool)
writer = store_instance.store
write_params.update(writer_id=output_config.writer_id, **output_config.write_params)
else:
writer = new_data_writer(output_config.writer_id)
write_params.update(**output_config.store_params, **output_config.write_params)
# TODO: develop an adapter from Dask callback to ProgressObserver and use it here.
data_id = writer.write_data(cube,
data_id=output_config.data_id,
replace=output_config.replace or False,
**write_params)
progress.worked(1)
return data_id
