feat(xrpc): sh.tangled.search.query route · ptr.pet/bobbin@9a275c1

+1

crates/xrpc/Cargo.toml

··· 9 9 bobbin-types = { workspace = true } 10 10 bobbin-edge-index = { workspace = true } 11 11 bobbin-record-lru = { workspace = true } 12 + bobbin-search = { workspace = true } 12 13 bobbin-slingshot-client = { workspace = true } 13 14 bobbin-knot-proxy = { workspace = true } 14 15 jacquard-common = { workspace = true }

+130 -12

crates/xrpc/src/lib.rs

··· 22 22 }; 23 23 use bobbin_knot_proxy::{KnotHost, KnotProxy, KnotProxyError, ProxyResponse, RepoSlug}; 24 24 use bobbin_record_lru::RecordStore; 25 + use bobbin_search::{SearchCursor, SearchError, SearchHit, SearchIndex, SearchOffset}; 25 26 use bobbin_slingshot_client::{SlingshotClient, SlingshotError}; 26 27 use bobbin_types::ids::{EdgeKey, nsid_static}; 27 28 use bobbin_types::record::RecordBody; 29 + use bobbin_types::search::SearchableRecord; 28 30 use bobbin_types::sh_tangled::actor::profile::{Profile, ProfileGetRecordOutput, ProfileRecord}; 29 31 use bobbin_types::sh_tangled::feed::star::{Star, StarRecord}; 30 32 use bobbin_types::sh_tangled::graph::follow::{Follow, FollowRecord}; ··· 37 39 use futures::stream::{self, StreamExt, TryStreamExt}; 38 40 use jacquard_common::types::did::Did; 39 41 use jacquard_common::types::ident::AtIdentifier; 42 + use jacquard_common::types::nsid::Nsid; 40 43 use jacquard_common::types::string::{AtUri, Cid}; 41 44 use jacquard_common::xrpc::XrpcResp; 42 45 use jacquard_common::{DefaultStr, IntoStatic}; ··· 53 56 pub edges: Arc<EdgeStore>, 54 57 pub coverage: Arc<CoverageWatch>, 55 58 pub knots: Arc<KnotProxy>, 59 + pub search: Arc<SearchIndex>, 56 60 } 57 61 58 62 impl AppState { ··· 62 66 edges: Arc<EdgeStore>, 63 67 coverage: Arc<CoverageWatch>, 64 68 knots: Arc<KnotProxy>, 69 + search: Arc<SearchIndex>, 65 70 ) -> Self { 66 71 Self { 67 72 records, ··· 69 74 edges, 70 75 coverage, 71 76 knots, 77 + search, 72 78 } 73 79 } 74 80 } ··· 95 101 "/xrpc/sh.tangled.repo.issue.countComments", 96 102 get(count_issue_comments), 97 103 ) 104 + .route("/xrpc/sh.tangled.search.query", get(search_query)) 98 105 .merge(knot_proxied_routes()) 99 106 .with_state(state) 100 107 } ··· 182 189 } 183 190 184 191 #[derive(Clone, Copy, Debug, Eq, PartialEq)] 185 - pub struct ExpectedNsid(&'static str); 192 + pub struct ExpectedNsid<'a>(&'a str); 186 193 187 - impl ExpectedNsid { 188 - pub const fn new(nsid: &'static str) -> Self { 194 + impl<'a> ExpectedNsid<'a> { 195 + pub const fn new(nsid: &'a str) -> Self { 189 196 Self(nsid) 190 197 } 191 198 192 - pub const fn as_str(self) -> &'static str { 199 + pub const fn as_str(self) -> &'a str { 193 200 self.0 194 201 } 195 202 } ··· 226 233 subject: RawAtUriParam, 227 234 } 228 235 236 + #[derive(Debug, Deserialize)] 237 + struct SearchQueryParams { 238 + q: String, 239 + nsid: Option<String>, 240 + cursor: Option<String>, 241 + limit: Option<u32>, 242 + } 243 + 229 244 pub struct XrpcQuery<T>(pub T); 230 245 231 246 impl<S, T> FromRequestParts<S> for XrpcQuery<T> ··· 255 270 UpstreamGone(String), 256 271 #[error("invalid record: {0}")] 257 272 InvalidRecord(String), 273 + #[error("internal: {0}")] 274 + Internal(String), 258 275 } 259 276 260 277 #[derive(Serialize)] ··· 271 288 Self::UpstreamUnavailable(_) => (StatusCode::BAD_GATEWAY, "UpstreamFailed"), 272 289 Self::UpstreamGone(_) => (StatusCode::BAD_GATEWAY, "UpstreamGone"), 273 290 Self::InvalidRecord(_) => (StatusCode::BAD_GATEWAY, "InvalidRecord"), 291 + Self::Internal(_) => (StatusCode::INTERNAL_SERVER_ERROR, "InternalError"), 274 292 }; 275 293 let body = ErrorBody { 276 294 error, ··· 322 340 distinct_authors: u64, 323 341 } 324 342 343 + #[derive(Serialize)] 344 + #[serde(rename_all = "camelCase")] 345 + struct SearchHitView { 346 + uri: AtUri<DefaultStr>, 347 + cid: Option<Cid<DefaultStr>>, 348 + nsid: Nsid<DefaultStr>, 349 + score: f32, 350 + value: SearchableRecord, 351 + } 352 + 353 + #[derive(Serialize)] 354 + #[serde(rename_all = "camelCase")] 355 + struct SearchResponse { 356 + coverage: CoverageEnvelope, 357 + hits: Vec<SearchHitView>, 358 + cursor: Option<String>, 359 + } 360 + 325 361 fn map_slingshot(err: SlingshotError) -> XrpcError { 326 362 use SlingshotError as E; 327 363 match err { ··· 419 455 .map_err(|e| XrpcError::InvalidParams(format!("limit: {e}"))) 420 456 } 421 457 458 + async fn resolve_for_view( 459 + state: &AppState, 460 + expected_nsid: &str, 461 + uri: AtUri<DefaultStr>, 462 + ) -> Result<Arc<RecordBody>, XrpcError> { 463 + let raw = uri.as_ref().to_owned(); 464 + resolve(state, ExpectedNsid::new(expected_nsid), uri) 465 + .await 466 + .map(|(body, _did)| body) 467 + .map_err(|e| match e { 468 + XrpcError::NotFound => XrpcError::UpstreamGone(raw), 469 + other => other, 470 + }) 471 + } 472 + 422 473 async fn resolve( 423 474 state: &AppState, 424 - expected: ExpectedNsid, 475 + expected: ExpectedNsid<'_>, 425 476 uri: AtUri<DefaultStr>, 426 477 ) -> Result<(Arc<RecordBody>, Did<DefaultStr>), XrpcError> { 427 478 let collection = uri ··· 466 517 ty: &'a str, 467 518 } 468 519 469 - fn verify_type_tag(body: &RecordBody, expected: ExpectedNsid) -> Result<(), XrpcError> { 520 + fn verify_type_tag(body: &RecordBody, expected: ExpectedNsid<'_>) -> Result<(), XrpcError> { 470 521 let tag: TypeTag<'_> = serde_json::from_slice(&body.value) 471 522 .map_err(|e| XrpcError::InvalidRecord(format!("$type peek: {e}")))?; 472 523 if tag.ty != expected.as_str() { ··· 555 606 let EdgePage { items, next } = state.edges.list(&key, cursor, limit); 556 607 let items = stream::iter(items) 557 608 .map(|uri| async move { 558 - let (body, _did) = resolve(state, ExpectedNsid::new(R::NSID), uri.clone()) 559 - .await 560 - .map_err(|e| match e { 561 - XrpcError::NotFound => XrpcError::UpstreamGone(uri.as_ref().to_owned()), 562 - other => other, 563 - })?; 609 + let body = resolve_for_view(state, R::NSID, uri).await?; 564 610 let value: V = serde_json::from_slice(&body.value) 565 611 .map_err(|e| XrpcError::InvalidRecord(e.to_string()))?; 566 612 Ok::<_, XrpcError>(RecordView { ··· 663 709 XrpcQuery(q): XrpcQuery<CountQuery>, 664 710 ) -> Result<Json<CountResponse>, XrpcError> { 665 711 count_for::<IssueCommentRecord>(&state, q).map(Json) 712 + } 713 + 714 + async fn search_query( 715 + State(state): State<AppState>, 716 + XrpcQuery(q): XrpcQuery<SearchQueryParams>, 717 + ) -> Result<Json<SearchResponse>, XrpcError> { 718 + if q.q.trim().is_empty() { 719 + return Err(XrpcError::InvalidParams("q must not be empty".into())); 720 + } 721 + let cursor = SearchCursor::from_token(q.cursor.as_deref()) 722 + .map_err(|e| XrpcError::InvalidParams(format!("cursor: {e}")))?; 723 + let limit = parse_limit(q.limit)?; 724 + let nsid_filter = q 725 + .nsid 726 + .as_deref() 727 + .map(|s| { 728 + Nsid::<DefaultStr>::new_owned(s) 729 + .map_err(|e| XrpcError::InvalidParams(format!("nsid: {e}"))) 730 + }) 731 + .transpose()?; 732 + let coverage = state.coverage.snapshot(); 733 + let page = state 734 + .search 735 + .search(&q.q, nsid_filter.as_ref(), cursor, limit.get()) 736 + .await 737 + .map_err(map_search_err)?; 738 + let state_ref = &state; 739 + let hits: Vec<SearchHitView> = stream::iter(page.hits) 740 + .map(|hit| hydrate_search_hit(state_ref, hit)) 741 + .buffered(FETCH_CONCURRENCY) 742 + .try_filter_map(|opt| async move { Ok(opt) }) 743 + .try_collect() 744 + .await?; 745 + Ok(Json(SearchResponse { 746 + coverage: coverage.into(), 747 + hits, 748 + cursor: page.next.map(SearchOffset::encode_token), 749 + })) 750 + } 751 + 752 + async fn hydrate_search_hit( 753 + state: &AppState, 754 + hit: SearchHit, 755 + ) -> Result<Option<SearchHitView>, XrpcError> { 756 + let SearchHit { uri, nsid, score } = hit; 757 + let body = match resolve_for_view(state, nsid.as_ref(), uri).await { 758 + Ok(b) => b, 759 + Err(XrpcError::UpstreamGone(_)) => return Ok(None), 760 + Err(other) => return Err(other), 761 + }; 762 + let value = SearchableRecord::from_json_bytes(nsid.as_ref(), &body.value) 763 + .map_err(|e| XrpcError::InvalidRecord(e.to_string()))?; 764 + Ok(Some(SearchHitView { 765 + uri: body.uri.clone(), 766 + cid: Some(body.cid.clone()), 767 + nsid, 768 + score, 769 + value, 770 + })) 771 + } 772 + 773 + fn map_search_err(err: SearchError) -> XrpcError { 774 + use SearchError as E; 775 + match err { 776 + E::Query(e) => XrpcError::InvalidParams(format!("query: {e}")), 777 + e @ (E::Tantivy(_) 778 + | E::InvalidUri(_) 779 + | E::InvalidNsid(_) 780 + | E::MissingField(_) 781 + | E::ThreadSpawn(_) 782 + | E::Cancelled(_)) => XrpcError::Internal(format!("search: {e}")), 783 + } 666 784 } 667 785 668 786 fn map_proxy_error(err: KnotProxyError) -> XrpcError {

+2

crates/xrpc/tests/aggregation.rs

··· 4 4 use bobbin_edge_index::{Coverage, CoverageWatch, EdgeStore, HydrantCursor, SourceId}; 5 5 use bobbin_knot_proxy::{KnotProxy, KnotProxyConfig}; 6 6 use bobbin_record_lru::{CacheCapacity, LruRecordStore}; 7 + use bobbin_search::{DEFAULT_WRITER_HEAP_BYTES, SearchIndex}; 7 8 use bobbin_slingshot_client::SlingshotClient; 8 9 use bobbin_types::edges::Edge; 9 10 use bobbin_xrpc::{AppState, router}; ··· 47 48 edges.clone(), 48 49 coverage.clone(), 49 50 Arc::new(KnotProxy::new(KnotProxyConfig::default()).unwrap()), 51 + Arc::new(SearchIndex::new(DEFAULT_WRITER_HEAP_BYTES).unwrap()), 50 52 ); 51 53 Self { 52 54 server,

+2

crates/xrpc/tests/cold_start.rs

··· 4 4 use bobbin_edge_index::{CoverageWatch, EdgeStore}; 5 5 use bobbin_knot_proxy::{KnotProxy, KnotProxyConfig}; 6 6 use bobbin_record_lru::{CacheCapacity, LruRecordStore}; 7 + use bobbin_search::{DEFAULT_WRITER_HEAP_BYTES, SearchIndex}; 7 8 use bobbin_slingshot_client::SlingshotClient; 8 9 use bobbin_xrpc::{AppState, router}; 9 10 use futures::stream::{self, StreamExt}; ··· 24 25 Arc::new(EdgeStore::new()), 25 26 Arc::new(CoverageWatch::new()), 26 27 Arc::new(KnotProxy::new(KnotProxyConfig::default()).unwrap()), 28 + Arc::new(SearchIndex::new(DEFAULT_WRITER_HEAP_BYTES).unwrap()), 27 29 ) 28 30 } 29 31

+2

crates/xrpc/tests/knot_proxy.rs

··· 5 5 use bobbin_edge_index::{CoverageWatch, EdgeStore}; 6 6 use bobbin_knot_proxy::{FailureThreshold, KnotProxy, KnotProxyConfig}; 7 7 use bobbin_record_lru::{CacheCapacity, LruRecordStore}; 8 + use bobbin_search::{DEFAULT_WRITER_HEAP_BYTES, SearchIndex}; 8 9 use bobbin_slingshot_client::SlingshotClient; 9 10 use bobbin_xrpc::{AppState, router}; 10 11 use http::{Request, StatusCode}; ··· 48 49 Arc::new(EdgeStore::new()), 49 50 Arc::new(CoverageWatch::new()), 50 51 Arc::new(KnotProxy::new(config).unwrap()), 52 + Arc::new(SearchIndex::new(DEFAULT_WRITER_HEAP_BYTES).unwrap()), 51 53 ); 52 54 Self { 53 55 slingshot: slingshot_server,

Configure Feed

Configure Feed