From 1c076f0f52d7b75a1b559b272a2d4e1592901942 Mon Sep 17 00:00:00 2001
From: sophia <scastellarin@quansight.com>
Date: Thu, 7 Nov 2024 15:00:59 -0800
Subject: [PATCH] Run lock from plugin

---
 .../_internal/worker/build.py                 | 37 +++++++++----------
 conda-store-server/conda_store_server/app.py  |  4 +-
 .../conda_store_server/plugins/__init__.py    |  2 +-
 .../{registry.py => plugin_registry.py}       |  0
 4 files changed, 20 insertions(+), 23 deletions(-)
 rename conda-store-server/conda_store_server/plugins/{registry.py => plugin_registry.py} (100%)

diff --git a/conda-store-server/conda_store_server/_internal/worker/build.py b/conda-store-server/conda_store_server/_internal/worker/build.py
index 50d4be82..aed4280c 100644
--- a/conda-store-server/conda_store_server/_internal/worker/build.py
+++ b/conda-store-server/conda_store_server/_internal/worker/build.py
@@ -17,6 +17,7 @@
 from sqlalchemy.orm import Session
 
 from conda_store_server import api
+from conda_store_server.plugins import plugin_context
 from conda_store_server._internal import action, conda_utils, orm, schema, utils
 
 
@@ -222,20 +223,22 @@ 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(
+                conda_lock_spec = conda_store.plugin_manager.hook.lock_environment(
+                    context=plugin_context.PluginContext(
+                        conda_store=conda_store,
+                        stdout=LoggedStream(
+                            db=db,
+                            conda_store=conda_store,
+                            build=build,
+                            prefix="hook-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,
-                        conda_store=conda_store,
-                        build=build,
-                        prefix="action_solve_lockfile: ",
-                    ),
                 )
 
             conda_store.storage.set(
@@ -243,14 +246,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 +336,14 @@ 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),
+    conda_lock_spec = conda_store.plugin_manager.hook.lock_environment(
+        context=plugin_context.PluginContext(),
+        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,
diff --git a/conda-store-server/conda_store_server/app.py b/conda-store-server/conda_store_server/app.py
index 15a4a4aa..43d94f5a 100644
--- a/conda-store-server/conda_store_server/app.py
+++ b/conda-store-server/conda_store_server/app.py
@@ -28,7 +28,7 @@
 from traitlets.config import LoggingConfigurable
 
 from conda_store_server import CONDA_STORE_DIR, BuildKey, api, registry, storage
-from conda_store_server.plugins import hookspec, registry
+from conda_store_server.plugins import hookspec, plugin_registry
 from conda_store_server.exception import CondaStorePluginNotFoundError
 from conda_store_server._internal import conda_utils, environment, orm, schema, utils
 
@@ -479,7 +479,7 @@ def plugin_registry(self):
         if hasattr(self, "_plugin_registry"):
             return self._plugin_registry
 
-        self._plugin_registry = registry.PluginRegistry()
+        self._plugin_registry = plugin_registry.PluginRegistry()
         self._plugin_registry.collect_plugins()
         return self._plugin_registry
 
diff --git a/conda-store-server/conda_store_server/plugins/__init__.py b/conda-store-server/conda_store_server/plugins/__init__.py
index 896f9127..9cb4f144 100644
--- a/conda-store-server/conda_store_server/plugins/__init__.py
+++ b/conda-store-server/conda_store_server/plugins/__init__.py
@@ -7,4 +7,4 @@
 
 BUILTIN_PLUGINS = [
     conda_lock.CondaLock
-]
\ No newline at end of file
+]
diff --git a/conda-store-server/conda_store_server/plugins/registry.py b/conda-store-server/conda_store_server/plugins/plugin_registry.py
similarity index 100%
rename from conda-store-server/conda_store_server/plugins/registry.py
rename to conda-store-server/conda_store_server/plugins/plugin_registry.py