use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};
use scryer_db::{CodeGraphEdge, Scope, ScryerDb, SourceFile, Symbol, SymbolReference};
use tokio::sync::mpsc;
use tokio::task::JoinHandle;
use crate::payload::{ParsedFilePayload, RawCodeGraphEdge, RawSymbolReference};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum ExistingFileStatus {
#[default]
Unknown,
New,
Existing(u64),
}
#[derive(Debug)]
pub enum IngestionMessage {
UpsertFile {
payload: ParsedFilePayload,
status: ExistingFileStatus,
dependency_package_id: Option<u64>,
},
DeleteFile(PathBuf),
}
pub struct IngestionHandle {
sender: mpsc::Sender<IngestionMessage>,
join_handle: JoinHandle<anyhow::Result<usize>>,
}
impl IngestionHandle {
pub fn sender(&self) -> mpsc::Sender<IngestionMessage> {
self.sender.clone()
}
pub async fn finish(self) -> anyhow::Result<usize> {
drop(self.sender);
self.join_handle.await?
}
}
pub struct BatchIngestionActor {
db: ScryerDb,
project_id: u64,
}
impl BatchIngestionActor {
pub fn new(db: ScryerDb, project_id: u64) -> Self {
Self { db, project_id }
}
pub fn spawn(db: ScryerDb, project_id: u64, buffer_capacity: usize) -> IngestionHandle {
let (sender, mut receiver) = mpsc::channel::<IngestionMessage>(buffer_capacity);
let actor = Self::new(db, project_id);
let join_handle = tokio::spawn(async move {
let mut processed = 0;
let mut batch = Vec::new();
while let Some(msg) = receiver.recv().await {
batch.push(msg);
while batch.len() < 500 {
match receiver.try_recv() {
Ok(m) => batch.push(m),
Err(_) => break,
}
}
let count = batch.len();
actor.ingest_batch(std::mem::take(&mut batch)).await?;
processed += count;
}
Ok(processed)
});
IngestionHandle {
sender,
join_handle,
}
}
pub async fn ingest_batch(&self, messages: Vec<IngestionMessage>) -> anyhow::Result<()> {
if messages.is_empty() {
return Ok(());
}
let batch_timer = std::time::Instant::now();
let n_messages = messages.len();
let mut db_guard = self.db.lock().await;
let mut tx = db_guard.transaction().await?;
let mut pending = Vec::new();
for msg in messages {
match msg {
IngestionMessage::UpsertFile {
payload,
status,
dependency_package_id,
} => {
pending.push(
self.insert_file_entities(&mut tx, payload, status, dependency_package_id)
.await?,
);
}
IngestionMessage::DeleteFile(rel_path) => {
self.delete_file_in_tx(&mut tx, &rel_path).await?;
}
}
}
let insert_ms = batch_timer.elapsed().as_millis();
let mut cache = HashMap::new();
self.prefetch_targets(&mut tx, &pending, &mut cache).await?;
for links in pending {
self.link_in_tx(&mut tx, links, &mut cache).await?;
}
let link_ms = batch_timer.elapsed().as_millis() - insert_ms;
tx.commit().await?;
tracing::info!(
"ingest_batch project={} messages={n_messages}: entities={insert_ms}ms link={link_ms}ms commit={}ms",
self.project_id,
batch_timer.elapsed().as_millis() - insert_ms - link_ms
);
Ok(())
}
pub async fn ingest_payload(&self, payload: ParsedFilePayload) -> anyhow::Result<()> {
self.ingest_batch(vec![IngestionMessage::UpsertFile {
payload,
status: ExistingFileStatus::Unknown,
dependency_package_id: None,
}])
.await
}
pub async fn ingest_payload_in_tx(
&self,
tx: &mut toasty::db::Transaction<'_>,
payload: ParsedFilePayload,
status: ExistingFileStatus,
) -> anyhow::Result<()> {
let links = self.insert_file_entities(tx, payload, status, None).await?;
self.link_in_tx(tx, links, &mut HashMap::new()).await
}
async fn insert_file_entities(
&self,
tx: &mut toasty::db::Transaction<'_>,
payload: ParsedFilePayload,
status: ExistingFileStatus,
dependency_package_id: Option<u64>,
) -> anyhow::Result<PendingLinks> {
let rel_path_str = payload.relative_path.to_string_lossy().to_string();
let effective_project_id = if dependency_package_id.is_some() {
0
} else {
self.project_id
};
let existing = match status {
ExistingFileStatus::New => None,
ExistingFileStatus::Existing(existing_id) => SourceFile::filter(
SourceFile::fields()
.project_id()
.eq(effective_project_id)
.and(SourceFile::fields().id().eq(existing_id)),
)
.first()
.exec(&mut *tx)
.await?
.map(|f| (existing_id, Some(f)))
.or(Some((existing_id, None))),
ExistingFileStatus::Unknown => {
if let Some(dep_pkg_id) = dependency_package_id {
SourceFile::filter(
SourceFile::fields()
.project_id()
.eq(0)
.and(
SourceFile::fields()
.dependency_package_id()
.eq(Some(dep_pkg_id)),
)
.and(SourceFile::fields().path().eq(&rel_path_str)),
)
.first()
.exec(&mut *tx)
.await?
.map(|f| (f.id, Some(f)))
} else {
SourceFile::filter(
SourceFile::fields()
.project_id()
.eq(effective_project_id)
.and(SourceFile::fields().path().eq(&rel_path_str)),
)
.first()
.exec(&mut *tx)
.await?
.map(|f| (f.id, Some(f)))
}
}
};
let mut old_symbols = Vec::new();
let file_id = match existing {
Some((existing_id, file)) => {
old_symbols = Symbol::filter(
Symbol::fields()
.project_id()
.eq(effective_project_id)
.and(Symbol::fields().file_id().eq(existing_id)),
)
.exec(&mut *tx)
.await?;
purge_file_records(
tx,
effective_project_id,
existing_id,
PurgeMode::KeepIncoming,
)
.await?;
match file {
Some(mut file) => {
file.update()
.content_hash(payload.content_hash.clone())
.line_count(payload.line_count)
.byte_size(payload.byte_size)
.exec(&mut *tx)
.await?;
Some(file.id)
}
None => None,
}
}
None => None,
};
let file_id = match file_id {
Some(id) => id,
None => {
let mut builder = SourceFile::create();
builder = builder
.project_id(effective_project_id)
.path(rel_path_str)
.content_hash(payload.content_hash.clone())
.language(payload.language.clone())
.line_count(payload.line_count)
.byte_size(payload.byte_size);
if let Some(dep_pkg_id) = dependency_package_id {
builder = builder.dependency_package_id(dep_pkg_id);
}
builder.exec(&mut *tx).await?.id
}
};
let scopes = payload.scopes;
let mut depth = Vec::with_capacity(scopes.len());
let mut local_depth = HashMap::new();
for raw_scope in &scopes {
let d = raw_scope
.parent_local_id
.and_then(|p| local_depth.get(&p).copied())
.map_or(0, |parent_depth| parent_depth + 1);
local_depth.insert(raw_scope.local_id, d);
depth.push(d);
}
let max_depth = depth.iter().copied().max().unwrap_or(0);
let mut local_to_db_id: HashMap<_, u64> = HashMap::with_capacity(scopes.len());
for level in 0..=max_depth {
let level_scopes: Vec<_> = scopes
.iter()
.zip(&depth)
.filter(|&(_, &d)| d == level)
.map(|(scope, _)| scope)
.collect();
for chunk in level_scopes.chunks(INSERT_CHUNK) {
let mut create = Scope::create_many();
for raw_scope in chunk {
let mut builder = Scope::create()
.project_id(effective_project_id)
.file_id(file_id)
.scope_kind(raw_scope.scope_kind.clone())
.start_byte(raw_scope.start_byte)
.end_byte(raw_scope.end_byte)
.start_line(raw_scope.start_line)
.end_line(raw_scope.end_line);
if let Some(parent_id) = raw_scope
.parent_local_id
.and_then(|p| local_to_db_id.get(&p).copied())
{
builder = builder.parent_scope_id(parent_id);
}
create = create.item(builder);
}
let inserted = create.exec(&mut *tx).await?;
anyhow::ensure!(
inserted.len() == chunk.len(),
"scope batch insert returned {} rows for {}",
inserted.len(),
chunk.len()
);
for (raw_scope, row) in chunk.iter().zip(inserted) {
anyhow::ensure!(
row.start_byte == raw_scope.start_byte
&& row.end_byte == raw_scope.end_byte,
"scope batch insert returned rows out of order"
);
local_to_db_id.insert(raw_scope.local_id, row.id);
}
}
}
let mut symbol_name_to_id = HashMap::new();
for chunk in payload.symbols.chunks(INSERT_CHUNK) {
let mut create = Symbol::create_many();
for raw_symbol in chunk {
let mut builder = Symbol::create()
.project_id(effective_project_id)
.file_id(file_id)
.name(raw_symbol.name.clone())
.qualified_name(raw_symbol.qualified_name.clone())
.kind(raw_symbol.kind.clone())
.visibility(raw_symbol.visibility.clone())
.signature(raw_symbol.signature.clone())
.start_byte(raw_symbol.start_byte)
.end_byte(raw_symbol.end_byte)
.start_line(raw_symbol.start_line)
.end_line(raw_symbol.end_line);
if let Some(dep_pkg_id) = dependency_package_id {
builder = builder.dependency_package_id(dep_pkg_id);
}
if let Some(scope_id) = raw_symbol
.scope_local_id
.and_then(|s| local_to_db_id.get(&s).copied())
{
builder = builder.scope_id(scope_id);
}
if let Some(doc) = &raw_symbol.docstring {
builder = builder.docstring(doc.clone());
}
create = create.item(builder);
}
let inserted = create.exec(&mut *tx).await?;
anyhow::ensure!(
inserted.len() == chunk.len(),
"symbol batch insert returned {} rows for {}",
inserted.len(),
chunk.len()
);
for (raw_symbol, row) in chunk.iter().zip(inserted) {
anyhow::ensure!(
row.qualified_name == raw_symbol.qualified_name
&& row.start_byte == raw_symbol.start_byte,
"symbol batch insert returned rows out of order"
);
symbol_name_to_id
.entry(raw_symbol.qualified_name.clone())
.or_insert(row.id);
symbol_name_to_id
.entry(raw_symbol.name.clone())
.or_insert(row.id);
}
}
let p = effective_project_id;
for old in old_symbols {
let stmts = match symbol_name_to_id.get(&old.qualified_name) {
Some(&new_id) => [
format!(
"UPDATE symbol_reference SET symbol_id = {new_id} WHERE project_id = {p} AND symbol_id = {};",
old.id
),
format!(
"UPDATE code_graph_edge SET target_symbol_id = {new_id} WHERE project_id = {p} AND target_symbol_id = {};",
old.id
),
],
None => [
format!(
"DELETE FROM symbol_reference WHERE project_id = {p} AND symbol_id = {};",
old.id
),
format!(
"DELETE FROM code_graph_edge WHERE project_id = {p} AND target_symbol_id = {};",
old.id
),
],
};
for stmt in &stmts {
toasty::sql::statement(stmt).exec(&mut *tx).await?;
}
}
Ok(PendingLinks {
file_id,
local_symbols: symbol_name_to_id,
references: payload.references,
edges: payload.edges,
})
}
async fn link_in_tx(
&self,
tx: &mut toasty::db::Transaction<'_>,
links: PendingLinks,
cache: &mut HashMap<String, Option<u64>>,
) -> anyhow::Result<()> {
let PendingLinks {
file_id,
local_symbols,
references,
edges,
} = links;
let mut ref_rows = Vec::with_capacity(references.len());
for raw_ref in references {
let target = self
.resolve_name(tx, &local_symbols, cache, &raw_ref.target_symbol_name)
.await?;
if let Some(sym_id) = target {
ref_rows.push(
SymbolReference::create()
.project_id(self.project_id)
.symbol_id(sym_id)
.file_id(file_id)
.role(raw_ref.role)
.start_byte(raw_ref.start_byte)
.end_byte(raw_ref.end_byte)
.line_number(raw_ref.line_number),
);
}
}
insert_all(tx, ref_rows, SymbolReference::create_many).await?;
let mut seen = HashSet::new();
let mut edge_rows = Vec::with_capacity(edges.len());
for raw_edge in edges {
let source = local_symbols.get(&raw_edge.source_symbol_name).copied();
let target = self
.resolve_name(tx, &local_symbols, cache, &raw_edge.target_symbol_name)
.await?;
if let (Some(src), Some(tgt)) = (source, target) {
if !seen.insert((src, tgt, raw_edge.edge_type.clone())) {
continue;
}
edge_rows.push(
CodeGraphEdge::create()
.project_id(self.project_id)
.source_symbol_id(src)
.target_symbol_id(tgt)
.edge_type(raw_edge.edge_type),
);
}
}
insert_all(tx, edge_rows, CodeGraphEdge::create_many).await?;
Ok(())
}
async fn resolve_name(
&self,
tx: &mut toasty::db::Transaction<'_>,
local_symbols: &HashMap<String, u64>,
cache: &mut HashMap<String, Option<u64>>,
target: &str,
) -> anyhow::Result<Option<u64>> {
if let Some(&id) = local_symbols.get(target) {
return Ok(Some(id));
}
if let Some(&cached) = cache.get(target) {
return Ok(cached);
}
let found = find_symbol_for_target(&mut *tx, self.project_id, target, None)
.await?
.map(|s| s.id);
cache.insert(target.to_string(), found);
Ok(found)
}
async fn prefetch_targets(
&self,
tx: &mut toasty::db::Transaction<'_>,
pending: &[PendingLinks],
cache: &mut HashMap<String, Option<u64>>,
) -> anyhow::Result<()> {
let mut targets: HashSet<&str> = HashSet::new();
for links in pending {
let names = links
.references
.iter()
.map(|r| r.target_symbol_name.as_str())
.chain(links.edges.iter().map(|e| e.target_symbol_name.as_str()));
for name in names {
if !links.local_symbols.contains_key(name) {
targets.insert(name);
}
}
}
let bare = |t: &str| bare_name(t).to_string();
let names: Vec<String> = targets
.iter()
.map(|t| bare(t))
.filter(|n| !n.is_empty())
.collect::<HashSet<_>>()
.into_iter()
.collect();
let mut by_name = fetch_candidates(&mut *tx, self.project_id, &names).await?;
if self.project_id != 0 {
let missing: Vec<String> = names
.iter()
.filter(|n| !by_name.contains_key(*n))
.cloned()
.collect();
by_name.extend(fetch_candidates(&mut *tx, 0, &missing).await?);
}
for target in targets {
let candidates = by_name.get(&bare(target)).cloned().unwrap_or_default();
let found = pick_candidate(candidates, target, None).map(|s| s.id);
cache.insert(target.to_string(), found);
}
Ok(())
}
pub async fn delete_file_in_tx(
&self,
tx: &mut toasty::db::Transaction<'_>,
rel_path: &Path,
) -> anyhow::Result<()> {
let rel_path_str = rel_path.to_string_lossy().to_string();
let existing_file = SourceFile::filter(
SourceFile::fields()
.project_id()
.eq(self.project_id)
.and(SourceFile::fields().path().eq(&rel_path_str)),
)
.first()
.exec(&mut *tx)
.await?;
if let Some(file) = existing_file {
purge_file_records(tx, self.project_id, file.id, PurgeMode::All).await?;
let del_source_file = format!(
"DELETE FROM source_file WHERE project_id = {} AND id = {};",
self.project_id, file.id
);
toasty::sql::statement(&del_source_file)
.exec(&mut *tx)
.await?;
}
Ok(())
}
pub async fn delete_file(&self, rel_path: &Path) -> anyhow::Result<()> {
self.ingest_batch(vec![IngestionMessage::DeleteFile(rel_path.to_path_buf())])
.await
}
}
const INSERT_CHUNK: usize = 500;
async fn insert_all<M, C>(
tx: &mut toasty::db::Transaction<'_>,
items: Vec<C>,
create_many: fn() -> toasty::stmt::CreateMany<M>,
) -> anyhow::Result<()>
where
M: toasty::schema::Model,
C: toasty::stmt::IntoInsert<Model = M>,
{
let mut items = items.into_iter().peekable();
while items.peek().is_some() {
let mut create = create_many();
for item in items.by_ref().take(INSERT_CHUNK) {
create = create.item(item);
}
create.exec(&mut *tx).await?;
}
Ok(())
}
struct PendingLinks {
file_id: u64,
local_symbols: HashMap<String, u64>,
references: Vec<RawSymbolReference>,
edges: Vec<RawCodeGraphEdge>,
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum PurgeMode {
All,
KeepIncoming,
}
async fn purge_file_records(
tx: &mut toasty::db::Transaction<'_>,
project_id: u64,
file_id: u64,
mode: PurgeMode,
) -> anyhow::Result<()> {
let file_symbols =
format!("SELECT id FROM symbol WHERE project_id = {project_id} AND file_id = {file_id}");
let mut statements = vec![
format!(
"DELETE FROM symbol_reference WHERE project_id = {project_id} AND file_id = {file_id};"
),
format!(
"DELETE FROM code_graph_edge WHERE project_id = {project_id} AND source_symbol_id IN ({file_symbols});"
),
];
if mode == PurgeMode::All {
statements.push(format!(
"DELETE FROM symbol_reference WHERE project_id = {project_id} AND symbol_id IN ({file_symbols});"
));
statements.push(format!(
"DELETE FROM code_graph_edge WHERE project_id = {project_id} AND target_symbol_id IN ({file_symbols});"
));
}
statements.push(format!(
"DELETE FROM symbol WHERE project_id = {project_id} AND file_id = {file_id};"
));
statements.push(format!(
"DELETE FROM scope WHERE project_id = {project_id} AND file_id = {file_id};"
));
for stmt in &statements {
toasty::sql::statement(stmt).exec(&mut *tx).await?;
}
Ok(())
}
async fn fetch_candidates(
executor: &mut dyn toasty::Executor,
project_id: u64,
names: &[String],
) -> anyhow::Result<HashMap<String, Vec<Symbol>>> {
let mut by_name: HashMap<String, Vec<Symbol>> = HashMap::new();
for chunk in names.chunks(200) {
let rows = Symbol::filter(
Symbol::fields()
.project_id()
.eq(project_id)
.and(Symbol::fields().name().in_list(chunk.to_vec())),
)
.exec(&mut *executor)
.await?;
for sym in rows {
by_name.entry(sym.name.clone()).or_default().push(sym);
}
}
for list in by_name.values_mut() {
list.sort_by_key(|s| s.id);
}
Ok(by_name)
}
fn bare_name(target: &str) -> &str {
let last = target.rsplit("::").next().unwrap_or(target);
last.rsplit('.').next().unwrap_or(last)
}
fn pick_candidate(
mut candidates: Vec<Symbol>,
target: &str,
prefer_file: Option<u64>,
) -> Option<Symbol> {
if target.contains("::") || target.contains('.') {
let rust_suffix = format!("::{target}");
let dotted_suffix = format!(".{target}");
candidates.retain(|s| {
s.qualified_name == target
|| s.qualified_name.ends_with(&rust_suffix)
|| s.qualified_name.ends_with(&dotted_suffix)
});
if let Some(pos) = candidates.iter().position(|s| s.qualified_name == target) {
return Some(candidates.swap_remove(pos));
}
}
let preferred = prefer_file.and_then(|f| candidates.iter().position(|s| s.file_id == f));
match preferred {
Some(pos) => Some(candidates.swap_remove(pos)),
None => candidates.into_iter().next(),
}
}
pub async fn find_symbol_for_target(
executor: &mut dyn toasty::Executor,
project_id: u64,
target: &str,
prefer_file: Option<u64>,
) -> anyhow::Result<Option<Symbol>> {
let name = bare_name(target);
if name.is_empty() {
return Ok(None);
}
let names = [name.to_string()];
let mut by_name = fetch_candidates(&mut *executor, project_id, &names).await?;
if by_name.is_empty() && project_id != 0 {
by_name = fetch_candidates(&mut *executor, 0, &names).await?;
}
let candidates = by_name.remove(name).unwrap_or_default();
Ok(pick_candidate(candidates, target, prefer_file))
}