use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use async_trait::async_trait;
use parking_lot::RwLock;
use crate::error::StorageResult;
use crate::search::{IndexValue, SearchParameterExtractor, SearchParameterRegistry};
use crate::tenant::TenantContext;
use crate::types::{
IncludeDirective, IncludeType, Page, ReverseChainedParameter, SearchBundle, SearchParamType,
SearchParameter, SearchQuery, SearchValue, StoredResource,
};
use super::storage::ResourceStorage;
#[derive(Debug, Clone)]
pub struct SearchResult {
pub resources: Page<StoredResource>,
pub included: Vec<StoredResource>,
pub total: Option<u64>,
pub scores: HashMap<String, f64>,
}
impl SearchResult {
pub fn new(resources: Page<StoredResource>) -> Self {
Self {
resources,
included: Vec::new(),
total: None,
scores: HashMap::new(),
}
}
pub fn with_included(mut self, included: Vec<StoredResource>) -> Self {
self.included = included;
self
}
pub fn with_total(mut self, total: u64) -> Self {
self.total = Some(total);
self
}
pub fn with_scores(mut self, scores: HashMap<String, f64>) -> Self {
self.scores = scores;
self
}
pub fn len(&self) -> usize {
self.resources.len()
}
pub fn is_empty(&self) -> bool {
self.resources.is_empty()
}
pub fn next_cursor(&self) -> Option<&String> {
self.resources.page_info.next_cursor.as_ref()
}
pub fn previous_cursor(&self) -> Option<&String> {
self.resources.page_info.previous_cursor.as_ref()
}
pub fn has_next(&self) -> bool {
self.resources.page_info.has_next
}
pub fn has_previous(&self) -> bool {
self.resources.page_info.has_previous
}
pub fn to_bundle(&self, base_url: &str, self_link: &str) -> SearchBundle {
use crate::types::{BundleEntry, SearchBundle};
let mut bundle = SearchBundle::new().with_self_link(self_link);
if let Some(total) = self.total {
bundle = bundle.with_total(total);
}
if let Some(ref cursor) = self.resources.page_info.next_cursor {
bundle = bundle.with_next_link(replace_cursor_param(self_link, cursor));
}
if let Some(ref cursor) = self.resources.page_info.previous_cursor {
bundle = bundle.with_previous_link(replace_cursor_param(self_link, cursor));
}
if self.resources.page_info.next_cursor.is_some()
|| self.resources.page_info.previous_cursor.is_some()
{
bundle = bundle.with_link("first", strip_paging_params(self_link));
}
for resource in &self.resources.items {
let full_url = format!("{}/{}", base_url, resource.url());
let entry = BundleEntry::match_entry(full_url, resource.content().clone())
.with_score(self.scores.get(&resource.url()).copied());
bundle = bundle.with_entry(entry);
}
for resource in &self.included {
let full_url = format!("{}/{}", base_url, resource.url());
bundle = bundle.with_entry(BundleEntry::include_entry(
full_url,
resource.content().clone(),
));
}
bundle
}
}
fn strip_paging_params(url: &str) -> String {
let (base, query) = match url.find('?') {
Some(pos) => (&url[..pos], &url[pos + 1..]),
None => return url.to_string(),
};
let parts: Vec<String> = query
.split('&')
.filter(|p| !p.is_empty() && !p.starts_with("_cursor=") && !p.starts_with("_offset="))
.map(str::to_string)
.collect();
if parts.is_empty() {
base.to_string()
} else {
format!("{}?{}", base, parts.join("&"))
}
}
fn replace_cursor_param(url: &str, cursor: &str) -> String {
let (base, query) = match url.find('?') {
Some(pos) => (&url[..pos], &url[pos + 1..]),
None => (url, ""),
};
let mut parts: Vec<String> = query
.split('&')
.filter(|p| !p.is_empty() && !p.starts_with("_cursor="))
.map(str::to_string)
.collect();
parts.push(format!("_cursor={}", cursor));
format!("{}?{}", base, parts.join("&"))
}
#[async_trait]
pub trait SearchProvider: ResourceStorage {
async fn search(
&self,
tenant: &TenantContext,
query: &SearchQuery,
) -> StorageResult<SearchResult>;
async fn search_count(&self, tenant: &TenantContext, query: &SearchQuery)
-> StorageResult<u64>;
fn search_param_registry(&self) -> &Arc<RwLock<SearchParameterRegistry>>;
fn supports_contained_search(&self) -> bool {
false
}
fn modifiers_for_param_type(&self, param_type: SearchParamType) -> Vec<&'static str> {
let _ = param_type;
Vec::new()
}
}
#[async_trait]
pub trait MultiTypeSearchProvider: SearchProvider {
async fn search_multi(
&self,
tenant: &TenantContext,
resource_types: &[&str],
query: &SearchQuery,
) -> StorageResult<SearchResult>;
}
#[async_trait]
pub trait IncludeProvider: SearchProvider {
async fn resolve_includes(
&self,
tenant: &TenantContext,
resources: &[StoredResource],
includes: &[IncludeDirective],
) -> StorageResult<Vec<StoredResource>>;
}
#[async_trait]
pub trait RevincludeProvider: SearchProvider {
async fn resolve_revincludes(
&self,
tenant: &TenantContext,
resources: &[StoredResource],
revincludes: &[IncludeDirective],
) -> StorageResult<Vec<StoredResource>>;
}
const MAX_INCLUDE_ITERATE_DEPTH: usize = 5;
const INCLUDE_FETCH_LIMIT: u32 = 10_000;
pub async fn resolve_includes_iterative<S>(
provider: &S,
tenant: &TenantContext,
matches: &[StoredResource],
includes: &[IncludeDirective],
) -> StorageResult<Vec<StoredResource>>
where
S: SearchProvider + ?Sized,
{
if matches.is_empty() || includes.is_empty() {
return Ok(Vec::new());
}
let extractor = SearchParameterExtractor::new(provider.search_param_registry().clone());
let key = |r: &StoredResource| format!("{}/{}", r.resource_type(), r.id());
let mut seen: HashSet<String> = matches.iter().map(&key).collect();
let mut included: Vec<StoredResource> = Vec::new();
let mut frontier: Vec<StoredResource> = matches.to_vec();
let mut first_hop = true;
let mut depth = 0;
loop {
let active: Vec<&IncludeDirective> =
includes.iter().filter(|d| first_hop || d.iterate).collect();
if active.is_empty() {
break;
}
let mut fetched: Vec<StoredResource> = Vec::new();
for directive in active {
match directive.include_type {
IncludeType::Include => {
let mut wanted: Vec<(String, String)> = Vec::new();
for res in &frontier {
if res.resource_type() != directive.source_type {
continue;
}
let def = provider
.search_param_registry()
.read()
.get_param(res.resource_type(), &directive.search_param);
let Some(def) = def else { continue };
if let Ok(values) = extractor.extract_for_param(res.content(), &def) {
for v in values {
if let IndexValue::Reference { reference, .. } = v.value {
if let Some((t, i)) = reference.split_once('/') {
if let Some(target) = &directive.target_type {
if t != target {
continue;
}
}
wanted.push((t.to_string(), i.to_string()));
}
}
}
}
}
let mut by_type: std::collections::HashMap<String, Vec<String>> =
std::collections::HashMap::new();
for (t, i) in wanted {
by_type.entry(t).or_default().push(i);
}
for (rtype, ids) in by_type {
let mut q = SearchQuery::new(&rtype).with_parameter(SearchParameter {
name: "_id".to_string(),
param_type: SearchParamType::Token,
modifier: None,
values: ids.iter().map(SearchValue::eq).collect(),
chain: vec![],
components: vec![],
});
q.count = Some(INCLUDE_FETCH_LIMIT);
let result = provider.search(tenant, &q).await?;
fetched.extend(result.resources.items);
}
}
IncludeType::Revinclude => {
let refs: Vec<SearchValue> =
frontier.iter().map(|r| SearchValue::eq(key(r))).collect();
if refs.is_empty() {
continue;
}
let mut q =
SearchQuery::new(&directive.source_type).with_parameter(SearchParameter {
name: directive.search_param.clone(),
param_type: SearchParamType::Reference,
modifier: None,
values: refs,
chain: vec![],
components: vec![],
});
q.count = Some(INCLUDE_FETCH_LIMIT);
let result = provider.search(tenant, &q).await?;
fetched.extend(result.resources.items);
}
}
}
let mut next = Vec::new();
for r in fetched {
if seen.insert(key(&r)) {
next.push(r.clone());
included.push(r);
}
}
first_hop = false;
depth += 1;
frontier = next;
if frontier.is_empty() || depth >= MAX_INCLUDE_ITERATE_DEPTH {
break;
}
}
Ok(included)
}
#[async_trait]
pub trait ChainedSearchProvider: SearchProvider {
async fn resolve_chain(
&self,
tenant: &TenantContext,
base_type: &str,
chain: &str,
value: &str,
) -> StorageResult<Vec<String>>;
async fn resolve_reverse_chain(
&self,
tenant: &TenantContext,
base_type: &str,
reverse_chain: &ReverseChainedParameter,
) -> StorageResult<Vec<String>>;
}
#[async_trait]
pub trait TerminologySearchProvider: SearchProvider {
async fn expand_value_set(&self, value_set_url: &str) -> StorageResult<Vec<(String, String)>>;
async fn codes_above(&self, system: &str, code: &str) -> StorageResult<Vec<String>>;
async fn codes_below(&self, system: &str, code: &str) -> StorageResult<Vec<String>>;
}
#[async_trait]
pub trait TextSearchProvider: SearchProvider {
async fn search_text(
&self,
tenant: &TenantContext,
resource_type: &str,
text: &str,
pagination: &crate::types::Pagination,
) -> StorageResult<SearchResult>;
async fn search_content(
&self,
tenant: &TenantContext,
resource_type: &str,
content: &str,
pagination: &crate::types::Pagination,
) -> StorageResult<SearchResult>;
}
pub trait FullSearchProvider:
SearchProvider
+ MultiTypeSearchProvider
+ IncludeProvider
+ RevincludeProvider
+ ChainedSearchProvider
{
}
impl<T> FullSearchProvider for T where
T: SearchProvider
+ MultiTypeSearchProvider
+ IncludeProvider
+ RevincludeProvider
+ ChainedSearchProvider
{
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::PageInfo;
use helios_fhir::FhirVersion;
#[test]
fn test_search_result_creation() {
let page = Page::new(Vec::new(), PageInfo::end());
let result = SearchResult::new(page);
assert!(result.included.is_empty());
assert!(result.total.is_none());
}
#[test]
fn test_search_result_with_included() {
let page = Page::new(Vec::new(), PageInfo::end());
let result = SearchResult::new(page)
.with_included(vec![StoredResource::new(
"Patient",
"123",
crate::tenant::TenantId::new("t1"),
serde_json::json!({}),
FhirVersion::default(),
)])
.with_total(100);
assert_eq!(result.included.len(), 1);
assert_eq!(result.total, Some(100));
}
#[test]
fn test_search_result_to_bundle() {
let resource = StoredResource::new(
"Patient",
"123",
crate::tenant::TenantId::new("t1"),
serde_json::json!({"resourceType": "Patient", "id": "123"}),
FhirVersion::default(),
);
let page = Page::new(vec![resource], PageInfo::end());
let result = SearchResult::new(page).with_total(1);
let bundle = result.to_bundle("http://example.com/fhir", "http://example.com/fhir/Patient");
assert_eq!(bundle.total, Some(1));
assert_eq!(bundle.entry.len(), 1);
assert!(bundle.entry[0].search.as_ref().unwrap().score.is_none());
}
#[test]
fn test_search_result_to_bundle_attaches_score() {
let resource = StoredResource::new(
"Patient",
"123",
crate::tenant::TenantId::new("t1"),
serde_json::json!({"resourceType": "Patient", "id": "123"}),
FhirVersion::default(),
);
let page = Page::new(vec![resource], PageInfo::end());
let mut scores = HashMap::new();
scores.insert("Patient/123".to_string(), 0.875);
let result = SearchResult::new(page).with_scores(scores);
let bundle = result.to_bundle("http://example.com/fhir", "http://example.com/fhir/Patient");
assert_eq!(bundle.entry.len(), 1);
assert_eq!(
bundle.entry[0].search.as_ref().unwrap().score,
Some(0.875),
"the matched entry carries Bundle.entry.search.score"
);
}
#[test]
fn test_replace_cursor_param_no_query() {
let url = replace_cursor_param("http://example.com/fhir/Patient", "abc");
assert_eq!(url, "http://example.com/fhir/Patient?_cursor=abc");
}
#[test]
fn test_replace_cursor_param_with_existing_params() {
let url = replace_cursor_param(
"http://example.com/fhir/Patient?_count=3&_elements=id",
"abc",
);
assert_eq!(
url,
"http://example.com/fhir/Patient?_count=3&_elements=id&_cursor=abc"
);
}
#[test]
fn test_replace_cursor_param_replaces_existing_cursor() {
let url = replace_cursor_param(
"http://example.com/fhir/Patient?_count=3&_cursor=old&_elements=id",
"new",
);
assert!(url.starts_with("http://example.com/fhir/Patient?"));
assert!(url.contains("_count=3"));
assert!(url.contains("_elements=id"));
assert!(url.contains("_cursor=new"));
assert!(!url.contains("_cursor=old"));
assert_eq!(url.matches("_cursor=").count(), 1);
}
#[test]
fn test_to_bundle_next_link_format() {
let page = Page::new(
Vec::<StoredResource>::new(),
PageInfo {
next_cursor: Some("CURSOR_VALUE".to_string()),
previous_cursor: None,
total: None,
has_next: true,
has_previous: false,
},
);
let result = SearchResult::new(page);
let bundle = result.to_bundle(
"http://example.com/fhir",
"http://example.com/fhir/Patient?_count=3&_elements=id",
);
let next = bundle
.link
.iter()
.find(|l| l.relation == "next")
.expect("next link present");
assert_eq!(
next.url,
"http://example.com/fhir/Patient?_count=3&_elements=id&_cursor=CURSOR_VALUE"
);
assert_eq!(
next.url.matches('?').count(),
1,
"exactly one '?' delimiter"
);
}
}