def __init__(self, settings):
self.__ESCAPE_PREFILTER = {
bool: int,
datetime: lambda dt: self.escape(self.datetime(dt)),
"iterable": self._prefilter_iterate,
}
self._transaction_state = self.__TRANSACTION_STATE_NONE
stpl = tuple((k, v) for k, v in settings.items())
self._db = preludedb.SQL(settings)
self._version = self._db.getServerVersion()
self._dbhash = hash(stpl)
self._dbtype = settings["type"]
def init_idmef_database(config):
"""
Create IDMEF database.
:return: prelude SQL object
:rtype: preludedb.SQL
"""
sql = preludedb.SQL(dict(config.idmef_database))
libpreludedb_sql_path = config.libpreludedb.get(
'sql_path', _LIBPRELUDEDB_DEFAULT_SQL_PATH)
# create database structure
sql_file_path = os.path.join(libpreludedb_sql_path,
'%s.sql' % config.idmef_database.type)
with open(sql_file_path, 'r') as sql_file:
sql.query(sql_file.read())
def clean_database(database):
"""
Remove all tables in a database.
:param database: database information from config
"""
db_type = database.type
sql = preludedb.SQL(dict(database))
if db_type == 'pgsql':
# https://stackoverflow.com/a/36023359
sql_query = """
DO $$ DECLARE
r RECORD;
BEGIN
FOR r IN (SELECT tablename FROM pg_tables WHERE schemaname = current_schema()) LOOP
EXECUTE 'DROP TABLE IF EXISTS ' || quote_ident(r.tablename) || ' CASCADE';
END LOOP;
END $$;"""
elif db_type == 'mysql':
# http://stackoverflow.com/a/18625545
sql_query = """
SET FOREIGN_KEY_CHECKS = 0;
SET GROUP_CONCAT_MAX_LEN=32768;
SET @tables = NULL;
SELECT GROUP_CONCAT('`', table_name, '`')
INTO @tables
FROM information_schema.tables
WHERE table_schema = (SELECT DATABASE());
SELECT IFNULL(@tables,'dummy') INTO @tables;
SET @tables = CONCAT('DROP TABLE IF EXISTS ', @tables);
PREPARE stmt FROM @tables;
EXECUTE stmt;
DEALLOCATE PREPARE stmt;
SET FOREIGN_KEY_CHECKS = 1;"""
else:
raise AttributeError('%s is not a valid database type for test suite' %
db_type)
sql.query(sql_query)
def __init__(self, config):
preludedb.DB.__init__(self, preludedb.SQL(dict(config)))