flake-tracker/src/bin/scan-flake.rs

use anyhow::{
    anyhow,
    Context,
    Result,
};
use chrono::{
    Utc,
};
use clap::{
    Parser,
};
use serde::{
    Deserialize,
};
use sqlx::{
    sqlite::SqlitePoolOptions,
};
use std::collections::HashMap;
use std::process::Command;

#[derive(Deserialize, Debug)]
struct FlakeMetadata {
    lastModified: i64,
    locked: FlakeLockedInfo,
    locks: FlakeLocks,
    original: FlakeSource,
    path: String,
    resolved: FlakeSource,
    revision: String,
}

#[derive(Deserialize, Debug, Clone)]
struct FlakeSource {
    owner: Option<String>,
    repo: Option<String>,
    r#ref: Option<String>,
    rev: Option<String>,
    r#type: String,
    url: Option<String>,
}

fn assemble_flake_git_uri(flake_source: &FlakeSource) -> Result<String> {
    let mut out = String::new();
    out.push_str("git+");
    out.push_str(&flake_source.url.clone()
        .context("Flake git uri does not contain an url")?);

    let mut query: Vec<String> = Vec::new();

    if let Some(r#ref) = &flake_source.r#ref {
        query.push(format!("ref={}", r#ref));
    }

    if let Some(rev) = &flake_source.rev {
        query.push(format!("rev={}", rev));
    }

    if query.len() > 0 {
        out.push_str("?");
        out.push_str(&query.join("&"));
    }

    Ok(out)
}

fn assemble_flake_github_uri(flake_source: &FlakeSource) -> Result<String> {
    let mut out = String::new();
    out.push_str("github:");
    out.push_str(&flake_source.owner.clone()
        .context("Flake github uri does not contain an owner")?);
    out.push_str("/");
    out.push_str(&flake_source.repo.clone()
        .context("Flake github uri does not contain an repo")?);

    if let Some(rev) = &flake_source.rev {
        out.push_str("/");
        out.push_str(&rev);
    }

    Ok(out)
}

fn assemble_flake_tarball_uri(flake_source: &FlakeSource) -> Result<String> {
    let mut out = String::new();
    out.push_str(&flake_source.url.clone()
        .context("Flake tarball uri does not contain an url")?);

    if let Some(rev) = &flake_source.rev {
        out.push_str("?rev=");
        out.push_str(rev);
    }

    Ok(out)
}

trait FlakeUri {
    fn flake_uri(&self) -> Result<String>;
}

impl FlakeUri for FlakeSource {
    fn flake_uri(&self) -> Result<String> {
        match self.r#type.as_str() {
            "git" => assemble_flake_git_uri(self),
            "github" => assemble_flake_github_uri(self),
            "tarball" => assemble_flake_tarball_uri(self),
            other => Err(anyhow!("Unsupported flake uri type {}", other)),
        }
    }
}

#[derive(Deserialize, Debug, Clone)]
struct FlakeLockedInfo {
    lastModified: i64,
    narHash: String,
    owner: Option<String>,
    repo: Option<String>,
    r#ref: Option<String>,
    rev: Option<String>,
    r#type: String,
    url: Option<String>,
}

impl FlakeUri for FlakeLockedInfo {
    fn flake_uri(&self) -> Result<String> {
        Into::<FlakeSource>::into(self).flake_uri()
    }
}

impl Into<FlakeSource> for &FlakeLockedInfo {
    fn into(self) -> FlakeSource {
        FlakeSource {
            owner: self.owner.clone(),
            repo: self.repo.clone(),
            r#ref: self.r#ref.clone(),
            rev: self.rev.clone(),
            r#type: self.r#type.clone(),
            url: self.url.clone(),
        }
    }
}

#[derive(Deserialize, Debug)]
struct FlakeLocks {
    nodes: HashMap<String, FlakeLocksNode>,
    root: String,
    version: u64,
}

#[derive(Deserialize, Debug, Clone)]
#[serde(untagged)]
enum FlakeLocksNodeInputs {
    String(String),
    Array(Vec<String>),
}

#[derive(Deserialize, Debug)]
struct FlakeLocksNode {
    inputs: Option<HashMap<String, FlakeLocksNodeInputs>>,
    locked: Option<FlakeLockedInfo>,
    original: Option<FlakeSource>,
}

#[derive(Parser)]
#[command(version, about, long_about = None)]
struct Cli {
    flake_uri: String,
}

#[tokio::main]
async fn main() -> Result<()> {
    let cli = Cli::parse();

    let scan_time = Utc::now().timestamp();

    let db = SqlitePoolOptions::new().connect("sqlite://flake-tracker.db").await?;

    let flake_metadata_raw = Command::new("nix")
        .arg("flake")
        .arg("metadata")
        .arg("--json")
        .arg(cli.flake_uri)
        .output()
        .context("Failed to fetch flake metadata")?;

    let flake_metadata: FlakeMetadata = serde_json::from_slice(&flake_metadata_raw.stdout)
        .context("Failed to parse flake metadata")?;

    sqlx::query("INSERT INTO revisions (revision_uri, flake_uri, nix_store_path, nar_hash, last_modified, tracker_last_scanned)
        VALUES (?, ?, ?, ?, ?, ?)
        ON CONFLICT(revision_uri) DO UPDATE SET
            flake_uri=excluded.flake_uri,
            nix_store_path=excluded.nix_store_path,
            nar_hash=excluded.nar_hash,
            last_modified=excluded.last_modified,
            tracker_last_scanned=excluded.tracker_last_scanned
        ")
        .bind(&flake_metadata.locked.flake_uri()?)
        .bind(&flake_metadata.resolved.flake_uri()?)
        .bind(&flake_metadata.path)
        .bind(&flake_metadata.locked.narHash)
        .bind(&flake_metadata.locked.lastModified)
        .bind(&scan_time)
        .execute(&db).await?;

    let locks_root_name = &flake_metadata.locks.root;
    let locks_root_node = flake_metadata.locks.nodes.get(locks_root_name)
        .context("Failed to get locks root node")?;

    for (input_name, locks_input_name) in locks_root_node.inputs.clone().context("No inputs found for flake")? {

        if let FlakeLocksNodeInputs::String(locks_input_name) = locks_input_name {
            let locks_input_node = flake_metadata.locks.nodes.get(&locks_input_name)
                .context("Failed to find lock of input")?;

            sqlx::query("INSERT INTO inputs (revision_uri, input_name, locked_revision_uri, locked_flake_uri, locked_nar_hash, last_modified)
                VALUES (?, ?, ?, ?, ?, ?)
                ON CONFLICT(revision_uri, input_name) DO UPDATE SET
                    locked_revision_uri=excluded.locked_revision_uri,
                    locked_flake_uri=excluded.locked_flake_uri,
                    locked_nar_hash=excluded.locked_nar_hash,
                    last_modified=excluded.last_modified
                ")
                .bind(flake_metadata.locked.flake_uri()?)
                .bind(input_name)
                .bind(locks_input_node.locked.clone().context("Unexpected missing lock")?.flake_uri()?)
                .bind(locks_input_node.original.clone().context("Unexpected missing lock")?.flake_uri()?)
                .bind(locks_input_node.locked.clone().context("Unexpected missing lock")?.narHash)
                .bind(locks_input_node.locked.clone().context("Unexpected missing lock")?.lastModified)
                .execute(&db).await?;

            sqlx::query("INSERT INTO revisions (revision_uri, flake_uri)
                VALUES (?, ?)
                ON CONFLICT(revision_uri) DO NOTHING
                ")
                .bind(locks_input_node.locked.clone().context("Unexpected missing lock")?.flake_uri()?)
                .bind(locks_input_node.original.clone().context("Unexpected missing lock")?.flake_uri()?)
                .execute(&db).await?;
        }
    }

    Ok(())
}