def test_dump_observations(db_connection):
with RdbApi(db_connection) as rdb_api, BlobApi(db_connection) as blob_api:
observations = pd.DataFrame.from_dict(observation_data)
_dump_observations(rdb_api=rdb_api,
blob_api=blob_api,
observations=observations)
blob_api.commit()
rdb_api.commit()
with RdbApi(db_connection) as rdb_api, BlobApi(db_connection) as blob_api:
poly_obs = rdb_api.get_observation("POLY_OBS")
assert poly_obs.id is not None
key_indexes = blob_api.get_blob(poly_obs.key_indexes_ref)
assert key_indexes.data == [0, 2, 4, 6, 8]
data_indexes = blob_api.get_blob(poly_obs.data_indexes_ref)
assert data_indexes.data == [10, 12, 14, 16, 18]
values = blob_api.get_blob(poly_obs.values_ref)
assert values.data == [2.0, 7.1, 21.1, 31.8, 53.2]
stds = blob_api.get_blob(poly_obs.stds_ref)
assert stds.data == [0.1, 1.1, 4.1, 9.1, 16.1]
test_obs = rdb_api.get_observation("TEST_OBS")
assert test_obs.id is not None
key_indexes = blob_api.get_blob(test_obs.key_indexes_ref)
assert key_indexes.data == [3, 6, 9]
data_indexes = blob_api.get_blob(test_obs.data_indexes_ref)
assert data_indexes.data == [3, 6, 9]
values = blob_api.get_blob(test_obs.values_ref)
assert values.data == [6, 12, 18]
stds = blob_api.get_blob(test_obs.stds_ref)
assert stds.data == [0.1, 0.2, 0.3]
def test_add_observation(db_connection):
observation_name = "test"
key_indexes = [0, 3]
data_indexes = [0, 3]
values = [22.1, 44.2]
stds = [1, 3]
with RdbApi(db_connection) as rdb_api, BlobApi(db_connection) as blob_api:
key_indexes_df = blob_api.add_blob(data=key_indexes)
data_indexes_df = blob_api.add_blob(data=data_indexes)
values_df = blob_api.add_blob(data=values)
stds_df = blob_api.add_blob(data=stds)
blob_api.commit()
observation = rdb_api.add_observation(
name=observation_name,
key_indexes_ref=key_indexes_df.id,
data_indexes_ref=data_indexes_df.id,
values_ref=values_df.id,
stds_ref=stds_df.id,
)
rdb_api.commit()
assert observation is not None
with RdbApi(db_connection) as rdb_api, BlobApi(db_connection) as blob_api:
observation = rdb_api.get_observation(observation_name)
assert observation is not None
assert blob_api.get_blob(
observation.key_indexes_ref).data == key_indexes
assert blob_api.get_blob(
observation.data_indexes_ref).data == data_indexes
assert blob_api.get_blob(observation.values_ref).data == values
assert blob_api.get_blob(observation.stds_ref).data == stds
def dump_to_new_storage(reference=None, rdb_session=None, blob_session=None):
start_time = time.time()
logger.debug("Starting extraction...")
if rdb_session is None:
rdb_session = ERT_STORAGE.RdbSession()
if blob_session is None:
blob_session = ERT_STORAGE.BlobSession()
rdb_api = RdbApi(session=rdb_session)
blob_api = BlobApi(session=blob_session)
try:
priors = _extract_and_dump_priors(
rdb_api=rdb_api) if reference is None else []
ensemble = _create_ensemble(rdb_api,
reference=reference,
priors=priors)
_extract_and_dump_observations(rdb_api=rdb_api, blob_api=blob_api)
_extract_and_dump_parameters(
rdb_api=rdb_api,
blob_api=blob_api,
ensemble_name=ensemble.name,
priors=priors,
)
_extract_and_dump_responses(rdb_api=rdb_api,
blob_api=blob_api,
ensemble_name=ensemble.name)
_extract_and_dump_update_data(ensemble.id, ensemble.name, rdb_api,
blob_api)
rdb_session.commit()
blob_session.commit()
ensemble_name = ensemble.name
end_time = time.time()
logger.debug(
"Extraction done... (Took {:.2f} seconds)".format(end_time -
start_time))
logger.debug("All ensembles in database: {}".format(", ".join(
[ensemble.name for ensemble in rdb_api.get_all_ensembles()])))
except:
rdb_session.rollback()
blob_session.rollback()
raise
finally:
rdb_session.close()
blob_session.close()
return ensemble_name
def test_add_reference_ensemble(db_connection):
reference_ensemble_name = "test_ensemble"
with RdbApi(db_connection) as rdb_api:
ensemble = rdb_api.add_ensemble(name=reference_ensemble_name)
rdb_api.commit()
with RdbApi(db_connection) as rdb_api:
result_ensemble = rdb_api.add_ensemble(
name="result_ensemble",
reference=(reference_ensemble_name, "es_mda"))
rdb_api.commit()
assert result_ensemble.parent.ensemble_reference.name == reference_ensemble_name
def dump_to_new_storage(reference=None,
rdb_connection=None,
blob_connection=None):
start_time = time.time()
logger.debug("Starting extraction...")
if rdb_connection is None:
rdb_url = "sqlite:///entities.db"
rdb_connection = connections.get_rdb_connection(rdb_url)
rdb_api = RdbApi(connection=rdb_connection)
if blob_connection is None:
blob_url = "sqlite:///blobs.db"
blob_connection = connections.get_blob_connection(blob_url)
blob_api = BlobApi(connection=blob_connection)
with rdb_api, blob_api:
priors = _extract_and_dump_priors(
rdb_api=rdb_api) if reference is None else []
ensemble = _create_ensemble(rdb_api,
reference=reference,
priors=priors)
_extract_and_dump_observations(rdb_api=rdb_api, blob_api=blob_api)
_extract_and_dump_parameters(
rdb_api=rdb_api,
blob_api=blob_api,
ensemble_name=ensemble.name,
priors=priors,
)
_extract_and_dump_responses(rdb_api=rdb_api,
blob_api=blob_api,
ensemble_name=ensemble.name)
_extract_and_dump_update_data(ensemble.id, ensemble.name, rdb_api,
blob_api)
blob_api.commit()
rdb_api.commit()
ensemble_name = ensemble.name
end_time = time.time()
logger.debug(
"Extraction done... (Took {:.2f} seconds)".format(end_time -
start_time))
logger.debug("All ensembles in database: {}".format(", ".join(
[ensemble.name for ensemble in rdb_api.get_all_ensembles()])))
rdb_connection.close()
blob_connection.close()
return ensemble_name
def test_get_parameter_definitions_by_ensemble_id(db_info):
populated_db, db_lookup = db_info
rdb_connection = connections.get_rdb_connection(populated_db)
with RdbApi(connection=rdb_connection) as rdb_api:
param_def = rdb_api.get_parameter_definitions_by_ensemble_id(
ensemble_id=db_lookup["ensemble"])
assert len(list(param_def)) == 3
def test_get_realizations_by_ensemble_id(db_info):
populated_db, db_lookup = db_info
rdb_connection = connections.get_rdb_connection(populated_db)
with RdbApi(connection=rdb_connection) as rdb_api:
reals = rdb_api.get_realizations_by_ensemble_id(
ensemble_id=db_lookup["ensemble"])
assert len(list(reals)) == 2
def test_get_all_ensembles(db_info):
populated_db, db_lookup = db_info
rdb_connection = connections.get_rdb_connection(populated_db)
with RdbApi(connection=rdb_connection) as rdb_api:
ens_list = rdb_api.get_all_ensembles()
assert len(list(ens_list)) == 1
assert ens_list[0].name == "ensemble_name"
def test_get_observation_attribute(db_info):
populated_db, _ = db_info
rdb_connection = connections.get_rdb_connection(populated_db)
with RdbApi(connection=rdb_connection) as rdb_api:
obs_attrib = rdb_api.get_observation_attribute(name="observation_one",
attribute="region")
assert obs_attrib == "1"
def test_dump_priors(db_connection):
priors = {
"COEFFS": [
{
"key": "COEFF_A",
"function": "UNIFORM",
"parameters": {
"MIN": 0.0,
"MAX": 1.0
},
},
{
"key": "COEFF_B",
"function": "UNIFORM",
"parameters": {
"MIN": 0.0,
"MAX": 2.0
},
},
{
"key": "COEFF_C",
"function": "UNIFORM",
"parameters": {
"MIN": 0.0,
"MAX": 5.0
},
},
]
}
with RdbApi(db_connection) as rdb_api:
_dump_priors(priors, rdb_api)
rdb_api.commit()
def test_get_observation(db_info):
populated_db, _ = db_info
rdb_connection = connections.get_rdb_connection(populated_db)
with RdbApi(connection=rdb_connection) as rdb_api:
obs = rdb_api.get_observation(name="observation_one")
assert obs is not None
assert obs.name == "observation_one"
def test_add_observation_response_definition_link(db_connection):
with RdbApi(db_connection) as rdb_api:
observation = rdb_api.add_observation(
name="test",
key_indexes_ref=None,
data_indexes_ref=None,
values_ref=None,
stds_ref=None,
)
ensemble = rdb_api.add_ensemble(name="test_ensemble")
response_definition = rdb_api.add_response_definition(
name="test_response_definition",
indexes_ref=0,
ensemble_name=ensemble.name)
rdb_api.flush()
link = rdb_api._add_observation_response_definition_link(
observation_id=observation.id,
response_definition_id=response_definition.id,
active_ref=1,
update_id=None,
)
rdb_api.commit()
assert link.id is not None
assert link.observation_id == observation.id
assert link.response_definition_id == response_definition.id
def test_add_parameter(db_connection):
value = 22.1
with RdbApi(db_connection) as rdb_api, BlobApi(db_connection) as blob_api:
value_df = blob_api.add_blob(data=value)
blob_api.commit()
ensemble = rdb_api.add_ensemble(name="test")
parameter_definition = rdb_api.add_parameter_definition(
name="test_param", group="test_group", ensemble_name=ensemble.name)
realization = rdb_api.add_realization(0, ensemble.name)
parameter = rdb_api.add_parameter(
name=parameter_definition.name,
group=parameter_definition.group,
value_ref=value_df.id,
realization_index=realization.index,
ensemble_name=ensemble.name,
)
rdb_api.commit()
with RdbApi(db_connection) as rdb_api, BlobApi(db_connection) as blob_api:
ensemble = rdb_api.get_ensemble(name="test")
assert ensemble.id is not None
parameter_definition = rdb_api._get_parameter_definition(
name="test_param", group="test_group", ensemble_id=ensemble.id)
assert parameter_definition.id is not None
assert parameter_definition.ensemble_id is not None
realization = rdb_api.get_realization(index=0,
ensemble_name=ensemble.name)
assert realization.id is not None
assert realization.ensemble_id is not None
parameter = rdb_api.get_parameter(
name="test_param",
group="test_group",
realization_index=0,
ensemble_name=ensemble.name,
)
assert parameter.id is not None
assert parameter.realization_id is not None
assert parameter.parameter_definition_id is not None
assert blob_api.get_blob(id=parameter.value_ref).data == value
def test_get_response_bundle(db_info):
populated_db, db_lookup = db_info
rdb_connection = connections.get_rdb_connection(populated_db)
with RdbApi(connection=rdb_connection) as rdb_api:
bundle = rdb_api.get_response_bundle(response_name="response_one",
ensemble_id=db_lookup["ensemble"])
assert bundle.name == "response_one"
assert len(bundle.responses) == 2
def test_add_response(db_connection):
indexes = [0, 2]
values = [22.1, 44.2]
with RdbApi(db_connection) as rdb_api, BlobApi(db_connection) as blob_api:
indexes_df = blob_api.add_blob(data=indexes)
values_df = blob_api.add_blob(data=values)
blob_api.commit()
ensemble = rdb_api.add_ensemble(name="test")
response_definition = rdb_api.add_response_definition(
name="test",
indexes_ref=indexes_df.id,
ensemble_name=ensemble.name)
realization = rdb_api.add_realization(0, ensemble.name)
response = rdb_api.add_response(
name=response_definition.name,
values_ref=values_df.id,
realization_index=realization.index,
ensemble_name=ensemble.name,
)
rdb_api.commit()
with RdbApi(db_connection) as rdb_api, BlobApi(db_connection) as blob_api:
ensemble = rdb_api.get_ensemble(name="test")
assert ensemble.id is not None
response_definition = rdb_api._get_response_definition(
name="test", ensemble_id=ensemble.id)
assert response_definition.id is not None
assert response_definition.ensemble_id is not None
assert blob_api.get_blob(
id=response_definition.indexes_ref).data == indexes
realization = rdb_api.get_realization(index=0,
ensemble_name=ensemble.name)
assert realization.id is not None
assert realization.ensemble_id is not None
response = rdb_api.get_response(response_definition.name,
realization.index, ensemble.name)
assert response.id is not None
assert response.realization_id is not None
assert response.response_definition_id is not None
assert blob_api.get_blob(id=response.values_ref).data == values
def api(initialize_databases):
session = ERT_STORAGE.Session()
api = RdbApi(session=session)
try:
yield api
finally:
session.rollback()
session.close()
def test_get_all_observation_keys(db_info):
populated_db, _ = db_info
rdb_connection = connections.get_rdb_connection(populated_db)
with RdbApi(connection=rdb_connection) as rdb_api:
obs_keys = rdb_api.get_all_observation_keys()
assert set([
"observation_one", "observation_two_first",
"observation_two_second"
]) == set(obs_keys)
def test_get_response_by_realization_id(db_info):
populated_db, db_lookup = db_info
rdb_connection = connections.get_rdb_connection(populated_db)
with RdbApi(connection=rdb_connection) as rdb_api:
response = rdb_api.get_response_by_realization_id(
response_definition_id=db_lookup["response_defition_one"],
realization_id=db_lookup["realization_0"],
)
assert response is not None
def db_api(populated_database):
db_lookup = populated_database
session = ERT_STORAGE.Session()
api = RdbApi(session=session)
try:
yield api, db_lookup
finally:
session.rollback()
session.close()
def test_add_realization(db_connection):
with RdbApi(db_connection) as rdb_api:
ensemble = rdb_api.add_ensemble(name="test_ensemble")
realizations = []
for i in range(5):
realization = rdb_api.add_realization(i, ensemble.name)
realizations.append(realization)
rdb_api.commit()
assert ensemble.id is not None
for realization in realizations:
assert realization.id is not None
with pytest.raises(sqlalchemy.exc.IntegrityError) as error, RdbApi(
connection=db_connection) as rdb_api:
rdb_api.add_realization(0, ensemble_name=ensemble.name)
rdb_api.commit()
def test_get_parameter_by_realization_id(db_info):
populated_db, db_lookup = db_info
rdb_connection = connections.get_rdb_connection(populated_db)
with RdbApi(connection=rdb_connection) as rdb_api:
param = rdb_api.get_parameter_by_realization_id(
parameter_definition_id=db_lookup["parameter_def_A_G"],
realization_id=db_lookup["realization_0"],
)
assert param.realization_id == db_lookup["realization_0"]
assert param.parameter_definition_id == db_lookup["parameter_def_A_G"]
def test_dump_parameters(db_connection):
ensemble_name = "default"
with RdbApi(db_connection) as rdb_api, BlobApi(db_connection) as blob_api:
ensemble = rdb_api.add_ensemble(name=ensemble_name)
for i in range(5):
rdb_api.add_realization(i, ensemble.name)
_dump_parameters(
rdb_api=rdb_api,
blob_api=blob_api,
parameters=parameters,
ensemble_name=ensemble.name,
priors=[],
)
blob_api.commit()
rdb_api.commit()
with RdbApi(db_connection) as rdb_api, BlobApi(db_connection) as blob_api:
parameter_0 = rdb_api.get_parameter("COEFF_A", "COEFFS", 0,
ensemble_name)
assert blob_api.get_blob(
parameter_0.value_ref).data == 0.7684484807065148
parameter_1 = rdb_api.get_parameter("COEFF_A", "COEFFS", 1,
ensemble_name)
assert blob_api.get_blob(
parameter_1.value_ref).data == 0.031542101926117616
parameter_2 = rdb_api.get_parameter("COEFF_A", "COEFFS", 2,
ensemble_name)
assert blob_api.get_blob(
parameter_2.value_ref).data == 0.9116906743615176
parameter_3 = rdb_api.get_parameter("COEFF_A", "COEFFS", 3,
ensemble_name)
assert blob_api.get_blob(
parameter_3.value_ref).data == 0.6985513230581486
parameter_4 = rdb_api.get_parameter("COEFF_A", "COEFFS", 4,
ensemble_name)
assert blob_api.get_blob(
parameter_4.value_ref).data == 0.5949261230249001
def test_get_parameter(db_info):
populated_db, db_lookup = db_info
rdb_connection = connections.get_rdb_connection(populated_db)
with RdbApi(connection=rdb_connection) as rdb_api:
param = rdb_api.get_parameter(name="A",
group="G",
realization_index=0,
ensemble_name="ensemble_name")
assert param.parameter_definition.name == "A"
assert param.parameter_definition.group == "G"
assert param.realization.index == 0
def test_get_parameter_bundle(db_info):
populated_db, db_lookup = db_info
rdb_connection = connections.get_rdb_connection(populated_db)
with RdbApi(connection=rdb_connection) as rdb_api:
bundle = rdb_api.get_parameter_bundle(
parameter_def_id=db_lookup["parameter_def_A_G"],
ensemble_id=db_lookup["ensemble"],
)
assert bundle is not None
assert bundle.name == "A"
assert bundle.group == "G"
assert len(bundle.parameters) == 2
def apis(initialize_databases):
rdb_session = ERT_STORAGE.RdbSession()
blob_session = ERT_STORAGE.BlobSession()
rdb_api = RdbApi(session=rdb_session)
blob_api = BlobApi(session=blob_session)
try:
yield rdb_api, blob_api
finally:
rdb_session.rollback()
blob_session.rollback()
rdb_session.close()
blob_session.close()
def test_observation_attribute(db_connection):
with RdbApi(db_connection) as rdb_api:
obs = rdb_api.add_observation(
name="test",
key_indexes_ref=1,
data_indexes_ref=1,
values_ref=1,
stds_ref=1,
)
obs.add_attribute("foo", "bar")
rdb_api.commit()
assert obs.get_attribute("foo") == "bar"
def test_dump_responses(db_connection):
ensemble_name = "default"
with RdbApi(db_connection) as rdb_api, BlobApi(db_connection) as blob_api:
ensemble = rdb_api.add_ensemble(name=ensemble_name)
observations = pd.DataFrame.from_dict(observation_data)
_dump_observations(rdb_api=rdb_api,
blob_api=blob_api,
observations=observations)
for i in range(5):
rdb_api.add_realization(i, ensemble.name)
_dump_response(
rdb_api=rdb_api,
blob_api=blob_api,
responses=responses,
ensemble_name=ensemble.name,
)
blob_api.commit()
rdb_api.commit()
with RdbApi(db_connection) as rdb_api, BlobApi(db_connection) as blob_api:
response_0 = rdb_api.get_response("POLY_RES", 0, ensemble_name)
response_values = blob_api.get_blob(response_0.values_ref).data
assert response_values == [
2.5995,
5.203511,
9.496884000000001,
15.479619,
23.151716,
32.513175000000004,
43.563995999999996,
56.304179,
70.73372400000001,
86.852631,
]
def test_two_ensembles_with_same_name(db_connection):
with RdbApi(db_connection) as rdb_api:
ensemble1 = rdb_api.add_ensemble(name="test_ensemble")
rdb_api.commit()
assert ensemble1.id is not None
time.sleep(1)
ensemble2 = rdb_api.add_ensemble(name="test_ensemble")
rdb_api.commit()
assert ensemble2.id is not None
ensemble = rdb_api.get_ensemble("test_ensemble")
assert ensemble.id == ensemble2.id
def db_apis(populated_database):
db_lookup = populated_database
rdb_session = ERT_STORAGE.RdbSession()
blob_session = ERT_STORAGE.BlobSession()
rdb_api = RdbApi(session=rdb_session)
blob_api = BlobApi(session=blob_session)
try:
yield rdb_api, blob_api, db_lookup
finally:
rdb_session.rollback()
blob_session.rollback()
rdb_session.close()
blob_session.close()
def session(self):
"""Provide a transactional scope around a series of operations."""
rdb_session = ERT_STORAGE.RdbSession()
blob_session = ERT_STORAGE.BlobSession()
rdb_api = RdbApi(session=rdb_session)
blob_api = BlobApi(session=blob_session)
try:
yield StorageApi(rdb_api, blob_api)
rdb_session.commit()
blob_session.commit()
except:
rdb_session.rollback()
blob_session.rollback()
raise
finally:
rdb_session.close()
blob_session.close()
