def get_engine(echo=True):
""" Creates a engine to a specific database.
:returns: engine
"""
global _ENGINE
if not _ENGINE:
sql_connection = config_get(DATABASE_SECTION, 'default')
config_params = [('pool_size', int), ('max_overflow', int), ('pool_timeout', int),
('pool_recycle', int), ('echo', int), ('echo_pool', str),
('pool_reset_on_return', str), ('use_threadlocal', int)]
# params = {'max_identifier_length': 128}
params = {}
for param, param_type in config_params:
try:
params[param] = param_type(config_get(DATABASE_SECTION, param))
except: # noqa: B901
pass
params['execution_options'] = {'schema_translate_map': {None: DEFAULT_SCHEMA_NAME}}
_ENGINE = create_engine(sql_connection, **params)
if 'mysql' in sql_connection:
event.listen(_ENGINE, 'checkout', mysql_ping_listener)
event.listen(_ENGINE, 'connect', mysql_convert_decimal_to_float)
elif 'postgresql' in sql_connection:
event.listen(_ENGINE, 'connect', psql_convert_decimal_to_float)
elif 'sqlite' in sql_connection:
event.listen(_ENGINE, 'connect', _fk_pragma_on_connect)
elif 'oracle' in sql_connection:
event.listen(_ENGINE, 'connect', my_on_connect)
assert _ENGINE
return _ENGINE
def get_session_pool():
sql_connection = config_get('database', 'default')
sql_connection = sql_connection.replace("oracle://", "")
user_pass, tns = sql_connection.split('@')
user, passwd = user_pass.split(':')
db_pool = cx_Oracle.SessionPool(user,
passwd,
tns,
min=12,
max=20,
increment=1)
schema = config_get('database', 'schema')
return db_pool, schema
def get_dump_engine(echo=False):
""" Creates a dump engine to a specific database.
:returns: engine """
statements = list()
def dump(sql, *multiparams, **params):
statement = str(sql.compile(dialect=engine.dialect))
if statement in statements:
return
statements.append(statement)
if statement.endswith(')\n\n'):
if engine.dialect.name == 'oracle':
print(statement.replace(')\n\n', ') PCTFREE 0;\n'))
else:
print(statement.replace(')\n\n', ');\n'))
elif statement.endswith(')'):
print(statement.replace(')', ');\n'))
else:
print(statement)
sql_connection = config_get(DATABASE_SECTION, 'default')
engine = create_engine(sql_connection,
echo=echo,
strategy='mock',
executor=dump)
return engine
def get_rest_cacher_dir():
cacher_dir = None
if config_has_section('rest') and config_has_option('rest', 'cacher_dir'):
cacher_dir = config_get('rest', 'cacher_dir')
if cacher_dir and os.path.exists(cacher_dir):
return cacher_dir
raise Exception("cacher_dir is not defined or it doesn't exist")
def setup_logging(name):
"""
Setup logging
"""
if config_has_section('common') and config_has_option('common', 'loglevel'):
loglevel = getattr(logging, config_get('common', 'loglevel').upper())
else:
loglevel = logging.INFO
if config_has_section('common') and config_has_option('common', 'logdir'):
logging.basicConfig(filename=os.path.join(config_get('common', 'logdir'), name),
level=loglevel,
format='%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s')
else:
logging.basicConfig(stream=sys.stdout, level=loglevel,
format='%(asctime)s\t%(name)s\t%(levelname)s\t%(message)s')
def load_config_agents():
if config_has_section(Sections.Main) and config_has_option(
Sections.Main, 'agents'):
agents = config_get(Sections.Main, 'agents')
agents = agents.split(',')
agents = [d.strip() for d in agents]
return agents
return []
def get_rest_host():
"""
Function to get rest host
"""
host = config_get('rest', 'host')
url_prefix = get_rest_url_prefix()
while host.endswith("/"):
host = host[:-1]
if url_prefix:
host = ''.join([host, url_prefix])
return host
def load_plugin_sequence(config_section, config_option='plugin_sequence'):
"""
load plugin sequence
"""
plugin_sequence = []
if config_has_section(config_section) and config_has_option(
config_section, config_option):
plugin_sequence = config_get(config_section, config_option)
plugin_sequence = plugin_sequence.split(",")
plugin_sequence = [plugin.strip() for plugin in plugin_sequence]
return plugin_sequence
def get_rest_url_prefix():
if config_has_section('rest') and config_has_option('rest', 'url_prefix'):
url_prefix = config_get('rest', 'url_prefix')
else:
url_prefix = None
if url_prefix:
while url_prefix.startswith('/'):
url_prefix = url_prefix[1:]
while url_prefix.endswith('/'):
url_prefix = url_prefix[:-1]
url_prefix = '/' + url_prefix
return url_prefix
def check_database():
"""
Function to check whether database is defined in config.
To be used to decide whether to skip some test functions.
:returns True: if database.default is available. Otherwise False.
"""
if config_has_option('database', 'default'):
database = config_get('database', 'default')
if database:
return True
return False
def check_rest_host():
"""
Function to check whether rest host is defined in config.
To be used to decide whether to skip some test functions.
:returns True: if rest host is available. Otherwise False.
"""
if config_has_option('rest', 'host'):
host = config_get('rest', 'host')
if host:
return True
return False
def build_database(echo=True, tests=False):
"""Build the database. """
engine = session.get_engine(echo=echo)
if config_has_option('database', 'schema'):
schema = config_get('database', 'schema')
if schema and not engine.dialect.has_schema(engine, schema):
print('Schema set in config, trying to create schema:', schema)
try:
engine.execute(CreateSchema(schema))
except Exception as e:
print(
'Cannot create schema, please validate manually if schema creation is needed, continuing:',
e)
print(traceback.format_exc())
models.register_models(engine)
def destroy_everything(echo=True):
""" Using metadata.reflect() to get all constraints and tables.
metadata.drop_all() as it handles cyclical constraints between tables.
Ref. http://www.sqlalchemy.org/trac/wiki/UsageRecipes/DropEverything
"""
engine = session.get_engine(echo=echo)
try:
# the transaction only applies if the DB supports
# transactional DDL, i.e. Postgresql, MS SQL Server
with engine.begin() as conn:
inspector = inspect(conn) # type: Union[Inspector, PGInspector]
for tname, fkcs in reversed(
inspector.get_sorted_table_and_fkc_names(schema='*')):
if tname:
drop_table_stmt = DropTable(
Table(tname, MetaData(), schema='*'))
conn.execute(drop_table_stmt)
elif fkcs:
if not engine.dialect.supports_alter:
continue
for tname, fkc in fkcs:
fk_constraint = ForeignKeyConstraint((), (), name=fkc)
Table(tname, MetaData(), fk_constraint)
drop_constraint_stmt = DropConstraint(fk_constraint)
conn.execute(drop_constraint_stmt)
if config_has_option('database', 'schema'):
schema = config_get('database', 'schema')
if schema:
conn.execute(DropSchema(schema, cascade=True))
if engine.dialect.name == 'postgresql':
assert isinstance(inspector,
PGInspector), 'expected a PGInspector'
for enum in inspector.get_enums(schema='*'):
sqlalchemy.Enum(**enum).drop(bind=conn)
except Exception as e:
print('Cannot destroy db:', e)
print(traceback.format_exc())
from sqlalchemy import create_engine, event
from sqlalchemy.exc import DatabaseError, DisconnectionError, OperationalError, TimeoutError
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, scoped_session
from idds.common.config import config_get, config_has_option
from idds.common.exceptions import IDDSException, DatabaseException
DATABASE_SECTION = 'database'
BASE = declarative_base()
DEFAULT_SCHEMA_NAME = None
if config_has_option(DATABASE_SECTION, 'schema'):
DEFAULT_SCHEMA_NAME = config_get(DATABASE_SECTION, 'schema')
if DEFAULT_SCHEMA_NAME:
BASE.metadata.schema = DEFAULT_SCHEMA_NAME
_MAKER, _ENGINE, _LOCK = None, None, Lock()
def _fk_pragma_on_connect(dbapi_con, con_record):
# Hack for previous versions of sqlite3
try:
dbapi_con.execute('pragma foreign_keys=ON')
except AttributeError:
pass
def mysql_ping_listener(dbapi_conn, connection_rec, connection_proxy):
def get_cache_url():
if config_has_section('cache') and config_has_option('cache', 'url'):
return config_get('cache', 'url')
return '127.0.0.1:11211'
