revng-revng/python/revng/cli/_commands/graphql/runner.py

#
# This file is distributed under the MIT License. See LICENSE.md for details.
#

import asyncio
import json
import sys
from graphlib import TopologicalSorter
from typing import Awaitable, Callable, Iterable, List, Tuple

import aiohttp
from aiohttp import ClientSession, ClientTimeout
from gql import Client, gql
from gql.client import AsyncClientSession
from gql.transport.aiohttp import AIOHTTPTransport

from .daemon_handler import DaemonHandler

Runner = Callable[[AsyncClientSession], Awaitable[None]]


async def run_on_daemon(handler: DaemonHandler, runners: Iterable[Runner]):
    await handler.wait_for_start()
    await check_server_up(handler.url)

    connector, address = get_connection(handler.url)
    transport = AIOHTTPTransport(
        f"http://{address}/graphql/",
        client_session_args={"connector": connector, "timeout": ClientTimeout()},
    )
    async with Client(
        transport=transport, fetch_schema_from_transport=True, execute_timeout=None
    ) as client:
        for runner in runners:
            await runner(client)


def upload_file(executable_path: str):
    async def runner(client: AsyncClientSession):
        upload_q = gql(
            """
            mutation upload($file: Upload!) {
                uploadFile(file: $file, container: "input")
            }"""
        )

        with open(executable_path, "rb") as binary_file:
            await client.execute(upload_q, variable_values={"file": binary_file}, upload_files=True)
        log("Upload complete")

    return runner


def run_analyses_lists(analyses_lists: List[str]):
    async def runner(client: AsyncClientSession):
        lists_q = gql("""{ info { analysesLists { name }}}""")
        available_analyses_lists = await client.execute(lists_q)
        list_names = [al["name"] for al in available_analyses_lists["info"]["analysesLists"]]

        for list_name in analyses_lists:
            assert list_name in list_names, f"Missing analyses list {list_name}"
            await client.execute(gql(f'mutation {{ runAnalysesList(name: "{list_name}") }}'))

        log("Autoanalysis complete")

    return runner


def produce_artifacts(filter_: List[str] | None = None):
    async def runner(client: AsyncClientSession):
        q = gql(
            """{ info { steps {
                name
                component
                parent
                artifacts {
                    kind { name }
                    container { name }
                }
            }}}"""
        )

        result = await client.execute(q)

        if filter_ is None:
            filtered_steps = list(result["info"]["steps"])
        else:
            filtered_steps = [
                step
                for step in result["info"]["steps"]
                if step["component"] in filter_ or step["name"] == "begin"
            ]

        steps = {step["name"]: step for step in filtered_steps}
        topo_sorter: TopologicalSorter = TopologicalSorter()
        for step in steps.values():
            if step["parent"] is not None:
                if step["parent"] in steps:
                    topo_sorter.add(step["name"], step["parent"])
                else:
                    topo_sorter.add(step["name"], "begin")

        for step_name in topo_sorter.static_order():
            step = steps[step_name]
            if step["artifacts"] is None:
                continue

            artifacts_container = step["artifacts"]["container"]["name"]
            artifacts_kind = step["artifacts"]["kind"]["name"]

            q = gql(
                """
            query cq($step: String!, $container: String!) {
                container(name: $container, step: $step) {
                    targets { serialized }
                }
            }"""
            )
            arguments = {"step": step_name, "container": artifacts_container}
            res = await client.execute(q, arguments)

            target_list = {
                target["serialized"]
                for target in res["container"]["targets"]
                if target["serialized"].endswith(f":{artifacts_kind}")
            }
            targets = ",".join(target_list)

            log(f"Producing {step_name}/{artifacts_container}/*:{artifacts_kind}")
            q = gql(
                """
            query($step: String!, $container: String!, $target: String!) {
                produce(step: $step, container: $container, targetList: $target)
            }"""
            )
            result = await client.execute(q, {**arguments, "target": targets})
            json_result = json.loads(result["produce"])
            assert target_list == set(json_result.keys()), "Some targets were not produced"

    return runner


async def check_server_up(url: str):
    connector, address = get_connection(url)
    session = ClientSession(connector=connector, timeout=ClientTimeout())
    for _ in range(10):
        try:
            async with session.get(f"http://{address}/status") as req:
                if req.status == 200:
                    connector.close()
                    return
                await asyncio.sleep(1.0)
        except aiohttp.ClientConnectionError:
            await asyncio.sleep(1.0)
    connector.close()
    raise ValueError()


def get_connection(url) -> Tuple[aiohttp.BaseConnector, str]:
    if url.startswith("unix:"):
        return (aiohttp.UnixConnector(url.replace("unix:", "", 1)), "dummy")
    return (aiohttp.TCPConnector(), url)


def log(string: str):
    sys.stderr.write(f"{string}\n")
    sys.stderr.flush()