def _get_inf_ids_from_db():
try:
db = DataBase(Config.DATA_DB)
if db.connect():
inf_list = []
if db.db_type == DataBase.MONGO:
res = db.find("inf_list", {"deleted": 0}, {"id": True},
[('id', -1)])
else:
res = db.select(
"select id from inf_list where deleted = 0 order by rowid desc"
)
for elem in res:
if db.db_type == DataBase.MONGO:
inf_list.append(elem['id'])
else:
inf_list.append(elem[0])
db.close()
return inf_list
else:
InfrastructureList.logger.error(
"ERROR connecting with the database!.")
return []
except Exception:
InfrastructureList.logger.exception(
"ERROR loading data. Correct or delete it!!")
return []
def _get_data_from_db(db_url, inf_id=None, auth=None):
"""
Get data from DB.
If no inf_id specified all Infrastructures are loaded.
If auth is specified only auth data will be loaded.
"""
if InfrastructureList.init_table():
db = DataBase(db_url)
if db.connect():
inf_list = {}
if inf_id:
if db.db_type == DataBase.MONGO:
res = db.find("inf_list", {"id": inf_id}, {"data": True})
else:
res = db.select("select data from inf_list where id = %s", (inf_id,))
else:
if db.db_type == DataBase.MONGO:
res = db.find("inf_list", {"deleted": 0}, {"data": True}, [('_id', -1)])
else:
res = db.select("select data from inf_list where deleted = 0 order by rowid desc")
if len(res) > 0:
for elem in res:
if db.db_type == DataBase.MONGO:
data = elem['data']
else:
data = elem[0]
try:
if auth:
inf = IM.InfrastructureInfo.InfrastructureInfo.deserialize_auth(data)
else:
inf = IM.InfrastructureInfo.InfrastructureInfo.deserialize(data)
inf_list[inf.id] = inf
except Exception:
InfrastructureList.logger.exception(
"ERROR reading infrastructure from database, ignoring it!.")
else:
InfrastructureList.logger.warn("No data in database!.")
db.close()
return inf_list
else:
InfrastructureList.logger.error("ERROR connecting with the database!.")
return {}
else:
InfrastructureList.logger.error("ERROR connecting with the database!.")
return {}
def test_mongo_db(self, mongo):
client = MagicMock()
mongo.return_value = client
database = MagicMock()
client.__getitem__.return_value = database
database.client = client
table = MagicMock()
database.__getitem__.return_value = table
table.database = database
db_url = "mongodb://*****:*****@server/db_name"
db = DataBase(db_url)
self.assertTrue(db.connect())
database.collection_names.return_value = ['table1']
res = db.table_exists("test")
self.assertFalse(res)
res = db.table_exists("table1")
self.assertTrue(res)
res = db.replace('table', {}, {'id': 1, 'data': 'test1'})
self.assertTrue(res)
self.assertEqual(table.replace_one.call_args_list[0][0], ({}, {'data': 'test1', 'id': 1}, True))
table.find.return_value = [{'id': 2, 'data': 'test2', '_id': 2}]
res = db.find('table', {'id': 2}, {'data': True})
self.assertEqual(len(res), 1)
self.assertEqual(table.find.call_args_list[0][0], ({'id': 2}, {'_id': False, 'data': True}))
del_res = MagicMock()
del_res.deleted_count = 1
table.delete_many.return_value = del_res
res = db.delete('table', {'id': 1})
self.assertEqual(res, 1)
self.assertEqual(table.delete_many.call_args_list[0][0], ({'id': 1},))
db.close()