paasta-tools 1.30.8__py3-none-any.whl → 1.30.9__py3-none-any.whl

paasta_tools/frameworks/task_store.py

@@ -1,245 +0,0 @@
-import copy
-import json
-from typing import Any
-from typing import Dict
-from typing import Tuple
-from typing import Type
-from typing import TypeVar
-from typing import Union
-
-from kazoo.client import KazooClient
-from kazoo.exceptions import BadVersionError
-from kazoo.exceptions import NodeExistsError
-from kazoo.exceptions import NoNodeError
-from kazoo.protocol.states import ZnodeStat
-
-from paasta_tools.utils import _log
-
-
-class MesosTaskParametersIsImmutableError(Exception):
-    pass
-
-
-_SelfT = TypeVar("_SelfT", bound="MesosTaskParameters")
-
-
-class MesosTaskParameters:
-    health: Any
-    mesos_task_state: str
-    is_draining: bool
-    is_healthy: bool
-    offer: Any
-    resources: Any
-
-    def __init__(
-        self,
-        health=None,
-        mesos_task_state=None,
-        is_draining=None,
-        is_healthy=None,
-        offer=None,
-        resources=None,
-    ):
-        self.__dict__["health"] = health
-        self.__dict__["mesos_task_state"] = mesos_task_state
-        self.__dict__["is_draining"] = is_draining
-        self.__dict__["is_healthy"] = is_healthy
-        self.__dict__["offer"] = offer
-        self.__dict__["resources"] = resources
-
-    def __eq__(self, other):
-        return self.__dict__ == other.__dict__
-
-    def __repr__(self):
-        return "{}(\n    {})".format(
-            type(self).__name__,
-            ",\n    ".join(["%s=%r" % kv for kv in self.__dict__.items()]),
-        )
-
-    def __setattr__(self, name, value):
-        raise MesosTaskParametersIsImmutableError()
-
-    def __delattr__(self, name):
-        raise MesosTaskParametersIsImmutableError()
-
-    def merge(self: _SelfT, **kwargs) -> "MesosTaskParameters":
-        """Return a merged MesosTaskParameters object, where attributes in other take precedence over self."""
-
-        new_dict = copy.deepcopy(self.__dict__)
-        new_dict.update(kwargs)
-
-        return MesosTaskParameters(**new_dict)
-
-    @classmethod
-    def deserialize(cls: Type[_SelfT], serialized_params: Union[str, bytes]) -> _SelfT:
-        return cls(**json.loads(serialized_params))
-
-    def serialize(self):
-        return json.dumps(self.__dict__).encode("utf-8")
-
-
-class TaskStore:
-    def __init__(self, service_name, instance_name, framework_id, system_paasta_config):
-        self.service_name = service_name
-        self.instance_name = instance_name
-        self.framework_id = framework_id
-        self.system_paasta_config = system_paasta_config
-
-    def get_task(self, task_id: str) -> MesosTaskParameters:
-        """Get task data for task_id. If we don't know about task_id, return None"""
-        raise NotImplementedError()
-
-    def get_all_tasks(self) -> Dict[str, MesosTaskParameters]:
-        """Returns a dictionary of task_id -> MesosTaskParameters for all known tasks."""
-        raise NotImplementedError()
-
-    def overwrite_task(self, task_id: str, params: MesosTaskParameters) -> None:
-        raise NotImplementedError()
-
-    def add_task_if_doesnt_exist(self, task_id: str, **kwargs) -> None:
-        """Add a task if it does not already exist. If it already exists, do nothing."""
-        if self.get_task(task_id) is not None:
-            return
-        else:
-            self.overwrite_task(task_id, MesosTaskParameters(**kwargs))
-
-    def update_task(self, task_id: str, **kwargs) -> MesosTaskParameters:
-        existing_task = self.get_task(task_id)
-        if existing_task:
-            merged_params = existing_task.merge(**kwargs)
-        else:
-            merged_params = MesosTaskParameters(**kwargs)
-
-        self.overwrite_task(task_id, merged_params)
-        return merged_params
-
-    def garbage_collect_old_tasks(self, max_dead_task_age: float) -> None:
-        # TODO: call me.
-        # TODO: implement in base class.
-        raise NotImplementedError()
-
-    def close(self):
-        pass
-
-
-class DictTaskStore(TaskStore):
-    def __init__(self, service_name, instance_name, framework_id, system_paasta_config):
-        self.tasks: Dict[str, MesosTaskParameters] = {}
-        super().__init__(
-            service_name, instance_name, framework_id, system_paasta_config
-        )
-
-    def get_task(self, task_id: str) -> MesosTaskParameters:
-        return self.tasks.get(task_id)
-
-    def get_all_tasks(self) -> Dict[str, MesosTaskParameters]:
-        """Returns a dictionary of task_id -> MesosTaskParameters for all known tasks."""
-        return dict(self.tasks)
-
-    def overwrite_task(self, task_id: str, params: MesosTaskParameters) -> None:
-        # serialize/deserialize to make sure the returned values are the same format as ZKTaskStore.
-        self.tasks[task_id] = MesosTaskParameters.deserialize(params.serialize())
-
-
-class ZKTaskStore(TaskStore):
-    def __init__(self, service_name, instance_name, framework_id, system_paasta_config):
-        super().__init__(
-            service_name, instance_name, framework_id, system_paasta_config
-        )
-        self.zk_hosts = system_paasta_config.get_zk_hosts()
-
-        # For some reason, I could not get the code suggested by this SO post to work to ensure_path on the chroot.
-        # https://stackoverflow.com/a/32785625/25327
-        # Plus, it just felt dirty to modify instance attributes of a running connection, especially given that
-        # KazooClient.set_hosts() doesn't allow you to change the chroot. Must be for a good reason.
-
-        chroot = f"task_store/{service_name}/{instance_name}/{framework_id}"
-
-        temp_zk_client = KazooClient(hosts=self.zk_hosts)
-        temp_zk_client.start()
-        temp_zk_client.ensure_path(chroot)
-        temp_zk_client.stop()
-        temp_zk_client.close()
-
-        self.zk_client = KazooClient(hosts=f"{self.zk_hosts}/{chroot}")
-        self.zk_client.start()
-        self.zk_client.ensure_path("/")
-
-    def close(self):
-        self.zk_client.stop()
-        self.zk_client.close()
-
-    def get_task(self, task_id: str) -> MesosTaskParameters:
-        params, stat = self._get_task(task_id)
-        return params
-
-    def _get_task(self, task_id: str) -> Tuple[MesosTaskParameters, ZnodeStat]:
-        """Like get_task, but also returns the ZnodeStat that self.zk_client.get() returns"""
-        try:
-            data, stat = self.zk_client.get("/%s" % task_id)
-            return MesosTaskParameters.deserialize(data), stat
-        except NoNodeError:
-            return None, None
-        except json.decoder.JSONDecodeError:
-            _log(
-                service=self.service_name,
-                instance=self.instance_name,
-                level="debug",
-                component="deploy",
-                line=f"Warning: found non-json-decodable value in zookeeper for task {task_id}: {data}",
-            )
-            return None, None
-
-    def get_all_tasks(self):
-        all_tasks = {}
-
-        for child_path in self.zk_client.get_children("/"):
-            task_id = self._task_id_from_zk_path(child_path)
-            params = self.get_task(task_id)
-            # sometimes there are bogus child ZK nodes. Ignore them.
-            if params is not None:
-                all_tasks[task_id] = params
-
-        return all_tasks
-
-    def update_task(self, task_id: str, **kwargs):
-        retry = True
-        while retry:
-            retry = False
-            existing_task, stat = self._get_task(task_id)
-
-            zk_path = self._zk_path_from_task_id(task_id)
-            if existing_task:
-                merged_params = existing_task.merge(**kwargs)
-                try:
-                    self.zk_client.set(
-                        zk_path, merged_params.serialize(), version=stat.version
-                    )
-                except BadVersionError:
-                    retry = True
-            else:
-                merged_params = MesosTaskParameters(**kwargs)
-                try:
-                    self.zk_client.create(zk_path, merged_params.serialize())
-                except NodeExistsError:
-                    retry = True
-
-        return merged_params
-
-    def overwrite_task(
-        self, task_id: str, params: MesosTaskParameters, version=-1
-    ) -> None:
-        try:
-            self.zk_client.set(
-                self._zk_path_from_task_id(task_id), params.serialize(), version=version
-            )
-        except NoNodeError:
-            self.zk_client.create(
-                self._zk_path_from_task_id(task_id), params.serialize()
-            )
-
-    def _zk_path_from_task_id(self, task_id: str) -> str:
-        return "/%s" % task_id
-
-    def _task_id_from_zk_path(self, zk_path: str) -> str:
-        return zk_path.lstrip("/")