1use std::collections::HashMap;
8use std::sync::Arc;
9
10use crate::facts::FileFacts;
11use crate::model::{Edge, Symbol};
12use anyhow::{Context, Result};
13use lancedb::arrow::arrow_array::{
14 ArrayRef, BooleanArray, FixedSizeListArray, Float32Array, RecordBatch, StringArray,
15 UInt32Array, UInt64Array,
16};
17use lancedb::arrow::arrow_schema::{DataType, Field, Schema};
18use lancedb::table::AddDataMode;
19
20type TextRow = (GlobalSymbolId, String, Vec<String>);
22type SearchRow = (GlobalSymbolId, String, Vec<f32>);
24type VectorRow = (GlobalSymbolId, Vec<f32>);
26
27fn symbols_schema() -> Arc<Schema> {
28 Arc::new(Schema::new(vec![
29 Field::new("symbol_id", DataType::Utf8, false),
30 Field::new("scheme", DataType::Utf8, false),
31 Field::new("package_manager", DataType::Utf8, false),
32 Field::new("package", DataType::Utf8, false),
33 Field::new("version", DataType::Utf8, false),
34 Field::new("descriptor", DataType::Utf8, false),
35 Field::new("name", DataType::Utf8, false),
36 Field::new("kind", DataType::Utf8, false),
37 Field::new("language", DataType::Utf8, false),
38 Field::new("repo", DataType::Utf8, false),
39 Field::new("file_path", DataType::Utf8, false),
40 Field::new("line_start", DataType::UInt32, false),
41 Field::new("col_start", DataType::UInt32, false),
42 Field::new("line_end", DataType::UInt32, false),
43 Field::new("col_end", DataType::UInt32, false),
44 Field::new("signature", DataType::Utf8, true),
45 Field::new("docstring", DataType::Utf8, true),
46 Field::new("fp_sig", DataType::UInt64, true),
47 Field::new("fp_body", DataType::UInt64, true),
48 Field::new("fp_shape", DataType::UInt64, true),
49 ]))
50}
51
52fn edges_schema() -> Arc<Schema> {
53 Arc::new(Schema::new(vec![
54 Field::new("src_id", DataType::Utf8, false),
55 Field::new("dst_id", DataType::Utf8, false),
56 Field::new("kind", DataType::Utf8, false),
57 Field::new("confidence", DataType::Float32, false),
58 Field::new("cross_repo", DataType::Boolean, false),
59 ]))
60}
61
62fn symbols_batch(symbols: &[Symbol]) -> Result<RecordBatch> {
63 let arrays: Vec<ArrayRef> = vec![
64 Arc::new(StringArray::from_iter_values(
65 symbols.iter().map(|s| s.id.render()),
66 )),
67 Arc::new(StringArray::from_iter_values(
68 symbols.iter().map(|s| s.id.scheme.clone()),
69 )),
70 Arc::new(StringArray::from_iter_values(
71 symbols.iter().map(|s| s.id.package_manager.clone()),
72 )),
73 Arc::new(StringArray::from_iter_values(
74 symbols.iter().map(|s| s.id.package.clone()),
75 )),
76 Arc::new(StringArray::from_iter_values(
77 symbols.iter().map(|s| s.id.version.clone()),
78 )),
79 Arc::new(StringArray::from_iter_values(
80 symbols.iter().map(|s| s.id.descriptor.clone()),
81 )),
82 Arc::new(StringArray::from_iter_values(
83 symbols.iter().map(|s| s.name.clone()),
84 )),
85 Arc::new(StringArray::from_iter_values(
86 symbols.iter().map(|s| format!("{:?}", s.kind)),
87 )),
88 Arc::new(StringArray::from_iter_values(
89 symbols.iter().map(|s| format!("{:?}", s.language)),
90 )),
91 Arc::new(StringArray::from_iter_values(
92 symbols.iter().map(|s| s.repo.0.clone()),
93 )),
94 Arc::new(StringArray::from_iter_values(
95 symbols.iter().map(|s| s.file_path.clone()),
96 )),
97 Arc::new(UInt32Array::from_iter_values(
98 symbols.iter().map(|s| s.range.start.line),
99 )),
100 Arc::new(UInt32Array::from_iter_values(
101 symbols.iter().map(|s| s.range.start.column),
102 )),
103 Arc::new(UInt32Array::from_iter_values(
104 symbols.iter().map(|s| s.range.end.line),
105 )),
106 Arc::new(UInt32Array::from_iter_values(
107 symbols.iter().map(|s| s.range.end.column),
108 )),
109 Arc::new(StringArray::from_iter(
110 symbols.iter().map(|s| s.signature.clone()),
111 )),
112 Arc::new(StringArray::from_iter(
113 symbols.iter().map(|s| s.docstring.clone()),
114 )),
115 Arc::new(UInt64Array::from_iter(
116 symbols.iter().map(|s| s.fingerprint.map(|f| f.sig)),
117 )),
118 Arc::new(UInt64Array::from_iter(
119 symbols.iter().map(|s| s.fingerprint.map(|f| f.body)),
120 )),
121 Arc::new(UInt64Array::from_iter(
122 symbols.iter().map(|s| s.fingerprint.map(|f| f.shape)),
123 )),
124 ];
125 RecordBatch::try_new(symbols_schema(), arrays).context("build symbols batch")
126}
127
128fn edges_batch(edges: &[Edge]) -> Result<RecordBatch> {
129 let arrays: Vec<ArrayRef> = vec![
130 Arc::new(StringArray::from_iter_values(
131 edges.iter().map(|e| e.src.render()),
132 )),
133 Arc::new(StringArray::from_iter_values(
134 edges.iter().map(|e| e.dst.render()),
135 )),
136 Arc::new(StringArray::from_iter_values(
137 edges.iter().map(|e| format!("{:?}", e.kind)),
138 )),
139 Arc::new(Float32Array::from_iter_values(
140 edges.iter().map(|e| e.confidence),
141 )),
142 Arc::new(BooleanArray::from(
143 edges.iter().map(|e| e.cross_repo).collect::<Vec<_>>(),
144 )),
145 ];
146 RecordBatch::try_new(edges_schema(), arrays).context("build edges batch")
147}
148
149pub(crate) async fn write_versioned(
154 db: &lancedb::Connection,
155 name: &str,
156 batch: RecordBatch,
157) -> Result<u64> {
158 let tbl = match db.open_table(name).execute().await {
160 Ok(t) => {
161 t.add(batch)
162 .mode(AddDataMode::Overwrite)
163 .execute()
164 .await
165 .with_context(|| format!("overwrite table `{name}`"))?;
166 t
167 }
168 Err(_) => db
169 .create_table(name, batch)
170 .execute()
171 .await
172 .with_context(|| format!("create table `{name}`"))?,
173 };
174 tbl.version()
175 .await
176 .with_context(|| format!("read version of `{name}`"))
177}
178
179pub async fn write_graph(uri: &str, symbols: &[Symbol], edges: &[Edge]) -> Result<(u64, u64)> {
182 let db = crate::storage::connect(uri)
183 .execute()
184 .await
185 .with_context(|| format!("connect lancedb at {uri}"))?;
186 let sv = write_versioned(&db, "symbols", symbols_batch(symbols)?).await?;
187 let ev = write_versioned(&db, "edges", edges_batch(edges)?).await?;
188 Ok((sv, ev))
189}
190
191pub async fn latest_versions(uri: &str) -> Result<(u64, u64)> {
193 let db = crate::storage::connect(uri).execute().await?;
194 let sv = db.open_table("symbols").execute().await?.version().await?;
195 let ev = db.open_table("edges").execute().await?.version().await?;
196 Ok((sv, ev))
197}
198
199use crate::model::{
202 EdgeKind, Fingerprint, GlobalSymbolId, Language, Position, Range, RepoId, SymbolKind,
203};
204use futures::TryStreamExt;
205use lancedb::arrow::arrow_array::Array;
206use lancedb::query::ExecutableQuery;
207
208fn kind_from_str(s: &str) -> SymbolKind {
209 match s {
210 "File" => SymbolKind::File,
211 "Module" => SymbolKind::Module,
212 "Namespace" => SymbolKind::Namespace,
213 "Class" => SymbolKind::Class,
214 "Interface" => SymbolKind::Interface,
215 "Trait" => SymbolKind::Trait,
216 "Enum" => SymbolKind::Enum,
217 "Struct" => SymbolKind::Struct,
218 "Function" => SymbolKind::Function,
219 "Method" => SymbolKind::Method,
220 "Field" => SymbolKind::Field,
221 "Variable" => SymbolKind::Variable,
222 "Constant" => SymbolKind::Constant,
223 "Import" => SymbolKind::Import,
224 "TypeAlias" => SymbolKind::TypeAlias,
225 "Process" => SymbolKind::Process,
226 "Contract" => SymbolKind::Contract,
227 _ => SymbolKind::Variable,
228 }
229}
230
231fn lang_from_str(s: &str) -> Language {
232 match s {
233 "Python" => Language::Python,
234 "TypeScript" => Language::TypeScript,
235 "JavaScript" => Language::JavaScript,
236 "Go" => Language::Go,
237 "Rust" => Language::Rust,
238 "Java" => Language::Java,
239 "Kotlin" => Language::Kotlin,
240 "Ruby" => Language::Ruby,
241 "CSharp" => Language::CSharp,
242 "Cpp" => Language::Cpp,
243 "C" => Language::C,
244 "Php" => Language::Php,
245 "Scala" => Language::Scala,
246 "Swift" => Language::Swift,
247 other => Language::Other(
249 other
250 .strip_prefix("Other(\"")
251 .and_then(|s| s.strip_suffix("\")"))
252 .unwrap_or(other)
253 .to_string(),
254 ),
255 }
256}
257
258fn id_from_rendered(s: &str) -> GlobalSymbolId {
260 let mut it = s.splitn(5, ' ');
261 let mut next = || it.next().unwrap_or(".").to_string();
262 GlobalSymbolId {
263 scheme: next(),
264 package_manager: next(),
265 package: next(),
266 version: next(),
267 descriptor: next(),
268 }
269}
270
271pub(crate) fn str_col<'a>(
272 batch: &'a lancedb::arrow::arrow_array::RecordBatch,
273 name: &str,
274) -> Result<&'a StringArray> {
275 batch
276 .column_by_name(name)
277 .and_then(|c| c.as_any().downcast_ref::<StringArray>())
278 .with_context(|| format!("column `{name}` missing or not Utf8"))
279}
280
281fn u32_col<'a>(
282 batch: &'a lancedb::arrow::arrow_array::RecordBatch,
283 name: &str,
284) -> Result<&'a UInt32Array> {
285 batch
286 .column_by_name(name)
287 .and_then(|c| c.as_any().downcast_ref::<UInt32Array>())
288 .with_context(|| format!("column `{name}` missing or not UInt32"))
289}
290
291pub async fn read_graph(uri: &str) -> Result<(Vec<Symbol>, Vec<Edge>)> {
294 read_graph_at(uri, None).await
295}
296
297pub async fn read_graph_at(uri: &str, at: Option<(u64, u64)>) -> Result<(Vec<Symbol>, Vec<Edge>)> {
299 let db = crate::storage::connect(uri).execute().await?;
300 read_graph_in(&db, at).await
301}
302
303async fn read_graph_in(
304 db: &lancedb::Connection,
305 at: Option<(u64, u64)>,
306) -> Result<(Vec<Symbol>, Vec<Edge>)> {
307 let open = |name: &'static str, v: Option<u64>| async move {
308 let t = db.open_table(name).execute().await?;
309 if let Some(v) = v {
310 t.checkout(v).await?;
311 }
312 anyhow::Ok(t)
313 };
314
315 let mut symbols = Vec::new();
317 let mut st = open("symbols", at.map(|a| a.0))
318 .await?
319 .query()
320 .execute()
321 .await?;
322 while let Some(b) = st.try_next().await? {
323 let (scheme, pm, pkg, ver, desc) = (
324 str_col(&b, "scheme")?,
325 str_col(&b, "package_manager")?,
326 str_col(&b, "package")?,
327 str_col(&b, "version")?,
328 str_col(&b, "descriptor")?,
329 );
330 let (name, kind, lang, repo, file) = (
331 str_col(&b, "name")?,
332 str_col(&b, "kind")?,
333 str_col(&b, "language")?,
334 str_col(&b, "repo")?,
335 str_col(&b, "file_path")?,
336 );
337 let (ls, cs, le, ce) = (
338 u32_col(&b, "line_start")?,
339 u32_col(&b, "col_start")?,
340 u32_col(&b, "line_end")?,
341 u32_col(&b, "col_end")?,
342 );
343 let sig = str_col(&b, "signature")?;
344 let doc = str_col(&b, "docstring")?;
345 let fp = |name: &str| {
347 b.column_by_name(name)
348 .and_then(|c| c.as_any().downcast_ref::<UInt64Array>())
349 };
350 let (fs, fb, fsh) = (fp("fp_sig"), fp("fp_body"), fp("fp_shape"));
351 for i in 0..b.num_rows() {
352 symbols.push(Symbol {
353 id: GlobalSymbolId {
354 scheme: scheme.value(i).to_string(),
355 package_manager: pm.value(i).to_string(),
356 package: pkg.value(i).to_string(),
357 version: ver.value(i).to_string(),
358 descriptor: desc.value(i).to_string(),
359 },
360 name: name.value(i).to_string(),
361 kind: kind_from_str(kind.value(i)),
362 language: lang_from_str(lang.value(i)),
363 repo: RepoId(repo.value(i).to_string()),
364 file_path: file.value(i).to_string(),
365 range: Range {
366 start: Position {
367 line: ls.value(i),
368 column: cs.value(i),
369 },
370 end: Position {
371 line: le.value(i),
372 column: ce.value(i),
373 },
374 },
375 signature: (!sig.is_null(i)).then(|| sig.value(i).to_string()),
376 docstring: (!doc.is_null(i)).then(|| doc.value(i).to_string()),
377 fingerprint: match (fs, fb, fsh) {
378 (Some(a), Some(c), Some(d)) if !a.is_null(i) => Some(Fingerprint {
379 sig: a.value(i),
380 body: c.value(i),
381 shape: d.value(i),
382 }),
383 _ => None,
384 },
385 });
386 }
387 }
388
389 let mut edges = Vec::new();
391 let mut st = open("edges", at.map(|a| a.1))
392 .await?
393 .query()
394 .execute()
395 .await?;
396 while let Some(b) = st.try_next().await? {
397 let src = str_col(&b, "src_id")?;
398 let dst = str_col(&b, "dst_id")?;
399 let kind = str_col(&b, "kind")?;
400 let conf = b
401 .column_by_name("confidence")
402 .and_then(|c| c.as_any().downcast_ref::<Float32Array>())
403 .context("confidence column")?;
404 let cross = b
405 .column_by_name("cross_repo")
406 .and_then(|c| c.as_any().downcast_ref::<BooleanArray>())
407 .context("cross_repo column")?;
408 for i in 0..b.num_rows() {
409 edges.push(Edge {
410 src: id_from_rendered(src.value(i)),
411 dst: id_from_rendered(dst.value(i)),
412 kind: edge_kind_from_str(kind.value(i)),
413 confidence: conf.value(i),
414 cross_repo: cross.value(i),
415 });
416 }
417 }
418
419 Ok((symbols, edges))
420}
421
422fn edge_kind_from_str(s: &str) -> EdgeKind {
423 match s {
424 "Contains" => EdgeKind::Contains,
425 "Imports" => EdgeKind::Imports,
426 "Calls" => EdgeKind::Calls,
427 "Instantiates" => EdgeKind::Instantiates,
428 "Inherits" => EdgeKind::Inherits,
429 "Implements" => EdgeKind::Implements,
430 "References" => EdgeKind::References,
431 "Defines" => EdgeKind::Defines,
432 "Uses" => EdgeKind::Uses,
433 "ParticipatesIn" => EdgeKind::ParticipatesIn,
434 "Serves" => EdgeKind::Serves,
435 "Consumes" => EdgeKind::Consumes,
436 "Publishes" => EdgeKind::Publishes,
437 "Subscribes" => EdgeKind::Subscribes,
438 "Reads" => EdgeKind::Reads,
439 "Writes" => EdgeKind::Writes,
440 _ => EdgeKind::References,
441 }
442}
443
444pub async fn graph_versions(uri: &str) -> Result<Vec<(u64, u64, String)>> {
448 let db = crate::storage::connect(uri).execute().await?;
449 let s = db.open_table("symbols").execute().await?;
450 let e = db.open_table("edges").execute().await?;
451 versions_of(&s, &e).await
452}
453
454async fn versions_of(s: &lancedb::Table, e: &lancedb::Table) -> Result<Vec<(u64, u64, String)>> {
455 let (s, e) = (s.list_versions().await?, e.list_versions().await?);
456 Ok(s.iter()
457 .zip(e.iter())
458 .map(|(a, b)| (a.version, b.version, a.timestamp.to_rfc3339()))
459 .collect())
460}
461
462pub fn graph_versions_blocking(uri: &str) -> Result<Vec<(u64, u64, String)>> {
463 crate::rt::block_on(graph_versions(uri))
464}
465
466pub fn read_graph_blocking(uri: &str) -> Result<(Vec<Symbol>, Vec<Edge>)> {
467 crate::rt::block_on(read_graph(uri))
468}
469
470pub async fn count_rows(uri: &str) -> Result<(usize, usize)> {
472 let db = crate::storage::connect(uri).execute().await?;
473 let symbols = db
474 .open_table("symbols")
475 .execute()
476 .await?
477 .count_rows(None)
478 .await?;
479 let edges = db
480 .open_table("edges")
481 .execute()
482 .await?
483 .count_rows(None)
484 .await?;
485 Ok((symbols, edges))
486}
487
488pub fn write_graph_blocking(uri: &str, symbols: &[Symbol], edges: &[Edge]) -> Result<(u64, u64)> {
491 crate::rt::block_on(write_graph(uri, symbols, edges))
492}
493
494pub fn count_rows_blocking(uri: &str) -> Result<(usize, usize)> {
496 crate::rt::block_on(count_rows(uri))
497}
498
499pub fn latest_versions_blocking(uri: &str) -> Result<(u64, u64)> {
501 crate::rt::block_on(latest_versions(uri))
502}
503
504pub async fn write_repo_meta(uri: &str, repo: &str, commit: &str) -> Result<()> {
510 let schema = Arc::new(Schema::new(vec![
511 Field::new("repo", DataType::Utf8, false),
512 Field::new("commit", DataType::Utf8, false),
513 ]));
514 let batch = RecordBatch::try_new(
515 schema,
516 vec![
517 Arc::new(StringArray::from(vec![repo])) as ArrayRef,
518 Arc::new(StringArray::from(vec![commit])) as ArrayRef,
519 ],
520 )
521 .context("build repo_meta batch")?;
522 let db = crate::storage::connect(uri).execute().await?;
523 write_versioned(&db, "repo_meta", batch).await?;
524 Ok(())
525}
526
527pub async fn read_repo_meta(uri: &str) -> Result<Option<String>> {
529 let db = crate::storage::connect(uri).execute().await?;
530 let Ok(tbl) = db.open_table("repo_meta").execute().await else {
531 return Ok(None);
532 };
533 let mut st = tbl.query().execute().await?;
534 while let Some(b) = st.try_next().await? {
535 if b.num_rows() > 0 {
536 return Ok(Some(str_col(&b, "commit")?.value(0).to_string()));
537 }
538 }
539 Ok(None)
540}
541
542pub fn write_repo_meta_blocking(uri: &str, repo: &str, commit: &str) -> Result<()> {
544 crate::rt::block_on(write_repo_meta(uri, repo, commit))
545}
546
547pub async fn write_repo_commits(uri: &str, repos: &[IndexedRepo]) -> Result<()> {
550 if repos.is_empty() {
551 return Ok(());
552 }
553 let schema = Arc::new(Schema::new(vec![
554 Field::new("repo", DataType::Utf8, false),
555 Field::new("commit", DataType::Utf8, false),
556 Field::new("remote", DataType::Utf8, true),
557 Field::new("branch", DataType::Utf8, true),
558 ]));
559 let batch = RecordBatch::try_new(
560 schema,
561 vec![
562 Arc::new(StringArray::from_iter_values(
563 repos.iter().map(|r| r.name.clone()),
564 )) as ArrayRef,
565 Arc::new(StringArray::from_iter_values(
566 repos.iter().map(|r| r.commit.clone()),
567 )) as ArrayRef,
568 Arc::new(StringArray::from_iter(
569 repos.iter().map(|r| r.remote.clone()),
570 )) as ArrayRef,
571 Arc::new(StringArray::from_iter(
572 repos.iter().map(|r| r.branch.clone()),
573 )) as ArrayRef,
574 ],
575 )
576 .context("build repo_commits batch")?;
577 let db = crate::storage::connect(uri).execute().await?;
578 write_versioned(&db, "repo_meta", batch).await?;
579 Ok(())
580}
581
582#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
584pub struct LoggedCommit {
585 pub sha: String,
586 pub time: i64,
588 #[serde(default, skip_serializing_if = "Option::is_none")]
589 pub mr: Option<u64>,
590 pub files: Vec<String>,
592}
593
594#[derive(Debug, Clone, Default, PartialEq, serde::Serialize, serde::Deserialize)]
596pub struct RepoLog {
597 pub commits: Vec<LoggedCommit>,
599 #[serde(default, skip_serializing_if = "std::collections::BTreeMap::is_empty")]
601 pub landed: std::collections::BTreeMap<String, Vec<String>>,
602}
603
604pub type RepoLogs = HashMap<String, RepoLog>;
606
607pub const LOG_DEPTH: usize = 1000;
609
610impl From<crate::git::CommitFiles> for LoggedCommit {
611 fn from(c: crate::git::CommitFiles) -> Self {
612 LoggedCommit {
613 sha: c.sha,
614 time: c.time,
615 mr: c.mr,
616 files: c.changed,
617 }
618 }
619}
620
621pub async fn write_repo_logs(uri: &str, logs: &RepoLogs) -> Result<()> {
623 if logs.is_empty() {
624 return Ok(());
625 }
626 let schema = Arc::new(Schema::new(vec![
627 Field::new("repo", DataType::Utf8, false),
628 Field::new("log", DataType::Utf8, false),
629 ]));
630 let (names, json): (Vec<String>, Vec<String>) = logs
631 .iter()
632 .map(|(r, l)| Ok((r.clone(), serde_json::to_string(l)?)))
633 .collect::<Result<Vec<_>>>()?
634 .into_iter()
635 .unzip();
636 let batch = RecordBatch::try_new(
637 schema,
638 vec![
639 Arc::new(StringArray::from(names)) as ArrayRef,
640 Arc::new(StringArray::from(json)) as ArrayRef,
641 ],
642 )
643 .context("build repo_log batch")?;
644 let db = crate::storage::connect(uri).execute().await?;
645 write_versioned(&db, "repo_log", batch).await?;
646 Ok(())
647}
648
649pub async fn read_repo_logs(uri: &str) -> Result<RepoLogs> {
651 let db = crate::storage::connect(uri).execute().await?;
652 let Ok(tbl) = db.open_table("repo_log").execute().await else {
653 return Ok(RepoLogs::new());
654 };
655 let mut out = RepoLogs::new();
656 let mut st = tbl.query().execute().await?;
657 while let Some(b) = st.try_next().await? {
658 let (repo, log) = (str_col(&b, "repo")?, str_col(&b, "log")?);
659 for i in 0..b.num_rows() {
660 if let Ok(l) = serde_json::from_str(log.value(i)) {
661 out.insert(repo.value(i).to_string(), l);
662 }
663 }
664 }
665 Ok(out)
666}
667
668pub fn read_repo_logs_blocking(uri: &str) -> Result<RepoLogs> {
670 crate::rt::block_on(read_repo_logs(uri))
671}
672
673#[derive(Debug, Clone, PartialEq, serde::Serialize)]
675pub struct IndexedRepo {
676 pub name: String,
677 pub commit: String,
678 pub remote: Option<String>,
679 pub branch: Option<String>,
680}
681
682pub async fn read_indexed_repos(uri: &str) -> Result<Vec<IndexedRepo>> {
684 let db = crate::storage::connect(uri).execute().await?;
685 let Ok(tbl) = db.open_table("repo_meta").execute().await else {
686 return Ok(Vec::new());
687 };
688 let mut out = Vec::new();
689 let mut st = tbl.query().execute().await?;
690 while let Some(b) = st.try_next().await? {
691 let repo = str_col(&b, "repo")?;
692 let commit = str_col(&b, "commit")?;
693 let remote = str_col(&b, "remote").ok();
694 let branch = str_col(&b, "branch").ok();
695 for i in 0..b.num_rows() {
696 out.push(IndexedRepo {
697 name: repo.value(i).to_string(),
698 commit: commit.value(i).to_string(),
699 remote: remote
700 .filter(|r| r.is_valid(i))
701 .map(|r| r.value(i).to_string()),
702 branch: branch
703 .filter(|r| r.is_valid(i))
704 .map(|r| r.value(i).to_string()),
705 });
706 }
707 }
708 out.sort_by(|a, b| a.name.cmp(&b.name));
709 Ok(out)
710}
711
712pub async fn read_repo_commits(uri: &str) -> Result<std::collections::HashMap<String, String>> {
714 Ok(read_indexed_repos(uri)
715 .await?
716 .into_iter()
717 .map(|r| (r.name, r.commit))
718 .collect())
719}
720
721pub fn write_repo_commits_blocking(uri: &str, repos: &[IndexedRepo]) -> Result<()> {
723 crate::rt::block_on(write_repo_commits(uri, repos))
724}
725
726pub fn read_repo_commits_blocking(uri: &str) -> Result<std::collections::HashMap<String, String>> {
728 crate::rt::block_on(read_repo_commits(uri))
729}
730
731pub fn read_repo_meta_blocking(uri: &str) -> Result<Option<String>> {
733 crate::rt::block_on(read_repo_meta(uri))
734}
735
736use lance_index::scalar::FullTextSearchQuery;
739use lancedb::index::scalar::FtsIndexBuilder;
740use lancedb::index::vector::IvfHnswFlatIndexBuilder;
741use lancedb::index::Index;
742use lancedb::query::{QueryBase, QueryExecutionOptions, Select};
743
744fn search_schema(dim: usize) -> Arc<Schema> {
745 let item = Arc::new(Field::new("item", DataType::Float32, true));
746 Arc::new(Schema::new(vec![
747 Field::new("symbol_id", DataType::Utf8, false),
748 Field::new("text", DataType::Utf8, false),
749 Field::new("vector", DataType::FixedSizeList(item, dim as i32), false),
750 ]))
751}
752
753pub async fn write_search_table(uri: &str, rows: &[SearchRow]) -> Result<u64> {
757 let Some((_, _, first)) = rows.first() else {
758 return Ok(0);
759 };
760 let dim = first.len();
761 let ids = StringArray::from_iter_values(rows.iter().map(|(id, _, _)| id.render()));
762 let texts = StringArray::from_iter_values(rows.iter().map(|(_, t, _)| t.clone()));
763 let flat: Vec<f32> = rows
764 .iter()
765 .flat_map(|(_, _, v)| v.iter().copied())
766 .collect();
767 let item = Arc::new(Field::new("item", DataType::Float32, true));
768 let vectors =
769 FixedSizeListArray::try_new(item, dim as i32, Arc::new(Float32Array::from(flat)), None)
770 .context("build vector column")?;
771 let batch = RecordBatch::try_new(
772 search_schema(dim),
773 vec![
774 Arc::new(ids) as ArrayRef,
775 Arc::new(texts) as ArrayRef,
776 Arc::new(vectors) as ArrayRef,
777 ],
778 )
779 .context("build search batch")?;
780
781 let db = crate::storage::connect(uri).execute().await?;
782 let version = write_versioned(&db, "search", batch).await?;
783
784 let tbl = db.open_table("search").execute().await?;
786 tbl.create_index(&["text"], Index::FTS(FtsIndexBuilder::default()))
787 .execute()
788 .await
789 .context("create FTS index")?;
790
791 let parts = (rows.len() / 20_000).max(1) as u32;
794 let ann = IvfHnswFlatIndexBuilder::default().num_partitions(parts);
795 if let Err(e) = tbl
796 .create_index(&["vector"], Index::IvfHnswFlat(ann))
797 .execute()
798 .await
799 {
800 eprintln!("comindly: vector ANN index skipped (non-fatal): {e}");
801 }
802 Ok(version)
803}
804
805pub async fn read_search_vectors(uri: &str) -> Result<Option<Vec<VectorRow>>> {
807 let db = crate::storage::connect(uri).execute().await?;
808 let Ok(tbl) = db.open_table("search").execute().await else {
809 return Ok(None);
810 };
811 let mut out = Vec::new();
812 let mut st = tbl
813 .query()
814 .select(Select::columns(&["symbol_id", "vector"]))
815 .execute()
816 .await?;
817 while let Some(b) = st.try_next().await? {
818 let ids = str_col(&b, "symbol_id")?;
819 let vecs = b
820 .column_by_name("vector")
821 .and_then(|c| c.as_any().downcast_ref::<FixedSizeListArray>())
822 .context("vector column")?;
823 for i in 0..b.num_rows() {
824 let sub = vecs.value(i);
825 let f = sub.as_any().downcast_ref::<Float32Array>().context("f32")?;
826 out.push((id_from_rendered(ids.value(i)), f.values().to_vec()));
827 }
828 }
829 Ok(Some(out))
830}
831
832pub async fn hybrid_search(
835 uri: &str,
836 query_text: &str,
837 query_vector: Vec<f32>,
838 k: usize,
839) -> Result<Vec<(String, f32)>> {
840 let db = crate::storage::connect(uri).execute().await?;
841 let tbl = db.open_table("search").execute().await?;
842 hybrid_in(&tbl, query_text, query_vector, k).await
843}
844
845async fn hybrid_in(
846 tbl: &lancedb::Table,
847 query_text: &str,
848 query_vector: Vec<f32>,
849 k: usize,
850) -> Result<Vec<(String, f32)>> {
851 let mut st = tbl
852 .query()
853 .select(lancedb::query::Select::columns(&["symbol_id"]))
854 .full_text_search(FullTextSearchQuery::new(query_text.to_owned()))
855 .nearest_to(query_vector)?
856 .limit(k)
857 .execute_hybrid(QueryExecutionOptions::default())
859 .await?;
860
861 let mut out = Vec::new();
862 let mut rank = 0usize;
863 while let Some(b) = st.try_next().await? {
864 let ids = str_col(&b, "symbol_id")?;
865 let scores = b
867 .column_by_name("_relevance_score")
868 .and_then(|c| c.as_any().downcast_ref::<Float32Array>());
869 for i in 0..b.num_rows() {
870 let score = scores
871 .map(|s| s.value(i))
872 .unwrap_or(1.0 / (rank as f32 + 1.0));
873 out.push((ids.value(i).to_string(), score));
874 rank += 1;
875 }
876 }
877 Ok(out)
878}
879
880pub struct Tables {
883 db: lancedb::Connection,
884 open: tokio::sync::Mutex<std::collections::HashMap<&'static str, lancedb::Table>>,
885}
886
887impl Tables {
888 pub async fn connect(uri: &str, fresh: std::time::Duration) -> Result<Tables> {
889 let db = crate::storage::connect(uri)
890 .read_consistency_interval(fresh)
891 .execute()
892 .await
893 .with_context(|| format!("connect lancedb at {uri}"))?;
894 Ok(Tables {
895 db,
896 open: Default::default(),
897 })
898 }
899
900 async fn table(&self, name: &'static str) -> Result<lancedb::Table> {
901 let mut open = self.open.lock().await;
902 if let Some(t) = open.get(name) {
903 return Ok(t.clone());
904 }
905 let t = self.db.open_table(name).execute().await?;
906 open.insert(name, t.clone());
907 Ok(t)
908 }
909
910 pub async fn hybrid_search(
911 &self,
912 query_text: &str,
913 query_vector: Vec<f32>,
914 k: usize,
915 ) -> Result<Vec<(String, f32)>> {
916 hybrid_in(&self.table("search").await?, query_text, query_vector, k).await
917 }
918
919 pub async fn versions(&self) -> Result<Vec<(u64, u64, String)>> {
920 versions_of(&self.table("symbols").await?, &self.table("edges").await?).await
921 }
922
923 pub async fn graph_at(&self, at: (u64, u64)) -> Result<(Vec<Symbol>, Vec<Edge>)> {
926 read_graph_in(&self.db, Some(at)).await
927 }
928}
929
930pub fn write_search_table_blocking(uri: &str, rows: &[SearchRow]) -> Result<u64> {
932 crate::rt::block_on(write_search_table(uri, rows))
933}
934pub fn read_search_vectors_blocking(uri: &str) -> Result<Option<Vec<VectorRow>>> {
935 crate::rt::block_on(read_search_vectors(uri))
936}
937pub fn hybrid_search_blocking(
938 uri: &str,
939 text: &str,
940 vector: Vec<f32>,
941 k: usize,
942) -> Result<Vec<(String, f32)>> {
943 crate::rt::block_on(hybrid_search(uri, text, vector, k))
944}
945
946fn enrichment_schema() -> Arc<Schema> {
949 Arc::new(Schema::new(vec![
950 Field::new("symbol_id", DataType::Utf8, false),
951 Field::new("summary", DataType::Utf8, false),
952 Field::new("queries", DataType::Utf8, false), ]))
954}
955
956pub async fn write_enrichment(uri: &str, rows: &[TextRow]) -> Result<u64> {
959 if rows.is_empty() {
960 return Ok(0);
961 }
962 let ids = StringArray::from_iter_values(rows.iter().map(|(id, _, _)| id.render()));
963 let summaries = StringArray::from_iter_values(rows.iter().map(|(_, s, _)| s.clone()));
964 let queries = StringArray::from_iter_values(
965 rows.iter()
966 .map(|(_, _, q)| serde_json::to_string(q).unwrap_or_else(|_| "[]".into())),
967 );
968 let batch = RecordBatch::try_new(
969 enrichment_schema(),
970 vec![
971 Arc::new(ids) as ArrayRef,
972 Arc::new(summaries) as ArrayRef,
973 Arc::new(queries) as ArrayRef,
974 ],
975 )
976 .context("build enrichment batch")?;
977 let db = crate::storage::connect(uri).execute().await?;
978 write_versioned(&db, "enrichment", batch).await
979}
980
981pub async fn read_enrichment(uri: &str) -> Result<Option<Vec<TextRow>>> {
983 let db = crate::storage::connect(uri).execute().await?;
984 let Ok(tbl) = db.open_table("enrichment").execute().await else {
985 return Ok(None);
986 };
987 let mut out = Vec::new();
988 let mut st = tbl.query().execute().await?;
989 while let Some(b) = st.try_next().await? {
990 let ids = str_col(&b, "symbol_id")?;
991 let summaries = str_col(&b, "summary")?;
992 let queries = str_col(&b, "queries")?;
993 for i in 0..b.num_rows() {
994 let q: Vec<String> = serde_json::from_str(queries.value(i)).unwrap_or_default();
995 out.push((
996 id_from_rendered(ids.value(i)),
997 summaries.value(i).to_string(),
998 q,
999 ));
1000 }
1001 }
1002 Ok(Some(out))
1003}
1004
1005pub fn write_enrichment_blocking(uri: &str, rows: &[TextRow]) -> Result<u64> {
1007 crate::rt::block_on(write_enrichment(uri, rows))
1008}
1009
1010pub fn read_enrichment_blocking(uri: &str) -> Result<Option<Vec<TextRow>>> {
1012 crate::rt::block_on(read_enrichment(uri))
1013}
1014
1015pub async fn write_keys(uri: &str, table: &str, keys: &[(String, String)]) -> Result<u64> {
1018 let schema = Arc::new(Schema::new(vec![
1019 Field::new("symbol_id", DataType::Utf8, false),
1020 Field::new("key", DataType::Utf8, false),
1021 ]));
1022 let batch = RecordBatch::try_new(
1023 schema,
1024 vec![
1025 Arc::new(StringArray::from_iter_values(
1026 keys.iter().map(|(i, _)| i.as_str()),
1027 )) as ArrayRef,
1028 Arc::new(StringArray::from_iter_values(
1029 keys.iter().map(|(_, k)| k.as_str()),
1030 )) as ArrayRef,
1031 ],
1032 )?;
1033 let db = crate::storage::connect(uri).execute().await?;
1034 write_versioned(&db, table, batch).await
1035}
1036
1037pub async fn read_keys(uri: &str, table: &str) -> Result<HashMap<String, String>> {
1039 let db = crate::storage::connect(uri).execute().await?;
1040 let tbl = db.open_table(table).execute().await?;
1041 let mut out = HashMap::new();
1042 let mut st = tbl.query().execute().await?;
1043 while let Some(b) = st.try_next().await? {
1044 let (ids, keys) = (str_col(&b, "symbol_id")?, str_col(&b, "key")?);
1045 for i in 0..b.num_rows() {
1046 out.insert(ids.value(i).to_string(), keys.value(i).to_string());
1047 }
1048 }
1049 Ok(out)
1050}
1051
1052pub fn write_keys_blocking(uri: &str, table: &str, keys: &[(String, String)]) -> Result<u64> {
1053 crate::rt::block_on(write_keys(uri, table, keys))
1054}
1055
1056pub fn read_keys_blocking(uri: &str, table: &str) -> HashMap<String, String> {
1057 crate::rt::block_on(read_keys(uri, table)).unwrap_or_default()
1058}
1059
1060fn flows_schema() -> Arc<Schema> {
1061 Arc::new(Schema::new(vec![
1062 Field::new("symbol_id", DataType::Utf8, false),
1063 Field::new("narration", DataType::Utf8, false),
1064 Field::new("queries", DataType::Utf8, false), ]))
1066}
1067
1068pub async fn write_flows(uri: &str, rows: &[TextRow]) -> Result<u64> {
1070 if rows.is_empty() {
1071 return Ok(0);
1072 }
1073 let ids = StringArray::from_iter_values(rows.iter().map(|(id, _, _)| id.render()));
1074 let narr = StringArray::from_iter_values(rows.iter().map(|(_, n, _)| n.clone()));
1075 let queries = StringArray::from_iter_values(
1076 rows.iter()
1077 .map(|(_, _, q)| serde_json::to_string(q).unwrap_or_else(|_| "[]".into())),
1078 );
1079 let batch = RecordBatch::try_new(
1080 flows_schema(),
1081 vec![
1082 Arc::new(ids) as ArrayRef,
1083 Arc::new(narr) as ArrayRef,
1084 Arc::new(queries) as ArrayRef,
1085 ],
1086 )
1087 .context("build flows batch")?;
1088 let db = crate::storage::connect(uri).execute().await?;
1089 write_versioned(&db, "flows", batch).await
1090}
1091
1092pub async fn read_flows(uri: &str) -> Result<Option<Vec<TextRow>>> {
1094 let db = crate::storage::connect(uri).execute().await?;
1095 let Ok(tbl) = db.open_table("flows").execute().await else {
1096 return Ok(None);
1097 };
1098 let mut out = Vec::new();
1099 let mut st = tbl.query().execute().await?;
1100 while let Some(b) = st.try_next().await? {
1101 let ids = str_col(&b, "symbol_id")?;
1102 let narr = str_col(&b, "narration")?;
1103 let queries = str_col(&b, "queries")?;
1104 for i in 0..b.num_rows() {
1105 let q: Vec<String> = serde_json::from_str(queries.value(i)).unwrap_or_default();
1106 out.push((id_from_rendered(ids.value(i)), narr.value(i).to_string(), q));
1107 }
1108 }
1109 Ok(Some(out))
1110}
1111
1112pub fn write_flows_blocking(uri: &str, rows: &[TextRow]) -> Result<u64> {
1114 crate::rt::block_on(write_flows(uri, rows))
1115}
1116pub fn read_flows_blocking(uri: &str) -> Result<Option<Vec<TextRow>>> {
1117 crate::rt::block_on(read_flows(uri))
1118}
1119
1120pub async fn write_style_guide(uri: &str, rows: &[(String, String)]) -> Result<()> {
1123 let schema = Arc::new(Schema::new(vec![
1124 Field::new("repo", DataType::Utf8, false),
1125 Field::new("content", DataType::Utf8, false),
1126 ]));
1127 let repos: Vec<&str> = rows.iter().map(|(r, _)| r.as_str()).collect();
1128 let contents: Vec<&str> = rows.iter().map(|(_, c)| c.as_str()).collect();
1129 let batch = RecordBatch::try_new(
1130 schema,
1131 vec![
1132 Arc::new(StringArray::from(repos)) as ArrayRef,
1133 Arc::new(StringArray::from(contents)) as ArrayRef,
1134 ],
1135 )
1136 .context("build style_guide batch")?;
1137 let db = crate::storage::connect(uri).execute().await?;
1138 write_versioned(&db, "style_guide", batch).await?;
1139 Ok(())
1140}
1141
1142pub async fn read_style_guide(uri: &str) -> Result<Vec<(String, String)>> {
1144 let db = crate::storage::connect(uri).execute().await?;
1145 let Ok(tbl) = db.open_table("style_guide").execute().await else {
1146 return Ok(vec![]);
1147 };
1148 let mut st = tbl.query().execute().await?;
1149 let mut out = Vec::new();
1150 while let Some(b) = st.try_next().await? {
1151 let repo = str_col(&b, "repo")?;
1152 let content = str_col(&b, "content")?;
1153 for i in 0..b.num_rows() {
1154 out.push((repo.value(i).to_string(), content.value(i).to_string()));
1155 }
1156 }
1157 Ok(out)
1158}
1159
1160pub fn write_style_guide_blocking(uri: &str, rows: &[(String, String)]) -> Result<()> {
1162 crate::rt::block_on(write_style_guide(uri, rows))
1163}
1164
1165pub fn read_style_guide_blocking(uri: &str) -> Result<Vec<(String, String)>> {
1167 crate::rt::block_on(read_style_guide(uri))
1168}
1169
1170pub async fn write_facts(uri: &str, facts: &[FileFacts]) -> Result<()> {
1172 let schema = Arc::new(Schema::new(vec![
1173 Field::new("repo", DataType::Utf8, false),
1174 Field::new("path", DataType::Utf8, false),
1175 Field::new("facts", DataType::Utf8, false),
1176 ]));
1177 let json: Vec<String> = facts
1178 .iter()
1179 .map(serde_json::to_string)
1180 .collect::<Result<_, _>>()
1181 .context("serialize facts")?;
1182 let batch = RecordBatch::try_new(
1183 schema,
1184 vec![
1185 Arc::new(StringArray::from_iter_values(
1186 facts.iter().map(|f| f.repo.as_str()),
1187 )) as ArrayRef,
1188 Arc::new(StringArray::from_iter_values(
1189 facts.iter().map(|f| f.path.as_str()),
1190 )),
1191 Arc::new(StringArray::from(json)),
1192 ],
1193 )
1194 .context("build file_facts batch")?;
1195 let db = crate::storage::connect(uri).execute().await?;
1196 write_versioned(&db, "file_facts", batch).await?;
1197 Ok(())
1198}
1199
1200pub async fn read_facts(uri: &str) -> Result<Option<Vec<FileFacts>>> {
1202 let db = crate::storage::connect(uri).execute().await?;
1203 let Ok(tbl) = db.open_table("file_facts").execute().await else {
1204 return Ok(None);
1205 };
1206 let mut st = tbl.query().execute().await?;
1207 let mut out = Vec::new();
1208 while let Some(b) = st.try_next().await? {
1209 let col = str_col(&b, "facts")?;
1210 for i in 0..b.num_rows() {
1211 out.push(serde_json::from_str(col.value(i)).context("parse facts row")?);
1212 }
1213 }
1214 Ok(Some(out))
1215}
1216
1217pub async fn write_recipes(uri: &str, rows: &[crate::conventions::Convention]) -> Result<()> {
1219 let schema = Arc::new(Schema::new(vec![
1220 Field::new("repo", DataType::Utf8, false),
1221 Field::new("row", DataType::Utf8, false),
1222 ]));
1223 let json: Vec<String> = rows
1224 .iter()
1225 .map(serde_json::to_string)
1226 .collect::<Result<_, _>>()
1227 .context("serialize recipes")?;
1228 let batch = RecordBatch::try_new(
1229 schema,
1230 vec![
1231 Arc::new(StringArray::from_iter_values(
1232 rows.iter().map(|r| r.repo.as_str()),
1233 )) as ArrayRef,
1234 Arc::new(StringArray::from(json)),
1235 ],
1236 )
1237 .context("build recipes batch")?;
1238 let db = crate::storage::connect(uri).execute().await?;
1239 write_versioned(&db, "recipes", batch).await?;
1240 Ok(())
1241}
1242
1243pub async fn read_recipes(uri: &str) -> Result<Vec<crate::conventions::Convention>> {
1244 let db = crate::storage::connect(uri).execute().await?;
1245 let Ok(tbl) = db.open_table("recipes").execute().await else {
1246 return Ok(Vec::new());
1247 };
1248 let mut st = tbl.query().execute().await?;
1249 let mut out = Vec::new();
1250 while let Some(b) = st.try_next().await? {
1251 let col = str_col(&b, "row")?;
1252 for i in 0..b.num_rows() {
1253 out.push(serde_json::from_str(col.value(i)).context("parse recipe row")?);
1254 }
1255 }
1256 Ok(out)
1257}
1258
1259pub fn write_recipes_blocking(uri: &str, rows: &[crate::conventions::Convention]) -> Result<()> {
1260 crate::rt::block_on(write_recipes(uri, rows))
1261}
1262
1263pub fn read_recipes_blocking(uri: &str) -> Result<Vec<crate::conventions::Convention>> {
1264 crate::rt::block_on(read_recipes(uri))
1265}
1266
1267pub async fn write_packages(uri: &str, rows: &[crate::manifests::PackageRow]) -> Result<()> {
1269 let schema = Arc::new(Schema::new(vec![
1270 Field::new("repo", DataType::Utf8, false),
1271 Field::new("row", DataType::Utf8, false),
1272 ]));
1273 let json: Vec<String> = rows
1274 .iter()
1275 .map(serde_json::to_string)
1276 .collect::<Result<_, _>>()
1277 .context("serialize packages")?;
1278 let batch = RecordBatch::try_new(
1279 schema,
1280 vec![
1281 Arc::new(StringArray::from_iter_values(
1282 rows.iter().map(|r| r.repo.as_str()),
1283 )) as ArrayRef,
1284 Arc::new(StringArray::from(json)),
1285 ],
1286 )
1287 .context("build packages batch")?;
1288 let db = crate::storage::connect(uri).execute().await?;
1289 write_versioned(&db, "packages", batch).await?;
1290 Ok(())
1291}
1292
1293pub async fn read_packages(uri: &str) -> Result<Vec<crate::manifests::PackageRow>> {
1294 let db = crate::storage::connect(uri).execute().await?;
1295 let Ok(tbl) = db.open_table("packages").execute().await else {
1296 return Ok(Vec::new());
1297 };
1298 let mut st = tbl.query().execute().await?;
1299 let mut out = Vec::new();
1300 while let Some(b) = st.try_next().await? {
1301 let col = str_col(&b, "row")?;
1302 for i in 0..b.num_rows() {
1303 out.push(serde_json::from_str(col.value(i)).context("parse package row")?);
1304 }
1305 }
1306 Ok(out)
1307}
1308
1309pub fn write_packages_blocking(uri: &str, rows: &[crate::manifests::PackageRow]) -> Result<()> {
1310 crate::rt::block_on(write_packages(uri, rows))
1311}
1312
1313pub fn read_packages_blocking(uri: &str) -> Result<Vec<crate::manifests::PackageRow>> {
1314 crate::rt::block_on(read_packages(uri))
1315}
1316
1317pub fn write_facts_blocking(uri: &str, facts: &[FileFacts]) -> Result<()> {
1318 crate::rt::block_on(write_facts(uri, facts))
1319}
1320
1321pub fn read_facts_blocking(uri: &str) -> Result<Option<Vec<FileFacts>>> {
1322 crate::rt::block_on(read_facts(uri))
1323}
1324
1325pub async fn write_sources(
1331 uri: &str,
1332 files: &[crate::sources::SourceFile],
1333 merge: bool,
1334) -> Result<(usize, usize)> {
1335 use std::collections::HashSet;
1336 let db = crate::storage::connect(uri).execute().await?;
1337 let blobs = db.open_table("source_blobs").execute().await.ok();
1338 let mut have: HashSet<String> = HashSet::new();
1339 if let Some(t) = &blobs {
1340 let mut st = t
1341 .query()
1342 .select(Select::columns(&["hash"]))
1343 .execute()
1344 .await?;
1345 while let Some(b) = st.try_next().await? {
1346 let h = str_col(&b, "hash")?;
1347 have.extend((0..b.num_rows()).map(|i| h.value(i).to_string()));
1348 }
1349 }
1350 let mut fresh: Vec<&crate::sources::SourceFile> = Vec::new();
1351 for f in files {
1352 if have.insert(f.hash.clone()) {
1353 fresh.push(f);
1354 }
1355 }
1356 if !fresh.is_empty() {
1357 let batch = RecordBatch::try_new(
1358 Arc::new(Schema::new(vec![
1359 Field::new("hash", DataType::Utf8, false),
1360 Field::new("content", DataType::Utf8, false),
1361 ])),
1362 vec![
1363 Arc::new(StringArray::from_iter_values(fresh.iter().map(|f| &f.hash))) as ArrayRef,
1364 Arc::new(StringArray::from_iter_values(
1365 fresh.iter().map(|f| &f.content),
1366 )) as ArrayRef,
1367 ],
1368 )?;
1369 let tbl = match blobs {
1370 Some(t) => {
1371 t.add(batch).mode(AddDataMode::Append).execute().await?;
1372 t
1373 }
1374 None => db.create_table("source_blobs", batch).execute().await?,
1375 };
1376 tbl.create_index(&["hash"], Index::BTree(Default::default()))
1377 .replace(true)
1378 .execute()
1379 .await
1380 .context("index source_blobs")?;
1381 }
1382
1383 let covered: HashSet<&str> = files.iter().map(|f| f.repo.as_str()).collect();
1384 let mut rows: Vec<(String, String, String)> = if merge {
1385 read_source_files(uri)
1386 .await
1387 .unwrap_or_default()
1388 .into_iter()
1389 .filter(|(r, _, _)| !covered.contains(r.as_str()))
1390 .collect()
1391 } else {
1392 Vec::new()
1393 };
1394 rows.extend(
1395 files
1396 .iter()
1397 .map(|f| (f.repo.clone(), f.path.clone(), f.hash.clone())),
1398 );
1399 let batch = RecordBatch::try_new(
1400 Arc::new(Schema::new(vec![
1401 Field::new("repo", DataType::Utf8, false),
1402 Field::new("path", DataType::Utf8, false),
1403 Field::new("hash", DataType::Utf8, false),
1404 ])),
1405 vec![
1406 Arc::new(StringArray::from_iter_values(rows.iter().map(|r| &r.0))) as ArrayRef,
1407 Arc::new(StringArray::from_iter_values(rows.iter().map(|r| &r.1))) as ArrayRef,
1408 Arc::new(StringArray::from_iter_values(rows.iter().map(|r| &r.2))) as ArrayRef,
1409 ],
1410 )?;
1411 write_versioned(&db, "source_files", batch).await?;
1412 Ok((files.len(), fresh.len()))
1413}
1414
1415pub async fn read_source_files(uri: &str) -> Result<Vec<(String, String, String)>> {
1417 let db = crate::storage::connect(uri).execute().await?;
1418 let tbl = db.open_table("source_files").execute().await?;
1419 let mut out = Vec::new();
1420 let mut st = tbl.query().execute().await?;
1421 while let Some(b) = st.try_next().await? {
1422 let (r, p, h) = (
1423 str_col(&b, "repo")?,
1424 str_col(&b, "path")?,
1425 str_col(&b, "hash")?,
1426 );
1427 for i in 0..b.num_rows() {
1428 out.push((
1429 r.value(i).to_string(),
1430 p.value(i).to_string(),
1431 h.value(i).to_string(),
1432 ));
1433 }
1434 }
1435 Ok(out)
1436}
1437
1438pub async fn read_source(uri: &str, repo: &str, path: &str) -> Result<Option<String>> {
1440 let quote = |s: &str| format!("'{}'", s.replace('\'', "''"));
1441 let db = crate::storage::connect(uri).execute().await?;
1442 let Ok(files) = db.open_table("source_files").execute().await else {
1443 return Ok(None);
1444 };
1445 let mut st = files
1446 .query()
1447 .only_if(format!("repo = {} AND path = {}", quote(repo), quote(path)))
1448 .select(Select::columns(&["hash"]))
1449 .limit(1)
1450 .execute()
1451 .await?;
1452 let Some(hash) = first_str(&mut st, "hash").await? else {
1453 return Ok(None);
1454 };
1455 let blobs = db.open_table("source_blobs").execute().await?;
1456 let mut st = blobs
1457 .query()
1458 .only_if(format!("hash = {}", quote(&hash)))
1459 .select(Select::columns(&["content"]))
1460 .limit(1)
1461 .execute()
1462 .await?;
1463 first_str(&mut st, "content").await
1464}
1465
1466async fn first_str(
1467 st: &mut lancedb::arrow::SendableRecordBatchStream,
1468 col: &str,
1469) -> Result<Option<String>> {
1470 while let Some(b) = st.try_next().await? {
1471 if b.num_rows() > 0 {
1472 return Ok(Some(str_col(&b, col)?.value(0).to_string()));
1473 }
1474 }
1475 Ok(None)
1476}
1477
1478#[cfg(test)]
1479mod tests {
1480 use super::*;
1481
1482 #[tokio::test]
1483 async fn repo_logs_round_trip() {
1484 let dir = std::env::temp_dir().join(format!("comindly-log-{}", std::process::id()));
1485 let uri = dir.to_str().unwrap();
1486 assert!(
1487 read_repo_logs(uri).await.unwrap().is_empty(),
1488 "older indexes have none"
1489 );
1490 let logs = RepoLogs::from([(
1491 "orders".to_string(),
1492 RepoLog {
1493 commits: vec![LoggedCommit {
1494 sha: "abc".into(),
1495 time: 1,
1496 mr: Some(7),
1497 files: vec!["app.py".into()],
1498 }],
1499 landed: [("mr:!7".to_string(), vec!["abc".to_string()])].into(),
1500 },
1501 )]);
1502 write_repo_logs(uri, &logs).await.unwrap();
1503 assert_eq!(read_repo_logs(uri).await.unwrap(), logs);
1504 std::fs::remove_dir_all(&dir).ok();
1505 }
1506
1507 #[tokio::test]
1508 async fn repo_meta_gains_remotes_and_old_tables_still_read() {
1509 let dir = std::env::temp_dir().join(format!("comindly-meta-{}", std::process::id()));
1510 let uri = dir.to_str().unwrap();
1511 let old = RecordBatch::try_new(
1512 Arc::new(Schema::new(vec![
1513 Field::new("repo", DataType::Utf8, false),
1514 Field::new("commit", DataType::Utf8, false),
1515 ])),
1516 vec![
1517 Arc::new(StringArray::from(vec!["orders"])) as ArrayRef,
1518 Arc::new(StringArray::from(vec!["abc"])) as ArrayRef,
1519 ],
1520 )
1521 .unwrap();
1522 let db = crate::storage::connect(uri).execute().await.unwrap();
1523 write_versioned(&db, "repo_meta", old).await.unwrap();
1524 let read = read_indexed_repos(uri).await.unwrap();
1525 assert_eq!(read[0].remote, None);
1526 let next = IndexedRepo {
1527 name: "orders".into(),
1528 commit: "def".into(),
1529 remote: Some("https://example.com/orders.git".into()),
1530 branch: Some("main".into()),
1531 };
1532 write_repo_commits(uri, std::slice::from_ref(&next))
1533 .await
1534 .unwrap();
1535 assert_eq!(read_indexed_repos(uri).await.unwrap(), vec![next]);
1536 std::fs::remove_dir_all(&dir).ok();
1537 }
1538
1539 #[tokio::test]
1540 async fn sources_store_each_content_once_and_merge_by_repo() {
1541 use crate::sources::SourceFile;
1542 let dir = std::env::temp_dir().join(format!("comindly-sources-idx-{}", std::process::id()));
1543 let _ = std::fs::remove_dir_all(&dir);
1544 let uri = dir.to_str().unwrap();
1545 let file = |repo: &str, path: &str, content: &str| SourceFile {
1546 repo: repo.into(),
1547 path: path.into(),
1548 hash: blake3::hash(content.as_bytes()).to_hex()[..32].to_string(),
1549 content: content.into(),
1550 };
1551 let a = file("orders", "src/it's.py", "x = 1\n");
1552 let b = file("billing", "lib/b.py", "x = 1\n");
1553 assert_eq!(
1554 write_sources(uri, &[a.clone(), b], false).await.unwrap(),
1555 (2, 1)
1556 );
1557 assert_eq!(
1558 read_source(uri, "orders", "src/it's.py")
1559 .await
1560 .unwrap()
1561 .as_deref(),
1562 Some("x = 1\n")
1563 );
1564
1565 let changed = file("orders", "src/it's.py", "x = 2\n");
1566 assert_eq!(write_sources(uri, &[changed], true).await.unwrap(), (1, 1));
1567 assert_eq!(
1568 read_source(uri, "orders", "src/it's.py")
1569 .await
1570 .unwrap()
1571 .as_deref(),
1572 Some("x = 2\n")
1573 );
1574 assert_eq!(
1575 read_source(uri, "billing", "lib/b.py")
1576 .await
1577 .unwrap()
1578 .as_deref(),
1579 Some("x = 1\n")
1580 );
1581 assert_eq!(
1582 read_source(uri, "billing", "missing.py").await.unwrap(),
1583 None
1584 );
1585 std::fs::remove_dir_all(&dir).ok();
1586 }
1587}