def test_query_sort(self, field, direction):
"""
Note: MongoDB 3.6 has changed how they compare array fields in a
sort. It used to compare the arrays element-by-element,
continuing until any "ties" were broken. Now it only
compares the highest/lowest valued element, apparently
ignoring the rest.
For this reason, a MongoDB query will not properly sort our
status and labels array fields, at least not in a simple way.
"""
session = connect_mongodb(self.settings)
recs, _total = session.query(sort=[(field, direction)],
projection=['metadata.labels'])
assert len(recs) == 26
for rec1, rec2 in zip(recs, recs[1:]):
value1 = self.deep_get(rec1, field, default=None)
value2 = self.deep_get(rec2, field, default=None)
if direction == 'desc':
if value1 is not None and value2 is not None:
assert value1 >= value2
elif value1 is None and value2 is not None:
# None value is less than a real value
assert False
else:
if value1 is not None and value2 is not None:
assert value1 <= value2
elif value2 is None and value1 is not None:
# None value is less than a real value
assert False
def test_replace_one(self):
session = connect_mongodb(self.settings)
ID = session.new_oil_id()
orig_name = 'original name'
new_name = 'new name'
# create a minimal oil
oil = Oil(ID)
oil.metadata.name = orig_name
# add it:
session.insert_one(oil)
oil_json = session.find_one(ID)
assert oil_json['oil_id'] == ID
assert oil_json['metadata']['name'] == orig_name
# update it:
oil.metadata.name = new_name
res = session.replace_one(oil)
assert res.matched_count == 1
assert res.modified_count == 1
oil_json = session.find_one(ID)
assert oil_json['oil_id'] == ID
assert oil_json['metadata']['name'] == new_name
def backup_db(settings, base_path):
"""
Here is where we backup our database. This is what we want to do:
- If the database does not exist, we flag an error and exit:
- If the database is already there:
- Gather the collections by name
- We want to start with an empty folder, so clear the base path
- For each collection:
- Create a subfolder under our path using the collection name
- Iterate the objects in the collection
- Save the objects as <basepath>/<collection>/<object>
"""
logger.info('connect_mongodb()...')
client = connect_mongodb(settings)
if settings['mongodb.database'] not in client.list_database_names():
print(f'The {settings["mongodb.database"]} database does not exist!')
return
db = client.get_database(settings['mongodb.database'])
collections = db.list_collection_names()
cleanup_folder(base_path)
for collection_name in collections:
add_folder(base_path, collection_name)
collection = getattr(db, collection_name)
for rec in collection.find({}):
export_to_file(base_path, collection_name, rec)
print('\nDatabase backup done!\n')
def restore_db(settings, base_path):
"""
Here is where we restore our database. This is what we want to do:
If the restore path does not exist, we flag an error and exit.
Otherwise:
- If the database does not exist, create it
- If the database is already there, initialize it
- Gather the collection names by directory name
- For each collection name:
- create the collection
- for each object in the collection directory
- Save the objects
"""
if not os.path.exists(base_path):
print(f'No path named {base_path}!')
return
logger.info('connect_mongodb()...')
client = connect_mongodb(settings)
drop_db(client, settings['mongodb.database'])
db = client.get_database(settings['mongodb.database'])
create_indices(db)
# load the database
for collection_name in os.listdir(base_path):
# filter out dotfiles
if not collection_name.startswith("."):
load_collection(db, base_path, collection_name)
print('\nDatabase restore done!\n')
def test_query(self):
session = connect_mongodb(self.settings)
recs, total = session.query()
assert len(recs) == 26 # our test set size
assert total == 26
def test_query_by_id(self):
session = connect_mongodb(self.settings)
recs, _total = session.query(oil_id='AD00020')
assert len(recs) == 1
assert recs[0]['oil_id'] == 'AD00020'
def test_connect_mongodb(mongodb_settings):
"""
Test a successful MongoDB connection. We need to have a running
MongoDB server for this to pass.
"""
mongodb_client = connect_mongodb(mongodb_settings)
assert mongodb_client.address == ('localhost', 27017)
def test_new_oil_id(self):
session = connect_mongodb(self.settings)
oil_id = session.new_oil_id()
assert isinstance(oil_id, str)
assert oil_id[:2] == 'XX'
assert oil_id[2:].isdigit()
def test_one_label_good(self):
session = connect_mongodb(self.settings)
for _id in (0, '0'):
rec = session.get_labels(_id)
for attr in ('name', 'product_types'):
assert attr in rec
def test_insert_one(self):
session = connect_mongodb(self.settings)
# create a minimal oil
ID = session.new_oil_id()
oil = Oil(ID)
# add it:
inserted_id = session.insert_one(oil)
assert inserted_id == ID
def test_query_by_labels(self, labels, expected):
session = connect_mongodb(self.settings)
recs, _total = session.query(labels=labels)
for rec in recs:
print(rec['metadata']['labels'])
assert len(recs) > 4 # so it's not too fragile if the data changes
for rec in recs:
assert rec['metadata']['labels'] == expected
def test_find_one(self):
session = connect_mongodb(self.settings)
# create a minimal oil
ID = session.new_oil_id()
oil = Oil(ID)
# add it:
session.insert_one(oil)
new_oil = session.find_one(ID)
assert new_oil['oil_id'] == ID
def test_one_label_bad(self):
session = connect_mongodb(self.settings)
# test non-existent, but valid IDs
assert session.get_labels(20000) is None
assert session.get_labels('20000') is None
# test an id that can't even be used
with pytest.raises(ValueError):
session.get_labels('bogus')
# test negative id that can't even be used
with pytest.raises(ValueError):
session.get_labels(-1)
def test_query_by_api(self, api, len_results, expected):
session = connect_mongodb(self.settings)
recs, _total = session.query(api=api)
assert len(recs) == len_results
for rec in recs:
_min, _max = expected
if _min is not None:
assert rec['metadata']['API'] >= _min
if _max is not None:
assert rec['metadata']['API'] <= _max
def db_setup():
# imported here so it's not needed if these tests aren't run
from adios_db.util.db_connection import connect_mongodb
# called for running each test in 'a' directory
config_path = os.path.join(os.path.dirname(__file__), '..', '..', '..',
'config-example.ini')
config = ConfigParser()
config.read(config_path)
print(config)
settings = dict(config['app:adios_db'].items())
return connect_mongodb(settings)
def test_query_with_projection(self):
session = connect_mongodb(self.settings)
recs, _total = session.query(projection=['metadata.name'])
assert len(recs) == 26 # our test set size
for rec in recs:
# We should only get the oil_id plus one field
# That field should be in the proper context however
assert len(rec) == 2
assert 'oil_id' in rec
assert 'metadata' in rec
assert len(rec['metadata']) == 1
assert 'name' in rec['metadata']
def test_query_by_name_location(self):
session = connect_mongodb(self.settings)
q_text = 'Alaska North Slope'
recs, _total = session.query(text=q_text)
assert len(recs) == 3
for rec in recs:
assert q_text.lower() in rec['metadata']['name'].lower()
q_text = 'Saudi Arabia'
recs, _total = session.query(text=q_text)
assert len(recs) == 4
for rec in recs:
assert q_text.lower() in rec['metadata']['location'].lower()
def test_delete_one(self):
session = connect_mongodb(self.settings)
ID = session.new_oil_id()
# create a minimal oil
oil = Oil(ID)
# add it:
session.insert_one(oil)
oil_json = session.find_one(ID)
assert oil_json['oil_id'] == ID
# delete it:
session.delete_one(oil.oil_id)
oil_json = session.find_one(ID)
assert oil_json is None
def init_db(settings, show_prompt=True):
"""
Here is where we create and initialize our database. This is what
we want to do:
- if the database is already there, we prompt for recreation
- if the database does not exist, or we are prompted for recreation:
- drop the database
- create the tables (if necessary)
- load the basic infrastructure data:
- Labels
"""
logger.info('connect_mongodb()...')
client = connect_mongodb(settings)
if show_prompt:
if settings['mongodb.database'] in client.list_database_names():
if prompt_drop_db():
print('Alright then...continuing on...')
else:
print('Ok, quitting the database initialization now...')
return
drop_db(client, settings['mongodb.database'])
db = client.get_database(settings['mongodb.database'])
print()
load_labels(db)
print_all_labels(db)
create_indices(db)
print('\nDatabase initialization done!\n')
def import_db_cmd(argv=sys.argv):
# Python 3 has made stderr buffered, so we have to fix it
sys.stderr = io.TextIOWrapper(sys.stderr.detach().detach(),
write_through=True)
logging.basicConfig(level=logging.INFO)
args = argp.parse_args(argv[1:])
if args.config is not None:
settings = file_settings(args.config)
else:
print('Using default settings')
settings = default_settings()
_add_datafiles(settings)
if args.overwrite:
settings['overwrite'] = True
else:
settings['overwrite'] = False
logger.info('connect_mongodb()...')
client = connect_mongodb(settings)
init_menu_item_collections(client, settings)
if args.all:
try:
add_all(settings)
except Exception:
print("{0}() FAILED\n".format(add_all.__name__))
raise
exit(0)
try:
import_db(settings)
except Exception:
print("{0}() FAILED\n".format(import_db.__name__))
raise
def oil_query(settings, export_file):
logger.info('connect_mongodb()...')
connect_mongodb(settings)
def test_init(self):
session = connect_mongodb(self.settings)
assert hasattr(session, 'get_labels')
def test_query_with_paging(self, page, expected):
session = connect_mongodb(self.settings)
recs, _total = session.query(page=page)
assert len(recs) == expected
# mockup of what we would like in order to separate
# the mongodb API from our code
from adios_db.util.db_connection import connect_mongodb
from adios_db.util.settings import file_settings
from pprint import pprint
# The session will be an instance of our oil_db session object, not mongodb
# our session is already connected to the database specified in settings
# We will be able to construct multiple session objects, each one connected to
# a mongodb database
settings = file_settings('settings_default.ini')
session = connect_mongodb(settings)
#session2 = connect_mongodb(settings2)
# The general purpose query function (this will be very verbose)
#print('Open query...')
#for rec in session.query():
# pprint(rec)
print('Query by id...')
print('Return only the name & location fields...')
for rec in session.query(oil_id='AD00020',
projection=['metadata.name',
'metadata.location']):
pprint(rec)
print('\n\nQuery by name and location, inclusive and case-sensitive.')
def test_all_labels(self):
session = connect_mongodb(self.settings)
recs = session.get_labels()
assert len(recs) == len(types_to_labels.product_types)
def test_init(self):
session = connect_mongodb(self.settings)
assert hasattr(session, 'query') # our object, not mongodb
import sys
from pymongo import MongoClient
from adios_db.util.db_connection import connect_mongodb
from adios_db.data_sources.oil.estimations import OilEstimation
from pprint import PrettyPrinter
pp = PrettyPrinter(indent=2, width=120)
client = connect_mongodb({'mongodb.host': 'localhost',
'mongodb.port': 27017,
'mongodb.database': 'oil_database',
'mongodb.alias': 'oil-db-app'})
db = client.oil_database
records = db.imported_record
oils = db.oil
def deep_get(obj, attr_path, default=None):
if isinstance(attr_path, str):
attr_path = attr_path.split('.')
attrs, current = attr_path, obj
try:
for p in attrs:
current = current[p]
def test_mongodb_connect_exceptions():
with pytest.raises(KeyError):
connect_mongodb({})
with pytest.raises(KeyError):
connect_mongodb({'mongodb.host': 'localhost'}) # port setting missing
