Setup pluggy with locking plugin (#965)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: jaimergp <[email protected]>

conda-store-server/conda_store_server/_internal/action/__init__.py

@@ -26,7 +26,6 @@
 )
 from conda_store_server._internal.action.generate_lockfile import (  # noqa
     action_save_lockfile,
-    action_solve_lockfile,
 )
 from conda_store_server._internal.action.get_conda_prefix_stats import (  # noqa
     action_get_conda_prefix_stats,

conda-store-server/conda_store_server/_internal/action/generate_lockfile.py

@@ -3,76 +3,9 @@
 # license that can be found in the LICENSE file.
 
 import json
-import os
 import pathlib
-import typing
 
-import yaml
-from conda_lock.conda_lock import run_lock
-
-from conda_store_server._internal import action, conda_utils, schema, utils
-from conda_store_server._internal.action.utils import logged_command
-
-
-@action.action
-def action_solve_lockfile(
-    context,
-    conda_command: str,
-    specification: schema.CondaSpecification,
-    platforms: typing.List[str] = [conda_utils.conda_platform()],
-    # Avoids package compatibility issues, see:
-    # https://docs.conda.io/projects/conda/en/latest/user-guide/tasks/manage-channels.html
-    conda_flags: str = "--strict-channel-priority",
-):
-    environment_filename = pathlib.Path.cwd() / "environment.yaml"
-    lockfile_filename = pathlib.Path.cwd() / "conda-lock.yaml"
-
-    with environment_filename.open("w") as f:
-        json.dump(specification.model_dump(), f)
-
-    context.log.info(
-        "Note that the output of `conda config --show` displayed below only reflects "
-        "settings in the conda configuration file, which might be overridden by "
-        "variables required to be set by conda-store via the environment. Overridden "
-        f"settings: CONDA_FLAGS={conda_flags}"
-    )
-
-    # The info command can be used with either mamba or conda
-    logged_command(context, [conda_command, "info"])
-    # The config command is not supported by mamba
-    logged_command(context, ["conda", "config", "--show"])
-    logged_command(context, ["conda", "config", "--show-sources"])
-
-    # conda-lock ignores variables defined in the specification, so this code
-    # gets the value of CONDA_OVERRIDE_CUDA and passes it to conda-lock via
-    # the with_cuda parameter, see:
-    # https://github.com/conda-incubator/conda-store/issues/719
-    # https://docs.conda.io/projects/conda/en/latest/user-guide/tasks/manage-virtual.html#overriding-detected-packages
-    # TODO: Support all variables once upstream fixes are made to conda-lock,
-    # see the discussion in issue 719.
-    if specification.variables is not None:
-        cuda_version = specification.variables.get("CONDA_OVERRIDE_CUDA")
-    else:
-        cuda_version = None
-
-    # CONDA_FLAGS is used by conda-lock in conda_solver.solve_specs_for_arch
-    try:
-        conda_flags_name = "CONDA_FLAGS"
-        print(f"{conda_flags_name}={conda_flags}")
-        os.environ[conda_flags_name] = conda_flags
-
-        run_lock(
-            environment_files=[environment_filename],
-            platforms=platforms,
-            lockfile_path=lockfile_filename,
-            conda_exe=conda_command,
-            with_cuda=cuda_version,
-        )
-    finally:
-        os.environ.pop(conda_flags_name, None)
-
-    with lockfile_filename.open() as f:
-        return yaml.safe_load(f)
+from conda_store_server._internal import action, schema, utils
 
 
 @action.action

conda-store-server/conda_store_server/_internal/plugins/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) conda-store development team. All rights reserved.
+# Use of this source code is governed by a BSD-style
+# license that can be found in the LICENSE file.

conda-store-server/conda_store_server/_internal/plugins/lock/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) conda-store development team. All rights reserved.
+# Use of this source code is governed by a BSD-style
+# license that can be found in the LICENSE file.

conda-store-server/conda_store_server/_internal/plugins/lock/conda_lock/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) conda-store development team. All rights reserved.
+# Use of this source code is governed by a BSD-style
+# license that can be found in the LICENSE file.

conda-store-server/conda_store_server/_internal/plugins/lock/conda_lock/conda_lock.py

@@ -0,0 +1,97 @@
+# Copyright (c) conda-store development team. All rights reserved.
+# Use of this source code is governed by a BSD-style
+# license that can be found in the LICENSE file.
+
+import json
+import os
+import pathlib
+import typing
+
+import yaml
+from conda_lock.conda_lock import run_lock
+
+from conda_store_server._internal import conda_utils, schema, utils
+from conda_store_server.plugins.hookspec import hookimpl
+from conda_store_server.plugins.plugin_context import PluginContext
+from conda_store_server.plugins.types import lock, types
+
+
+class CondaLock(lock.LockPlugin):
+    def _conda_command(self, conda_store) -> str:
+        with conda_store.session_factory() as db:
+            settings = conda_store.get_settings(db=db)
+        return settings.conda_command
+
+    def _conda_flags(self, conda_store) -> str:
+        return conda_store.conda_flags
+
+    @utils.run_in_tempdir
+    def lock_environment(
+        self,
+        context: PluginContext,
+        spec: schema.CondaSpecification,
+        platforms: typing.List[str] = [conda_utils.conda_platform()],
+    ) -> str:
+        context.log.info("lock_environment entrypoint for conda-lock")
+        conda_command = self._conda_command(context.conda_store)
+        conda_flags = self._conda_flags(context.conda_store)
+
+        environment_filename = pathlib.Path.cwd() / "environment.yaml"
+        lockfile_filename = pathlib.Path.cwd() / "conda-lock.yaml"
+
+        with environment_filename.open("w") as f:
+            json.dump(spec.dict(), f)
+
+        context.log.info(
+            "Note that the output of `conda config --show` displayed below only reflects "
+            "settings in the conda configuration file, which might be overridden by "
+            "variables required to be set by conda-store via the environment. Overridden "
+            f"settings: CONDA_FLAGS={conda_flags}"
+        )
+
+        # The info command can be used with either mamba or conda
+        context.run_command([conda_command, "info"])
+        # The config command is not supported by mamba
+        context.run_command(["conda", "config", "--show"])
+        context.run_command(["conda", "config", "--show-sources"])
+
+        # conda-lock ignores variables defined in the specification, so this code
+        # gets the value of CONDA_OVERRIDE_CUDA and passes it to conda-lock via
+        # the with_cuda parameter, see:
+        # https://github.com/conda-incubator/conda-store/issues/719
+        # https://docs.conda.io/projects/conda/en/latest/user-guide/tasks/manage-virtual.html#overriding-detected-packages
+        # TODO: Support all variables once upstream fixes are made to conda-lock,
+        # see the discussion in issue 719.
+        if spec.variables is not None:
+            cuda_version = spec.variables.get("CONDA_OVERRIDE_CUDA")
+        else:
+            cuda_version = None
+
+        # CONDA_FLAGS is used by conda-lock in conda_solver.solve_specs_for_arch
+        try:
+            conda_flags_name = "CONDA_FLAGS"
+            print(f"{conda_flags_name}={conda_flags}")
+            os.environ[conda_flags_name] = conda_flags
+
+            run_lock(
+                environment_files=[environment_filename],
+                platforms=platforms,
+                lockfile_path=lockfile_filename,
+                conda_exe=conda_command,
+                with_cuda=cuda_version,
+            )
+        finally:
+            os.environ.pop(conda_flags_name, None)
+
+        with lockfile_filename.open() as f:
+            return yaml.safe_load(f)
+
+
+@hookimpl
+def lock_plugins():
+    """conda-lock locking plugin"""
+    yield types.TypeLockPlugin(
+        name="conda-lock",
+        synopsis="Generate a lockfile using conda-lock",
+        backend=CondaLock,
+    )

conda-store-server/conda_store_server/_internal/utils.py

@@ -3,15 +3,17 @@
 # license that can be found in the LICENSE file.
 
 import contextlib
+import functools
 import hashlib
 import json
 import os
 import pathlib
 import re
 import subprocess
 import sys
+import tempfile
 import time
-from typing import AnyStr
+from typing import AnyStr, Callable
 
 from filelock import FileLock
 
@@ -187,3 +189,20 @@ def compile_arn_sql_like(
         re.sub(r"\*", "%", match.group(1)),
         re.sub(r"\*", "%", match.group(2)),
     )
+
+
+def run_in_tempdir(f: Callable):
+    @functools.wraps(f)
+    def wrapper(*args, **kwargs):
+        with contextlib.ExitStack() as stack:
+            # create a temporary directory
+            tmpdir = stack.enter_context(tempfile.TemporaryDirectory())
+
+            # enter temporary directory
+            stack.enter_context(chdir(tmpdir))
+
+            # run function and store result
+            result = f(*args, **kwargs)
+        return result
+
+    return wrapper

conda-store-server/conda_store_server/_internal/worker/build.py

@@ -18,6 +18,7 @@
 
 from conda_store_server import api
 from conda_store_server._internal import action, conda_utils, orm, schema, utils
+from conda_store_server.plugins import plugin_context
 
 
 class LoggedStream:
@@ -222,19 +223,20 @@ def build_conda_environment(db: Session, conda_store, build):
                         prefix="action_save_lockfile: ",
                     ),
                 )
+                conda_lock_spec = context.result
             else:
-                context = action.action_solve_lockfile(
-                    settings.conda_command,
-                    specification=schema.CondaSpecification.parse_obj(
-                        build.specification.spec
-                    ),
+                lock_backend, locker = conda_store.lock_plugin()
+                conda_lock_spec = locker.lock_environment(
+                    spec=schema.CondaSpecification.parse_obj(build.specification.spec),
                     platforms=settings.conda_solve_platforms,
-                    conda_flags=conda_store.conda_flags,
-                    stdout=LoggedStream(
-                        db=db,
+                    context=plugin_context.PluginContext(
                         conda_store=conda_store,
-                        build=build,
-                        prefix="action_solve_lockfile: ",
+                        stdout=LoggedStream(
+                            db=db,
+                            conda_store=conda_store,
+                            build=build,
+                            prefix=f"plugin-{lock_backend}: ",
+                        ),
                     ),
                 )
 
@@ -243,14 +245,12 @@ def build_conda_environment(db: Session, conda_store, build):
                 build.id,
                 build.conda_lock_key,
                 json.dumps(
-                    context.result, indent=4, cls=utils.CustomJSONEncoder
+                    conda_lock_spec, indent=4, cls=utils.CustomJSONEncoder
                 ).encode("utf-8"),
                 content_type="application/json",
                 artifact_type=schema.BuildArtifactType.LOCKFILE,
             )
 
-            conda_lock_spec = context.result
-
             context = action.action_fetch_and_extract_conda_packages(
                 conda_lock_spec=conda_lock_spec,
                 pkgs_dir=conda_utils.conda_root_package_dir(),
@@ -335,18 +335,15 @@ def build_conda_environment(db: Session, conda_store, build):
 
 
 def solve_conda_environment(db: Session, conda_store, solve: orm.Solve):
-    settings = conda_store.get_settings(db=db)
-
     solve.started_on = datetime.datetime.utcnow()
     db.commit()
 
-    context = action.action_solve_lockfile(
-        conda_command=settings.conda_command,
-        specification=schema.CondaSpecification.parse_obj(solve.specification.spec),
+    _, locker = conda_store.lock_plugin()
+    conda_lock_spec = locker.lock_environment(
+        context=plugin_context.PluginContext(conda_store=conda_store),
+        spec=schema.CondaSpecification.parse_obj(solve.specification.spec),
         platforms=[conda_utils.conda_platform()],
-        conda_flags=conda_store.conda_flags,
     )
-    conda_lock_spec = context.result
 
     action.action_add_lockfile_packages(
         db=db,

conda-store-server/conda_store_server/app.py

@@ -28,6 +28,8 @@
 
 from conda_store_server import CONDA_STORE_DIR, BuildKey, api, registry, storage
 from conda_store_server._internal import conda_utils, environment, orm, schema, utils
+from conda_store_server.plugins import hookspec, plugin_manager
+from conda_store_server.plugins.types import lock
 
 
 def conda_store_validate_specification(
@@ -192,6 +194,12 @@ def _check_build_key_version(self, proposal):
         config=True,
     )
 
+    lock_backend = Unicode(
+        default_value="conda-lock",
+        allow_none=False,
+        config=True,
+    )
+
     pypi_default_packages = List(
         [],
         help="PyPi packages that included by default if none are included",
@@ -471,6 +479,27 @@ def celery_app(self):
         self._celery_app.config_from_object(self.celery_config)
         return self._celery_app
 
+    @property
+    def plugin_manager(self):
+        """Creates a plugin manager (if it doesn't already exist) and registers all plugins"""
+        if hasattr(self, "_plugin_manager"):
+            return self._plugin_manager
+
+        self._plugin_manager = plugin_manager.PluginManager(hookspec.spec_name)
+        self._plugin_manager.add_hookspecs(hookspec.CondaStoreSpecs)
+
+        # Register all available plugins
+        self._plugin_manager.collect_plugins()
+
+        return self._plugin_manager
+
+    def lock_plugin(self) -> tuple[str, lock.LockPlugin]:
+        """Returns the configured lock plugin"""
+        # TODO: get configured lock plugin name from settings
+        lock_plugin = self.plugin_manager.get_lock_plugin(name=self.lock_backend)
+        locker = lock_plugin.backend()
+        return lock_plugin.name, locker
+
     def ensure_settings(self, db: Session):
         """Ensure that conda-store traitlets settings are applied"""
         settings = schema.Settings(

conda-store-server/conda_store_server/exception.py

@@ -0,0 +1,19 @@
+# Copyright (c) conda-store development team. All rights reserved.
+# Use of this source code is governed by a BSD-style
+# license that can be found in the LICENSE file.
+
+
+class CondaStoreError(Exception):
+    pass
+
+
+class CondaStorePluginNotFoundError(CondaStoreError):
+    """Exception raised by conda store when a specified plugin is not found
+    Attributes:
+        plugin -- plugin that was not found
+        available_plugins -- list of registered plugins
+    """
+
+    def __init__(self, plugin, available_plugins):
+        self.message = f"Plugin {plugin} was requested but not found! The following plugins are available: {', '.join(available_plugins)}"
+        super().__init__(self.message)

conda-store-server/conda_store_server/plugins/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) conda-store development team. All rights reserved.
+# Use of this source code is governed by a BSD-style
+# license that can be found in the LICENSE file.
+
+from conda_store_server._internal.plugins.lock.conda_lock import conda_lock
+
+BUILTIN_PLUGINS = [
+    conda_lock,
+]