class GlobalInit:
def __init__(self, config_path='conf/config.json'):
from sm.engine.db import ConnectionPool # pylint: disable=import-outside-toplevel
self.sm_config = on_startup(config_path)
self.pool = ConnectionPool(self.sm_config['db'])
def __enter__(self):
return self.sm_config
def __exit__(self, ext_type, ext_value, traceback):
self.pool.close()
def empty_test_db(sm_config):
db_name = sm_config['db']['database']
db_owner = sm_config['db']['user']
_autocommit_execute(
{
**sm_config['db'], 'database': 'postgres'
},
f'DROP DATABASE IF EXISTS {db_name}',
f'CREATE DATABASE {db_name} OWNER {db_owner}',
)
conn_pool = ConnectionPool(sm_config['db'])
yield
conn_pool.close()
_autocommit_execute({
**sm_config['db'], 'database': 'postgres'
}, f'DROP DATABASE IF EXISTS {db_name}')
def save_additional_info_to_db(db_id, user_id, input_path):
conf = SMConfig.get_conf()
with ConnectionPool(conf['db']):
db = DB()
if db.select_one('SELECT * FROM molecular_db WHERE id = %s',
(db_id, )):
print(f'Updating existing molecular database {db_id}')
DB().alter(
'UPDATE molecular_db SET user_id = %s, input_path = %s WHERE id = %s',
(user_id, input_path, db_id),
)
else:
print(f'Specified molecular database {db_id} does not exist.')
def main():
parser = argparse.ArgumentParser(
description='Merge mol_dbs and adducts into config')
parser.add_argument('--config',
default='conf/config.json',
help='SM config path')
args = parser.parse_args()
SMConfig.set_path(args.config)
init_loggers(SMConfig.get_conf()['logs'])
conf = SMConfig.get_conf()
with ConnectionPool(conf['db']):
db = DB()
populate_ion_formula(db)
populate_ions(db)
populate_ion_id(db)
def migrate_dataset(ds, i, n, force=False):
output.print(f'Migrating dataset {ds["id"]} ({i}/{n})')
with ConnectionPool(sm_config['db']):
try:
with timeit('Dataset'):
if force or ds['status'] == DatasetStatus.FINISHED:
migrate_isotopic_images(ds['id'])
migrate_ion_thumbnail(ds['id'])
migrate_optical_images(ds['id'])
except Exception:
output.print(f'Migration of {ds["id"]} failed:\n{traceback.format_exc()}')
lock.acquire()
with open('FAILED_DATASETS.txt', 'a') as f:
f.write(',' + ds['id'])
lock.release()
else:
lock.acquire()
with open('SUCCEEDED_DATASETS.txt', 'a') as f:
f.write(',' + ds['id'])
lock.release()
return output.flush()
def __init__(self, config_path='conf/config.json'):
from sm.engine.db import ConnectionPool # pylint: disable=import-outside-toplevel
self.sm_config = on_startup(config_path)
self.pool = ConnectionPool(self.sm_config['db'])
args = parser.parse_args()
image_base_path = args.image_base_path
sm_config = on_startup(args.config)
bucket_name = sm_config['image_storage']['bucket']
logging.getLogger('engine').setLevel(logging.WARNING)
logging.getLogger('engine.db').setLevel(logging.WARNING)
output = Output()
lock = Lock()
if args.ds_ids:
ds_ids = list({ds_id for ds_id in args.ds_ids.split(',') if ds_id})
with ConnectionPool(sm_config['db']):
dss_to_process = DB().select_with_fields(SEL_SPEC_DSS, params=(ds_ids,))
n = len(dss_to_process)
for i, ds in enumerate(dss_to_process, 1):
logs = migrate_dataset(ds, i, n, force=True)
print(logs)
else:
with ConnectionPool(sm_config['db']):
dss = DB().select_with_fields(SEL_ALL_DSS)
processed_ds_ids = read_ds_list('SUCCEEDED_DATASETS.txt') | read_ds_list(
'FAILED_DATASETS.txt'
)
dss_to_process = [ds for ds in dss if ds['id'] not in processed_ds_ids]
n = len(dss_to_process)
tasks = [(ds, i, n) for i, ds in enumerate(dss_to_process, 1)]
import pandas as pd
from sm.engine.db import DB, ConnectionPool
from sm.engine.config import init_loggers
if __name__ == '__main__':
parser = argparse.ArgumentParser(
description='Update molecular database molecule names')
parser.add_argument('--config',
default='conf/config.json',
help='SM config path')
parser.add_argument('file_path', help='Path to file with new names')
args = parser.parse_args()
init_loggers()
logger = logging.getLogger('engine')
logger.info(f'Importing new names from {args.file_path}')
db_config = {"host": "localhost", "database": "mol_db", "user": "******"}
with ConnectionPool(db_config):
db = DB()
names_df = pd.read_csv(args.file_path, sep='\t')[['id', 'name']]
sql = (
'WITH molecule_name AS (SELECT UNNEST(%s::text[]) as id_, UNNEST(%s::text[]) as name_) '
'UPDATE molecule SET mol_name = molecule_name.name_ '
'FROM molecule_name WHERE molecule.mol_id = molecule_name.id_')
db.alter(sql,
[names_df.id.values.tolist(),
names_df.name.values.tolist()])
def main(daemon_name, exit_after):
logger.info(f'Starting {daemon_name}-daemon')
conn_pool = ConnectionPool(sm_config['db'])
# Ensure the S3 bucket for image storage exists and is configured correctly.
image_storage.configure_bucket(sm_config)
daemons = []
if daemon_name == 'annotate':
daemons.append(
SMAnnotateDaemon(manager=get_manager(),
annot_qdesc=SM_ANNOTATE,
upd_qdesc=SM_UPDATE))
elif daemon_name == 'update':
make_update_queue_cons = partial(
QueueConsumer,
config=sm_config['rabbitmq'],
qdesc=SM_UPDATE,
logger=logger,
poll_interval=1,
)
for _ in range(sm_config['services']['update_daemon_threads']):
daemon = SMUpdateDaemon(get_manager(), make_update_queue_cons)
daemons.append(daemon)
elif daemon_name == 'lithops':
try:
# Raise the soft limit of open files, as Lithops sometimes makes many network
# connections in parallel, exceeding the default limit of 1024.
# The hard limit cannot be changed by non-sudo users.
soft_rlimit, hard_rlimit = resource.getrlimit(
resource.RLIMIT_NOFILE)
new_rlimit = 10000
if soft_rlimit < new_rlimit:
resource.setrlimit(resource.RLIMIT_NOFILE,
(new_rlimit, hard_rlimit))
logger.debug(
f'Raised open file limit from {soft_rlimit} to {new_rlimit}'
)
except Exception:
logger.warning('Failed to set the open file limit (non-critical)',
exc_info=True)
daemon = LithopsDaemon(get_manager(),
lit_qdesc=SM_LITHOPS,
annot_qdesc=SM_ANNOTATE,
upd_qdesc=SM_UPDATE)
daemons.append(daemon)
else:
raise Exception(f'Wrong SM daemon name: {daemon_name}')
def cb_stop_daemons(*args): # pylint: disable=redefined-outer-name
if parent_process() is not None:
# Multiprocessing worker processes (used by Lithops) inherit this signal handler.
# Avoid interacting with the queues from a worker process as they aren't functional
return
logger.info(f'Stopping {daemon_name}-daemon')
for d in daemons: # pylint: disable=invalid-name
d.stop()
conn_pool.close()
sys.exit(1)
def timer_stop_daemons():
# Raise a signal to stop the daemons. Only the main thread is able to set an exit code,
# so a signal is raised instead of calling cb_stop_daemons directly from the timer thread.
os.kill(os.getpid(), signal.SIGINT)
signal.signal(signal.SIGINT, cb_stop_daemons)
signal.signal(signal.SIGTERM, cb_stop_daemons)
if exit_after is not None:
exit_timer = Timer(exit_after, timer_stop_daemons)
exit_timer.setDaemon(
True) # Don't prevent shutdown if the timer is still running
exit_timer.start()
for daemon in daemons:
daemon.start()
for daemon in daemons:
daemon.join()