From 9b731f61bfa19680247cee0076b395652e367ecf Mon Sep 17 00:00:00 2001
From: Nikita Karetnikov <nkaretnikov@quansight.com>
Date: Fri, 17 Nov 2023 09:05:04 +0000
Subject: [PATCH] Allow to set `build_key_version` via the config

---
 .../conda_store_server/__init__.py            |  4 ++++
 conda-store-server/conda_store_server/app.py  | 21 +++++++++++++++++++
 conda-store-server/conda_store_server/orm.py  | 14 ++++++++++++-
 3 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/conda-store-server/conda_store_server/__init__.py b/conda-store-server/conda_store_server/__init__.py
index d3b075936..eca3025f1 100644
--- a/conda-store-server/conda_store_server/__init__.py
+++ b/conda-store-server/conda_store_server/__init__.py
@@ -4,3 +4,7 @@
 
 
 CONDA_STORE_DIR = Path.home() / ".conda-store"
+
+
+# Default build_key_version. Must be None here, initialized from the config file
+_BUILD_KEY_VERSION = None
diff --git a/conda-store-server/conda_store_server/app.py b/conda-store-server/conda_store_server/app.py
index a3aa62564..623953297 100644
--- a/conda-store-server/conda_store_server/app.py
+++ b/conda-store-server/conda_store_server/app.py
@@ -106,6 +106,22 @@ class CondaStore(LoggingConfigurable):
         config=True,
     )
 
+    build_key_version = Integer(
+        2,
+        help="Build key version to use: 1 (long, legacy), 2 (short, default)",
+        config=True,
+    )
+
+    @validate("build_key_version")
+    def _check_build_key_version(self, proposal):
+        expected = [1, 2]
+        if proposal.value not in expected:
+            raise TraitError(
+                f"c.CondaStore.build_key_version: invalid build key version: "
+                f"{proposal.value}, expected: {expected}"
+            )
+        return proposal.value
+
     conda_command = Unicode(
         "mamba",
         help="conda executable to use for solves",
@@ -364,6 +380,11 @@ def session_factory(self):
             url=self.database_url,
             poolclass=QueuePool,
         )
+
+        # Sets the default build_key_version value in the DB based on the config
+        import conda_store_server
+
+        conda_store_server._BUILD_KEY_VERSION = self.build_key_version
         return self._session_factory
 
     # Do not define this as a FastAPI dependency! That would cause Sessions
diff --git a/conda-store-server/conda_store_server/orm.py b/conda-store-server/conda_store_server/orm.py
index b1549be67..681ea58eb 100644
--- a/conda-store-server/conda_store_server/orm.py
+++ b/conda-store-server/conda_store_server/orm.py
@@ -158,6 +158,18 @@ class Solve(Base):
 _BUILD_KEY_V2_HASH_SIZE = 8
 
 
+# Avoids a cyclic dependency between the orm module and the module defining
+# CondaStore.build_key_version. Because the orm module is loaded early on
+# startup, we want to delay initialization of the build_key_version field until
+# it's been read from the config.
+def _get_build_key_version():
+    from conda_store_server import _BUILD_KEY_VERSION
+
+    # None means the value is not set, likely due to an import error
+    assert _BUILD_KEY_VERSION is not None
+    return _BUILD_KEY_VERSION
+
+
 class Build(Base):
     """The state of a build of a given specification"""
 
@@ -185,7 +197,7 @@ class Build(Base):
     started_on = Column(DateTime, default=None)
     ended_on = Column(DateTime, default=None)
     deleted_on = Column(DateTime, default=None)
-    build_key_version = Column(Integer, default=2, nullable=False)
+    build_key_version = Column(Integer, default=_get_build_key_version, nullable=False)
 
     @validates("build_key_version")
     def validate_build_key_version(self, key, build_key_version):