Speed up ddb data collection

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Victor Fuentes
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

packages/libsnow-generators/default.nix

@@ -29,4 +29,9 @@ rustPlatform.buildRustPackage rec {
     openssl
     sqlite
   ];
+
+  patchPhase = ''
+    substituteInPlace src/ddb/mod.rs --replace "const REGISTRY: &str = \"./registry.nix\"" \
+      "const REGISTRY: &str = \"${../../registry.nix}\""
+  '';
 }

registry.nix

@@ -0,0 +1,88 @@
+# Adapted from https://github.com/replit/rippkgs/blob/main/lib/default.nix
+{lib, ...}: let
+in {
+  genRegistry = platform: pkgs: let
+    inherit (builtins) deepSeq filter listToAttrs map parseDrvName seq tryEval;
+    inherit (lib) filterAttrs flatten foldl isDerivation mapAttrsToList optional optionals removePrefix traceVal;
+
+    registerPackage = name: value: let
+      safeValue = tryEval value;
+      safeVal = safeValue.value;
+
+      safeRegistryValue = tryEval (deepSeq registryValue registryValue);
+      registryValue = {
+        version = safeVal.version or null;
+        storePaths = let
+          getOutput = out: {
+            name = out;
+            value = let
+              outPath = tryEval safeVal.${out}.outPath;
+            in if outPath.success then removePrefix "/nix/store/" outPath.value else "<broken>";
+          };
+
+          outputs-list = map getOutput (safeVal.outputs or []);
+          relevant-outputs = filter ({name, ...}: name == "out") outputs-list;
+        in
+          listToAttrs relevant-outputs;
+      };
+
+      platformForAvailability = {system = platform;};
+      isAvailableOn = tryEval (lib.meta.availableOn platformForAvailability safeValue.value);
+      available = safeValue.success && isDerivation value && isAvailableOn.success && isAvailableOn.value;
+
+      checkRegistryCondition = prev: {
+        reason,
+        ok,
+      }: let
+        isOk =
+          if !ok
+          then seq (traceVal "${name}: ${reason}") false
+          else true;
+      in
+        # change to `prev && isOk` to debug why a value isn't included
+        prev && ok;
+
+      shouldBeInRegistry = foldl checkRegistryCondition true [
+        {
+          reason = "not available on ${platformForAvailability.system}";
+          ok = available;
+        }
+        {
+          reason = "failed eval";
+          ok = safeRegistryValue.success;
+        }
+        {
+          reason = "broken outpath";
+          ok = safeRegistryValue.value.storePaths != { out = "<broken>"; } && safeRegistryValue.value.storePaths != {};
+        }
+      ];
+    in
+      optional shouldBeInRegistry {
+        inherit name;
+        value = filterAttrs (_: v: v != null) safeRegistryValue.value;
+      };
+
+    registerScope = scope-name: scope: let
+      safeScope = tryEval scope;
+
+      list-of-scope-packages = mapAttrsToList registerPackage safeScope.value;
+      scope-registry-inner = flatten list-of-scope-packages;
+      scope-registry =
+        map (item: {
+          name = "${scope-name}.${item.name}";
+          value = item.value;
+        })
+        scope-registry-inner;
+
+      shouldBeInRegistry = safeScope.success && safeScope.value ? recurseForDerivations && safeScope.value.recurseForDerivations;
+    in
+      optionals shouldBeInRegistry scope-registry;
+
+    list-of-registry-packages = mapAttrsToList registerPackage pkgs;
+    registry-items = flatten list-of-registry-packages;
+
+    scoped-registries = flatten (mapAttrsToList registerScope pkgs);
+    registry = listToAttrs (registry-items ++ scoped-registries);
+  in
+    registry;
+}

src/ddb/mod.rs

@@ -8,4 +8,7 @@ pub mod batch_put;
 pub struct Store {
     pub attribute: String,
     pub store: String,
+    pub version: Option<String>,
 }
+
+pub const REGISTRY: &str = "./registry.nix";

src/ddb/nix.rs

@@ -1,3 +1,5 @@
+use crate::ddb::REGISTRY;
+
 use super::Store;
 use log::{debug, error, info};
 use serde::Deserialize;
@@ -6,10 +8,11 @@ use tokio::process::Command;
 
 #[derive(Deserialize, Debug, Clone)]
 struct Package {
-    outputs: Option<HashMap<String, String>>,
+    #[serde(rename = "storePaths")]
+    outputs: HashMap<String, String>,
+    version: Option<String>,
 }
 
-//nix-env -qa --meta --json --out-path -f https://github.com/NixOS/nixpkgs/archive/5f64a12a728902226210bf01d25ec6cbb9d9265b.tar.gz
 pub async fn get_store(rev: &str) -> HashMap<String, Store> {
     let nixpath = Command::new("nix-instantiate")
         .arg("--eval")
@@ -30,40 +33,37 @@ pub async fn get_store(rev: &str) -> HashMap<String, Store> {
     let nixpath = String::from_utf8_lossy(&nixpath.stdout).trim().to_string();
     debug!("nixpath: {}", nixpath);
 
-    let output = Command::new("nix-env")
+    let output = Command::new("nix-instantiate")
         .env("NIXPKGS_ALLOW_UNFREE", "1")
         .env("NIXPKGS_ALLOW_INSECURE", "1")
         // .env("NIXPKGS_ALLOW_BROKEN", "0")
         // .env("NIXPKGS_ALLOW_UNSUPPORTED_SYSTEM", "0")
-        .arg("-f")
-        .arg(&nixpath)
+        .arg("--eval")
+        .arg("-E")
+        .arg(format!("with import {} {{}}; (import {} {{ inherit lib; }}).genRegistry \"x86_64-linux\" pkgs", nixpath, REGISTRY))
         .arg("-I")
         .arg(format!("nixpkgs={}", nixpath))
-        .arg("-qa")
-        // .arg("--meta")
         .arg("--json")
-        .arg("--out-path")
-        .arg("--arg")
-        .arg("config")
-        .arg(format!(
-            "import {}/pkgs/top-level/packages-config.nix",
-            nixpath
-        ))
+        .arg("--strict")
         .output()
         .await
         .expect("failed to execute process");
 
-    let output: HashMap<String, Package> = serde_json::from_slice(&output.stdout).expect("failed to parse nix-env output");
+    println!("{}", String::from_utf8_lossy(&output.stderr));
+
+
+    let output: HashMap<String, Package> = serde_json::from_slice(&output.stdout).expect("failed to parse nix-instantiate output");
 
-    info!("nix-env: got {} packages", output.len());
+    info!("nix-instantiate: got {} packages", output.len());
 
     let store = output
         .iter()
         .filter_map(|(attr, pkg)| {
-            if let Some(outpath) = pkg.outputs.as_ref().and_then(|x| x.get("out")) {
+            if let Some(outpath) = pkg.outputs.get("out") {
                 let store = Store {
                     attribute: attr.to_string(),
                     store: outpath.split("/").last().unwrap().to_string(),
+                    version: pkg.version.clone()
                 };
 
                 Some((attr.to_string(), store))
@@ -73,7 +73,7 @@ pub async fn get_store(rev: &str) -> HashMap<String, Store> {
         })
         .collect::<HashMap<String, Store>>();
 
-    info!("nix-env: got {} store paths", store.len());
+    info!("nix-instantiate: got {} store paths", store.len());
 
     return store;
 }