From 093ed4763615a73edb1e8567d2f41d55ff5aaa0c Mon Sep 17 00:00:00 2001
From: Alexander Bantyev <balsoft@balsoft.ru>
Date: Wed, 30 Jun 2021 22:13:32 +0300
Subject: [PATCH] nix registry: add --registry flag

---
 src/nix/registry-add.md    |  7 ++++
 src/nix/registry-pin.md    |  7 ++++
 src/nix/registry-remove.md |  6 ++++
 src/nix/registry.cc        | 67 ++++++++++++++++++++++++++++++--------
 4 files changed, 73 insertions(+), 14 deletions(-)

diff --git a/src/nix/registry-add.md b/src/nix/registry-add.md
index 80a31996a..a947fa0b3 100644
--- a/src/nix/registry-add.md
+++ b/src/nix/registry-add.md
@@ -21,6 +21,13 @@ R""(
   # nix registry add nixpkgs/nixos-20.03 ~/Dev/nixpkgs
   ```
 
+* Add `nixpkgs` pointing to `github:nixos/nixpkgs` to your custom flake
+  registry:
+
+  ```console
+  nix registry add --registry ./custom-flake-registry.json nixpkgs github:nixos/nixpkgs
+  ```
+
 # Description
 
 This command adds an entry to the user registry that maps flake
diff --git a/src/nix/registry-pin.md b/src/nix/registry-pin.md
index 6e97e003e..7b163c463 100644
--- a/src/nix/registry-pin.md
+++ b/src/nix/registry-pin.md
@@ -24,6 +24,13 @@ R""(
   …
   ```
 
+* Pin `nixpkgs` in a custom registry to its most recent Git revision
+
+  ```console
+  # nix registry pin --registry ./custom-flake-registry.json nixpkgs
+  ```
+
+
 # Description
 
 This command adds an entry to the user registry that maps flake
diff --git a/src/nix/registry-remove.md b/src/nix/registry-remove.md
index 4c0eb4947..eecd4c6e7 100644
--- a/src/nix/registry-remove.md
+++ b/src/nix/registry-remove.md
@@ -8,6 +8,12 @@ R""(
   # nix registry remove nixpkgs
   ```
 
+* Remove the entry `nixpkgs` from a custom registry:
+
+  ```console
+  # nix registry remove --registry ./custom-flake-registry.json nixpkgs
+  ```
+
 # Description
 
 This command removes from the user registry any entry for flake
diff --git a/src/nix/registry.cc b/src/nix/registry.cc
index f9719600f..4400cc8dc 100644
--- a/src/nix/registry.cc
+++ b/src/nix/registry.cc
@@ -10,6 +10,45 @@
 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 +84,7 @@ struct CmdRegistryList : StoreCommand
     }
 };
 
-struct CmdRegistryAdd : MixEvalArgs, Command
+struct CmdRegistryAdd : MixEvalArgs, Command, RegistryCommand
 {
     std::string fromUrl, toUrl;
 
@@ -71,16 +110,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,13 +142,13 @@ 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;
 
@@ -132,14 +171,14 @@ struct CmdRegistryPin : virtual Args, EvalCommand
 
     void run(nix::ref<nix::Store> store) override
     {
+        auto registry = getRegistry();
         auto ref = parseFlakeRef(url);
-        auto userRegistry = fetchers::getUserRegistry();
-        userRegistry->remove(ref.input);
+        registry->remove(ref.input);
         auto [tree, resolved] = 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());
     }
 };