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