parakeet / consumer / src / backfill / repo.rs

use super::{
    types::{CarCommitEntry, CarEntry, CarRecordEntry},
    CopyStore,
};
use crate::indexer::records;
use crate::indexer::types::{AggregateDeltaStore, RecordTypes};
use crate::{db, indexer};
use deadpool_postgres::Transaction;
use ipld_core::cid::Cid;
use iroh_car::CarReader;
use metrics::counter;
use parakeet_index::AggregateType;
use redis::aio::MultiplexedConnection;
use std::collections::HashMap;
use std::path::Path;
use tokio::io::BufReader;

type BackfillDeltaStore = HashMap<(String, i32), i32>;

pub async fn insert_repo(
    t: &mut Transaction<'_>,
    rc: &mut MultiplexedConnection,
    tmp_dir: &Path,
    repo: &str,
) -> eyre::Result<(CarCommitEntry, BackfillDeltaStore, CopyStore)> {
    let car = tokio::fs::File::open(tmp_dir.join(repo)).await?;
    let mut car_stream = CarReader::new(BufReader::new(car)).await?;

    // the root should be the commit block
    let root = car_stream.header().roots().first().cloned().unwrap();

    let mut commit = None;
    let mut mst_nodes: HashMap<Cid, String> = HashMap::new();
    let mut records: HashMap<Cid, RecordTypes> = HashMap::new();
    let mut deltas = HashMap::new();
    let mut copies = CopyStore::default();

    while let Some((cid, block)) = car_stream.next_block().await? {
        let Ok(block) = serde_ipld_dagcbor::from_slice::<CarEntry>(&block) else {
            tracing::warn!("failed to parse block {cid}");
            continue;
        };

        if root == cid {
            if let CarEntry::Commit(commit_entry) = block {
                commit = Some(commit_entry);
            } else {
                tracing::warn!("root did not point to a commit entry");
            }
            continue;
        }

        match block {
            CarEntry::Commit(_) => {
                tracing::warn!("got commit entry that was not in root")
            }
            CarEntry::Record(CarRecordEntry::Known(record)) => {
                if let Some(path) = mst_nodes.remove(&cid) {
                    record_index(t, rc, &mut copies, &mut deltas, repo, &path, cid, record).await?;
                } else {
                    records.insert(cid, record);
                }
            }
            CarEntry::Record(CarRecordEntry::Other { ty }) => {
                tracing::debug!("repo contains unknown record type: {ty} ({cid})");
            }
            CarEntry::Mst(mst) => {
                let mut out = Vec::with_capacity(mst.e.len());

                for node in mst.e {
                    let ks = String::from_utf8_lossy(&node.k);

                    let key = if node.p == 0 {
                        ks.to_string()
                    } else {
                        let (_, prev): &(Cid, String) = out.last().unwrap();
                        let prefix = &prev[..node.p as usize];

                        format!("{prefix}{ks}")
                    };

                    out.push((node.v, key.to_string()));
                }

                mst_nodes.extend(out);
            }
        }
    }

    for (cid, record) in records {
        if let Some(path) = mst_nodes.remove(&cid) {
            record_index(t, rc, &mut copies, &mut deltas, repo, &path, cid, record).await?;
        } else {
            tracing::warn!("couldn't find MST node for record {cid}")
        }
    }

    let commit = commit.unwrap();

    Ok((commit, deltas, copies))
}

async fn record_index(
    t: &mut Transaction<'_>,
    rc: &mut MultiplexedConnection,
    copies: &mut CopyStore,
    deltas: &mut BackfillDeltaStore,
    did: &str,
    path: &str,
    cid: Cid,
    record: RecordTypes,
) -> eyre::Result<()> {
    let Some((collection_raw, rkey)) = path.split_once("/") else {
        tracing::warn!("op contained invalid path {path}");
        return Ok(());
    };

    counter!("backfilled_commits", "collection" => collection_raw.to_string()).increment(1);

    let at_uri = format!("at://{did}/{path}");

    match record {
        RecordTypes::AppBskyFeedLike(rec) => {
            deltas.incr(&rec.subject.uri, AggregateType::Like).await;

            copies.push_record(&at_uri, cid);
            copies
                .likes
                .push((rkey.to_string(), rec.subject, rec.via, rec.created_at));
        }
        RecordTypes::AppBskyFeedPost(rec) => {
            let maybe_reply = rec.reply.as_ref().map(|v| v.parent.uri.clone());
            let maybe_embed = rec
                .embed
                .as_ref()
                .and_then(|v| v.as_bsky())
                .and_then(|v| match v {
                    records::AppBskyEmbed::Record(r) => Some(r.record.uri.clone()),
                    records::AppBskyEmbed::RecordWithMedia(r) => Some(r.record.record.uri.clone()),
                    _ => None,
                });

            if let Some(labels) = rec.labels.clone() {
                db::maintain_self_labels(t, did, Some(cid), &at_uri, labels).await?;
            }
            if let Some(embed) = rec.embed.clone().and_then(|embed| embed.into_bsky()) {
                db::post_embed_insert(t, &at_uri, embed, rec.created_at).await?;
            }

            deltas.incr(did, AggregateType::ProfilePost).await;
            if let Some(reply) = maybe_reply {
                deltas.incr(&reply, AggregateType::Reply).await;
            }
            if let Some(embed) = maybe_embed {
                deltas.incr(&embed, AggregateType::Embed).await;
            }

            copies.push_record(&at_uri, cid);
            copies.posts.push((at_uri, cid, rec));
        }
        RecordTypes::AppBskyFeedRepost(rec) => {
            deltas.incr(&rec.subject.uri, AggregateType::Repost).await;

            copies.push_record(&at_uri, cid);
            copies
                .reposts
                .push((rkey.to_string(), rec.subject, rec.via, rec.created_at));
        }
        RecordTypes::AppBskyGraphBlock(rec) => {
            copies.push_record(&at_uri, cid);
            copies
                .blocks
                .push((rkey.to_string(), rec.subject, rec.created_at));
        }
        RecordTypes::AppBskyGraphFollow(rec) => {
            deltas.incr(did, AggregateType::Follow).await;
            deltas.incr(&rec.subject, AggregateType::Follower).await;

            copies.push_record(&at_uri, cid);
            copies
                .follows
                .push((rkey.to_string(), rec.subject, rec.created_at));
        }
        RecordTypes::AppBskyGraphListItem(rec) => {
            let split_aturi = rec.list.rsplitn(4, '/').collect::<Vec<_>>();
            if did != split_aturi[2] {
                // it's also probably a bad idea to log *all* the attempts to do this...
                tracing::warn!("tried to create a listitem on a list we don't control!");
                return Ok(());
            }

            copies.push_record(&at_uri, cid);
            copies.list_items.push((at_uri, rec));
        }
        RecordTypes::AppBskyGraphVerification(rec) => {
            copies.push_record(&at_uri, cid);
            copies.verifications.push((at_uri, cid, rec));
        }
        _ => indexer::index_op(t, rc, deltas, did, cid, record, &at_uri, rkey).await?,
    }

    Ok(())
}