mccode-plumber 0.6.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

mccode_plumber/file_writer_control/JobHandler.py

@@ -0,0 +1,102 @@
+from datetime import datetime
+
+from file_writer_control.CommandHandler import CommandHandler
+from file_writer_control.JobStatus import JobState
+from file_writer_control.WorkerFinder import WorkerFinder
+from file_writer_control.WriteJob import WriteJob
+
+
+class JobHandler:
+    """
+    A stand in for controlling and checking the state of a job running on a file-writer instance.
+    """
+
+    def __init__(self, worker_finder: WorkerFinder, job_id=""):
+        """
+        Constructor.
+        :param worker_finder: An instance of a class that inherits from WorkerFinder and implements the member function
+        try_start_job of that class.
+        :param job_id: (Optional) The job identifier of an existing job.
+        """
+        self.worker_finder = worker_finder
+        self._job_id = job_id
+
+    def start_job(self, job: WriteJob) -> CommandHandler:
+        """
+        Start a write job. This call is not blocking. It does not guarantee that the write job will actually be started.
+        :param job: The write to be started.
+        .. note:: Starting a new job will cause the current instance of this class to no longer being able to track or
+        control previous jobs.
+        :return: A CommandHandler instance that can be used to monitor the outcome of the attempt to start a write job.
+        """
+        self._job_id = job.job_id
+        return self.worker_finder.try_start_job(job)
+
+    def get_state(self) -> JobState:
+        """
+        Get the state of the job.
+        """
+        return self.worker_finder.get_job_state(self._job_id)
+
+    def is_done(self) -> bool:
+        """
+        :return: True if job was completed without errors. False otherwise.
+        .. note:: If the job was completed with errors, this call will return False.
+        """
+        current_job_state = self.worker_finder.get_job_state(self._job_id)
+        if current_job_state == JobState.ERROR:
+            raise RuntimeError(f'Job failed with error message "{self.get_message()}".')
+        if current_job_state == JobState.TIMEOUT:
+            raise RuntimeError("Timed out while trying to start write job.")
+        return current_job_state == JobState.DONE
+
+    def get_message(self) -> str:
+        """
+        Get a string describing the error that was encountered when running the job. (If there was an error.)
+        """
+        current_status = self.worker_finder.get_job_status(self._job_id)
+        if current_status is None:
+            return ""
+        return current_status.message
+
+    def set_stop_time(self, stop_time: datetime) -> CommandHandler:
+        """
+        Set a new stop time for the file-writing job. There is no guarantee that the stop time will actually be changed.
+        This call is not blocking. Calling this member function will have no effect on the stop-time before the write
+        job has started.
+
+        :param stop_time: The new stop time of the job.
+        :return: A CommandHandler instance that can be used to monitor the outcome of the attempt to set a new stop time.
+        """
+        current_status = self.worker_finder.get_job_status(self._job_id)
+        return self.worker_finder.try_send_stop_time(
+            current_status.service_id if current_status else None,
+            self._job_id,
+            stop_time,
+        )
+
+    def stop_now(self) -> CommandHandler:
+        """
+        See the documentation for abort_write_job().
+        """
+        return self.abort_write_job()
+
+    def abort_write_job(self) -> CommandHandler:
+        """
+        Tell the file-writing to abort writing. There is no guarantee that will actually happen though.
+        This call is not blocking. Calling this member function will have no effect if done before a write job has
+        actually started.
+
+        :return: A CommandHandler instance that can be used to monitor the outcome of the attempt to set a new stop time.
+        """
+        current_status = self.worker_finder.get_job_status(self._job_id)
+        return self.worker_finder.try_send_abort(
+            current_status.service_id if current_status else None, self._job_id
+        )
+
+    @property
+    def job_id(self) -> str:
+        """
+        The job identifier of the job that this instance of the JobHandler class represent.
+        """
+        return self._job_id

mccode_plumber/file_writer_control/JobStatus.py

@@ -0,0 +1,147 @@
+from datetime import datetime, timedelta
+from enum import Enum, auto
+from typing import Dict, Optional
+
+DEFAULT_TIMEOUT = timedelta(seconds=15)
+
+
+class JobState(Enum):
+    """
+    The state of a job.
+    """
+
+    NO_JOB = auto()
+    WAITING = auto()
+    WRITING = auto()
+    TIMEOUT = auto()
+    ERROR = auto()
+    DONE = auto()
+    UNAVAILABLE = auto()
+
+
+class JobStatus:
+    """
+    Contains general information about the (execution) of a job.
+    """
+
+    def __init__(self, job_id: str, timeout: Optional[timedelta] = DEFAULT_TIMEOUT):
+        self._job_id = job_id
+        self._timeout = timeout
+        self._service_id = ""
+        self._file_name = ""
+        self._last_update = datetime.now()
+        self._state = JobState.WAITING
+        self._metadata: Optional[Dict] = None
+        self._message = ""
+
+    def update_status(self, new_status: "JobStatus") -> None:
+        """
+        Updates the status/state of a this instance of the JobStatus class, using another instance.
+        .. note:: The job identifier of this instance and the other must be identical.
+        :param new_status: The other instance of the JobStatus class.
+        """
+        if new_status.job_id != self.job_id:
+            raise RuntimeError(
+                f"Job id of status update is not correct ({self.job_id} vs {new_status.job_id})"
+            )
+        self._state = new_status.state
+        if new_status.message:
+            self._message = new_status.message
+        self._service_id = new_status.service_id
+        self._file_name = new_status.file_name
+        self._last_update = new_status.last_update
+        self._metadata = new_status.metadata
+
+    def check_if_outdated(self, current_time: datetime):
+        """
+        Given the current time, state and the time of the last update: Have we lost the connection?
+        :param current_time: The current time
+        """
+        if (
+            self.state != JobState.DONE
+            and self.state != JobState.ERROR
+            and self.state != JobState.TIMEOUT
+            and current_time - self.last_update > self._timeout
+        ):
+            self._state = JobState.TIMEOUT
+            self._last_update = current_time
+
+    @property
+    def job_id(self) -> str:
+        """
+        The (unique) job identifier.
+        """
+        return self._job_id
+
+    @property
+    def service_id(self) -> str:
+        """
+        The (unique) service identifier of the instance of the file-writer that executes the current job.
+        """
+        return self._service_id
+
+    @service_id.setter
+    def service_id(self, new_service_id: str) -> None:
+        if not self._service_id:
+            self._service_id = new_service_id
+            self._last_update = datetime.now()
+        elif self._service_id == new_service_id:
+            return
+        else:
+            raise RuntimeError(
+                f'Can not set service_id of job with id "{self._job_id}" to "{new_service_id}" as it has already been set to "{self._service_id}".'
+            )
+
+    @property
+    def last_update(self) -> datetime:
+        """
+        The local time stamp of the last update of the status of the job.
+        """
+        return self._last_update
+
+    @property
+    def state(self) -> JobState:
+        """
+        The current state of the job.
+        """
+        return self._state
+
+    @property
+    def file_name(self) -> str:
+        """
+        The file name of the job. None if the file name is not known.
+        """
+        if self._file_name == "":
+            return None
+        return self._file_name
+
+    @file_name.setter
+    def file_name(self, new_file_name: str) -> None:
+        self._file_name = new_file_name
+        self._last_update = datetime.now()
+
+    @state.setter
+    def state(self, new_state: JobState) -> None:
+        self._state = new_state
+        self._last_update = datetime.now()
+
+    @property
+    def message(self) -> str:
+        """
+        Status/state message of the job as received from the file-writer.
+        """
+        return self._message
+
+    @property
+    def metadata(self) -> Optional[Dict]:
+        return self._metadata
+
+    @metadata.setter
+    def metadata(self, metadata: Dict) -> None:
+        self._metadata = metadata
+
+    @message.setter
+    def message(self, new_message: str) -> None:
+        if new_message:
+            self._message = new_message
+            self._last_update = datetime.now()

mccode_plumber/file_writer_control/KafkaTopicUrl.py

@@ -0,0 +1,22 @@
+import re
+
+
+class KafkaTopicUrl:
+    """
+    Class for extracting address, port and topic name from a Kafka topic url.
+    """
+
+    test_regexp = re.compile(
+        r"^\s*(?:kafka://)?(?:(?P<host>[^/?#:]+)(?::(?P<port>\d+){1,5})?)/(?P<topic>[a-zA-Z0-9._-]+)\s*$"
+    )
+
+    def __init__(self, url: str):
+        result = re.match(KafkaTopicUrl.test_regexp, url)
+        if result is None:
+            raise RuntimeError("Unable to match kafka url.")
+        self.port = 9092  # Default Kafka broker port
+        if result.group("port") is not None:
+            self.port = int(result.group("port"))
+        self.host = result.group("host")
+        self.host_port = f"{self.host}:{self.port}"
+        self.topic = result.group("topic")

mccode_plumber/file_writer_control/StateExtractor.py

@@ -0,0 +1,58 @@
+from json import loads
+from typing import Optional
+
+from streaming_data_types.action_response_answ import (
+    ActionOutcome,
+    ActionResponse,
+    ActionType,
+)
+from streaming_data_types.status_x5f2 import StatusMessage
+
+from file_writer_control.CommandStatus import CommandState
+from file_writer_control.JobStatus import JobState
+from file_writer_control.WorkerStatus import WorkerState
+
+
+def extract_worker_state_from_status(status: StatusMessage) -> WorkerState:
+    """
+    Determine the worker state (i.e. file-writer state) based on a file-writer status message.
+    :param status: A status update message from a file-writer.
+    :return: The extracted worker state.
+    """
+    json_struct = loads(status.status_json)
+    status_map = {"writing": WorkerState.WRITING, "idle": WorkerState.IDLE}
+    try:
+        status_string = json_struct["state"]
+        return status_map[status_string]
+    except KeyError:
+        return WorkerState.UNKNOWN
+
+
+def extract_state_from_command_answer(answer: ActionResponse) -> CommandState:
+    """
+    Determine the command state from a action response message.
+    :param answer: The action (either "start a job" or "set top time") response from a file-writer.
+    :return: The extracted command state/response.
+    """
+    status_map = {
+        ActionOutcome.Failure: CommandState.ERROR,
+        ActionOutcome.Success: CommandState.SUCCESS,
+    }
+    try:
+        return status_map[answer.outcome]
+    except KeyError:
+        return CommandState.ERROR
+
+
+def extract_job_state_from_answer(answer: ActionResponse) -> Optional[JobState]:
+    """
+    Determine the file writing job state from a action response message.
+    :param answer: The action (either "start a job" or "set top time") response from a file-writer.
+    :return: The extracted job state, None if job state can not be determined from this answer.
+    """
+    if answer.action == ActionType.StartJob:
+        if answer.outcome == ActionOutcome.Success:
+            return JobState.WRITING
+        else:
+            return JobState.ERROR
+    return None

mccode_plumber/file_writer_control/WorkerFinder.py

@@ -0,0 +1,139 @@
+import uuid
+from datetime import datetime
+from typing import Dict, List, Optional
+
+from kafka import KafkaProducer
+from kafka.errors import NoBrokersAvailable
+from streaming_data_types.run_stop_6s4t import serialise_6s4t as serialise_stop
+
+from file_writer_control.CommandChannel import CommandChannel
+from file_writer_control.CommandHandler import CommandHandler
+from file_writer_control.CommandStatus import CommandStatus
+from file_writer_control.JobStatus import JobState, JobStatus
+from file_writer_control.KafkaTopicUrl import KafkaTopicUrl
+from file_writer_control.WorkerStatus import WorkerStatus
+from file_writer_control.WriteJob import WriteJob
+
+
+class WorkerFinderBase:
+    def __init__(
+        self,
+        command_topic: str,
+        command_channel: CommandChannel,
+        message_producer: KafkaProducer,
+    ):
+        """
+        Constructor.
+        """
+        self.command_channel = command_channel
+        self.command_topic = command_topic
+        self.message_producer = message_producer
+
+    def send_command(self, message: bytes):
+        """
+        Send a message (command) to the file-writer "command"-topic.
+        :param message: The command/message as binary data.
+        """
+        self.message_producer.send(self.command_topic, message)
+
+    def try_start_job(self, job: WriteJob) -> CommandHandler:
+        """
+        Attempts to start a file-writing job. This function is not blocking. No guarantees are given that the job will
+        be successfully started.
+        .. note:: This class must be implemented by the classes inheriting from this one.
+        :param job: The file-writing job to be started.
+        :return: A CommandHandler instance for (more) easily checking the outcome of attempting to start a write job.
+        """
+        raise NotImplementedError("Not implemented in base class.")
+
+    def try_send_stop_time(
+        self, service_id: Optional[str], job_id: str, stop_time: datetime
+    ) -> CommandHandler:
+        """
+        Sends a "set stop time" message to a file-writer running a job as identified by the parameters.
+        This function is not blocking. No guarantees are given that this command will be followed.
+        :param service_id: The (optional) service identifier of the file-writer to receive the command.
+        :param job_id: The job identifier of the currently running file-writer job.
+        :param stop_time: The new stop time.
+        :return: A CommandHandler instance for (more) easily checking the outcome of setting a new stop time.
+        """
+        command_id = str(uuid.uuid1())
+        message = serialise_stop(
+            job_id=job_id,
+            service_id=service_id,
+            command_id=command_id,
+            stop_time=stop_time,
+        )
+        self.command_channel.add_command_id(job_id=job_id, command_id=command_id)
+        self.send_command(message)
+        return CommandHandler(self.command_channel, command_id)
+
+    def try_send_stop_now(
+        self, service_id: Optional[str], job_id: str
+    ) -> CommandHandler:
+        """
+        See documentation for `try_send_abort()`.
+        """
+        return self.try_send_abort(service_id, job_id)
+
+    def try_send_abort(self, service_id: Optional[str], job_id: str) -> CommandHandler:
+        """
+        Sends a "abort" message to a file-writer running a job as identified by the parameters of this function.
+        This function is not blocking. No guarantees are given that this command will be followed.
+        :param service_id: The (optional) service identifier of the file-writer to receive the command.
+        :param job_id: The job identifier of the currently running file-writer job.
+        :return: A CommandHandler instance for (more) easily checking the outcome of the "abort" command.
+        """
+        return self.try_send_stop_time(service_id, job_id, 0)
+
+    def list_known_workers(self) -> List[WorkerStatus]:
+        """
+        :return: A list of the (known) status of the workers publishing status updates to the configured command topic.
+        """
+        return self.command_channel.list_workers()
+
+    def list_known_jobs(self) -> List[JobStatus]:
+        """
+        :return: A list of the (known) jobs and their status as published on the configured command topic.
+        """
+        return self.command_channel.list_jobs()
+
+    def list_known_commands(self) -> List[CommandStatus]:
+        """
+        :return: A list of the (known) commands and their outcomes as published on the configured command topic.
+        """
+        return self.command_channel.list_commands()
+
+    def get_job_state(self, job_id: str) -> JobState:
+        """
+        Get the state of a specific job.
+        :param job_id: The (unique) identifier of the job that we are trying to find the state of.
+        :return: The state of the job if known, JobState.UNAVAILABLE if job is not known.
+        """
+        current_job = self.command_channel.get_job(job_id)
+        if current_job is None:
+            return JobState.UNAVAILABLE
+        return current_job.state
+
+    def get_job_status(self, job_id: str) -> JobStatus:
+        """
+        Get the full (known) status of a specific job.
+        :param job_id: The (unique) identifier of the job that we are trying to find the status of.
+        :return: The status of the job if known. None if it is not.
+        """
+        return self.command_channel.get_job(job_id)
+
+
+class WorkerFinder(WorkerFinderBase):
+    def __init__(self, command_topic_url: str, kafka_config: Dict[str, str] = {}):
+        temp_cmd_ch = CommandChannel(command_topic_url, kafka_config=kafka_config)
+        command_url = KafkaTopicUrl(command_topic_url)
+        try:
+            temp_producer = KafkaProducer(
+                bootstrap_servers=[command_url.host_port], **kafka_config
+            )
+        except NoBrokersAvailable as e:
+            raise NoBrokersAvailable(
+                f'Unable to find brokers (or connect to brokers) on address: "{command_url.host_port}"'
+            ) from e
+        super().__init__(command_url.topic, temp_cmd_ch, temp_producer)

mccode_plumber/file_writer_control/WorkerJobPool.py

@@ -0,0 +1,70 @@
+from typing import Dict
+
+from kafka import KafkaProducer
+from kafka.errors import NoBrokersAvailable
+
+from file_writer_control.CommandHandler import CommandHandler
+from file_writer_control.CommandStatus import CommandState
+from file_writer_control.KafkaTopicUrl import KafkaTopicUrl
+from file_writer_control.WorkerFinder import WorkerFinder
+from file_writer_control.WriteJob import WriteJob
+
+
+class WorkerJobPool(WorkerFinder):
+    """
+    A child of WorkerFinder intended for use with "worker pool" style of starting a file-writing job.
+    """
+
+    def __init__(
+        self,
+        job_topic_url: str,
+        command_topic_url: str,
+        max_message_size: int = 1048576 * 200,
+        kafka_config: Dict[str, str] = {},
+    ):
+        """
+        :param job_topic_url: The Kafka topic that the available file-writers are listening to for write jobs.
+        :param command_topic_url: The Kafka topic that a file-writer uses to send status updates to and receive direct
+        commands from.
+        :param max_message_size: The maximum message (actually "request") size.
+        """
+        super().__init__(command_topic_url, kafka_config=kafka_config)
+        self._job_pool = KafkaTopicUrl(job_topic_url)
+        self._max_message_size = max_message_size
+        try:
+            self._pool_producer = KafkaProducer(
+                bootstrap_servers=[self._job_pool.host_port],
+                max_request_size=max_message_size,
+                buffer_memory=max_message_size,
+                **kafka_config,
+            )
+        except NoBrokersAvailable as e:
+            raise NoBrokersAvailable(
+                f'Unable to find brokers (or connect to brokers) on address: "{self._job_pool.host_port}"'
+            ) from e
+
+    def _send_pool_message(self, message: bytes):
+        """
+        Send a message to the Kafka topic that is configured as the job-pool topic.
+        .. note:: If the file-writer has been configured properly, it will only accept start-job messages to this topic.
+        :param message: The binary data of the message.
+        """
+        if len(message) >= self._max_message_size:
+            raise RuntimeError(
+                f"Unable to send Kafka message as message size is too large ({len(message)} vs"
+                f"{self._max_message_size} bytes). Increase max message size with the 'max_message_size'"
+                f"constructor argument."
+            )
+        self._pool_producer.send(self._job_pool.topic, message)
+
+    def try_start_job(self, job: WriteJob) -> CommandHandler:
+        """
+        See base class for documentation.
+        """
+        self.command_channel.add_job_id(job.job_id)
+        self.command_channel.add_command_id(job.job_id, job.job_id)
+        self.command_channel.get_command(
+            job.job_id
+        ).state = CommandState.WAITING_RESPONSE
+        self._send_pool_message(job.get_start_message())
+        return CommandHandler(self.command_channel, job.job_id)

mccode_plumber/file_writer_control/WorkerStatus.py

@@ -0,0 +1,88 @@
+from datetime import datetime, timedelta
+from enum import Enum, auto
+from typing import Optional
+
+DEFAULT_TIMEOUT = timedelta(seconds=15)
+
+
+class WorkerState(Enum):
+    """
+    The state of a worker (i.e. a file-writer instance).
+    """
+
+    IDLE = auto()
+    WRITING = auto()
+    UNKNOWN = auto()
+    UNAVAILABLE = auto()
+
+
+class WorkerStatus(object):
+    """
+    Contains general status information about a worker.
+    """
+
+    def __init__(self, service_id: str, timeout: Optional[timedelta] = DEFAULT_TIMEOUT):
+        self._last_update = datetime.now()
+        self._service_id = service_id
+        self._timeout = timeout
+        self._state = WorkerState.UNAVAILABLE
+
+    def __eq__(self, other_status: "WorkerStatus") -> bool:
+        if not isinstance(other_status, WorkerStatus):
+            raise NotImplementedError
+        return (
+            self.service_id == other_status.service_id
+            and self.state == other_status.state
+        )
+
+    def update_status(self, new_status: "WorkerStatus"):
+        """
+        Updates the status/state of this instance of the WorkerStatus class using another instance.
+        .. note:: The service identifier of both this instance and the other one must be identical.
+        :param new_status: The other instance of the WorkerStatus class.
+        """
+        if new_status.service_id != self.service_id:
+            raise RuntimeError(
+                f"Service id of status update is not correct ({self.service_id} vs {new_status.service_id})"
+            )
+        self._state = new_status.state
+        self._last_update = new_status.last_update
+
+    def check_if_outdated(self, current_time: datetime):
+        """
+        Given the current time, state and the time of the last update: Have we lost the connection?
+        :param current_time: The current time
+        """
+        if (
+            self.state != WorkerState.UNAVAILABLE
+            and current_time - self.last_update > self._timeout
+        ):
+            self._state = WorkerState.UNAVAILABLE
+            self._last_update = current_time
+
+    @property
+    def state(self) -> WorkerState:
+        """
+        The current state of the worker.
+        """
+        return self._state
+
+    @property
+    def service_id(self) -> str:
+        """
+        The service identifier of the worker that this instance of the WorkerState class represent.
+        """
+        return self._service_id
+
+    @property
+    def last_update(self) -> datetime:
+        """
+        The local time stamp of the last update of the status of the file-writer instance that this instance of the
+        WorkerStatus class represents.
+        """
+        return self._last_update
+
+    @state.setter
+    def state(self, new_state: WorkerState):
+        self._last_update = datetime.now()
+        self._state = new_state