1 files changed, 71 insertions, 16 deletions

diff --git a/src/nix/registry.cc b/src/nix/registry.cc
index f9719600f..6a92576c7 100644
--- a/src/nix/registry.cc
+++ b/src/nix/registry.cc
@@ -10,6 +10,46 @@
 using namespace nix;
 using namespace nix::flake;
 
+
+class RegistryCommand : virtual Args
+{
+    std::string registry_path;
+
+    std::shared_ptr<fetchers::Registry> registry;
+
+public:
+
+    RegistryCommand()
+    {
+        addFlag({
+            .longName = "registry",
+            .description = "The registry to operate on.",
+            .labels = {"registry"},
+            .handler = {&registry_path},
+        });
+    }
+
+    std::shared_ptr<fetchers::Registry> getRegistry()
+    {
+        if (registry) return registry;
+        if (registry_path.empty()) {
+            registry = fetchers::getUserRegistry();
+        } else {
+            registry = fetchers::getCustomRegistry(registry_path);
+        }
+        return registry;
+    }
+
+    Path getRegistryPath()
+    {
+        if (registry_path.empty()) {
+            return fetchers::getUserRegistryPath();
+        } else {
+            return registry_path;
+        }
+    }
+};
+
 struct CmdRegistryList : StoreCommand
 {
     std::string description() override
@@ -45,7 +85,7 @@ struct CmdRegistryList : StoreCommand
     }
 };
 
-struct CmdRegistryAdd : MixEvalArgs, Command
+struct CmdRegistryAdd : MixEvalArgs, Command, RegistryCommand
 {
     std::string fromUrl, toUrl;
 
@@ -71,16 +111,16 @@ struct CmdRegistryAdd : MixEvalArgs, Command
     {
         auto fromRef = parseFlakeRef(fromUrl);
         auto toRef = parseFlakeRef(toUrl);
+        auto registry = getRegistry();
         fetchers::Attrs extraAttrs;
         if (toRef.subdir != "") extraAttrs["dir"] = toRef.subdir;
-        auto userRegistry = fetchers::getUserRegistry();
-        userRegistry->remove(fromRef.input);
-        userRegistry->add(fromRef.input, toRef.input, extraAttrs);
-        userRegistry->write(fetchers::getUserRegistryPath());
+        registry->remove(fromRef.input);
+        registry->add(fromRef.input, toRef.input, extraAttrs);
+        registry->write(getRegistryPath());
     }
 };
 
-struct CmdRegistryRemove : virtual Args, MixEvalArgs, Command
+struct CmdRegistryRemove : RegistryCommand, Command
 {
     std::string url;
 
@@ -103,19 +143,21 @@ struct CmdRegistryRemove : virtual Args, MixEvalArgs, Command
 
     void run() override
     {
-        auto userRegistry = fetchers::getUserRegistry();
-        userRegistry->remove(parseFlakeRef(url).input);
-        userRegistry->write(fetchers::getUserRegistryPath());
+        auto registry = getRegistry();
+        registry->remove(parseFlakeRef(url).input);
+        registry->write(getRegistryPath());
     }
 };
 
-struct CmdRegistryPin : virtual Args, EvalCommand
+struct CmdRegistryPin : RegistryCommand, EvalCommand
 {
     std::string url;
 
+    std::string locked;
+
     std::string description() override
     {
-        return "pin a flake to its current version in user flake registry";
+        return "pin a flake to its current version or to the current version of a flake URL";
     }
 
     std::string doc() override
@@ -128,18 +170,31 @@ struct CmdRegistryPin : virtual Args, EvalCommand
     CmdRegistryPin()
     {
         expectArg("url", &url);
+
+        expectArgs({
+            .label = "locked",
+            .optional = true,
+            .handler = {&locked},
+            .completer = {[&](size_t, std::string_view prefix) {
+                completeFlakeRef(getStore(), prefix);
+            }}
+        });
     }
 
     void run(nix::ref<nix::Store> store) override
     {
+        if (locked.empty()) {
+            locked = url;
+        }
+        auto registry = getRegistry();
         auto ref = parseFlakeRef(url);
-        auto userRegistry = fetchers::getUserRegistry();
-        userRegistry->remove(ref.input);
-        auto [tree, resolved] = ref.resolve(store).input.fetch(store);
+        auto locked_ref = parseFlakeRef(locked);
+        registry->remove(ref.input);
+        auto [tree, resolved] = locked_ref.resolve(store).input.fetch(store);
         fetchers::Attrs extraAttrs;
         if (ref.subdir != "") extraAttrs["dir"] = ref.subdir;
-        userRegistry->add(ref.input, resolved, extraAttrs);
-        userRegistry->write(fetchers::getUserRegistryPath());
+        registry->add(ref.input, resolved, extraAttrs);
+        registry->write(getRegistryPath());
     }
 };