quickpeep/quickpeep_raker/src/bin/qp-raker-db.rs

use clap::Parser;
use std::borrow::Cow;

use std::fmt::Debug;

use env_logger::Env;

use anyhow::{bail, Context};

use colour::{dark_yellow_ln, red_ln};
use libmdbx::{Database, TableObject, RO};

use std::path::PathBuf;

use quickpeep_raker::config;

use quickpeep_raker::storage::mdbx_helper_types::MdbxBare;
use quickpeep_raker::storage::records::{
    ActiveDomainRecord, BackingOffDomainRecord, DomainRecord, OnHoldUrlRecord, QueueUrlRecord,
    UrlVisitedRecord,
};
use quickpeep_raker::storage::{RakerStore, RakerTxn};

/// Seeds a raker's queue with URLs
#[derive(Clone, Debug, Parser)]
pub struct Opts {
    #[clap(long = "config")]
    config: Option<PathBuf>,

    /// Table name
    table: String,

    /// Key name to look up
    key_name: String,

    /// Search for any prefix, not an exact match.
    #[clap(long = "prefix", short = 'p')]
    prefix: bool,
}

#[tokio::main]
pub async fn main() -> anyhow::Result<()> {
    env_logger::Builder::from_env(Env::default().default_filter_or("info,quickpeep=debug")).init();

    let opts: Opts = Opts::parse();

    let config_path = opts
        .config
        .unwrap_or_else(|| PathBuf::from("quickpeep.ron"));
    let config = config::RakerConfig::load(&config_path).context("Failed to load config")?;

    if !config.raker.workbench_dir.exists() {
        bail!(
            "Workbench directory ({:?}) doesn't exist.",
            config.raker.workbench_dir
        );
    }
    if !config.seed_dir.exists() {
        bail!("Seed directory ({:?}) doesn't exist.", config.seed_dir);
    }

    let store = RakerStore::open(&config.raker.workbench_dir.join("raker.mdbx"))?;

    let txn = store.ro_txn()?;
    match opts.table.as_ref() {
        "queue_urls" | "urls_queue" => {
            inspect::<MdbxBare<QueueUrlRecord>>(
                opts.key_name.as_ref(),
                opts.prefix,
                &txn.mdbx.borrow_dbs().queue_urls,
                &txn,
            )?;
        }
        "active_domains" => {
            inspect::<MdbxBare<ActiveDomainRecord>>(
                opts.key_name.as_ref(),
                opts.prefix,
                &txn.mdbx.borrow_dbs().active_domains,
                &txn,
            )?;
        }
        "active_domains_raffle" => {
            inspect::<MdbxBare<String>>(
                opts.key_name.as_ref(),
                opts.prefix,
                &txn.mdbx.borrow_dbs().active_domain_raffle,
                &txn,
            )?;
        }
        "backing_off_reinstatements" => {
            inspect::<MdbxBare<String>>(
                opts.key_name.as_ref(),
                opts.prefix,
                &txn.mdbx.borrow_dbs().backing_off_reinstatements,
                &txn,
            )?;
        }
        "backing_off_domains" => {
            inspect::<MdbxBare<BackingOffDomainRecord>>(
                opts.key_name.as_ref(),
                opts.prefix,
                &txn.mdbx.borrow_dbs().backing_off_domains,
                &txn,
            )?;
        }
        "visited_urls" => {
            inspect::<MdbxBare<UrlVisitedRecord>>(
                opts.key_name.as_ref(),
                opts.prefix,
                &txn.mdbx.borrow_dbs().visited_urls,
                &txn,
            )?;
        }
        "domains" => {
            inspect::<MdbxBare<DomainRecord>>(
                opts.key_name.as_ref(),
                opts.prefix,
                &txn.mdbx.borrow_dbs().domains,
                &txn,
            )?;
        }
        "urls_on_hold" => {
            inspect::<MdbxBare<OnHoldUrlRecord>>(
                opts.key_name.as_ref(),
                opts.prefix,
                &txn.mdbx.borrow_dbs().urls_on_hold,
                &txn,
            )?;
        }
        other => {
            dark_yellow_ln!("Unknown database {:?}", other);
        }
    }

    Ok(())
}

trait Inspectable {
    fn inspect(&self) -> String;
}

impl<T: Debug> Inspectable for MdbxBare<T> {
    fn inspect(&self) -> String {
        format!("{:?}", &self.0)
    }
}

fn inspect<'a, IV: Inspectable + TableObject<'a> + 'static>(
    key: &str,
    prefix: bool,
    database: &Database<'a>,
    txn: &'a RakerTxn<'a, RO>,
) -> anyhow::Result<()> {
    if prefix {
        let mut cur = txn.mdbx_txn.cursor(database)?;
        for item in cur.iter_from::<Cow<'_, [u8]>, IV>(key.as_bytes()) {
            let (k, v) = item?;
            if !k.starts_with(key.as_bytes()) {
                break;
            }
            println!("• {}", std::str::from_utf8(&k).unwrap_or("<Not UTF-8>"));
            println!("  = {}", v.inspect());
        }
    } else {
        if let Some(entry) = txn.mdbx_txn.get::<IV>(database, key.as_bytes())? {
            println!("{}", entry.inspect());
        } else {
            red_ln!("no value");
        }
    }

    Ok(())
}