Add ability to create environment from lockfile

Fixes conda-incubator#676.

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

@@ -1,6 +1,7 @@
 from conda_store_server.action.base import action  # noqa
 
 from conda_store_server.action.generate_lockfile import action_solve_lockfile  # noqa
+from conda_store_server.action.generate_lockfile import action_save_lockfile  # noqa
 from conda_store_server.action.download_packages import (
     action_fetch_and_extract_conda_packages,  # noqa
 )

conda-store-server/conda_store_server/action/generate_constructor_installer.py

@@ -3,6 +3,7 @@
 import sys
 import tempfile
 import warnings
+from typing import Union
 
 import yaml
 from conda_store_server import action, schema
@@ -20,13 +21,20 @@ def get_installer_platform():
     return context.subdir
 
 
+# TODO: Ideally, this would just use the lockfile instead of reading data
+# from the specification. But we currently also need to support legacy lockfiles
+# generated by conda-store, which were not generated by conda-lock, and those
+# don't have enough information for constructor to generate an installer. In
+# particular, the channels are missing in that format. For more information on
+# lockfile formats, see: api_get_build_lockfile and test_api_get_build_lockfile
 @action.action
 def action_generate_constructor_installer(
     context,
     conda_command: str,
-    specification: schema.CondaSpecification,
+    specification: Union[schema.CondaSpecification, schema.LockfileSpecification],
     installer_dir: pathlib.Path,
     version: str,
+    is_lockfile: bool = False,
 ):
     def write_file(filename, s):
         with open(filename, "w") as f:
@@ -52,11 +60,28 @@ def write_file(filename, s):
     # conda and pip need to be in dependencies for the post_install script
     dependencies = ["conda", "pip"]
     pip_dependencies = []
-    for d in specification.dependencies:
-        if type(d) is schema.CondaSpecificationPip:
-            pip_dependencies.extend(d.pip)
-        else:
-            dependencies.append(d)
+    channels = []
+
+    if is_lockfile:
+        # Adds channels
+        channels = [c.url for c in specification.lockfile.metadata.channels]
+
+        # Adds dependencies
+        for p in specification.lockfile.package:
+            if p.manager == "pip":
+                pip_dependencies.append(f"{p.name}=={p.version}")
+            else:
+                dependencies.append(f"{p.name}=={p.version}")
+    else:
+        # Adds channels
+        channels = specification.channels
+
+        # Adds dependencies
+        for d in specification.dependencies:
+            if type(d) is schema.CondaSpecificationPip:
+                pip_dependencies.extend(d.pip)
+            else:
+                dependencies.append(d)
 
     # Creates the construct.yaml file and post_install script
     ext = ".exe" if sys.platform == "win32" else ".sh"
@@ -78,7 +103,7 @@ def write_file(filename, s):
             "installer_filename": str(installer_filename),
             "post_install": str(post_install_file),
             "name": specification.name,
-            "channels": specification.channels,
+            "channels": channels,
             "specs": dependencies,
             "version": version,
         }

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

@@ -5,7 +5,7 @@
 
 import yaml
 from conda_lock.conda_lock import run_lock
-from conda_store_server import action, conda_utils, schema
+from conda_store_server import action, conda_utils, schema, utils
 from conda_store_server.action.utils import logged_command
 
 
@@ -61,3 +61,17 @@ def action_solve_lockfile(
 
     with lockfile_filename.open() as f:
         return yaml.safe_load(f)
+
+
+@action.action
+def action_save_lockfile(
+    context,
+    specification: schema.LockfileSpecification,
+):
+    lockfile = specification.dict()["lockfile"]
+    lockfile_filename = pathlib.Path.cwd() / "conda-lock.yaml"
+
+    with lockfile_filename.open("w") as f:
+        json.dump(lockfile, f, cls=utils.CustomJSONEncoder)
+
+    return lockfile

conda-store-server/conda_store_server/alembic/versions/cd5b48ff57b5_lockfile_spec.py

@@ -0,0 +1,26 @@
+"""lockfile spec
+
+Revision ID: cd5b48ff57b5
+Revises: 03c839888c82
+Create Date: 2024-03-02 09:21:02.519805
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "cd5b48ff57b5"
+down_revision = "03c839888c82"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.add_column(
+        "specification",
+        sa.Column("is_lockfile", sa.Boolean(), nullable=False, server_default="False"),
+    )
+
+
+def downgrade():
+    op.drop_column("specification", "is_lockfile")

conda-store-server/conda_store_server/api.py

@@ -1,5 +1,5 @@
 import re
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Union
 
 from conda_store_server import conda_utils, orm, schema, utils
 from sqlalchemy import distinct, func, null, or_
@@ -378,19 +378,29 @@ def get_environment(
     return db.query(orm.Environment).join(orm.Namespace).filter(*filters).first()
 
 
-def ensure_specification(db, specification: schema.CondaSpecification):
+def ensure_specification(
+    db,
+    specification: Union[schema.CondaSpecification, schema.LockfileSpecification],
+    is_lockfile: bool = False,
+):
     specification_sha256 = utils.datastructure_hash(specification.dict())
     specification_orm = get_specification(db, sha256=specification_sha256)
 
     if specification_orm is None:
-        specification_orm = create_speficication(db, specification)
+        specification_orm = create_speficication(
+            db, specification, is_lockfile=is_lockfile
+        )
         db.commit()
 
     return specification_orm
 
 
-def create_speficication(db, specification: schema.CondaSpecification):
-    specification_orm = orm.Specification(specification.dict())
+def create_speficication(
+    db,
+    specification: Union[schema.CondaSpecification, schema.LockfileSpecification],
+    is_lockfile: bool = False,
+):
+    specification_orm = orm.Specification(specification.dict(), is_lockfile=is_lockfile)
     db.add(specification_orm)
     return specification_orm
 

conda-store-server/conda_store_server/app.py

@@ -618,8 +618,10 @@ def register_environment(
         specification: dict,
         namespace: str = None,
         force: bool = True,
+        is_lockfile: bool = False,
     ):
         """Register a given specification to conda store with given namespace/name."""
+
         settings = self.get_settings(db)
 
         namespace = namespace or settings.default_namespace
@@ -632,12 +634,18 @@ def register_environment(
             action=schema.Permissions.ENVIRONMENT_CREATE,
         )
 
-        specification_model = self.validate_specification(
-            db=db,
-            conda_store=self,
-            namespace=namespace.name,
-            specification=schema.CondaSpecification.parse_obj(specification),
-        )
+        if is_lockfile:
+            # It's a lockfile, do not do any validation in this case. If there
+            # are problems, these would be caught earlier during parsing or
+            # later when conda-lock attempts to install it.
+            specification_model = specification
+        else:
+            specification_model = self.validate_specification(
+                db=db,
+                conda_store=self,
+                namespace=namespace.name,
+                specification=schema.CondaSpecification.parse_obj(specification),
+            )
 
         spec_sha256 = utils.datastructure_hash(specification_model.dict())
         matching_specification = api.get_specification(db, sha256=spec_sha256)
@@ -651,7 +659,9 @@ def register_environment(
         ):
             return None
 
-        specification = api.ensure_specification(db, specification_model)
+        specification = api.ensure_specification(
+            db, specification_model, is_lockfile=is_lockfile
+        )
         environment_was_empty = (
             api.get_environment(db, name=specification.name, namespace_id=namespace.id)
             is None
@@ -663,15 +673,23 @@ def register_environment(
             description=specification.spec["description"],
         )
 
-        build = self.create_build(db, environment.id, specification.sha256)
+        build = self.create_build(
+            db, environment.id, specification.sha256, is_lockfile=is_lockfile
+        )
 
         if environment_was_empty:
             environment.current_build = build
             db.commit()
 
         return build.id
 
-    def create_build(self, db: Session, environment_id: int, specification_sha256: str):
+    def create_build(
+        self,
+        db: Session,
+        environment_id: int,
+        specification_sha256: str,
+        is_lockfile: bool = False,
+    ):
         environment = api.get_environment(db, id=environment_id)
         self.validate_action(
             db=db,

conda-store-server/conda_store_server/build.py

@@ -169,29 +169,45 @@ def build_conda_environment(db: Session, conda_store, build):
         environment_prefix = build.environment_path(conda_store)
         environment_prefix.parent.mkdir(parents=True, exist_ok=True)
 
+        is_lockfile = build.specification.is_lockfile
+
         with utils.timer(conda_store.log, f"building conda_prefix={conda_prefix}"):
-            context = action.action_solve_lockfile(
-                settings.conda_command,
-                specification=schema.CondaSpecification.parse_obj(
-                    build.specification.spec
-                ),
-                platforms=settings.conda_solve_platforms,
-            )
+            if is_lockfile:
+                context = action.action_save_lockfile(
+                    specification=schema.LockfileSpecification.parse_obj(
+                        build.specification.spec
+                    ),
+                )
+            else:
+                context = action.action_solve_lockfile(
+                    settings.conda_command,
+                    specification=schema.CondaSpecification.parse_obj(
+                        build.specification.spec
+                    ),
+                    platforms=settings.conda_solve_platforms,
+                )
 
             conda_store.storage.set(
                 db,
                 build.id,
                 build.conda_lock_key,
-                json.dumps(context.result, indent=4).encode("utf-8"),
+                json.dumps(
+                    context.result, indent=4, cls=utils.CustomJSONEncoder
+                ).encode("utf-8"),
                 content_type="application/json",
                 artifact_type=schema.BuildArtifactType.LOCKFILE,
             )
 
+            if is_lockfile:
+                lockfile_action_name = "action_save_lockfile"
+            else:
+                lockfile_action_name = "action_solve_lockfile"
+
             append_to_logs(
                 db,
                 conda_store,
                 build,
-                "::group::action_solve_lockfile\n"
+                f"::group::{lockfile_action_name}\n"
                 + context.stdout.getvalue()
                 + "\n::endgroup::\n",
             )
@@ -414,13 +430,23 @@ def build_constructor_installer(db: Session, conda_store, build: orm.Build):
         conda_store.log, f"creating installer for conda environment={conda_prefix}"
     ):
         with tempfile.TemporaryDirectory() as tmpdir:
+            is_lockfile = build.specification.is_lockfile
+
+            if is_lockfile:
+                specification = schema.LockfileSpecification.parse_obj(
+                    build.specification.spec
+                )
+            else:
+                specification = schema.CondaSpecification.parse_obj(
+                    build.specification.spec
+                )
+
             context = action.action_generate_constructor_installer(
                 conda_command=settings.conda_command,
-                specification=schema.CondaSpecification.parse_obj(
-                    build.specification.spec
-                ),
+                specification=specification,
                 installer_dir=pathlib.Path(tmpdir),
                 version=build.build_key,
+                is_lockfile=is_lockfile,
             )
             output_filename = context.result
             append_to_logs(

conda-store-server/conda_store_server/dbutil.py

@@ -1,9 +1,11 @@
+import json
 import os
 from contextlib import contextmanager
 from tempfile import TemporaryDirectory
 
 from alembic import command
 from alembic.config import Config
+from conda_store_server import utils
 from sqlalchemy import create_engine, inspect
 
 _here = os.path.abspath(os.path.dirname(__file__))
@@ -71,7 +73,11 @@ def upgrade(db_url, revision="head"):
     revision: str [default: head]
         The alembic revision to upgrade to.
     """
-    engine = create_engine(db_url)
+    engine = create_engine(
+        db_url,
+        # See the comment on the CustomJSONEncoder class on why this is needed
+        json_serializer=lambda x: json.dumps(x, cls=utils.CustomJSONEncoder),
+    )
 
     # retrieves the names of tables in the DB
     current_table_names = set(inspect(engine).get_table_names())