class BLASTData:
    """
    A data layer domain for NuVs BLAST data.
    """

    def __init__(self, db, pg: AsyncEngine, tasks: TasksClient):
        self._db = db
        self._pg = pg
        self._tasks = tasks

    async def create_nuvs_blast(
        self, analysis_id: str, sequence_index: int
    ) -> Document:
        """
        Create a NuVs BLAST record for the sequence associated with a specific analysis
        ID and sequence index.

        A BLAST task will be spawned that runs a BLAST search against NCBI and
        populates the database with result. The analysis and BLAST records are
        updated as the task proceeds.

        :param analysis_id: the ID for the analysis to create BLAST for
        :param sequence_index: the index of the sequence being BLASTed.
        :return: the dictionary representation of the BLAST record.
        """
        timestamp = virtool.utils.timestamp()

        async with AsyncSession(self._pg) as session:
            await self._remove_nuvs_blast(session, analysis_id, sequence_index)
            await session.flush()

            blast = NuVsBlast(
                analysis_id=analysis_id,
                created_at=timestamp,
                last_checked_at=timestamp,
                ready=False,
                sequence_index=sequence_index,
                updated_at=timestamp,
            )

            session.add(blast)
            await session.flush()

            await self._tasks.add(
                BLASTTask,
                {"analysis_id": analysis_id, "sequence_index": sequence_index},
            )

            blast_data = blast.to_dict()
            await session.commit()

        return blast_data

    @staticmethod
    async def _get_nuvs_blast(session, analysis_id: str, sequence_index: int):
        result = await session.execute(
            select(NuVsBlast)
            .filter(NuVsBlast.analysis_id == analysis_id)
            .filter(NuVsBlast.sequence_index == sequence_index)
        )

        return result.scalar_one()

    async def get_nuvs_blast(self, analysis_id: str, sequence_index: int) -> Document:
        """
        Get a NuVs BLAST record by its analysis ID and sequence index.

        :param analysis_id:
        :param sequence_index:
        :return: the dictionary representation of the BLAST record
        """
        async with AsyncSession(self._pg) as session:
            return await self._get_nuvs_blast(session, analysis_id, sequence_index)

    async def update_nuvs_blast(
        self,
        analysis_id: str,
        sequence_index: int,
        error: Optional[str],
        last_checked_at: Optional[datetime],
        rid: Optional[str],
        ready: bool,
        result: Optional[dict],
    ) -> Document:
        """
        Update an existing NuVsBlast as it progresses in Virtool and on NCBI.

        Returns a `dict` representation of the BLAST record.

        :param analysis_id: the ID of the NuVs analysis
        :param sequence_index: the index of the NuVs nucleotide sequence
        :param error: an optional error string
        :param last_checked_at: when the NCBI API was last checked for results
        :param rid: the NCBI ID for the BLAST request
        :param ready: whether the NCBI BLAST search has completed
        :param result: the BLAST search result from NCBI
        :return: the dictionary representation of the BLAST search
        """
        timestamp = virtool.utils.timestamp()

        async with AsyncSession(self._pg) as session:
            blast = await self._get_nuvs_blast(session, analysis_id, sequence_index)

            blast.updated_at = timestamp
            blast.error = error
            blast.last_checked_at = last_checked_at
            blast.rid = rid
            blast.ready = ready
            blast.result = result

            session.add(blast)
            await session.flush()
            document = blast.to_dict()
            await session.commit()

        await self._db.analyses.update_one(
            {"_id": analysis_id}, {"$set": {"updated_at": virtool.utils.timestamp()}}
        )

        return document

    async def _remove_nuvs_blast(
        self, session: AsyncSession, analysis_id: str, sequence_index: int
    ) -> int:
        result = await session.execute(
            delete(NuVsBlast)
            .where(NuVsBlast.analysis_id == analysis_id)
            .where(NuVsBlast.sequence_index == sequence_index)
        )
        await session.commit()

        await self._db.analyses.update_one(
            {"_id": analysis_id}, {"$set": {"updated_at": virtool.utils.timestamp()}}
        )

        return result.rowcount

    async def remove_nuvs_blast(self, analysis_id: str, sequence_index: int) -> int:
        """
        Remove a NuVs BLAST record.

        :param analysis_id: the analysis the BLAST belongs to
        :param sequence_index: the index of the BLASTed NuVs sequence
        :return: the number of deleted records
        """
        async with AsyncSession(self._pg) as session:
            return await self._remove_nuvs_blast(session, analysis_id, sequence_index)

    async def list_by_analysis(self, analysis_id: str) -> List[Document]:
        """
        Get all BLAST records associated with an analysis.

        :param analysis_id: the ID of the analysis to list BLASTs for
        :return: a list of BLAST records
        """
        async with AsyncSession(self._pg) as session:
            result = await session.execute(
                select(NuVsBlast).where(NuVsBlast.analysis_id == analysis_id)
            )

            return [
                {
                    "id": blast.id,
                    "created_at": blast.created_at,
                    "updated_at": blast.updated_at,
                    "last_checked_at": blast.last_checked_at,
                    "interval": blast.interval,
                    "rid": blast.rid,
                    "ready": blast.ready,
                    "result": blast.result,
                }
                for blast in result.scalars()
            ] class DataLayer:

    """
    Provides access to Virtool application data through an abstract interface over
    database and storage.

    """

    analyses: AnalysisData
    blast: BLASTData class DispatcherClient:
    def __init__(self, redis: Redis):
        self._redis = redis
        self._q = asyncio.Queue()

    async def run(self):
        """
        Run the dispatcher.

        Continually publishes enqueued changes to a Redis channel.

        """
        try:
            while True:
                json_string = await self._q.get()
                await self._redis.publish("channel:dispatch", json_string)
        except CancelledError:
            pass

    def enqueue_change(
        self, interface: str, operation: Operation, id_list: Sequence[Union[str, int]]
    ):
        """
        Enqueue the description of a change to be processed by a dispatcher.

        :param interface: the interface the change is on (eg. samples)
        :param operation: the operation the change represents (eg. update)
        :param id_list: the IDs of the resources changed

        """
        json_string = json.dumps(
            {"interface": interface, "operation": operation, "id_list": id_list}
        )

        self._q.put_nowait(json_string) class DispatcherSQLEvents:
    """
    Listens on SQLAlchemy session events to know what models have changes on transaction commits.

    Calls enqueue change with the correct dispatcher interface and list of modified IDs when a
    transaction is committed. Gracefully handles rollbacks.

    Inspired by signalling in [Flask-SQLAlchemy](https://github.com/pallets/flask-sqlalchemy).

    """

    def __init__(
        self, enqueue_change: Callable[[str, Operation, List[Union[str, int]]], None]
    ):
        self._enqueue_change = enqueue_change
        self._changes = defaultdict(dict)

        event.listen(Session, "before_flush", self._record_ops)
        event.listen(Session, "before_commit", self._record_ops)
        event.listen(Session, "after_commit", self._after_commit)
        event.listen(Session, "after_rollback", self._after_rollback)

    def _after_commit(self, session: Session):
        """
        Executed after a commit.

        Processes recorded changes and uses them to call :meth:`enqueue_change` as needed. Clears
        the record changes when complete.

        :param session: the session that was committed

        """
        changes = list(self._changes[id(session)].values())

        for data, change_type in changes:
            interface = get_interface_from_model(data)
            self._enqueue_change(interface, change_type, [data.id])

        del self._changes[id(session)]

    def _after_rollback(self, session: Session):
        """
        Executed after session rollback.

        Clears the records changes for the passed session.

        :param session: the session that was rolled back

        """
        del self._changes[id(session)]

    def _record_ops(self, session: Session, flush_context=None, instances=None):
        """
        Called when a session fires the `after_flush` or `before_commit` events.

        Records the operations that have occurred in the session and stores them so they can be
        sent via `enqueue_change` when the transaction is committed.

        The `flush_context` and `instances` arguments are included to match the callback function
        signature expected by SQLAlchemy for the `after_flush` or `before_commit` events. They are
        not used.

        :param session: the session the event was fired on
        :param flush_context: unused
        :param instances: unused

        """
        session_changes = self._changes[id(session)]

        for targets, operation in (
            (session.new, "insert"),
            (session.dirty, "update"),
            (session.deleted, "delete"),
        ):
            for target in targets:
                state = inspect(target)
                key = state.identity_key if state.has_identity else id(target)
                session_changes[key] = (target, operation) class RedisDispatcherListener(AsyncIterable):
    """
    Asynchronously iterates through messages on a Redis Pub/Sub channel returning :class:`.Change`
    objects.

    """

    def __init__(self, channel: Channel):
        self._channel = channel

    def __aiter__(self):
        return self

    async def __anext__(self) -> Change:
        """
        Get the next JSON-encoded message from the channel and return it as a :class:`Change`
        object.

        :return: the change derived from the incoming JSON object

        """
        change = await self._channel.get_json()

        return Change(change["interface"], change["operation"], change["id_list"]) class FakerWrapper:
    def __init__(self):
        self.fake = Faker()
        self.fake.seed_instance(0)
        self.fake.add_provider(address)
        self.fake.add_provider(misc)
        self.fake.add_provider(lorem)
        self.fake.add_provider(python)
        self.fake.add_provider(date_time)
        self.fake.add_provider(profile)

        self.country = self.fake.country

        self.text = self.fake.text
        self.word = self.fake.word
        self.words = self.fake.words

        self.integer = self.fake.pyint
        self.list = self.fake.pylist
        self.boolean = self.fake.pybool
        self.profile = self.fake.profile
        self.date_time = self.fake.date_time
        self.random_element = self.fake.random_element

    def get_mongo_id(self) -> str:
        """
        Create a predictable, fake ID for MongoDB that imitates Virtool IDs.

        :return: a fake MongoDB ID

        """
        return self.fake.password(length=8, special_chars=False) class JWKArgs:
    kty: str
    kid: str
    use: str
    n: str
    e: str async def create_test_database(connection_string: str, name: str):
    """
    Create a testing database with the passed name.

    Used for creating disposable databases during test runs and when the ``--fake`` option is in
    use on the backend.

    :param connection_string: the base connection string
    :type name: the database name

    """
    engine = create_async_engine(
        f"{connection_string}",
        isolation_level="AUTOCOMMIT",
        json_serializer=virtool.api.json.dumps,
    )

    async with engine.connect() as conn:
        try:
            await conn.execute(text(f"CREATE DATABASE {name}"))
        except ProgrammingError as exc:
            if "DuplicateDatabaseError" not in str(exc):
                raise

    await engine.dispose() class CleanReferencesTask(Task):
    task_type = "clean_references"

    def __init__(self, app, task_id):
        super().__init__(app, task_id)

        self.steps = [self.clean_timed_out_updates]

    async def _remove_latest_update(self, ref_id):
        await self.db.references.update_one(
            {"_id": ref_id}, {"$pop": {"updates": -1}, "$set": {"updating": False}}
        )

    async def clean_timed_out_updates(self):
        async for reference in self.db.references.find(
            {"remotes_from": {"$exists": True}}, ["installed", "task", "updates"]
        ):
            if len(reference["updates"]) == 0:
                continue

            latest_update = reference["updates"][-1]

            if latest_update["ready"]:
                continue

            installed_version = VersionInfo.parse(
                reference["installed"]["name"].lstrip("v")
            )

            latest_update_version = VersionInfo.parse(latest_update["name"].lstrip("v"))

            if latest_update_version <= installed_version:
                continue

            if arrow.utcnow() - arrow.get(latest_update["created_at"]) > timedelta(
                minutes=15
            ):
                await self._remove_latest_update(reference["_id"]) class CreateIndexJSONTask(Task):
    task_type = "create_index_json"

    def __init__(self, app, task_id):
        super().__init__(app, task_id)

        self.steps = [self.create_index_json_files]

    async def create_index_json_files(self):
        tracker = await self.get_tracker()
        async for index in self.db.indexes.find({"has_json": {"$ne": True}}):
            index_id = index["_id"]
            ref_id = index["reference"]["id"]

            document = await self.db.references.find_one(
                ref_id, ["data_type", "organism", "targets"]
            )

            otu_list = await virtool.references.db.export(self.app, ref_id)

            data = {
                "otus": otu_list,
                "data_type": document["data_type"],
                "organism": document["organism"],
            }

            try:
                data["targets"] = document["targets"]
            except KeyError:
                pass

            file_path = (
                self.app["config"].data_path
                / "references"
                / ref_id
                / index_id
                / "reference.json.gz"
            )

            # Convert the list of OTUs to a JSON-formatted string.
            json_string = json.dumps(data, cls=CustomEncoder)

            # Compress the JSON string to a gzip file.
            await self.run_in_thread(compress_json_with_gzip, json_string, file_path)

            await self.db.indexes.find_one_and_update(
                {"_id": index_id}, {"$set": {"has_json": True}}
            )

            await virtool.tasks.pg.update(
                self.pg,
                self.id,
                progress=tracker.step_completed,
                step="create_index_json_files",
            ) class DeleteReferenceTask(Task):
    task_type = "delete_reference"

    def __init__(self, app, task_id):
        super().__init__(app, task_id)

        self.steps = [
            self.remove_directory,
            self.remove_indexes,
            self.remove_unreferenced_otus,
            self.remove_referenced_otus,
        ]

        self.non_existent_references = []

    async def remove_directory(self):
        tracker = await self.get_tracker()

        path = self.app["config"].data_path / "references"

        reference_ids = os.listdir(path)
        existent_references = await self.db.references.distinct(
            "_id", {"_id": {"$in": reference_ids}}
        )
        self.non_existent_references = [
            ref_id for ref_id in reference_ids if ref_id not in existent_references
        ]

        for dir_name in self.non_existent_references:
            await self.app["run_in_thread"](shutil.rmtree, path / dir_name, True)

        await virtool.tasks.pg.update(
            self.pg, self.id, progress=tracker.step_completed, step="remove_directory"
        )

    async def remove_indexes(self):
        tracker = await self.get_tracker()

        await self.db.indexes.delete_many(
            {"reference.id": {"$in": self.non_existent_references}}
        )

        await virtool.tasks.pg.update(
            self.pg, self.id, progress=tracker.step_completed, step="remove_indexes"
        )

    async def remove_unreferenced_otus(self):
        tracker = await self.get_tracker()

        for ref_id in self.non_existent_references:
            referenced_otu_ids = await self.db.analyses.distinct(
                "results.otu.id", {"reference.id": ref_id}
            )

            unreferenced_otu_ids = await self.db.otus.distinct(
                "_id",
                {"reference.id": ref_id, "_id": {"$not": {"$in": referenced_otu_ids}}},
            )

            diff_file_change_ids = await self.db.history.distinct(
                "_id", {"diff": "file", "otu.id": {"$in": unreferenced_otu_ids}}
            )

            await asyncio.gather(
                self.db.otus.delete_many({"_id": {"$in": unreferenced_otu_ids}}),
                self.db.history.delete_many({"otu.id": {"$in": unreferenced_otu_ids}}),
                self.db.sequences.delete_many(
                    {"otu_id": {"$in": unreferenced_otu_ids}}
                ),
                remove_diff_files(self.app, diff_file_change_ids),
            )

            await virtool.tasks.pg.update(
                self.pg,
                self.id,
                progress=tracker.step_completed,
                step="remove_unreferenced_otus",
            )

    async def remove_referenced_otus(self):
        tracker = await self.get_tracker()

        user_id = self.context["user_id"]

        for ref_id in self.non_existent_references:
            async for document in self.db.otus.find({"reference.id": ref_id}):
                await virtool.otus.db.remove(
                    self.app, document["_id"], user_id, document=document, silent=True
                )

        await virtool.tasks.pg.update(
            self.pg,
            self.id,
            progress=tracker.step_completed,
            step="remove_referenced_otus",
        ) def setup_routes(app):
    app.router.add_get("/ws", virtool.http.ws.root)

    if app["config"].dev:
        logger.info("Enabling development API")
        app.router.add_routes(virtool.dev.api.routes)

    for routes in ROUTES:
        app.router.add_routes(routes) class CompressSamplesTask(Task):
    """
    Compress the legacy FASTQ files for all uncompressed samples.

    """

    task_type = "compress_samples"

    def __init__(self, app, process_id):
        super().__init__(app, process_id)
        self.steps = [self.compress_samples]

    async def compress_samples(self):
        query = {"is_legacy": True, "is_compressed": {"$exists": False}}

        count = await self.db.samples.count_documents(query)

        tracker = await self.get_tracker(count)

        while True:
            sample = await self.db.samples.find_one(query)

            if sample is None:
                break

            await virtool.samples.db.compress_sample_reads(self.app, sample)
            await tracker.add(1)

            logger.info(
                f"Compressed legacy sample {sample['_id']} ({tracker.progress}%)"
            )

        await virtool.tasks.pg.update(self.pg, self.id, step="compress_samples") class MoveSampleFilesTask(Task):
    """
    Move pre-SQL samples' file information to new `sample_reads` and `uploads` tables.

    """

    task_type = "move_sample_files"

    def __init__(self, app, task_id):
        super().__init__(app, task_id)

        self.steps = [self.move_sample_files]

    async def move_sample_files(self):
        query = {
            "files": {"$exists": True},
            "$or": [{"is_legacy": False}, {"is_legacy": True, "is_compressed": True}],
        }

        count = await self.db.samples.count_documents(query)

        tracker = await self.get_tracker(count)

        while True:
            sample = await self.db.samples.find_one(query)

            if sample is None:
                break

            await virtool.samples.db.move_sample_files_to_pg(self.app, sample)
            await tracker.add(1)

            logger.info(
                f"Moved files in sample {sample['_id']} to SQL ({tracker.progress}%)"
            )

        await virtool.tasks.pg.update(self.pg, self.id, step="move_sample_files") class TaskRunner:
    def __init__(self, channel, app):
        self._channel = channel
        self.app = app

    async def run(self):
        logging.info("Started task runner")

        try:
            while True:
                logging.info("Waiting for next task")
                await asyncio.sleep(0.3)

                task_id = await self._channel.get_json()

                await self.run_task(task_id)

                logging.info(f"Task finished: {task_id}")

        except asyncio.CancelledError:
            logging.info("Stopped task runner")

    async def run_task(self, task_id: int):
        """
        Run task with given `task_id`.

        :param task_id: ID of the task

        """
        task: Task = await get_row_by_id(self.app["pg"], Task, task_id)

        logging.info(f"Task starting: {task.id} {task.type}")

        loop = asyncio.get_event_loop()

        for task_class in virtool.tasks.task.Task.__subclasses__():
            if task.type == task_class.task_type:
                current_task = task_class(self.app, task_id)
                await loop.create_task(current_task.run()) class UpdateUserDocumentsTask(Task):
    """
    Update user documents that don't contain a "handle" field.

    For B2C users with b2c_given_name and b2c_family_name included in their user document, generate handle.

    For all other users use existing _id values as handle.

    """

    task_type = "update_user_documents"

    def __init__(self, app: App, task_id: int):
        super().__init__(app, task_id)

        self.steps = [self.update_user_documents]

    async def update_user_documents(self):
        async for document in self.db.users.find({"handle": {"$exists": False}}):
            user_id = document["_id"]

            if "b2c_given_name" in document and "b2c_family_name" in document:
                handle = await virtool.users.db.generate_handle(
                    self.db, document["b2c_given_name"], document["b2c_family_name"]
                )
            else:
                handle = user_id

            await self.db.users.find_one_and_update(
                {"_id": user_id}, {"$set": {"handle": handle}}, projection=PROJECTION
            )
import asyncio
import concurrent.futures
import logging
import signal
import sys
import typing
import aiohttp.client
import aiojobs
import aiojobs.aiohttp
import pymongo.errors
import virtool.db.mongo
import virtool.pg.utils
import virtool.redis
from dataclasses import dataclass
from pathlib import Path
from typing import Dict
from urllib.parse import urlparse, urlunparse
from aiohttp.web import Application
from msal import ClientApplication
from virtool.analyses.data import AnalysisData
from virtool.analyses.tasks import StoreNuvsFilesTask
from virtool.blast.data import BLASTData
from virtool.data.layer import DataLayer
from virtool.db.migrate import migrate
from virtool.dev.fake import create_fake_data_path, populate
from virtool.dispatcher.client import DispatcherClient
from virtool.dispatcher.dispatcher import Dispatcher
from virtool.dispatcher.events import DispatcherSQLEvents
from virtool.dispatcher.listener import RedisDispatcherListener
from virtool.fake.wrapper import FakerWrapper
from virtool.hmm.db import refresh
from virtool.indexes.tasks import AddIndexFilesTask
from virtool.jobs.client import JobsClient
from virtool.oidc.utils import JWKArgs
from virtool.pg.testing import create_test_database
from virtool.redis import periodically_ping_redis
from virtool.references.db import refresh_remotes
from virtool.references.tasks import (
    CleanReferencesTask,
    CreateIndexJSONTask,
    DeleteReferenceTask,
)
from virtool.routes import setup_routes
from virtool.samples.tasks import CompressSamplesTask, MoveSampleFilesTask
from virtool.sentry import setup
from virtool.settings.db import ensure
from virtool.subtractions.db import check_subtraction_fasta_files
from virtool.subtractions.tasks import (
    AddSubtractionFilesTask,
    WriteSubtractionFASTATask,
)
from virtool.tasks.client import TasksClient
from virtool.tasks.runner import TaskRunner
from virtool.types import App
from virtool.uploads.tasks import MigrateFilesTask
from virtool.users.tasks import UpdateUserDocumentsTask
from virtool.utils import ensure_data_dir, random_alphanumeric
from virtool.version import determine_server_version
        not app["config"].no_sentry
        and app["settings"].enable_sentry is not False
        and not app["config"].dev
    ):
        logger.info("Configuring Sentry")
        setup(app["version"], app["config"].sentry_dsn)

    else:
        logger.info("Skipped configuring Sentry")


async def startup_settings(app: typing.Union[dict, Application]):
    """
    Draws settings from the settings database collection and populates `app["settings"`.

    Performs migration of old settings style to `v3.3.0` if necessary.

    :param app: the app object

    """
    app["settings"] = await ensure(app["db"])


async def startup_version(app: typing.Union[dict, Application]):
    """
    Bind the application version to the application state `dict`.

    The value will come by checking `--force-version`, the `VERSION` file, or the
    current Git tag if the containing folder is a Git repository.

    :param app: the application object

    """

    force_version = app["config"].force_version

    if force_version:
        version = force_version
    else:
        version = await determine_server_version(Path(sys.path[0]))

    logger.info(f"Virtool {version}")
    logger.info(f"Mode: {app['mode']}")

    app["version"] = version


async def startup_b2c(app: Application):
    """
    Initiate connection to Azure AD B2C tenant.

    :param app: Application object
    """
    b2c_tenant = app["config"].b2c_tenant
    b2c_user_flow = app["config"].b2c_user_flow

    if not all(
        [
            app["config"].b2c_client_id,
            app["config"].b2c_client_secret,
            b2c_tenant,
            b2c_user_flow,
        ]
    ):
        logger.fatal(
            "Required B2C client information not provided for --use-b2c option"
        )
        sys.exit(1)

    authority = f"https://{b2c_tenant}.b2clogin.com/{b2c_tenant}.onmicrosoft.com/{b2c_user_flow}"

    msal = ClientApplication(
        client_id=app["config"].b2c_client_id,
        authority=authority,
        client_credential=app["config"].b2c_client_secret,
    )

    app["b2c"] = B2C(msal, authority)


async def startup_task_runner(app: Application):
    """
    An application `on_startup` callback that initializes a Virtool
    :class:`virtool.tasks.runner.TaskRunner` object and puts it in app state.

    :param app: the app object

    """
    scheduler = get_scheduler_from_app(app)
    (channel,) = await app["redis"].subscribe("channel:tasks")

    app["tasks"] = TasksClient(app["redis"], app["pg"])

    await scheduler.spawn(TaskRunner(channel, app).run())


async def startup_tasks(app: Application):
    if app["config"].no_check_db:
        return logger.info("Skipping subtraction FASTA files checks")

    scheduler = get_scheduler_from_app(app)

    logger.info("Checking subtraction FASTA files")
    subtractions_without_fasta = await check_subtraction_fasta_files(
        app["db"], app["config"]
    )
    for subtraction in subtractions_without_fasta:
        await app["tasks"].add(
            WriteSubtractionFASTATask, context={"subtraction": subtraction}
        )

    logger.info("Checking index JSON files")

    await app["tasks"].add(CreateIndexJSONTask)
    await app["tasks"].add(DeleteReferenceTask, context={"user_id": "virtool"})
    await app["tasks"].add(AddSubtractionFilesTask)
    await app["tasks"].add(AddIndexFilesTask)
    await app["tasks"].add(StoreNuvsFilesTask)
    await app["tasks"].add(CompressSamplesTask)
    await app["tasks"].add(MoveSampleFilesTask)

    await app["tasks"].add(UpdateUserDocumentsTask)