fluidattacks_batch_client 0.1.0__py3-none-any.whl

fluidattacks_batch_client/__init__.py

@@ -0,0 +1,9 @@
+from ._logger import (
+    setup_logger,
+)
+from fa_purity import (
+    Unsafe,
+)
+
+__version__ = "0.1.0"
+Unsafe.compute(setup_logger(__name__))

fluidattacks_batch_client/_cli.py

@@ -0,0 +1,171 @@
+import dataclasses
+
+from fluidattacks_batch_client.decode import JobDefDecoder
+from fluidattacks_batch_client.sender import new_sender
+from .api import (
+    new_client,
+)
+from .core import (
+    AllowDuplicates,
+    JobPipeline,
+    JobRequest,
+    JobName,
+)
+import click
+from fa_purity import (
+    Bool,
+    Cmd,
+    FrozenList,
+    NewFrozenList,
+    PureIterFactory,
+    PureIterTransform,
+    ResultE,
+    UnitType,
+    Unsafe,
+    unit,
+)
+from fa_purity.json import (
+    JsonObj,
+    JsonValueFactory,
+    UnfoldedFactory,
+    Unfolder,
+)
+from typing import (
+    IO,
+    Callable,
+    NoReturn,
+    TypeVar,
+)
+import logging
+from . import _utils
+
+LOG = logging.getLogger(__name__)
+
+_T = TypeVar("_T")
+
+
+def _read_file(file_path: str, transform: Callable[[IO[str]], _T]) -> Cmd[_T]:
+    def _action() -> _T:
+        with open(file_path, "r", encoding="utf-8") as file:
+            return transform(file)
+
+    return Cmd.wrap_impure(_action)
+
+
+def _decode_json(file_path: str) -> Cmd[ResultE[JsonObj]]:
+    return _read_file(file_path, UnfoldedFactory.load)
+
+
+def _decode_json_list(file_path: str) -> Cmd[ResultE[FrozenList[JsonObj]]]:
+    return _read_file(
+        file_path,
+        lambda file: JsonValueFactory.load(file).bind(
+            lambda v: Unfolder.to_list_of(v, Unfolder.to_json)
+        ),
+    )
+
+
+@click.command()
+@click.option("--job", required=True, help="json encoded str defining a JobRequest")
+@click.option("--allow-duplicates", is_flag=True)
+@click.option("--args-in-name", is_flag=True)
+@click.option("--dry-run", is_flag=True)
+@click.argument("args", nargs=-1)  # additional args to append into the job command
+def submit_job(
+    job: str,
+    allow_duplicates: bool,
+    args_in_name: bool,
+    dry_run: bool,
+    args: FrozenList[str],
+) -> NoReturn:
+    def _execute(job: JobRequest) -> Cmd[None]:
+        name = Bool.from_primitive(len(args) > 0 and args_in_name).map(
+            lambda _: JobName.normalize(job.name.raw + "-" + "-".join(args)),
+            lambda _: job.name,
+        )
+        overriden_job = dataclasses.replace(job, name=name)
+        if dry_run:
+            return Cmd.wrap_impure(
+                lambda: LOG.info(
+                    "[dry-run] the following batch job will be sent: %s", overriden_job
+                )
+            )
+        return (
+            new_client()
+            .map(new_sender)
+            .bind(
+                lambda s: s.send_single_job(
+                    overriden_job, AllowDuplicates(allow_duplicates)
+                ).map(lambda _: None)
+            )
+        )
+
+    _decoder = _utils.get_environment.map(JobDefDecoder)
+    cmd: Cmd[None] = _decoder.bind(
+        lambda decoder: _decode_json(job)
+        .map(lambda r: r.bind(decoder.decode_job))
+        .map(lambda r: r.alt(Unsafe.raise_exception).to_union())
+        .bind(_execute)
+    )
+    cmd.compute()
+
+
+@click.command()
+@click.option(
+    "--pipeline", required=True, help="json encoded str defining a JobRequest"
+)
+@click.option("--dry-run", is_flag=True)
+def submit_pipeline(
+    pipeline: str,
+    dry_run: bool,
+) -> NoReturn:
+    def _execute(job_pipeline: JobPipeline) -> Cmd[UnitType]:
+        if dry_run:
+            msg = Cmd.wrap_impure(
+                lambda: LOG.info(
+                    "[dry-run] A batch pipeline will be sent",
+                )
+            )
+
+            msgs = NewFrozenList(tuple(enumerate(job_pipeline.jobs, start=1))).map(
+                lambda j: Cmd.wrap_impure(
+                    lambda: LOG.info(
+                        "[dry-run] [pipeline-job-#%s] this job will be sent: %s",
+                        j[0],
+                        j[1],
+                    )
+                )
+            )
+            jobs_msgs = (
+                PureIterFactory.from_list(msgs.items)
+                .transform(PureIterTransform.consume)
+                .map(lambda _: unit)
+            )
+            return msg + jobs_msgs
+        return (
+            new_client()
+            .map(new_sender)
+            .bind(
+                lambda s: s.send_pipeline(
+                    PureIterFactory.from_list(job_pipeline.jobs.items)
+                )
+            )
+        )
+
+    _decoder = _utils.get_environment.map(JobDefDecoder)
+    cmd: Cmd[UnitType] = _decoder.bind(
+        lambda decoder: _decode_json_list(pipeline)
+        .map(lambda r: r.map(NewFrozenList).bind(decoder.decode_pipeline))
+        .map(lambda r: r.alt(Unsafe.raise_exception).to_union())
+        .bind(_execute)
+    )
+    cmd.compute()
+
+
+@click.group()
+def main() -> None:
+    pass
+
+
+main.add_command(submit_job)
+main.add_command(submit_pipeline)

fluidattacks_batch_client/_logger.py

@@ -0,0 +1,17 @@
+from fa_purity import (
+    Cmd,
+)
+import logging
+import sys
+
+
+def setup_logger(name: str) -> Cmd[None]:
+    def _action() -> None:
+        handler = logging.StreamHandler(sys.stderr)
+        formatter = logging.Formatter("[%(levelname)s] %(message)s")
+        handler.setFormatter(formatter)
+        log = logging.getLogger(name)
+        log.addHandler(handler)
+        log.setLevel(logging.INFO)
+
+    return Cmd.wrap_impure(_action)

fluidattacks_batch_client/_utils.py

@@ -0,0 +1,102 @@
+from __future__ import (
+    annotations,
+)
+
+from collections.abc import (
+    Callable,
+)
+from dataclasses import (
+    dataclass,
+)
+from fa_purity import (
+    Cmd,
+    FrozenDict,
+    FrozenList,
+    Maybe,
+    PureIter,
+    Result,
+    ResultE,
+    Coproduct,
+    CoproductFactory,
+    Unsafe,
+)
+from typing import (
+    NoReturn,
+    TypeVar,
+)
+import os
+
+_T = TypeVar("_T")
+
+
+@dataclass(frozen=True)
+class LibraryBug(Exception):
+    traceback: Exception
+
+    def __str__(self) -> str:
+        return (
+            "If raised then there is a bug in the `fluidattacks_batch_client` library"
+        )
+
+
+def str_to_int(raw: str) -> ResultE[int]:
+    try:
+        return Result.success(int(raw))
+    except ValueError as err:
+        return Result.failure(Exception(err))
+
+
+def int_to_str(item: int) -> str:
+    return str(item)
+
+
+def handle_value_error(transform: Callable[[], _T | NoReturn]) -> ResultE[_T]:
+    try:
+        return Result.success(transform())
+    except ValueError as err:
+        return Result.failure(Exception(err))
+
+
+def handle_key_error(transform: Callable[[], _T | NoReturn]) -> ResultE[_T]:
+    try:
+        return Result.success(transform())
+    except KeyError as err:
+        return Result.failure(Exception(err))
+
+
+def handle_index_error(transform: Callable[[], _T | NoReturn]) -> ResultE[_T]:
+    try:
+        return Result.success(transform())
+    except IndexError as err:
+        return Result.failure(Exception(err))
+
+
+def get_index(items: FrozenList[_T], index: int) -> Maybe[_T]:
+    return Maybe.from_result(
+        handle_index_error(lambda: items[index]).alt(lambda _: None)
+    )
+
+
+def extract_single(items: PureIter[_T]) -> Coproduct[_T, PureIter[_T]]:
+    _factory: CoproductFactory[_T, PureIter[_T]] = CoproductFactory()
+    single_element = (
+        items.enumerate(1)
+        .find_first(lambda t: t[0] >= 2)
+        .map(lambda _: False)
+        .value_or(True)
+    )
+    if single_element:
+        return _factory.inl(
+            items.enumerate(1)
+            .find_first(lambda t: t[0] == 1)
+            .to_result()
+            .alt(lambda _: LibraryBug(Exception("no first element")))
+            .alt(Unsafe.raise_exception)
+            .to_union()[1]
+        )
+    return _factory.inr(items)
+
+
+get_environment = Cmd[FrozenDict[str, str]].wrap_impure(
+    lambda: FrozenDict(dict(os.environ))
+)

fluidattacks_batch_client/api/__init__.py

@@ -0,0 +1,11 @@
+from ._client_1 import (
+    new_client,
+)
+from ._core import (
+    ApiClient,
+)
+
+__all__ = [
+    "ApiClient",
+    "new_client",
+]

fluidattacks_batch_client/api/_client_1/__init__.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+from . import (
+    _list_jobs,
+    _send_job,
+    _get_command,
+)
+from fluidattacks_batch_client.api._core import (
+    ApiClient,
+)
+import boto3
+from fa_purity import (
+    Cmd,
+)
+
+if TYPE_CHECKING:
+    from mypy_boto3_batch.client import (
+        BatchClient,
+    )
+
+
+def _new_batch_client() -> Cmd[BatchClient]:
+    def _action() -> BatchClient:
+        return boto3.client("batch")
+
+    return Cmd.wrap_impure(_action)
+
+
+def new_client() -> Cmd[ApiClient]:
+    return _new_batch_client().map(
+        lambda client: ApiClient(
+            lambda n, q, s: _list_jobs.list_jobs(client, n, q, s),
+            lambda j: _send_job.send_job(client, j),
+            lambda j: _get_command.get_command(client, j),
+        )
+    )

fluidattacks_batch_client/api/_client_1/_get_command.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING, TypeVar
+
+from mypy_boto3_batch.type_defs import (
+    ContainerPropertiesOutputTypeDef,
+    JobDefinitionTypeDef,
+)
+from fluidattacks_batch_client.core import (
+    Command,
+    JobDefinitionName,
+)
+from fa_purity import (
+    Cmd,
+    FrozenList,
+    Maybe,
+    Result,
+    ResultE,
+    cast_exception,
+)
+
+if TYPE_CHECKING:
+    from mypy_boto3_batch.client import (
+        BatchClient,
+    )
+
+_T = TypeVar("_T")
+
+
+def _decode_command(raw: JobDefinitionTypeDef) -> Maybe[Command]:
+    props: Maybe[ContainerPropertiesOutputTypeDef] = Maybe.from_optional(
+        raw.get("containerProperties")
+    )
+    return props.bind_optional(lambda c: c.get("command")).map(
+        lambda c: Command(tuple(c))
+    )
+
+
+def _assert_one(items: FrozenList[_T]) -> ResultE[_T]:
+    if len(items) == 0:
+        return Result.failure(ValueError("list does not have elements"))
+    if len(items) > 1:
+        return Result.failure(ValueError("list has more than one element"))
+    return Result.success(items[0])
+
+
+def get_command(
+    client: BatchClient,
+    job_def: JobDefinitionName,
+) -> Cmd[ResultE[Maybe[Command]]]:
+    def _action() -> ResultE[Maybe[Command]]:
+        try:
+            result = client.describe_job_definitions(
+                jobDefinitionName=job_def.raw, status="ACTIVE"
+            )
+        except Exception as e:
+            return Result.failure(e)
+        return (
+            _assert_one(tuple(result["jobDefinitions"]))
+            .alt(
+                lambda e: ValueError(
+                    f"Could not determine the current active job definition i.e. {e}"
+                )
+            )
+            .alt(cast_exception)
+            .map(_decode_command)
+        )
+
+    return Cmd.wrap_impure(_action)

fluidattacks_batch_client/api/_client_1/_list_jobs.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+from fluidattacks_batch_client import (
+    _utils,
+)
+from fluidattacks_batch_client.core import (
+    BatchJob,
+    BatchJobObj,
+    JobArn,
+    JobId,
+    JobName,
+    JobStatus,
+    QueueName,
+)
+from dataclasses import (
+    dataclass,
+)
+from fa_purity import (
+    Cmd,
+    FrozenList,
+    Maybe,
+    ResultE,
+    Stream,
+    PureIterFactory,
+    StreamFactory,
+    StreamTransform,
+    ResultTransform,
+    Unsafe,
+)
+
+if TYPE_CHECKING:
+    from mypy_boto3_batch.client import (
+        BatchClient,
+    )
+    from mypy_boto3_batch.type_defs import (
+        JobSummaryTypeDef,
+        KeyValuesPairTypeDef,
+        ListJobsResponseTypeDef,
+    )
+
+
+def _decode_job(raw: JobSummaryTypeDef) -> ResultE[BatchJob]:
+    def _inner() -> ResultE[BatchJob]:
+        return JobStatus.to_status(raw["status"]).map(
+            lambda status: BatchJob(
+                raw["createdAt"],
+                status,
+                Maybe.from_optional(raw.get("statusReason")),
+                Maybe.from_optional(raw.get("startedAt")),
+                Maybe.from_optional(raw.get("stoppedAt")),
+            )
+        )
+
+    return _utils.handle_key_error(_inner).bind(lambda x: x)
+
+
+def _decode_job_obj(
+    raw: JobSummaryTypeDef,
+) -> ResultE[BatchJobObj]:
+    def _inner() -> ResultE[BatchJobObj]:
+        _arn = JobArn(raw["jobArn"])
+        _id = JobId(raw["jobId"])
+        _name = JobName.new(raw["jobName"])
+        return _name.bind(
+            lambda name: _decode_job(raw).map(lambda j: BatchJobObj(_id, _arn, name, j))
+        )
+
+    return _utils.handle_key_error(_inner).bind(lambda x: x)
+
+
+@dataclass
+class JobsPage:
+    items: FrozenList[BatchJobObj]
+    next_item: Maybe[str]
+
+
+def _decode_respose(response: ListJobsResponseTypeDef) -> ResultE[JobsPage]:
+    def _inner() -> ResultE[JobsPage]:
+        items = PureIterFactory.from_list(response["jobSummaryList"]).map(
+            _decode_job_obj
+        )
+        _next = Maybe.from_optional(response.get("nextToken"))
+        return ResultTransform.all_ok(items.to_list()).map(lambda i: JobsPage(i, _next))
+
+    return _utils.handle_key_error(_inner).bind(lambda x: x)
+
+
+def _list_jobs_page(
+    client: BatchClient,
+    queue: QueueName,
+    name: JobName,
+    _next: Maybe[str],
+) -> Cmd[JobsPage]:
+    def _action() -> JobsPage:
+        _filter: FrozenList[KeyValuesPairTypeDef] = (
+            {"name": "JOB_NAME", "values": [name.raw]},
+        )
+        result = _next.map(
+            lambda n: client.list_jobs(jobQueue=queue.raw, filters=_filter, nextToken=n)
+        ).or_else_call(lambda: client.list_jobs(jobQueue=queue.raw, filters=_filter))
+        return _decode_respose(result).alt(Unsafe.raise_exception).to_union()
+
+    return Cmd.wrap_impure(_action)
+
+
+def list_jobs(
+    client: BatchClient,
+    name: JobName,
+    queue: QueueName,
+    status: frozenset[JobStatus],
+) -> Stream[BatchJobObj]:
+    def _extract(page: JobsPage) -> Maybe[Maybe[str]]:
+        return page.next_item.map(lambda n: Maybe.some(n))
+
+    def _cmd(index: Maybe[str]) -> Cmd[JobsPage]:
+        return _list_jobs_page(client, queue, name, index)
+
+    return (
+        StreamFactory.generate(_cmd, _extract, Maybe.empty(str))
+        .map(lambda j: PureIterFactory.from_list(j.items))
+        .transform(lambda x: StreamTransform.chain(x))
+        .filter(lambda j: j.job.status in status)
+    )