nostr-rs-relay/src/bin/bulkloader.rs

use nostr_rs_relay::config;
use nostr_rs_relay::error::{Error, Result};
use nostr_rs_relay::event::{single_char_tagname, Event};
use nostr_rs_relay::repo::sqlite::{build_pool, PooledConnection};
use nostr_rs_relay::repo::sqlite_migration::{curr_db_version, DB_VERSION};
use nostr_rs_relay::utils::is_lower_hex;
use rusqlite::params;
use rusqlite::{OpenFlags, Transaction};
use std::io;
use std::path::Path;
use std::sync::mpsc;
use std::thread;
use tracing::info;

/// Bulk load JSONL data from STDIN to the database specified in config.toml (or ./nostr.db as a default).
/// The database must already exist, this will not create a new one.
/// Tested against schema v13.

pub fn main() -> Result<()> {
    let _trace_sub = tracing_subscriber::fmt::try_init();
    println!("Nostr-rs-relay Bulk Loader");
    // check for a database file, or create one.
    let settings = config::Settings::new(&None)?;
    if !Path::new(&settings.database.data_directory).is_dir() {
        info!("Database directory does not exist");
        return Err(Error::DatabaseDirError);
    }
    // Get a database pool
    let pool = build_pool(
        "bulk-loader",
        &settings,
        OpenFlags::SQLITE_OPEN_READ_WRITE | OpenFlags::SQLITE_OPEN_CREATE,
        1,
        4,
        false,
    );
    {
        // check for database schema version
        let mut conn: PooledConnection = pool.get()?;
        let version = curr_db_version(&mut conn)?;
        info!("current version is: {:?}", version);
        // ensure the schema version is current.
        if version != DB_VERSION {
            info!("version is not current, exiting");
            panic!("cannot write to schema other than v{DB_VERSION}");
        }
    }
    // this channel will contain parsed events ready to be inserted
    let (event_tx, event_rx) = mpsc::sync_channel(100_000);
    // Thread for reading events
    let _stdin_reader_handler = thread::spawn(move || {
        let stdin = io::stdin();
        for readline in stdin.lines() {
            if let Ok(line) = readline {
                // try to parse a nostr event
                let eres: Result<Event, serde_json::Error> = serde_json::from_str(&line);
                if let Ok(mut e) = eres {
                    if let Ok(()) = e.validate() {
                        e.build_index();
                        //debug!("Event: {:?}", e);
                        event_tx.send(Some(e)).ok();
                    } else {
                        info!("could not validate event");
                    }
                } else {
                    info!("error reading event: {:?}", eres);
                }
            } else {
                // error reading
                info!("error reading: {:?}", readline);
            }
        }
        info!("finished parsing events");
        event_tx.send(None).ok();
        let ok: Result<()> = Ok(());
        ok
    });
    let mut conn: PooledConnection = pool.get()?;
    let mut events_read = 0;
    let event_batch_size = 50_000;
    let mut new_events = 0;
    let mut has_more_events = true;
    while has_more_events {
        // begin a transaction
        let tx = conn.transaction()?;
        // read in batch_size events and commit
        for _ in 0..event_batch_size {
            match event_rx.recv() {
                Ok(Some(e)) => {
                    events_read += 1;
                    // ignore ephemeral events
                    if !(e.kind >= 20000 && e.kind < 30000) {
                        match write_event(&tx, e) {
                            Ok(c) => {
                                new_events += c;
                            }
                            Err(e) => {
                                info!("error inserting event: {:?}", e);
                            }
                        }
                    }
                }
                Ok(None) => {
                    // signal that the sender will never produce more
                    // events
                    has_more_events = false;
                    break;
                }
                Err(_) => {
                    info!("sender is closed");
                    // sender is done
                }
            }
        }
        info!("committed {} events...", new_events);
        tx.commit()?;
        conn.execute_batch("pragma wal_checkpoint(truncate)")?;
    }
    info!("processed {} events", events_read);
    info!("stored {} new events", new_events);
    // get a connection for writing events
    // read standard in.
    info!("finished reading input");
    Ok(())
}

/// Write an event and update the tag table.
/// Assumes the event has its index built.
fn write_event(tx: &Transaction, e: Event) -> Result<usize> {
    let id_blob = hex::decode(&e.id).ok();
    let pubkey_blob: Option<Vec<u8>> = hex::decode(&e.pubkey).ok();
    let delegator_blob: Option<Vec<u8>> = e.delegated_by.as_ref().and_then(|d| hex::decode(d).ok());
    let event_str = serde_json::to_string(&e).ok();
    // ignore if the event hash is a duplicate.
    let ins_count = tx.execute(
	"INSERT OR IGNORE INTO event (event_hash, created_at, kind, author, delegated_by, content, first_seen, hidden) VALUES (?1, ?2, ?3, ?4, ?5, ?6, strftime('%s','now'), FALSE);",
	params![id_blob, e.created_at, e.kind, pubkey_blob, delegator_blob, event_str]
    )?;
    if ins_count == 0 {
        return Ok(0);
    }
    // we want to capture the event_id that had the tag, the tag name, and the tag hex value.
    let event_id = tx.last_insert_rowid();
    // look at each event, and each tag, creating new tag entries if appropriate.
    for t in e.tags.iter().filter(|x| x.len() > 1) {
        let tagname = t.get(0).unwrap();
        let tagnamechar_opt = single_char_tagname(tagname);
        if tagnamechar_opt.is_none() {
            continue;
        }
        // safe because len was > 1
        let tagval = t.get(1).unwrap();
        // insert as BLOB if we can restore it losslessly.
        // this means it needs to be even length and lowercase.
        if (tagval.len() % 2 == 0) && is_lower_hex(tagval) {
            tx.execute(
                "INSERT INTO tag (event_id, name, value_hex) VALUES (?1, ?2, ?3);",
                params![event_id, tagname, hex::decode(tagval).ok()],
            )?;
        } else {
            // otherwise, insert as text
            tx.execute(
                "INSERT INTO tag (event_id, name, value) VALUES (?1, ?2, ?3);",
                params![event_id, tagname, &tagval],
            )?;
        }
    }
    if e.is_replaceable() {
        //let query = "SELECT id FROM event WHERE kind=? AND author=? ORDER BY created_at DESC LIMIT 1;";
        //let count: usize = tx.query_row(query, params![e.kind, pubkey_blob], |row| row.get(0))?;
        //info!("found {} rows that /would/ be preserved", count);
        match tx.execute(
	    "DELETE FROM event WHERE kind=? and author=? and id NOT IN (SELECT id FROM event WHERE kind=? AND author=? ORDER BY created_at DESC LIMIT 1);",
	    params![e.kind, pubkey_blob, e.kind, pubkey_blob],
	) {
	    Ok(_) => {},
	    Err(x) => {info!("error deleting replaceable event: {:?}",x);}
	}
    }
    Ok(ins_count)
}