From 444ba45f30dad718e51b7043418cff5ef83583ea Mon Sep 17 00:00:00 2001
From: Joonas Rautiola <joonas.rautiola@unikie.com>
Date: Wed, 28 Aug 2024 14:28:21 +0300
Subject: [PATCH] Optimize provenance generation speed

Signed-off-by: Joonas Rautiola <joonas.rautiola@unikie.com>
---
 src/provenance/main.py | 38 ++++++++++++++++++--------------------
 1 file changed, 18 insertions(+), 20 deletions(-)

diff --git a/src/provenance/main.py b/src/provenance/main.py
index d8b78b7..4985a73 100755
--- a/src/provenance/main.py
+++ b/src/provenance/main.py
@@ -83,36 +83,34 @@ def get_subjects(outputs: dict) -> list[dict]:
 def get_dependencies(drv_path: str, recursive: bool = False) -> list[dict]:
     """Get dependencies of derivation and parse them into ResourceDescriptors"""
 
-    LOG.info("Querying derivation dependencies%s", "recursively" if recursive else "")
+    LOG.info("Querying derivation dependencies %s", "recursively" if recursive else "")
 
     depth = "--requisites" if recursive else "--references"
-    deps_drv = exec_cmd(["nix-store", "--query", depth, drv_path]).stdout.split()
+
+    references = exec_cmd(
+        ["nix-store", "--query", depth, "--include-outputs", drv_path]
+    ).stdout.split()
+    hashes = exec_cmd(["nix-store", "--query", "--hash"] + references).stdout.split()
+    infos = json.loads(exec_cmd(["nix", "derivation", "show", "-r", drv_path]).stdout)
 
     dependencies = []
-    for drv in deps_drv:
-        store_hash = exec_cmd(["nix-store", "--query", "--hash", drv]).stdout.strip()
-        hash_type, hash_value = store_hash.split(":")
+    for drv, output_hash in zip(references, hashes):
+        LOG.debug("Creating dependency entry for %s", drv)
+        hash_type, hash_value = output_hash.split(":")
 
-        dependency_json = {
+        package = {
+            "name": drv.split("-", 1)[-1].removesuffix(".drv"),
             "uri": drv,
             "digest": {hash_type: hash_value},
         }
 
-        annotations = {}
-
-        if drv.endswith(".drv"):
-            dep_json = json.loads(exec_cmd(["nix", "derivation", "show", drv]).stdout)
-            env = dep_json[drv]["env"]
-            dependency_json["name"] = dep_json[drv]["name"]
-
-            version = env.get("version")
-            if version:
-                annotations["version"] = version
-
-        if annotations:
-            dependency_json["annotations"] = annotations
+        info = infos.get(drv)
+        if info:
+            package["name"] = info["name"]
+            if version := info["env"].get("version"):
+                package["annotations"] = {"version": version}
 
-        dependencies.append(dependency_json)
+        dependencies.append(package)
 
     return dependencies