def test_sqlite_db(self):
filename = "/tmp/inf.dat"
if os.path.exists(filename):
os.unlink(filename)
db_url = "sqlite://" + filename
db = DataBase(db_url)
self.assertTrue(db.connect())
if not db.table_exists("test"):
db.execute("CREATE TABLE test(id int PRIMARY KEY, date TIMESTAMP, data LONGBLOB)")
self.assertTrue(db.table_exists("test"))
db.execute("insert into test (id, data, date) values (%s, %s, now())", (1, "Data"))
res = db.select("select data from test where id = %s", (1,))
self.assertEqual(res, [("Data",)])
db.close()
def get_data_from_db(db_url):
db = DataBase(db_url)
if db.connect():
if not db.table_exists("inf_list"):
db.execute("CREATE TABLE inf_list(id VARCHAR(255) PRIMARY KEY, date TIMESTAMP, data LONGBLOB)")
db.close()
return {}
else:
inf_list = {}
res = db.select("select * from inf_list order by id desc")
if len(res) > 0:
for elem in res:
#inf_id = elem[0]
#date = elem[1]
try:
inf = pickle.loads(elem[2])
if not inf.deleted:
inf_list[inf.id] = inf
except:
InfrastructureManager.logger.exception("ERROR reading infrastructure %d from database, ignoring it!." % inf.id)
else:
InfrastructureManager.logger.error("ERROR getting inf_list from database!.")
db.close()
return inf_list
else:
InfrastructureManager.logger.error("ERROR connecting with the database!.")
return {}
def get_data_from_db(db_url):
db = DataBase(db_url)
if db.connect():
if not db.table_exists("inf_list"):
return {}
else:
inf_list = {}
res = db.select("select * from inf_list order by id desc")
if len(res) > 0:
for elem in res:
# inf_id = elem[0]
# date = elem[1]
try:
inf = pickle.loads(elem[2])
inf_list[inf.id] = inf
except:
sys.stderr.write(
"ERROR reading infrastructure from database, ignoring it!."
)
sys.exit(-1)
else:
sys.stderr.write("ERROR getting inf_list from database!.")
sys.exit(-1)
db.close()
return inf_list
else:
sys.stderr.write("ERROR connecting with the database!.")
sys.exit(-1)
def test_db(self, mdb_conn):
filename = "/tmp/inf.dat"
if os.path.exists(filename):
os.unlink(filename)
db_url = "sqlite://" + filename
db = DataBase(db_url)
self.assertTrue(db.connect())
if not db.table_exists("test"):
db.execute(
"CREATE TABLE test(id int PRIMARY KEY, date TIMESTAMP, data LONGBLOB)"
)
self.assertTrue(db.table_exists("test"))
db.execute("insert into test (id, data, date) values (%s, %s, now())",
(1, "Data"))
res = db.select("select data from test where id = %s", (1, ))
self.assertEqual(res, [("Data", )])
db.close()
connection = MagicMock()
mdb_conn.return_value = connection
db_url = "mysql://*****:*****@server/db_name"
db = DataBase(db_url)
self.assertTrue(db.connect())
if not db.table_exists("test"):
db.execute(
"CREATE TABLE test(id int PRIMARY KEY, date TIMESTAMP, data LONGBLOB)"
)
db.execute("insert into test (id, data, date) values (%s, %s, now())",
(1, "Data"))
cursor = MagicMock()
cursor.fetchall.return_value = [("Data", )]
connection.cursor.return_value = cursor
res = db.select("select data from test where id = %s", (1, ))
self.assertEqual(res, [("Data", )])
db.close()
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()
def init_table():
""" Creates de database """
db = DataBase(Config.DATA_DB)
if db.connect():
if not db.table_exists("inf_list"):
db.execute(
"CREATE TABLE inf_list(id VARCHAR(255) PRIMARY KEY, deleted INTEGER,"
" date TIMESTAMP, data LONGBLOB)")
db.close()
return True
else:
InfrastructureList.logger.error(
"ERROR connecting with the database!.")
return False
def init_table():
""" Creates de database """
db = DataBase(Config.DATA_DB)
if db.connect():
if not db.table_exists("inf_list"):
InfrastructureList.logger.debug("Creating the IM database!.")
if db.db_type == DataBase.MYSQL:
db.execute("CREATE TABLE inf_list(rowid INTEGER NOT NULL AUTO_INCREMENT UNIQUE,"
" id VARCHAR(255) PRIMARY KEY, deleted INTEGER, date TIMESTAMP, data LONGBLOB)")
elif db.db_type == DataBase.SQLITE:
db.execute("CREATE TABLE inf_list(id VARCHAR(255) PRIMARY KEY, deleted INTEGER,"
" date TIMESTAMP, data LONGBLOB)")
db.close()
return True
else:
InfrastructureList.logger.error("ERROR connecting with the database!.")
return False
def rename_old_data():
db = DataBase(Config.DATA_DB)
if db.connect():
if db.table_exists("inf_list"):
now = str(int(time.time() * 100))
if db.db_type == DataBase.SQLITE:
db.execute('ALTER TABLE inf_list RENAME TO inf_list_%s;' % now)
db.close()
elif db.db_type == DataBase.MYSQL:
db.execute('RENAME TABLE inf_list TO inf_list_%s;' % now)
db.close()
else:
db.close()
sys.stderr.write("ERROR connecting with the database!.")
sys.exit(-1)
else:
db.close()
else:
sys.stderr.write("ERROR connecting with the database!.")
sys.exit(-1)
import sys
sys.path.append("..")
sys.path.append(".")
from IM.config import Config
from IM.db import DataBase
if __name__ == "__main__":
if not Config.DATA_DB:
sys.stderr.write("No DATA_DB defined in the im.cfg file!!")
sys.exit(-1)
db = DataBase(Config.DATA_DB)
if db.connect():
if db.table_exists("inf_list"):
if db.db_type == DataBase.MYSQL:
sys.stdout.write("Updating DB: %s.\n" % Config.DATA_DB)
db.execute(
"ALTER TABLE `inf_list` ADD COLUMN `rowid` INT AUTO_INCREMENT UNIQUE FIRST;"
)
else:
sys.stdout.write("SQLite DB does not need to be updated.")
db.close()
else:
sys.stdout.write(
"There are no inf_list table. Do not need to update.")
else:
sys.stderr.write("Error connecting with DB: %s\n" % Config.DATA_DB)
sys.exit(-1)
