<reponame>rueberger/observers<gh_stars>0
""" This module contains base classes and logic for observe backends, which are called Archivers
"""

import bson

from abc import ABC, abstractmethod

import pymongo

from observers.step import Step


# TODO: need to figure out if it's worth the extra complication to add abstractmethods
#  for traversals, to allow traversals to be optimized for each backend
class Archiver(ABC):
    """ Base class for an observer backend that supports branching experiments and mandates that
     all backends include logic for retrieving results.

    In order to support branching, some concessions must be made. It is no longer possible
     to assign well-defined sequential ids to entire experiments. Instead, experiments are
     tracked as a series of possibly branching steps. Each trajectory is uniquely identified
     by the hash of its endpoint. Same thing as git, steps are commits.

    """

    def __init__(self, archive_name):
        """ Archiver constructor

        Args:
          archive_name - str
        """
        self.archive_name = archive_name

    def commit_step(self, parent_id, data):
        """ Commit step to backend. parent_id may be None

        Args:
          parent_id - int | None
          data: BSON serializable data

        Returns:
          step_id - int
        """
        assert bson.is_valid(data)

        self._commit_step(parent_id, data)
        step = Step(parent_id, data)
        return hash(step)

    @abstractmethod
    def _commit_step(self, parent_id, data):
        """ Commit step to backend. parent_id may be None

        Args:
          parent_id - int | None
          data: BSON serializable data

        """
        pass

    def fetch_step(self, step_id):
        """ Fetch a step from the backend

        Args:
          step_id - str

        Returns:
          step - Step
        """
        parent_id, data = self._fetch_step(step_id)
        return Step(parent_id, data)

    @abstractmethod
    def _fetch_step(self, step_id):
        """ Fetch a step from the backend

        Args:
          step_id - str

        Returns:
          parent_id - int
          data
        """
        pass

    @abstractmethod
    def descendants(self, step_id):
        """ Fetch all descendants of a given step

        Args:
          step_id - int

        Returns:
          steps - [Step]
        """
        pass

    @abstractmethod
    def roots(self):
        """ Fetch all steps which have no parent

        Returns:
          steps - [Step]
        """
        pass

    @abstractmethod
    def leaves(self):
        """ Fetch all steps which have no descendants

        Returns:
          steps - [Step]
        """
        pass

    def pull(self, step_id):
        """ Pull trajectory ending at step_id

        Args:
          step_id - int

        Returns:
          trajectory: ending at step_id - [Step]
        """
        curr_step = self.fetch_step(step_id)
        trajectory = [curr_step]
        while curr_step.parent_id is not None:
            curr_step = self.fetch_step(curr_step.parent_id)
            trajectory.append(curr_step)
        return trajectory[::-1]


class MongoArchiver(Archiver):
    """ MongoDB backed archiver
    """

    def __init__(self, mongo_client, archive_name):
        """ MongoArchiver constructor

        Args:
          mongo_client - pymongo.MongoClient
          archive_name - str
        """
        assert isinstance(mongo_client, pymongo.MongoClient)

        self.client = mongo_client

        # TODO: broken in 2.7
        super().__init__(archive_name)

        self.archive = self.client['archive'][self.archive_name]

    # TODO: needs initialize method

    def database_initialized(self):
        """ Checks if the archive db has been initialized, by checking for
          the presence of indices

        Returns:
          initialized - bool
        """
        return len(self.archive.index_information()) > 1


    @abstractmethod
    def _commit_step(self, parent_id, data):
        """ Commit step to backend. parent_id may be None

        Args:
          parent_id - int | None
          data: BSON serializable data

        """
        pass

    @abstractmethod
    def _fetch_step(self, step_id):
        """ Fetch a step from the backend

        Args:
          step_id - str

        Returns:
          parent_id - int
          data
        """
        pass

    @abstractmethod
    def descendants(self, step_id):
        """ Fetch all descendants of a given step

        Args:
          step_id - int

        Returns:
          steps - [Step]
        """
        pass

    @abstractmethod
    def roots(self):
        """ Fetch all steps which have no parent

        Returns:
          steps - [Step]
        """
        pass

    @abstractmethod
    def leaves(self):
        """ Fetch all steps which have no descendants

        Returns:
          steps - [Step]
        """
        pass
