dkist-processing-common 10.2.2__py3-none-any.whl

dkist_processing_common/__init__.py

@@ -0,0 +1,9 @@
+"""Package providing support classes and methods used by all workflow tasks."""
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version
+
+try:
+    __version__ = version(distribution_name=__name__)
+except PackageNotFoundError:
+    # package is not installed
+    __version__ = "unknown"

dkist_processing_common/_util/__init__.py

@@ -0,0 +1 @@
+"""Utilities not intended for direct access outside the dkist_processing_common package."""

dkist_processing_common/_util/constants.py

@@ -0,0 +1,98 @@
+"""Wrapper for interactions with shared database that holds arbitrary data that persists across the entire recipe run."""
+import json
+from collections.abc import MutableMapping
+from enum import Enum
+from typing import Generator
+
+from dkist_processing_common._util.tags import TagDB
+
+
+class ConstantsDb(MutableMapping):
+    """
+    Base class defining the constants db.
+
+    Initialize a connection to the shared database.
+
+    Parameters
+    ----------
+    recipe_run_id
+        The resipe_run_id
+    task_name
+        The task name
+    """
+
+    def __init__(self, recipe_run_id: int, task_name: str):
+        self.store = TagDB(recipe_run_id, task_name, "constant")
+        self._audit_db = TagDB(
+            recipe_run_id=recipe_run_id, task_name=task_name, namespace="constants_audit"
+        )
+        self._audit_set_tag = f"SET_{task_name}"
+
+    @staticmethod
+    def extract_value(value: set) -> int | str | float:
+        """
+        Pop the first (and only) value from set and convert it from a json string.
+
+        Parameters
+        ----------
+        value
+            The set from which to pop the value
+
+        Returns
+        -------
+        The value popped from the set.
+        """
+        return json.loads(value.pop())
+
+    def __getitem__(self, key: str) -> int | str | float | tuple | list:
+        """Return the constant stored at a specific key. Raise and error if the key doesn't exist."""
+        if isinstance(key, Enum):
+            key = key.value
+        value = self.store.all(key)
+        if not value:
+            raise KeyError(f"Constant {key} does not exist")
+        return self.extract_value(value)
+
+    def __delitem__(self, key: str):
+        """'delete' a key by making it map to an empty set."""
+        self.store.clear_tag(key)
+
+    def __setitem__(self, key: str, value: str | int | float | tuple | list):
+        """Set a constant key with the specified value. Raise an error if the key already exists."""
+        if self.store.all(key):
+            raise ValueError(f"Constant {key} already exists")
+        self.store.add(key, json.dumps(value))
+        # audit which constant key had a value set
+        self._audit_db.add(tag=self._audit_set_tag, value=key)
+
+    def __iter__(self) -> Generator[str, None, None]:
+        """Yield the currently defined constants as strings."""
+        yield from self.store.tags
+
+    def __len__(self):
+        """Return the number of constants currently defined."""
+        return len(self.store.tags)
+
+    def close(self):
+        """Close the db connection.  Call on __exit__ of a Task."""
+        self.store.close()
+        self._audit_db.close()
+
+    def purge(self):
+        """Remove all constants associated with the instance recipe run id."""
+        self.store.purge()
+        self._audit_db.purge()
+
+    def rollback(self):
+        """Remove all constants associated with the instance recipe run id and task name."""
+        # remove constants
+        for key in self._audit_db.all(tags=self._audit_set_tag):
+            self.store.clear_tag(key)
+        # remove audit
+        self._audit_db.clear_tag(tag=self._audit_set_tag)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()

dkist_processing_common/_util/graphql.py

@@ -0,0 +1,93 @@
+"""Extension of the GraphQL supporting retries for data processing use cases."""
+import logging
+from typing import Any
+from typing import Callable
+
+import requests
+from gqlclient.base import DefaultParameters
+from gqlclient.base import GraphQLClientBase
+from gqlclient.request_wrap import wrap_request
+from requests.adapters import HTTPAdapter
+from requests.packages.urllib3.util.retry import Retry
+
+from dkist_processing_common.config import common_configurations
+
+logger = logging.getLogger(__name__)
+
+
+class GraphQLClient(GraphQLClientBase):
+    """Helper class for formatting and executing synchronous GraphQL queries and mutations."""
+
+    adapter = HTTPAdapter(
+        max_retries=Retry(
+            total=10,
+            backoff_factor=1,
+            status_forcelist=[502, 503, 404],
+            allowed_methods=["POST"],  # all graphql methods are POST
+        )
+    )
+
+    def execute_gql_call(self, query: dict, **kwargs) -> dict:
+        """
+        Execute a GraphQL query or mutation using requests.
+
+        :param query: Dictionary formatted graphql query
+
+        :param kwargs: Optional arguments that `requests` takes. e.g. headers
+
+        :return: Dictionary containing the response from the GraphQL endpoint
+        """
+        logger.debug(f"Executing graphql call: host={self.gql_uri}")
+        kwargs["headers"] = {
+            **kwargs.get("headers", {}),
+            "Service-Name": __name__,
+            "Authorization": common_configurations.gql_auth_token,
+        }
+        with requests.sessions.Session() as http:
+            http.mount("http://", self.adapter)
+            response = http.post(url=self.gql_uri, json=query, **kwargs)
+        try:
+            response.raise_for_status()
+        except requests.exceptions.HTTPError as e:
+            logger.error(
+                f"Error executing graphql call: status_code={e.response.status_code}, detail={e.response.text}"
+            )
+            raise e
+        return response.json()
+
+    def execute_gql_query(
+        self,
+        query_base: str,
+        query_response_cls: type,
+        query_parameters: object | None = DefaultParameters,
+        response_encoder: Callable[[str, list[dict] | dict, type], Any] | None = None,
+        **kwargs,
+    ) -> Any:
+        """Execute gql query with parameters dynamically wrapped."""
+        if query_parameters is not None and query_parameters is not DefaultParameters:
+            query_parameters = wrap_request(query_parameters)
+        return super().execute_gql_query(
+            query_base=query_base,
+            query_response_cls=query_response_cls,
+            query_parameters=query_parameters,
+            response_encoder=response_encoder,
+            **kwargs,
+        )
+
+    def execute_gql_mutation(
+        self,
+        mutation_base: str,
+        mutation_parameters: object,
+        mutation_response_cls: type | None = None,
+        response_encoder: Callable[[str, list[dict] | dict, type], Any] | None = None,
+        **kwargs,
+    ) -> Any:
+        """Execute gql mutation with parameters dynamically wrapped."""
+        mutation_parameters = wrap_request(mutation_parameters)
+        return super().execute_gql_mutation(
+            mutation_base=mutation_base,
+            mutation_parameters=mutation_parameters,
+            mutation_response_cls=mutation_response_cls,
+            response_encoder=response_encoder,
+            **kwargs,
+        )

dkist_processing_common/_util/scratch.py

@@ -0,0 +1,337 @@
+"""Scratch file system api."""
+import logging
+from contextlib import contextmanager
+from os import umask
+from pathlib import Path
+from shutil import rmtree
+from typing import Generator
+
+from dkist_processing_common._util.tags import TagDB
+from dkist_processing_common.config import common_configurations
+
+
+logger = logging.getLogger(__name__)
+
+
+class WorkflowFileSystem:
+    """
+    Wrapper for interactions with the shared file system "scratch" supporting recipe run id based namespaces and tagged data.
+
+    Create a workflow file system object.
+
+    Parameters
+    ----------
+    recipe_run_id
+        The recipe_run_id
+    task_name
+        The task_name
+    scratch_base_path
+        The base path at which to create the file system
+
+    """
+
+    def __init__(
+        self,
+        recipe_run_id: int = 0,
+        task_name: str = "dev_task",
+        scratch_base_path: Path | str | None = None,
+    ):
+        self.recipe_run_id = recipe_run_id
+        self.task_name = task_name
+        if not scratch_base_path:
+            scratch_base_path = common_configurations.scratch_base_path
+        self.scratch_base_path = scratch_base_path
+        self.workflow_base_path = Path(self.scratch_base_path) / str(recipe_run_id)
+        with self._mask():
+            self.workflow_base_path.mkdir(parents=True, exist_ok=True)
+        self._tag_db = TagDB(recipe_run_id=self.recipe_run_id, task_name=self.task_name)
+        self._audit_db = TagDB(
+            recipe_run_id=self.recipe_run_id, task_name=self.task_name, namespace="scratch_audit"
+        )
+        self._audit_write_tag = f"WRITE_{self.task_name}"
+        self._audit_tag_tag = f"TAG_{self.task_name}"
+        self._audit_new_tag_cache = dict()
+
+    @staticmethod
+    @contextmanager
+    def _mask():
+        """Set a permissive umask to allow other users (e.g. globus) to modify resources created by the scratch library."""
+        old_mask = umask(0)
+        try:
+            yield
+        finally:
+            umask(old_mask)
+
+    def absolute_path(self, relative_path: Path | str) -> Path:
+        """
+        Convert a relative path to an absolute path with the base directories for the that workflow instance.
+
+        Parameters
+        ----------
+        relative_path
+            The relative_path input
+
+        Returns
+        -------
+        The absolute path.
+        """
+        relative_path = Path(relative_path)
+        if relative_path.is_absolute():
+            raise ValueError("Relative path must be relative")
+
+        return self.workflow_base_path / relative_path
+
+    @staticmethod
+    def _parse_tags(tags: str | list | None) -> list:
+        """Parse tags to support an individual tag in the form of a string or an arbitrarily nested list of strings."""
+        if tags is None:
+            return []
+        if isinstance(tags, str):
+            return [tags]
+        return _flatten_list(tags)
+
+    def write(
+        self,
+        file_obj: bytes,
+        relative_path: Path | str,
+        tags: str | list | None = None,
+        overwrite: bool = False,
+    ) -> None:
+        """
+        Write a file object to the path specified and tagged with any tags listed in tags.
+
+        Parameters
+        ----------
+        file_obj
+            The file object to be written
+        relative_path
+            The relative path at which to write the file
+        tags
+            The tags to be associated with the file object
+        overwrite
+            Should the file be overwritten if it already exists?
+
+        Returns
+        -------
+        None
+        """
+        tags = self._parse_tags(tags)
+        path = self.absolute_path(relative_path)
+        # audit the path that was written to scratch
+        self._audit_db.add(tag=self._audit_write_tag, value=str(path))
+        with self._mask():
+            path.parent.mkdir(parents=True, exist_ok=True)
+            if overwrite:
+                mode = "wb"
+            else:
+                mode = "xb"
+            with path.open(mode=mode) as f:
+                f.write(file_obj)
+        self.tag(path, tags)
+
+    def delete(self, path: Path | str):
+        """
+        Delete the file or path.
+
+        Parameters
+        ----------
+        path
+            The path to be deleted
+
+        Returns
+        -------
+        None
+        """
+        path = Path(path)
+        path.unlink(missing_ok=True)
+        self._tag_db.clear_value(value=path)
+
+    def tag(self, path: Path | str, tags: list | str) -> None:
+        """
+        Tag existing paths.
+
+        The path must be relative to the WorkflowFileSystem base path and must exist.
+
+        Parameters
+        ----------
+        path
+            The path to tag
+        tags
+            The tags associated with the path.
+
+        Returns
+        -------
+        None
+        """
+        tags = self._parse_tags(tags)
+        path = Path(path)
+        if not (self.workflow_base_path in path.parents):
+            raise ValueError(
+                f"Cannot tag paths which are not children of the base path {self.workflow_base_path}"
+            )
+        if not path.exists():
+            raise FileNotFoundError(f"Cannot tag paths which do not exist. {path=}")
+
+        for tag in tags:
+            # audit the tag that was newly added to the scratch tag db
+            if self._tag_is_new(tag=tag):
+                self._audit_db.add(tag=self._audit_tag_tag, value=tag)
+            self._tag_db.add(tag, str(path))
+
+    def _tag_is_new(self, tag: str) -> bool:
+        if self._audit_new_tag_cache.get(tag, None) is None:
+            tag_is_new = not bool(self._tag_db.all(tags=tag))
+            self._audit_new_tag_cache[tag] = tag_is_new
+        return self._audit_new_tag_cache[tag]
+
+    def tags(self, path: Path | str):
+        """
+        Return the tags associated with the given file object.
+
+        Parameters
+        ----------
+        path
+            The input file object
+        Returns
+        -------
+        An iterable containing the tags associated with the file
+        """
+        value = str(path)
+        return self._tag_db.tags_for_value(value=value)
+
+    def remove_tags(self, path: Path | str, tags: list | str) -> None:
+        """Remove a tag or tags from a given path."""
+        tags = self._parse_tags(tags)
+        for tag in tags:
+            self._tag_db.remove(tag, str(path))
+
+    def find_any(self, tags: str | list) -> Generator[Path, None, None]:
+        """
+        Return a generator of Path objects that are tagged by the union of the input tags.
+
+        Parameters
+        ----------
+        tags
+            The tags to be used in the search
+
+        Returns
+        -------
+        A generator of path objects matching the union of the desired tags
+        """
+        tags = self._parse_tags(tags)
+        paths = self._tag_db.any(tags)
+        logger.debug(f"Found {len(paths)} files containing the set of {tags=}")
+        for path in paths:
+            yield Path(path)
+
+    def find_all(self, tags: str | list) -> Generator[Path, None, None]:
+        """
+        Return a generator of Path objects that are tagged by the intersection of the input tags.
+
+        Parameters
+        ----------
+        tags
+            The tags to be used in the search
+
+        Returns
+        -------
+        A generator of path objects matching the intersection of the desired tags
+        """
+        tags = self._parse_tags(tags)
+        paths = self._tag_db.all(tags)
+        logger.debug(f"Found {len(paths)} files containing the set of {tags=}")
+        for path in paths:
+            yield Path(path)
+
+    def count_any(self, tags: str | list) -> int:
+        """
+        Return the number of objects that are tagged by the union of the input tags.
+
+        Parameters
+        ----------
+        tags
+            The tags to be used in the search
+
+        Returns
+        -------
+        The number of objects tagged with the union of the input tags.
+        """
+        tags = self._parse_tags(tags)
+        return len(self._tag_db.any(tags))
+
+    def count_all(self, tags: str | list) -> int:
+        """
+        Return the number of objects that are tagged by the intersection of the input tags.
+
+        Parameters
+        ----------
+        tags
+            The tags to be used in the search
+
+        Returns
+        -------
+        The number of objects tagged with the intersection of the input tags.
+
+        """
+        tags = self._parse_tags(tags)
+        return len(self._tag_db.all(tags))
+
+    def close(self):
+        """Close the db connection.  Call on __exit__ of a Task."""
+        self._tag_db.close()
+        self._audit_db.close()
+
+    def purge(self, ignore_errors: bool = False):
+        """
+        Remove all data (tags, files, and folders) for the instance.
+
+        Call when tearing down a workflow
+
+        Parameters
+        ----------
+        ignore_errors
+            If set, errors will be ignored, otherwise stop at the first error
+        Returns
+        -------
+        None
+        """
+        rmtree(self.workflow_base_path, ignore_errors=ignore_errors)
+        self._tag_db.purge()
+        self._audit_db.purge()
+
+    def rollback(self):
+        """Remove all files and new tags associated with the instance recipe run id and task name."""
+        # remove files
+        for path in self._audit_db.all(tags=self._audit_write_tag):
+            path = Path(path)
+            path.unlink(missing_ok=True)
+            self._tag_db.clear_value(path)
+        # remove tags
+        for tag in self._audit_db.all(tags=self._audit_tag_tag):
+            self._tag_db.clear_tag(tag=tag)
+        # remove audit
+        self._audit_db.clear_tag(tag=self._audit_write_tag)
+        self._audit_db.clear_tag(tag=self._audit_tag_tag)
+
+    def __repr__(self):
+        return f"WorkflowFileSystem(recipe_run_id={self.recipe_run_id}, task_name={self.task_name}, scratch_base_path={self.scratch_base_path})"
+
+    def __str__(self):
+        return f"{self!r} connected to {self._tag_db}"
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+
+def _flatten_list(elements: list) -> list:
+    """Flatten an arbitrarily nested list."""
+    result = []
+    for element in elements:
+        if isinstance(element, list):
+            result.extend(_flatten_list(element))
+        else:
+            result.append(element)
+    return result