Refactor state directories

Author: mgax

opslib/__init__.py

@@ -5,14 +5,12 @@
 from .local import run
 from .places import Command, Directory, File, LocalHost, SshHost
 from .props import Prop
-from .state import JsonState
 
 __all__ = [
     "Command",
     "Component",
     "Directory",
     "File",
-    "JsonState",
     "Lazy",
     "LocalHost",
     "MaybeLazy",

opslib/ansible.py

@@ -11,6 +11,7 @@
 from ansible.playbook.play import Play
 from ansible.plugins.callback import CallbackBase
 from ansible.vars.manager import VariableManager
+from opslib.state import StatefulMixin
 
 from .callbacks import Callbacks
 from .components import Component
@@ -146,7 +147,7 @@ def run_ansible(hostname, ansible_variables, action, check=False):
     return result
 
 
-class AnsibleAction(Component):
+class AnsibleAction(StatefulMixin, Component):
     """
     The AnsibleAction component executes an Ansible module.
 

opslib/cli.py

@@ -7,9 +7,9 @@
 
 import click
 
+import opslib
 from .operations import apply, print_report
 from .results import OperationError
-from .state import run_gc
 
 logger = logging.getLogger(__name__)
 
@@ -89,7 +89,7 @@ def decorator(func):
         return decorator
 
 
-def get_cli(component) -> click.Group:
+def get_cli(component: "opslib.Component") -> click.Group:
     @click.group(cls=ComponentGroup)
     def cli():
         pass
@@ -110,7 +110,8 @@ def shell():
     @cli.command()
     @click.option("-n", "--dry-run", is_flag=True)
     def gc(dry_run):
-        run_gc(component, dry_run=dry_run)
+        provider = component._meta.stack._state_provider
+        provider.run_gc(component, dry_run=dry_run)
 
     @cli.forward_command("component")
     @click.pass_context

opslib/components.py

@@ -3,23 +3,20 @@
 import sys
 from functools import cached_property
 from pathlib import Path
-from typing import Any, Type, TypeVar
+from typing import Any, Type, TypeVar, cast
 
 from .props import get_instance_props
 from .results import Result
-from .state import StateDirectory
+from .state import FilesystemStateProvider
 
 logger = logging.getLogger(__name__)
 
 
 class Meta:
-    statedir = StateDirectory()
-
-    def __init__(self, component, name, parent, stateroot=None):
+    def __init__(self, component: "Component", name: str, parent: "Component | None"):
         self.component = component
         self.name = name
         self.parent = parent
-        self.stateroot = stateroot
 
     @cached_property
     def full_name(self):
@@ -29,8 +26,12 @@ def full_name(self):
         return f"{self.parent._meta.full_name}.{self.name}"
 
     @cached_property
-    def stack(self):
-        return self.component if self.parent is None else self.parent._meta.stack
+    def stack(self) -> "Stack":
+        return (
+            cast(Stack, self.component)
+            if self.parent is None
+            else self.parent._meta.stack
+        )
 
 
 class Component:
@@ -146,18 +147,17 @@ def __init__(self, import_name=None, stateroot=None, **kwargs):
         if import_name is None and stateroot is None:
             raise ValueError("Either `import_name` or `stateroot` must be set")
 
+        self._state_provider = FilesystemStateProvider(
+            stateroot or get_stateroot(import_name)
+        )
+
         super().__init__(**kwargs)
 
-        self._meta = self.Meta(
-            component=self,
-            name="__root__",
-            parent=None,
-            stateroot=stateroot or get_stateroot(import_name),
-        )
+        self._meta = self.Meta(component=self, name="__root__", parent=None)
         self.build()
 
 
-def walk(component):
+def walk(component) -> Iterator[Component]:
     """
     Iterate depth-first over all child components. The first item is
     ``component`` itself.

opslib/extras/restic.py

@@ -9,15 +9,15 @@
 from opslib.local import run
 from opslib.props import Prop
 from opslib.results import OperationError, Result
-from opslib.state import JsonState
+from opslib.state import JsonState, StatefulMixin
 
 BASH_PREAMBLE = """\
 #!/bin/bash
 set -euo pipefail
 """
 
 
-class ResticRepository(Component):
+class ResticRepository(StatefulMixin, Component):
     class Props:
         repository = Prop(str)
         password = Prop(str, lazy=True)

opslib/places.py

@@ -11,7 +11,7 @@
 from .local import LocalRunResult, run
 from .props import Prop
 from .results import Result
-from .state import JsonState
+from .state import JsonState, StatefulMixin
 from .utils import diff
 
 
@@ -404,7 +404,7 @@ def run(self, *args, **kwargs):
         return self.host.run(cwd=self.path, *args, **kwargs)
 
 
-class Command(Component):
+class Command(StatefulMixin, Component):
     """
     The Command component represents a command that should be run on the
     host during deployment.

opslib/state.py

@@ -1,75 +1,85 @@
+from collections.abc import Iterator
+from contextlib import contextmanager
 import json
 import logging
 from pathlib import Path
 import shutil
+from typing import cast
 
-logger = logging.getLogger(__name__)
+import opslib
 
+logger = logging.getLogger(__name__)
 
-class ComponentStateDirectory:
-    def __init__(self, meta):
-        self.meta = meta
 
-    @property
-    def prefix(self):
-        return self.get_prefix(create=True)
+class FilesystemStateProvider:
+    def __init__(self, stateroot: Path):
+        self.stateroot = stateroot
 
-    def get_prefix(self, create=False) -> Path:
-        if self.meta.parent is None:
-            prefix = self.meta.stateroot
+    def _get_directory(self, component: "opslib.Component") -> Path:
+        if component._meta.parent is None:
+            return self.stateroot
 
-        else:
-            parent_meta = self.meta.parent._meta
-            prefix = parent_meta.statedir.get_prefix(create=create) / self.meta.name
+        return self._get_directory(component._meta.parent) / component._meta.name
 
-        if create:
-            self._mkdir(prefix)
+    def _get_state_directory(self, component: "opslib.Component"):
+        return self._get_directory(component) / "_statedir"
 
-        return prefix
+    @contextmanager
+    def state_directory(self, component: "opslib.Component"):
+        statedir = self._get_state_directory(component)
+        if not statedir.exists():
+            statedir.mkdir(parents=True)
+        yield statedir
 
-    @property
-    def path(self):
-        return self.get_path(create=True)
+    def run_gc(self, component: "opslib.Component", dry_run=False):
+        child_names = {child._meta.name for child in component}
 
-    def get_path(self, create=False) -> Path:
-        path = self.get_prefix(create=create) / "_statedir"
+        def unexpected(item):
+            if isinstance(component, StatefulMixin) and item.name == "_statedir":
+                return False
 
-        if create:
-            self._mkdir(path)
+            if not item.is_dir():
+                return False
 
-        return path
+            return item.name not in child_names
 
-    def _mkdir(self, path):
-        if not path.is_dir():
-            logger.debug("ComponentState init %s", path)
-            path.mkdir(mode=0o700)
+        directory = self._get_directory(component)
+        if directory.exists():
+            for item in directory.iterdir():
+                if unexpected(item):
+                    print(item)
+                    if dry_run:
+                        continue
 
+                    shutil.rmtree(item)
 
-class StateDirectory:
-    def __get__(self, obj, objtype=None):
-        return ComponentStateDirectory(obj)
+        for child in component:
+            self.run_gc(child, dry_run=dry_run)
 
 
 class ComponentJsonState:
     def __init__(self, component):
         self.component = component
 
-    @property
-    def _path(self):
-        return self.component._meta.statedir.path / "state.json"
+    @contextmanager
+    def json_path(self) -> Iterator[Path]:
+        with self.component.state_directory() as statedir:
+            yield statedir / "state.json"
 
     @property
     def _data(self):
         try:
-            with self._path.open() as f:
-                return json.load(f)
+            with self.json_path() as json_path:
+                with json_path.open() as f:
+                    return json.load(f)
 
         except FileNotFoundError:
             return {}
 
     def save(self, data=(), **kwargs):
-        with self._path.open("w") as f:
-            json.dump(dict(data, **kwargs), f, indent=2)
+        with self.json_path() as json_path:
+            with json_path.open("w") as f:
+                json.dump(dict(data, **kwargs), f, indent=2)
 
         self.__dict__["_data"] = data
 
@@ -99,26 +109,11 @@ def __get__(self, obj, objtype=None):
         return ComponentJsonState(obj)
 
 
-def run_gc(component, dry_run=False):
-    child_names = {child._meta.name for child in component}
-    statedir_prefix = component._meta.statedir.prefix
-
-    def unexpected(item):
-        if item.name.startswith("_"):
-            return False
-
-        if not item.is_dir():
-            return False
-
-        return item.name not in child_names
-
-    for item in statedir_prefix.iterdir():
-        if unexpected(item):
-            print(item)
-            if dry_run:
-                continue
-
-            shutil.rmtree(item)
+class StatefulMixin:
+    _meta: "opslib.components.Meta"
 
-    for child in component:
-        run_gc(child, dry_run=dry_run)
+    @contextmanager
+    def state_directory(self):
+        provider = self._meta.stack._state_provider
+        with provider.state_directory(cast(opslib.Component, self)) as statedir:
+            yield statedir