tasks_api.py

"""
Process and analyze your data with tasks in the InfluxDB task engine.

Use tasks (scheduled Flux queries) to input a data stream and then analyze, modify, and act on the data accordingly.
"""

import datetime
from typing import List

from influxdb_client import TasksService, Task, TaskCreateRequest, TaskUpdateRequest, LabelResponse, LabelMapping, \
    AddResourceMemberRequestBody, RunManually, Run, LogEvent


class _Page:
    def __init__(self, values, has_next, next_after):
        self.has_next = has_next
        self.values = values
        self.next_after = next_after

    @staticmethod
    def empty():
        return _Page([], False, None)

    @staticmethod
    def initial(after):
        return _Page([], True, after)


class _PageIterator:
    def __init__(self, page: _Page, get_next_page):
        self.page = page
        self.get_next_page = get_next_page

    def __iter__(self):
        return self

    def __next__(self):
        if not self.page.values:
            if self.page.has_next:
                self.page = self.get_next_page(self.page)
            if not self.page.values:
                raise StopIteration
        return self.page.values.pop(0)


class TasksApi(object):
    """Implementation for '/api/v2/tasks' endpoint."""

    def __init__(self, influxdb_client):
        """Initialize defaults."""
        self._influxdb_client = influxdb_client
        self._service = TasksService(influxdb_client.api_client)

    def find_task_by_id(self, task_id) -> Task:
        """Retrieve a task."""
        task = self._service.get_tasks_id(task_id)
        return task

    def find_tasks(self, **kwargs):
        """List all tasks up to set limit (max 500).

        :key str name: only returns tasks with the specified name
        :key str after: returns tasks after specified ID
        :key str user: filter tasks to a specific user ID
        :key str org: filter tasks to a specific organization name
        :key str org_id: filter tasks to a specific organization ID
        :key int limit: the number of tasks to return
        :return: Tasks
        """
        return self._service.get_tasks(**kwargs).tasks

    def find_tasks_iter(self, **kwargs):
        """Iterate over all tasks with pagination.

        :key str name: only returns tasks with the specified name
        :key str after: returns tasks after specified ID
        :key str user: filter tasks to a specific user ID
        :key str org: filter tasks to a specific organization name
        :key str org_id: filter tasks to a specific organization ID
        :key int limit: the number of tasks in one page
        :return: Tasks iterator
        """

        def get_next_page(page: _Page):
            return self._find_tasks_next_page(page, **kwargs)

        return iter(_PageIterator(_Page.initial(kwargs.get('after')), get_next_page))

    def create_task(self, task: Task = None, task_create_request: TaskCreateRequest = None) -> Task:
        """Create a new task."""
        if task_create_request is not None:
            return self._service.post_tasks(task_create_request)

        if task is not None:
            request = TaskCreateRequest(flux=task.flux, org_id=task.org_id, org=task.org, description=task.description,
                                        status=task.status)

            return self.create_task(task_create_request=request)

        raise ValueError("task or task_create_request must be not None")

    @staticmethod
    def _create_task(name: str, flux: str, every, cron, org_id: str) -> Task:

        task = Task(id=0, name=name, org_id=org_id, status="active", flux=flux)

        repetition = ""
        if every is not None:
            repetition += "every: "
            repetition += every

        if cron is not None:
            repetition += "cron: "
            repetition += '"' + cron + '"'

        flux_with_options = '{} \n\noption task = {{name: "{}", {}}}'.format(flux, name, repetition)
        task.flux = flux_with_options

        return task

    def create_task_every(self, name, flux, every, organization) -> Task:
        """Create a new task with every repetition schedule."""
        task = self._create_task(name, flux, every, None, organization.id)
        return self.create_task(task)

    def create_task_cron(self, name: str, flux: str, cron: str, org_id: str) -> Task:
        """Create a new task with cron repetition schedule."""
        task = self._create_task(name=name, flux=flux, cron=cron, org_id=org_id, every=None)
        return self.create_task(task)

    def delete_task(self, task_id: str):
        """Delete a task."""
        if task_id is not None:
            return self._service.delete_tasks_id(task_id=task_id)

    def update_task(self, task: Task) -> Task:
        """Update a task."""
        req = TaskUpdateRequest(flux=task.flux, description=task.description, every=task.every, cron=task.cron,
                                status=task.status, offset=task.offset)

        return self.update_task_request(task_id=task.id, task_update_request=req)

    def update_task_request(self, task_id, task_update_request: TaskUpdateRequest) -> Task:
        """Update a task."""
        return self._service.patch_tasks_id(task_id=task_id, task_update_request=task_update_request)

    def clone_task(self, task: Task) -> Task:
        """Clone a task."""
        cloned = Task(id=0, name=task.name, org_id=task.org_id, org=task.org, flux=task.flux, status="active")

        created = self.create_task(cloned)
        if task.id:
            labels = self.get_labels(task.id)
            for label in labels.labels:
                self.add_label(label.id, created.id)
        return created

    def get_labels(self, task_id):
        """List all labels for a task."""
        return self._service.get_tasks_id_labels(task_id=task_id)

    def add_label(self, label_id: str, task_id: str) -> LabelResponse:
        """Add a label to a task."""
        label_mapping = LabelMapping(label_id=label_id)
        return self._service.post_tasks_id_labels(task_id=task_id, label_mapping=label_mapping)

    def delete_label(self, label_id: str, task_id: str):
        """Delete a label from a task."""
        return self._service.delete_tasks_id_labels_id(task_id=task_id, label_id=label_id)

    def get_members(self, task_id: str):
        """List all task members."""
        return self._service.get_tasks_id_members(task_id=task_id).users

    def add_member(self, member_id, task_id):
        """Add a member to a task."""
        user = AddResourceMemberRequestBody(id=member_id)
        return self._service.post_tasks_id_members(task_id=task_id, add_resource_member_request_body=user)

    def delete_member(self, member_id, task_id):
        """Remove a member from a task."""
        return self._service.delete_tasks_id_members_id(user_id=member_id, task_id=task_id)

    def get_owners(self, task_id):
        """List all owners of a task."""
        return self._service.get_tasks_id_owners(task_id=task_id).users

    def add_owner(self, owner_id, task_id):
        """Add an owner to a task."""
        user = AddResourceMemberRequestBody(id=owner_id)
        return self._service.post_tasks_id_owners(task_id=task_id, add_resource_member_request_body=user)

    def delete_owner(self, owner_id, task_id):
        """Remove an owner from a task."""
        return self._service.delete_tasks_id_owners_id(user_id=owner_id, task_id=task_id)

    def get_runs(self, task_id, **kwargs) -> List['Run']:
        """
        Retrieve list of run records for a task.

        :param task_id: task id
        :key str after: returns runs after specified ID
        :key int limit: the number of runs to return
        :key datetime after_time: filter runs to those scheduled after this time, RFC3339
        :key datetime before_time: filter runs to those scheduled before this time, RFC3339
        """
        return self._service.get_tasks_id_runs(task_id=task_id, **kwargs).runs

    def get_run(self, task_id: str, run_id: str) -> Run:
        """
        Get run record for specific task and run id.

        :param task_id: task id
        :param run_id: run id
        :return: Run for specified task and run id
        """
        return self._service.get_tasks_id_runs_id(task_id=task_id, run_id=run_id)

    def get_run_logs(self, task_id: str, run_id: str) -> List['LogEvent']:
        """Retrieve all logs for a run."""
        return self._service.get_tasks_id_runs_id_logs(task_id=task_id, run_id=run_id).events

    def run_manually(self, task_id: str, scheduled_for: datetime = None):
        """
        Manually start a run of the task now overriding the current schedule.

        :param task_id:
        :param scheduled_for: planned execution
        """
        r = RunManually(scheduled_for=scheduled_for)
        return self._service.post_tasks_id_runs(task_id=task_id, run_manually=r)

    def retry_run(self, task_id: str, run_id: str):
        """
        Retry a task run.

        :param task_id: task id
        :param run_id: run id
        """
        return self._service.post_tasks_id_runs_id_retry(task_id=task_id, run_id=run_id)

    def cancel_run(self, task_id: str, run_id: str):
        """
        Cancel a currently running run.

        :param task_id:
        :param run_id:
        """
        return self._service.delete_tasks_id_runs_id(task_id=task_id, run_id=run_id)

    def get_logs(self, task_id: str) -> List['LogEvent']:
        """
        Retrieve all logs for a task.

        :param task_id: task id
        """
        return self._service.get_tasks_id_logs(task_id=task_id).events

    def find_tasks_by_user(self, task_user_id):
        """List all tasks by user."""
        return self.find_tasks(user=task_user_id)

    def _find_tasks_next_page(self, page: _Page, **kwargs):
        if not page.has_next:
            return _Page.empty()

        args = {**kwargs, 'after': page.next_after} if page.next_after is not None else kwargs
        tasks_response = self._service.get_tasks(**args)

        tasks = tasks_response.tasks
        has_next = tasks_response.links.next is not None
        last_id = tasks[-1].id if tasks else None

        return _Page(tasks, has_next, last_id)