Add ability to create environment from lockfile (#772)

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

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

@@ -4,6 +4,8 @@
 import tempfile
 import warnings
 
+from typing import Union
+
 import yaml
 
 from conda_store_server import action, schema
@@ -25,9 +27,10 @@ def get_installer_platform():
 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:
@@ -53,11 +56,37 @@ 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)
+
+    if is_lockfile:
+        # Adds channels
+        channels = [c.url for c in specification.lockfile.metadata.channels]
+
+        # Adds dependencies
+        for p in specification.lockfile.package:
+            # Ignores packages not matching the current platform. Versions can
+            # be different between platforms or a package might not support all
+            # platforms. constructor is cross-friendly, but we're currently
+            # building only for the current architecture, see the comment in
+            # get_installer_platform
+            if p.platform not in ["noarch", get_installer_platform()]:
+                continue
+            if p.manager == "pip":
+                pip_dependencies.append(f"{p.name}=={p.version}")
+            else:
+                ext = ".tar.bz2" if p.url.endswith(".tar.bz2") else ".conda"
+                build_string = p.url[: -len(ext)].rsplit("-", maxsplit=1)[-1]
+                dependencies.append(f"{p.name}=={p.version}={build_string}")
+
+    else:
+        # Adds channels
+        channels = specification.channels
+
+        # Adds dependencies
+        for d in specification.dependencies:
+            if isinstance(d, 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"
@@ -79,7 +108,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

@@ -7,7 +7,7 @@
 
 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
 
 
@@ -63,3 +63,19 @@ 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,
+):
+    # Note: this calls dict on specification so that the version field is
+    # part of the output
+    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/bf065abf375b_lockfile_spec.py

@@ -0,0 +1,32 @@
+"""lockfile spec
+
+Revision ID: bf065abf375b
+Revises: e17b4cc6e086
+Create Date: 2024-03-02 09:21:02.519805
+
+"""
+
+import sqlalchemy as sa
+
+from alembic import op
+
+
+# revision identifiers, used by Alembic.
+revision = "bf065abf375b"
+down_revision = "e17b4cc6e086"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.add_column(
+        "specification",
+        # https://github.com/sqlalchemy/sqlalchemy/issues/1403#issuecomment-1698365595
+        sa.Column(
+            "is_lockfile", sa.Boolean(), nullable=False, server_default=sa.sql.false()
+        ),
+    )
+
+
+def downgrade():
+    op.drop_column("specification", "is_lockfile")

conda-store-server/conda_store_server/api.py

@@ -1,6 +1,6 @@
 import re
 
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Union
 
 from sqlalchemy import distinct, func, null, or_
 from sqlalchemy.orm import aliased
@@ -380,19 +380,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

@@ -627,8 +627,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
@@ -641,12 +643,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)
@@ -660,7 +668,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

conda-store-server/conda_store_server/build.py

@@ -199,27 +199,44 @@ def build_conda_environment(db: Session, conda_store, build):
         if environment_prefix is not None:
             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,
-                conda_flags=conda_store.conda_flags,
-                stdout=LoggedStream(
-                    db=db,
-                    conda_store=conda_store,
-                    build=build,
-                    prefix="action_solve_lockfile: ",
-                ),
-            )
+            if is_lockfile:
+                context = action.action_save_lockfile(
+                    specification=schema.LockfileSpecification.parse_obj(
+                        build.specification.spec
+                    ),
+                    stdout=LoggedStream(
+                        db=db,
+                        conda_store=conda_store,
+                        build=build,
+                        prefix="action_save_lockfile: ",
+                    ),
+                )
+            else:
+                context = action.action_solve_lockfile(
+                    settings.conda_command,
+                    specification=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(
                 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,
             )
@@ -459,11 +476,39 @@ 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:
+                try:
+                    # Tries to use the lockfile if it's available since it has
+                    # pinned dependencies. This code is wrapped into try/except
+                    # because the lockfile lookup might fail if the file is not
+                    # in external storage or on disk, or if parsing fails
+                    specification = schema.LockfileSpecification.parse_obj(
+                        {
+                            "name": build.specification.name,
+                            "lockfile": json.loads(
+                                conda_store.storage.get(build.conda_lock_key)
+                            ),
+                        }
+                    )
+                    is_lockfile = True
+                except Exception as e:
+                    conda_store.log.warning(
+                        "Exception while obtaining lockfile, using specification",
+                        exc_info=e,
+                    )
+                    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,
                 stdout=LoggedStream(
@@ -472,6 +517,7 @@ def build_constructor_installer(db: Session, conda_store, build: orm.Build):
                     build=build,
                     prefix="action_generate_constructor_installer: ",
                 ),
+                is_lockfile=is_lockfile,
             )
             output_filename = context.result
             if output_filename is None:

conda-store-server/conda_store_server/dbutil.py

@@ -1,12 +1,16 @@
+import json
 import os
 
 from contextlib import contextmanager
+from functools import partial
 from tempfile import TemporaryDirectory
 
 from alembic import command
 from alembic.config import Config
 from sqlalchemy import create_engine, inspect
 
+from conda_store_server import utils
+
 
 _here = os.path.abspath(os.path.dirname(__file__))
 
@@ -73,7 +77,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=partial(json.dumps, cls=utils.CustomJSONEncoder),
+    )
 
     # retrieves the names of tables in the DB
     current_table_names = set(inspect(engine).get_table_names())