From f6890b93e868db9dfe94499550ec5e91e9a0de0f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=B6rg=20Thalheim?= <joerg@thalheim.io>
Date: Fri, 9 Sep 2022 11:24:53 +0000
Subject: [PATCH] add check-cache-status option to query wether a binary cache
 has the build

---
 src/nix-eval-jobs.cc | 45 +++++++++++++++++++++++++++++++++++++-------
 tests/test_eval.py   | 20 ++++++++++++++++----
 2 files changed, 54 insertions(+), 11 deletions(-)

diff --git a/src/nix-eval-jobs.cc b/src/nix-eval-jobs.cc
index 5a335a0..8544b15 100644
--- a/src/nix-eval-jobs.cc
+++ b/src/nix-eval-jobs.cc
@@ -44,6 +44,7 @@ struct MyArgs : MixEvalArgs, MixCommonArgs {
     bool meta = false;
     bool showTrace = false;
     bool impure = false;
+    bool checkCacheStatus = false;
     size_t nrWorkers = 1;
     size_t maxMemorySize = 4096;
 
@@ -93,6 +94,15 @@ struct MyArgs : MixEvalArgs, MixCommonArgs {
                  .description = "include derivation meta field in output",
                  .handler = {&meta, true}});
 
+        addFlag(
+            {.longName = "check-cache-status",
+             .description =
+                 "Check if the derivations are present locally or in "
+                 "any configured substituters (i.e. binary cache). The "
+                 "information "
+                 "will be exposed in the `isCached` field of the JSON output.",
+             .handler = {&checkCacheStatus, true}});
+
         addFlag({.longName = "show-trace",
                  .description =
                      "print out a stack trace in case of evaluation errors",
@@ -171,11 +181,26 @@ Value *topLevelValue(EvalState &state, Bindings &autoArgs) {
                         : releaseExprTopLevelValue(state, autoArgs);
 }
 
+bool queryIsCached(Store &store, std::map<std::string, std::string> &outputs) {
+    uint64_t downloadSize, narSize;
+    StorePathSet willBuild, willSubstitute, unknown;
+
+    std::vector<StorePathWithOutputs> paths;
+    for (auto const &[key, val] : outputs) {
+        paths.push_back(followLinksToStorePathWithOutputs(store, val));
+    }
+
+    store.queryMissing(toDerivedPaths(paths), willBuild, willSubstitute,
+                       unknown, downloadSize, narSize);
+    return willBuild.empty() && unknown.empty();
+}
+
 /* The fields of a derivation that are printed in json form */
 struct Drv {
     std::string name;
     std::string system;
     std::string drvPath;
+    bool isCached;
     std::map<std::string, std::string> outputs;
     std::optional<nlohmann::json> meta;
 
@@ -209,6 +234,9 @@ struct Drv {
             }
             meta = meta_;
         }
+        if (myArgs.checkCacheStatus) {
+            isCached = queryIsCached(*localStore, outputs);
+        }
 
         name = drvInfo.queryName();
         system = drvInfo.querySystem();
@@ -217,15 +245,18 @@ struct Drv {
 };
 
 static void to_json(nlohmann::json &json, const Drv &drv) {
-    json = nlohmann::json{
-        {"name", drv.name},
-        {"system", drv.system},
-        {"drvPath", drv.drvPath},
-        {"outputs", drv.outputs},
-    };
+    json = nlohmann::json{{"name", drv.name},
+                          {"system", drv.system},
+                          {"drvPath", drv.drvPath},
+                          {"outputs", drv.outputs}};
 
-    if (drv.meta.has_value())
+    if (drv.meta.has_value()) {
         json["meta"] = drv.meta.value();
+    }
+
+    if (myArgs.checkCacheStatus) {
+        json["isCached"] = drv.isCached;
+    }
 }
 
 std::string attrPathJoin(json input) {
diff --git a/tests/test_eval.py b/tests/test_eval.py
index 8ab26b7..5cdad3b 100644
--- a/tests/test_eval.py
+++ b/tests/test_eval.py
@@ -4,14 +4,14 @@ import subprocess
 import json
 from tempfile import TemporaryDirectory
 from pathlib import Path
-from typing import List
+from typing import List, Dict, Any
 
 TEST_ROOT = Path(__file__).parent.resolve()
 PROJECT_ROOT = TEST_ROOT.parent
 BIN = PROJECT_ROOT.joinpath("build", "src", "nix-eval-jobs")
 
 
-def common_test(extra_args: List[str]) -> None:
+def common_test(extra_args: List[str]) -> List[Dict[str, Any]]:
     with TemporaryDirectory() as tempdir:
         cmd = [str(BIN), "--gc-roots-dir", tempdir, "--meta"] + extra_args
         res = subprocess.run(
@@ -47,14 +47,26 @@ def common_test(extra_args: List[str]) -> None:
         assert substituted_job["attr"] == "substitutedJob"
         assert substituted_job["name"].startswith("hello-")
         assert substituted_job["meta"]["broken"] is False
+        return results
 
 
 def test_flake() -> None:
-    common_test(["--flake", ".#hydraJobs"])
+    results = common_test(["--flake", ".#hydraJobs"])
+    for result in results:
+        assert "isCached" not in result
+
+
+def test_query_cache_status() -> None:
+    results = common_test(["--flake", ".#hydraJobs", "--check-cache-status"])
+    # FIXME in the nix sandbox we cannot query binary caches, this would need some local one
+    for result in results:
+        assert "isCached" in result
 
 
 def test_expression() -> None:
-    common_test(["ci.nix"])
+    results = common_test(["ci.nix"])
+    for result in results:
+        assert "isCached" not in result
 
     with open(TEST_ROOT.joinpath("assets/ci.nix"), "r") as ci_nix:
         common_test(["-E", ci_nix.read()])