10403778022

Committed 15 Aug 2024 12:29PM UTC coverage: 86.292%. First build

Build # 10403778022

Build Type

Pull #463

github

Committed by

seallard

Commit Message

Add integration test

Pull Request Pull Request #463: Clean up tower and slurm interactions

Run Details

156 of 221 new or added lines in 20 files covered. (70.59%)

2027 of 2349 relevant lines covered (86.29%)

0.86 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

72.73

/trailblazer/cli/core.py

import logging
import sys
from datetime import datetime
from pathlib import Path

import click
import coloredlogs
from sqlalchemy.orm import scoped_session
from dependency_injector.wiring import inject, Provide

import trailblazer
from trailblazer.cli.utils.ls_helper import _get_ls_analysis_message
from trailblazer.cli.utils.user_helper import is_existing_user, is_user_archived
from trailblazer.constants import TRAILBLAZER_TIME_STAMP, FileFormat, TrailblazerStatus
from trailblazer.containers import Container
from trailblazer.environ import environ_email
from trailblazer.io.controller import ReadFile
from trailblazer.models import Config
from trailblazer.server.wiring import setup_dependency_injection
from trailblazer.services.analysis_service.analysis_service import AnalysisService
from trailblazer.store.database import get_session, initialize_database
from trailblazer.store.models import Analysis, User
from trailblazer.store.store import Store

LOG = logging.getLogger(__name__)
LEVELS = ["DEBUG", "INFO", "WARNING", "ERROR"]


class DatabaseResource:
    """
    Setup the database and ensure resources are released when the
    CLI command has been processed.
    """

    def __init__(self, db_uri: str):
        self.db_uri = db_uri

    def __enter__(self):
        initialize_database(self.db_uri)

    def __exit__(self, _, __, ___):
        session: scoped_session = get_session()
        session.remove()


@click.group()
@click.option("-c", "--config", required=True, type=click.File())
@click.option(
    "-l", "--log-level", type=click.Choice(LEVELS), default="INFO", help="lowest level to log at"
)
@click.option("--verbose", is_flag=True, help="Show full log information, time stamp etc")
@click.version_option(trailblazer.__version__, prog_name=trailblazer.__title__)
@click.pass_context
def base(
    context: click.Context,
    config: click.File,
    log_level: str,
    verbose: bool,
):
    """Trailblazer - Monitor analyses"""
    if verbose:
        log_format = "%(asctime)s %(hostname)s %(name)s[%(process)d] %(levelname)s %(message)s"
    else:
        log_format = "%(message)s" if sys.stdout.isatty() else None

    coloredlogs.install(level=log_level, fmt=log_format)
    setup_dependency_injection()

    validated_config = Config(
        **ReadFile.get_content_from_file(file_format=FileFormat.YAML, file_path=Path(config.name))
    )
    context.obj = dict(validated_config)
    context.with_resource(DatabaseResource(validated_config.database_url))
    context.obj["trailblazer_db"] = Store()


@base.command()
@click.option("--reset", is_flag=True, help="reset database before setting up tables")
@click.option("--force", is_flag=True, help="bypass manual confirmations")
@click.pass_context
def init(context, reset, force):
    """Setup the database."""
    existing_tables = context.obj["trailblazer_db"].engine.table_names()
    if force or reset:
        if existing_tables and not force:
            message = f"Delete existing tables? [{', '.join(existing_tables)}]"
            click.confirm(click.style(message, fg="yellow"), abort=True)
        context.obj["trailblazer_db"].drop_all()
    elif existing_tables:
        LOG.warning("Database already exists, use '--reset'")
        context.abort()
    context.obj["trailblazer_db"].setup()
    LOG.info(f"Success! New tables: {', '.join(context.obj['trailblazer'].engine.table_names())}")


@base.command()
@inject
def scan(analysis_service: AnalysisService = Provide[Container.analysis_service]):
    """Scan ongoing analyses in SLURM"""
    analysis_service.update_ongoing_analyses()
    analysis_service.update_uploading_analyses()
    LOG.info("All analyses updated!")


@base.command("update-analysis")
@click.argument("analysis_id")
@click.pass_context
def update_analysis(
    analysis_id: int,
    analysis_service: AnalysisService = Provide[Container.analysis_service],
):
    """Update status of a single analysis."""
    analysis_service.update_analysis_meta_data(analysis_id)


@base.command("add-user")
@click.option("--name", help="Name of new user to add")
@click.option("--abbreviation", help="Abbreviation of new user to add")
@click.argument("email", default=environ_email())
@click.pass_context
def add_user(context, email: str, name: str, abbreviation: str) -> None:
    """Add a new user to the database."""
    trailblazer_db: Store = context.obj["trailblazer_db"]
    existing_user = trailblazer_db.get_user(email=email, exclude_archived=False)
    if is_existing_user(user=existing_user, email=email):
        return
    new_user = trailblazer_db.add_user(email=email, name=name, abbreviation=abbreviation)
    LOG.info(f"New user added: {new_user.email} ({new_user.id})")


@base.command("get-user")
@click.argument("email", default=environ_email())
@click.pass_context
def get_user_from_db(context, email: str) -> None:
    """Display information about an existing user."""
    trailblazer_db: Store = context.obj["trailblazer_db"]
    existing_user = trailblazer_db.get_user(email=email, exclude_archived=False)
    if not is_existing_user(user=existing_user, email=email):
        return
    LOG.info(f"Existing user found: {existing_user.to_dict()}")


@base.command("get-users")
@click.option("--name", type=click.types.STRING, help="Name of new users to list")
@click.option("--email", type=click.types.STRING, help="Name of new users to list")
@click.option("--exclude-archived", is_flag=True, help="Exclude archived users")
@click.pass_context
def get_users_from_db(context, name: str, email: str, exclude_archived: bool) -> None:
    """Display information about existing users."""
    trailblazer_db: Store = context.obj["trailblazer_db"]
    users: list[User] = trailblazer_db.get_users(
        email=email, exclude_archived=exclude_archived, name=name
    )
    LOG.info("Listing users in database:")
    for user in users:
        LOG.info(f"{user.to_dict()}")


@base.command("archive-user")
@click.argument("email", default=environ_email())
@click.pass_context
def archive_user(context, email: str) -> None:
    """Archive an existing user identified by email."""
    trailblazer_db: Store = context.obj["trailblazer_db"]
    existing_user: User = trailblazer_db.get_user(email=email, exclude_archived=False)

    if not is_existing_user(user=existing_user, email=email):
        return
    if is_user_archived(user=existing_user, email=email):
        return

    trailblazer_db.update_user_is_archived(user=existing_user, archive=True)
    LOG.info(f"User archived: {email}")


@base.command("unarchive-user")
@click.argument("email", default=environ_email())
@click.pass_context
def unarchive_user(context, email: str) -> None:
    """Unarchive an existing user identified by email."""
    trailblazer_db: Store = context.obj["trailblazer_db"]
    existing_user: User = trailblazer_db.get_user(email=email, exclude_archived=False)

    if not is_existing_user(user=existing_user, email=email):
        return
    if not is_user_archived(user=existing_user, email=email):
        return

    trailblazer_db.update_user_is_archived(user=existing_user, archive=False)
    LOG.info(f"User unarchived: {email}")


@base.command()
@click.argument("analysis_id", type=int)
@click.pass_context
def cancel(
    analysis_id: int,
    analysis_service: AnalysisService = Provide[Container.analysis_service],
):
    """Cancel all jobs in an analysis."""
    try:
        analysis_service.cancel_analysis(analysis_id)
    except Exception as error:
        LOG.error(error)


@base.command("set-completed")
@click.argument("analysis_id", type=int)
@click.pass_context
def set_analysis_completed(context, analysis_id):
    """Set status of an analysis to 'COMPLETED'."""
    trailblazer_db: Store = context.obj["trailblazer_db"]
    try:
        trailblazer_db.update_analysis_status_to_completed(analysis_id=analysis_id)
    except Exception as error:
        LOG.error(error)


@base.command("set-status")
@click.option(
    "--status",
    is_flag=False,
    type=str,
    help=f"Status to be set. Can take the values:{TrailblazerStatus.statuses()}",
)
@click.argument("case_id", type=str)
@click.pass_context
def set_analysis_status(
    context,
    case_id: str,
    status: str,
):
    """Set the status of the latest analysis for a given case id."""
    trailblazer_db: Store = context.obj["trailblazer_db"]
    try:
        trailblazer_db.update_analysis_status_by_case_id(case_id=case_id, status=status)
    except ValueError as error:
        LOG.error(error)
        raise click.Abort from error
    except Exception as error:
        LOG.error(error)


@base.command()
@click.option("--force", is_flag=True, help="Force delete if analysis is ongoing")
@click.option("--cancel-jobs", is_flag=True, help="Cancel all ongoing jobs before deleting")
@click.argument("analysis_id", type=int)
@click.pass_context
def delete(context, analysis_id: int, force: bool, cancel_jobs: bool):
    """Delete analysis completely from database, and optionally cancel all ongoing analysis jobs."""
    trailblazer_db: Store = context.obj["trailblazer_db"]
    try:
        if cancel_jobs:
            trailblazer_db.cancel_ongoing_analysis(analysis_id=analysis_id)
        trailblazer_db.delete_analysis(analysis_id=analysis_id, force=force)
    except Exception as error:
        LOG.error(error)


@base.command("ls")
@click.option(
    "-s",
    "--status",
    type=click.Choice(TrailblazerStatus.statuses()),
    help="Find analysis with specified status",
)
@click.option("-b", "--before", type=str, help="Find analyses started before date")
@click.option("-c", "--comment", type=str, help="Find analysis with comment")
@click.option("--limit", type=int, default=30, help="Limit the number of analysis returned")
@click.pass_context
def ls_cmd(context, before: str, status: TrailblazerStatus, comment: str, limit: int = 30):
    """Display recent logs for the latest analyses."""
    trailblazer_db: Store = context.obj["trailblazer_db"]
    analyses: list[Analysis] | None = trailblazer_db.get_analyses_by_status_started_at_and_comment(
        status=status,
        before=datetime.strptime(before, TRAILBLAZER_TIME_STAMP).date() if before else None,
        comment=comment,
    )
    if analyses:
        for analysis in analyses[:limit]:
            (message, color) = _get_ls_analysis_message(analysis=analysis)
            click.echo(click.style(message, fg=color))
    else:
        LOG.warning("No analyses matching search criteria")

1	import logging	1✔
2	import sys	1✔
3	from datetime import datetime	1✔
4	from pathlib import Path	1✔
5
6	import click	1✔
7	import coloredlogs	1✔
8	from sqlalchemy.orm import scoped_session	1✔
9	from dependency_injector.wiring import inject, Provide	1✔
10
11	import trailblazer	1✔
12	from trailblazer.cli.utils.ls_helper import _get_ls_analysis_message	1✔
13	from trailblazer.cli.utils.user_helper import is_existing_user, is_user_archived	1✔
14	from trailblazer.constants import TRAILBLAZER_TIME_STAMP, FileFormat, TrailblazerStatus	1✔
15	from trailblazer.containers import Container	1✔
16	from trailblazer.environ import environ_email	1✔
17	from trailblazer.io.controller import ReadFile	1✔
18	from trailblazer.models import Config	1✔
19	from trailblazer.server.wiring import setup_dependency_injection	1✔
20	from trailblazer.services.analysis_service.analysis_service import AnalysisService	1✔
21	from trailblazer.store.database import get_session, initialize_database	1✔
22	from trailblazer.store.models import Analysis, User	1✔
23	from trailblazer.store.store import Store	1✔
24
25	LOG = logging.getLogger(__name__)	1✔
26	LEVELS = ["DEBUG", "INFO", "WARNING", "ERROR"]	1✔
27
28
29	class DatabaseResource:	1✔
30	"""
31	Setup the database and ensure resources are released when the
32	CLI command has been processed.
33	"""
34
35	def __init__(self, db_uri: str):	1✔
36	self.db_uri = db_uri	×
37
38	def __enter__(self):	1✔
39	initialize_database(self.db_uri)	×
40
41	def __exit__(self, _, __, ___):	1✔
42	session: scoped_session = get_session()	×
43	session.remove()	×
44
45
46	@click.group()	1✔
47	@click.option("-c", "--config", required=True, type=click.File())	1✔
48	@click.option(	1✔
49	"-l", "--log-level", type=click.Choice(LEVELS), default="INFO", help="lowest level to log at"
50	)
51	@click.option("--verbose", is_flag=True, help="Show full log information, time stamp etc")	1✔
52	@click.version_option(trailblazer.__version__, prog_name=trailblazer.__title__)	1✔
53	@click.pass_context	1✔
54	def base(	1✔
55	context: click.Context,
56	config: click.File,
57	log_level: str,
58	verbose: bool,
59	):
60	"""Trailblazer - Monitor analyses"""
61	if verbose:	×
62	log_format = "%(asctime)s %(hostname)s %(name)s[%(process)d] %(levelname)s %(message)s"	×
63	else:
64	log_format = "%(message)s" if sys.stdout.isatty() else None	×
65
66	coloredlogs.install(level=log_level, fmt=log_format)	×
67	setup_dependency_injection()	×
68
69	validated_config = Config(	×
70	**ReadFile.get_content_from_file(file_format=FileFormat.YAML, file_path=Path(config.name))
71	)
72	context.obj = dict(validated_config)	×
73	context.with_resource(DatabaseResource(validated_config.database_url))	×
74	context.obj["trailblazer_db"] = Store()	×
75
76
77	@base.command()	1✔
78	@click.option("--reset", is_flag=True, help="reset database before setting up tables")	1✔
79	@click.option("--force", is_flag=True, help="bypass manual confirmations")	1✔
80	@click.pass_context	1✔
81	def init(context, reset, force):	1✔
82	"""Setup the database."""
83	existing_tables = context.obj["trailblazer_db"].engine.table_names()	×
84	if force or reset:	×
85	if existing_tables and not force:	×
86	message = f"Delete existing tables? [{', '.join(existing_tables)}]"	×
87	click.confirm(click.style(message, fg="yellow"), abort=True)	×
88	context.obj["trailblazer_db"].drop_all()	×
89	elif existing_tables:	×
90	LOG.warning("Database already exists, use '--reset'")	×
91	context.abort()	×
92	context.obj["trailblazer_db"].setup()	×
93	LOG.info(f"Success! New tables: {', '.join(context.obj['trailblazer'].engine.table_names())}")	×
94
95
96	@base.command()	1✔
97	@inject	1✔
98	def scan(analysis_service: AnalysisService = Provide[Container.analysis_service]):	1✔
99	"""Scan ongoing analyses in SLURM"""
100	analysis_service.update_ongoing_analyses()	×
101	analysis_service.update_uploading_analyses()	×
102	LOG.info("All analyses updated!")	×
103
104
105	@base.command("update-analysis")	1✔
106	@click.argument("analysis_id")	1✔
107	@click.pass_context	1✔
108	def update_analysis(	1✔
109	analysis_id: int,
110	analysis_service: AnalysisService = Provide[Container.analysis_service],
111	):
112	"""Update status of a single analysis."""
NEW 113	analysis_service.update_analysis_meta_data(analysis_id)	×
114
115
116	@base.command("add-user")	1✔
117	@click.option("--name", help="Name of new user to add")	1✔
118	@click.option("--abbreviation", help="Abbreviation of new user to add")	1✔
119	@click.argument("email", default=environ_email())	1✔
120	@click.pass_context	1✔
121	def add_user(context, email: str, name: str, abbreviation: str) -> None:	1✔
122	"""Add a new user to the database."""
123	trailblazer_db: Store = context.obj["trailblazer_db"]	1✔
124	existing_user = trailblazer_db.get_user(email=email, exclude_archived=False)	1✔
125	if is_existing_user(user=existing_user, email=email):	1✔
126	return	1✔
127	new_user = trailblazer_db.add_user(email=email, name=name, abbreviation=abbreviation)	1✔
128	LOG.info(f"New user added: {new_user.email} ({new_user.id})")	1✔
129
130
131	@base.command("get-user")	1✔
132	@click.argument("email", default=environ_email())	1✔
133	@click.pass_context	1✔
134	def get_user_from_db(context, email: str) -> None:	1✔
135	"""Display information about an existing user."""
136	trailblazer_db: Store = context.obj["trailblazer_db"]	1✔
137	existing_user = trailblazer_db.get_user(email=email, exclude_archived=False)	1✔
138	if not is_existing_user(user=existing_user, email=email):	1✔
139	return	1✔
140	LOG.info(f"Existing user found: {existing_user.to_dict()}")	1✔
141
142
143	@base.command("get-users")	1✔
144	@click.option("--name", type=click.types.STRING, help="Name of new users to list")	1✔
145	@click.option("--email", type=click.types.STRING, help="Name of new users to list")	1✔
146	@click.option("--exclude-archived", is_flag=True, help="Exclude archived users")	1✔
147	@click.pass_context	1✔
148	def get_users_from_db(context, name: str, email: str, exclude_archived: bool) -> None:	1✔
149	"""Display information about existing users."""
150	trailblazer_db: Store = context.obj["trailblazer_db"]	1✔
151	users: list[User] = trailblazer_db.get_users(	1✔
152	email=email, exclude_archived=exclude_archived, name=name
153	)
154	LOG.info("Listing users in database:")	1✔
155	for user in users:	1✔
156	LOG.info(f"{user.to_dict()}")	1✔
157
158
159	@base.command("archive-user")	1✔
160	@click.argument("email", default=environ_email())	1✔
161	@click.pass_context	1✔
162	def archive_user(context, email: str) -> None:	1✔
163	"""Archive an existing user identified by email."""
164	trailblazer_db: Store = context.obj["trailblazer_db"]	1✔
165	existing_user: User = trailblazer_db.get_user(email=email, exclude_archived=False)	1✔
166
167	if not is_existing_user(user=existing_user, email=email):	1✔
168	return	1✔
169	if is_user_archived(user=existing_user, email=email):	1✔
170	return	1✔
171
172	trailblazer_db.update_user_is_archived(user=existing_user, archive=True)	1✔
173	LOG.info(f"User archived: {email}")	1✔
174
175
176	@base.command("unarchive-user")	1✔
177	@click.argument("email", default=environ_email())	1✔
178	@click.pass_context	1✔
179	def unarchive_user(context, email: str) -> None:	1✔
180	"""Unarchive an existing user identified by email."""
181	trailblazer_db: Store = context.obj["trailblazer_db"]	1✔
182	existing_user: User = trailblazer_db.get_user(email=email, exclude_archived=False)	1✔
183
184	if not is_existing_user(user=existing_user, email=email):	1✔
185	return	×
186	if not is_user_archived(user=existing_user, email=email):	1✔
187	return	1✔
188
189	trailblazer_db.update_user_is_archived(user=existing_user, archive=False)	1✔
190	LOG.info(f"User unarchived: {email}")	1✔
191
192
193	@base.command()	1✔
194	@click.argument("analysis_id", type=int)	1✔
195	@click.pass_context	1✔
196	def cancel(	1✔
197	analysis_id: int,
198	analysis_service: AnalysisService = Provide[Container.analysis_service],
199	):
200	"""Cancel all jobs in an analysis."""
201	try:	×
NEW 202	analysis_service.cancel_analysis(analysis_id)	×
203	except Exception as error:	×
204	LOG.error(error)	×
205
206
207	@base.command("set-completed")	1✔
208	@click.argument("analysis_id", type=int)	1✔
209	@click.pass_context	1✔
210	def set_analysis_completed(context, analysis_id):	1✔
211	"""Set status of an analysis to 'COMPLETED'."""
212	trailblazer_db: Store = context.obj["trailblazer_db"]	1✔
213	try:	1✔
214	trailblazer_db.update_analysis_status_to_completed(analysis_id=analysis_id)	1✔
215	except Exception as error:	×
216	LOG.error(error)	×
217
218
219	@base.command("set-status")	1✔
220	@click.option(	1✔
221	"--status",
222	is_flag=False,
223	type=str,
224	help=f"Status to be set. Can take the values:{TrailblazerStatus.statuses()}",
225	)
226	@click.argument("case_id", type=str)	1✔
227	@click.pass_context	1✔
228	def set_analysis_status(	1✔
229	context,
230	case_id: str,
231	status: str,
232	):
233	"""Set the status of the latest analysis for a given case id."""
234	trailblazer_db: Store = context.obj["trailblazer_db"]	1✔
235	try:	1✔
236	trailblazer_db.update_analysis_status_by_case_id(case_id=case_id, status=status)	1✔
237	except ValueError as error:	1✔
238	LOG.error(error)	1✔
239	raise click.Abort from error	1✔
240	except Exception as error:	×
241	LOG.error(error)	×
242
243
244	@base.command()	1✔
245	@click.option("--force", is_flag=True, help="Force delete if analysis is ongoing")	1✔
246	@click.option("--cancel-jobs", is_flag=True, help="Cancel all ongoing jobs before deleting")	1✔
247	@click.argument("analysis_id", type=int)	1✔
248	@click.pass_context	1✔
249	def delete(context, analysis_id: int, force: bool, cancel_jobs: bool):	1✔
250	"""Delete analysis completely from database, and optionally cancel all ongoing analysis jobs."""
251	trailblazer_db: Store = context.obj["trailblazer_db"]	×
252	try:	×
253	if cancel_jobs:	×
254	trailblazer_db.cancel_ongoing_analysis(analysis_id=analysis_id)	×
255	trailblazer_db.delete_analysis(analysis_id=analysis_id, force=force)	×
256	except Exception as error:	×
257	LOG.error(error)	×
258
259
260	@base.command("ls")	1✔
261	@click.option(	1✔
262	"-s",
263	"--status",
264	type=click.Choice(TrailblazerStatus.statuses()),
265	help="Find analysis with specified status",
266	)
267	@click.option("-b", "--before", type=str, help="Find analyses started before date")	1✔
268	@click.option("-c", "--comment", type=str, help="Find analysis with comment")	1✔
269	@click.option("--limit", type=int, default=30, help="Limit the number of analysis returned")	1✔
270	@click.pass_context	1✔
271	def ls_cmd(context, before: str, status: TrailblazerStatus, comment: str, limit: int = 30):	1✔
272	"""Display recent logs for the latest analyses."""
273	trailblazer_db: Store = context.obj["trailblazer_db"]	×
274	analyses: list[Analysis] \| None = trailblazer_db.get_analyses_by_status_started_at_and_comment(	×
275	status=status,
276	before=datetime.strptime(before, TRAILBLAZER_TIME_STAMP).date() if before else None,
277	comment=comment,
278	)
279	if analyses:	×
280	for analysis in analyses[:limit]:	×
281	(message, color) = _get_ls_analysis_message(analysis=analysis)	×
282	click.echo(click.style(message, fg=color))	×
283	else:
284	LOG.warning("No analyses matching search criteria")	×

Clinical-Genomics / trailblazer / 10403778022

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous