1026 lines
32 KiB
Rust
1026 lines
32 KiB
Rust
mod query_builder;
|
|
|
|
use std::collections::{HashMap, HashSet};
|
|
use std::path::Path;
|
|
use std::sync::{Arc, Mutex};
|
|
|
|
use common::proto::komp_ac::search::searcher_server::Searcher;
|
|
pub use common::proto::komp_ac::search::searcher_server::SearcherServer;
|
|
use common::proto::komp_ac::search::{
|
|
SearchOrderDirection, SearchRequest, SearchResponse, search_response::Hit,
|
|
};
|
|
use common::search::{SchemaFields, register_tokenizers, search_index_path};
|
|
use query_builder::{
|
|
ConstraintMode, SearchConstraint, SearchConstraintTarget, build_master_query,
|
|
};
|
|
use sqlx::{AssertSqlSafe, PgPool, Row};
|
|
use tantivy::collector::TopDocs;
|
|
use tantivy::schema::Value;
|
|
use tantivy::{Index, IndexReader, ReloadPolicy, TantivyDocument};
|
|
use tonic::{Request, Response, Status};
|
|
use tracing::info;
|
|
|
|
const INDEX_ROOT: &str = "./tantivy_indexes";
|
|
const DEFAULT_RESULT_LIMIT: usize = 60;
|
|
const HARD_RESULT_LIMIT: usize = 200;
|
|
const DEFAULT_LIST_LIMIT: usize = 5;
|
|
const SEARCH_SCORE_RELATIVE_FLOOR: f32 = 0.25;
|
|
const SEARCH_SCORE_GROUP_WIDTH: f32 = 0.25;
|
|
|
|
pub struct SearcherService {
|
|
pub pool: PgPool,
|
|
profiles: Mutex<HashMap<String, Arc<ProfileIndex>>>,
|
|
}
|
|
|
|
impl SearcherService {
|
|
pub fn new(pool: PgPool) -> Self {
|
|
Self {
|
|
pool,
|
|
profiles: Mutex::new(HashMap::new()),
|
|
}
|
|
}
|
|
|
|
async fn run_rpc(
|
|
&self,
|
|
request: Request<SearchRequest>,
|
|
) -> Result<Response<SearchResponse>, Status> {
|
|
let req = request.into_inner();
|
|
let normalized = normalize_request(req)?;
|
|
|
|
if !profile_exists(&self.pool, &normalized.profile_name).await? {
|
|
return Err(Status::not_found(format!(
|
|
"Profile '{}' was not found",
|
|
normalized.profile_name
|
|
)));
|
|
}
|
|
|
|
if let Some(table_name) = normalized.table_name.as_deref() {
|
|
if !table_exists(&self.pool, &normalized.profile_name, table_name).await? {
|
|
return Err(Status::not_found(format!(
|
|
"Table '{}' was not found in profile '{}'",
|
|
table_name, normalized.profile_name
|
|
)));
|
|
}
|
|
}
|
|
|
|
if !normalized.has_input() {
|
|
let Some(table_name) = normalized.table_name.as_deref() else {
|
|
return Err(Status::invalid_argument(
|
|
"table_name is required when query is empty",
|
|
));
|
|
};
|
|
|
|
let hits = fetch_ordered_rows(
|
|
&self.pool,
|
|
&normalized.profile_name,
|
|
table_name,
|
|
normalized.order.as_ref(),
|
|
normalized.limit.unwrap_or(DEFAULT_LIST_LIMIT),
|
|
normalized.offset,
|
|
)
|
|
.await?;
|
|
return Ok(Response::new(SearchResponse { hits }));
|
|
}
|
|
|
|
let index_path = search_index_path(Path::new(INDEX_ROOT), &normalized.profile_name);
|
|
if !index_path.exists() {
|
|
return Err(Status::not_found(format!(
|
|
"No search index found for profile '{}'",
|
|
normalized.profile_name
|
|
)));
|
|
}
|
|
|
|
let resolved_must = resolve_constraints(
|
|
&self.pool,
|
|
&normalized.profile_name,
|
|
normalized.table_name.as_deref(),
|
|
&normalized.must,
|
|
)
|
|
.await?;
|
|
|
|
let profile = profile_index(&self.profiles, &normalized.profile_name, &index_path)?;
|
|
let mut hits = run_search(
|
|
&self.pool,
|
|
&profile,
|
|
&normalized.profile_name,
|
|
normalized.table_name.as_deref(),
|
|
&normalized.free_query,
|
|
&resolved_must,
|
|
normalized.limit.unwrap_or(DEFAULT_RESULT_LIMIT),
|
|
normalized.offset,
|
|
normalized.order.as_ref(),
|
|
)
|
|
.await?;
|
|
|
|
if normalized.order.is_none() {
|
|
hits.sort_by(|left, right| right.score.total_cmp(&left.score));
|
|
if let Some(limit) = normalized.limit {
|
|
if hits.len() > limit {
|
|
hits.truncate(limit);
|
|
}
|
|
}
|
|
}
|
|
|
|
info!(
|
|
"search: profile={} table={:?} free='{}' constraints={} hits={}",
|
|
normalized.profile_name,
|
|
normalized.table_name,
|
|
normalized.free_query,
|
|
resolved_must.len(),
|
|
hits.len()
|
|
);
|
|
|
|
Ok(Response::new(SearchResponse { hits }))
|
|
}
|
|
}
|
|
|
|
struct ProfileIndex {
|
|
index: Index,
|
|
reader: IndexReader,
|
|
fields: SchemaFields,
|
|
}
|
|
|
|
impl ProfileIndex {
|
|
fn open(path: &Path) -> Result<Self, Status> {
|
|
let index = Index::open_in_dir(path)
|
|
.map_err(|e| Status::internal(format!("Failed to open index: {}", e)))?;
|
|
register_tokenizers(&index)
|
|
.map_err(|e| Status::internal(format!("Failed to register tokenizers: {}", e)))?;
|
|
let reader = index
|
|
.reader_builder()
|
|
.reload_policy(ReloadPolicy::OnCommitWithDelay)
|
|
.try_into()
|
|
.map_err(|e| Status::internal(format!("Failed to build index reader: {}", e)))?;
|
|
let fields = SchemaFields::from(&index.schema()).map_err(|e| {
|
|
Status::internal(format!(
|
|
"Search index schema mismatch. Delete the stale index and create it again: {}",
|
|
e
|
|
))
|
|
})?;
|
|
|
|
Ok(Self {
|
|
index,
|
|
reader,
|
|
fields,
|
|
})
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct NormalizedSearchRequest {
|
|
profile_name: String,
|
|
table_name: Option<String>,
|
|
free_query: String,
|
|
must: Vec<NormalizedColumnConstraint>,
|
|
limit: Option<usize>,
|
|
offset: usize,
|
|
order: Option<NormalizedSearchOrder>,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct NormalizedSearchOrder {
|
|
column: String,
|
|
direction: SearchOrderDirection,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct NormalizedColumnConstraint {
|
|
column: String,
|
|
query: String,
|
|
mode: ConstraintMode,
|
|
}
|
|
|
|
impl NormalizedSearchRequest {
|
|
fn has_input(&self) -> bool {
|
|
!self.free_query.is_empty() || !self.must.is_empty()
|
|
}
|
|
}
|
|
|
|
fn profile_index(
|
|
cache: &Mutex<HashMap<String, Arc<ProfileIndex>>>,
|
|
profile_name: &str,
|
|
path: &Path,
|
|
) -> Result<Arc<ProfileIndex>, Status> {
|
|
{
|
|
let cache_guard = cache
|
|
.lock()
|
|
.map_err(|_| Status::internal("Profile index cache lock poisoned"))?;
|
|
if let Some(index) = cache_guard.get(profile_name) {
|
|
return Ok(index.clone());
|
|
}
|
|
}
|
|
|
|
let opened = Arc::new(ProfileIndex::open(path)?);
|
|
let mut cache_guard = cache
|
|
.lock()
|
|
.map_err(|_| Status::internal("Profile index cache lock poisoned"))?;
|
|
if let Some(index) = cache_guard.get(profile_name) {
|
|
return Ok(index.clone());
|
|
}
|
|
cache_guard.insert(profile_name.to_string(), opened.clone());
|
|
Ok(opened)
|
|
}
|
|
|
|
fn validate_identifier(value: &str, field_name: &str) -> Result<(), Status> {
|
|
let mut chars = value.chars();
|
|
let Some(first) = chars.next() else {
|
|
return Err(Status::invalid_argument(format!(
|
|
"{field_name} must not be empty"
|
|
)));
|
|
};
|
|
|
|
if !(first.is_ascii_alphabetic() || first == '_')
|
|
|| !chars.all(|ch| ch.is_ascii_alphanumeric() || ch == '_')
|
|
{
|
|
return Err(Status::invalid_argument(format!(
|
|
"{field_name} contains invalid characters"
|
|
)));
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn validate_search_column(value: &str) -> Result<(), Status> {
|
|
if value.is_empty() {
|
|
return Err(Status::invalid_argument(
|
|
"constraint.column must not be empty",
|
|
));
|
|
}
|
|
|
|
if value.chars().any(|ch| ch.is_control() || ch == '\0') {
|
|
return Err(Status::invalid_argument(
|
|
"constraint.column contains invalid characters",
|
|
));
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn qualify_profile_table(profile_name: &str, table_name: &str) -> String {
|
|
format!("\"{}\".\"{}\"", profile_name, table_name)
|
|
}
|
|
|
|
async fn profile_exists(pool: &PgPool, profile_name: &str) -> Result<bool, Status> {
|
|
let exists =
|
|
sqlx::query_scalar::<_, bool>("SELECT EXISTS(SELECT 1 FROM schemas WHERE name = $1)")
|
|
.bind(profile_name)
|
|
.fetch_one(pool)
|
|
.await
|
|
.map_err(|e| Status::internal(format!("Profile lookup failed: {}", e)))?;
|
|
Ok(exists)
|
|
}
|
|
|
|
async fn table_exists(pool: &PgPool, profile_name: &str, table_name: &str) -> Result<bool, Status> {
|
|
let exists = sqlx::query_scalar::<_, bool>(
|
|
r#"
|
|
SELECT EXISTS(
|
|
SELECT 1
|
|
FROM table_definitions td
|
|
JOIN schemas s ON td.schema_id = s.id
|
|
WHERE s.name = $1 AND td.table_name = $2
|
|
)
|
|
"#,
|
|
)
|
|
.bind(profile_name)
|
|
.bind(table_name)
|
|
.fetch_one(pool)
|
|
.await
|
|
.map_err(|e| Status::internal(format!("Table lookup failed: {}", e)))?;
|
|
Ok(exists)
|
|
}
|
|
|
|
fn normalize_request(req: SearchRequest) -> Result<NormalizedSearchRequest, Status> {
|
|
let profile_name = req.profile_name.trim();
|
|
if profile_name.is_empty() {
|
|
return Err(Status::invalid_argument("profile_name is required"));
|
|
}
|
|
validate_identifier(profile_name, "profile_name")?;
|
|
|
|
let table_name = match req.table_name.as_deref().map(str::trim) {
|
|
Some(table_name) if !table_name.is_empty() => {
|
|
validate_identifier(table_name, "table_name")?;
|
|
Some(table_name.to_string())
|
|
}
|
|
_ => None,
|
|
};
|
|
|
|
let free_query = req.free_query.trim().to_string();
|
|
let mut must = Vec::new();
|
|
|
|
for constraint in req.must {
|
|
let column = constraint.column.trim();
|
|
validate_search_column(column)?;
|
|
|
|
let query = constraint.query.trim();
|
|
if query.is_empty() {
|
|
return Err(Status::invalid_argument(
|
|
"constraint.query must not be empty",
|
|
));
|
|
}
|
|
|
|
must.push(NormalizedColumnConstraint {
|
|
column: column.to_string(),
|
|
query: query.to_string(),
|
|
mode: constraint_mode_from_proto(constraint.mode),
|
|
});
|
|
}
|
|
|
|
let limit = req
|
|
.limit
|
|
.map(|value| (value as usize).min(HARD_RESULT_LIMIT));
|
|
let offset = req.offset.unwrap_or_default() as usize;
|
|
let order = req
|
|
.order
|
|
.map(|order| {
|
|
let column = order.column.trim();
|
|
if column.is_empty() {
|
|
return Err(Status::invalid_argument("order.column must not be empty"));
|
|
}
|
|
validate_search_column(column)?;
|
|
let direction = SearchOrderDirection::try_from(order.direction).map_err(|_| {
|
|
Status::invalid_argument("order.direction must be ASC or DESC")
|
|
})?;
|
|
if direction == SearchOrderDirection::Unspecified {
|
|
return Err(Status::invalid_argument(
|
|
"order.direction must be ASC or DESC",
|
|
));
|
|
}
|
|
Ok(NormalizedSearchOrder {
|
|
column: column.to_string(),
|
|
direction,
|
|
})
|
|
})
|
|
.transpose()?;
|
|
|
|
if order.is_some() && table_name.is_none() {
|
|
return Err(Status::invalid_argument(
|
|
"table_name is required when order is specified",
|
|
));
|
|
}
|
|
|
|
Ok(NormalizedSearchRequest {
|
|
profile_name: profile_name.to_string(),
|
|
table_name,
|
|
free_query,
|
|
must,
|
|
limit,
|
|
offset,
|
|
order,
|
|
})
|
|
}
|
|
|
|
fn constraint_mode_from_proto(raw_mode: i32) -> ConstraintMode {
|
|
match raw_mode {
|
|
2 => ConstraintMode::Exact,
|
|
_ => ConstraintMode::Fuzzy,
|
|
}
|
|
}
|
|
|
|
async fn resolve_constraints(
|
|
pool: &PgPool,
|
|
profile_name: &str,
|
|
table_filter: Option<&str>,
|
|
must: &[NormalizedColumnConstraint],
|
|
) -> Result<Vec<SearchConstraint>, Status> {
|
|
let mut resolved = Vec::with_capacity(must.len());
|
|
for constraint in must {
|
|
let targets = resolve_constraint_targets(pool, profile_name, table_filter, &constraint.column)
|
|
.await?;
|
|
|
|
resolved.push(SearchConstraint {
|
|
targets,
|
|
query: constraint.query.clone(),
|
|
mode: constraint.mode,
|
|
});
|
|
}
|
|
Ok(resolved)
|
|
}
|
|
|
|
async fn resolve_constraint_targets(
|
|
pool: &PgPool,
|
|
profile_name: &str,
|
|
table_filter: Option<&str>,
|
|
column: &str,
|
|
) -> Result<Vec<SearchConstraintTarget>, Status> {
|
|
let rows = if let Some(table_name) = table_filter {
|
|
sqlx::query(
|
|
r#"
|
|
SELECT td.table_name, tdc.physical_name
|
|
FROM schemas s
|
|
JOIN table_definitions td ON td.schema_id = s.id
|
|
JOIN table_definition_columns tdc ON tdc.table_definition_id = td.id
|
|
WHERE s.name = $1
|
|
AND td.table_name = $2
|
|
AND tdc.display_name = $3
|
|
"#,
|
|
)
|
|
.bind(profile_name)
|
|
.bind(table_name)
|
|
.bind(column)
|
|
.fetch_all(pool)
|
|
.await
|
|
} else {
|
|
sqlx::query(
|
|
r#"
|
|
SELECT td.table_name, tdc.physical_name
|
|
FROM schemas s
|
|
JOIN table_definitions td ON td.schema_id = s.id
|
|
JOIN table_definition_columns tdc ON tdc.table_definition_id = td.id
|
|
WHERE s.name = $1
|
|
AND tdc.display_name = $2
|
|
"#,
|
|
)
|
|
.bind(profile_name)
|
|
.bind(column)
|
|
.fetch_all(pool)
|
|
.await
|
|
}
|
|
.map_err(|e| Status::internal(format!("Column mapping lookup failed: {}", e)))?;
|
|
|
|
if rows.is_empty() {
|
|
let target = SearchConstraintTarget {
|
|
table_name: table_filter.map(str::to_string),
|
|
column: column.to_string(),
|
|
};
|
|
return Ok(vec![target]);
|
|
}
|
|
|
|
let mut seen = HashSet::new();
|
|
let mut targets = Vec::new();
|
|
for row in rows {
|
|
let table_name: String = row
|
|
.try_get("table_name")
|
|
.map_err(|e| Status::internal(format!("Column mapping table read failed: {}", e)))?;
|
|
let physical_name: String = row
|
|
.try_get("physical_name")
|
|
.map_err(|e| Status::internal(format!("Column mapping physical read failed: {}", e)))?;
|
|
let key = (table_name.clone(), physical_name.clone());
|
|
if seen.insert(key) {
|
|
targets.push(SearchConstraintTarget {
|
|
table_name: if table_filter.is_some() {
|
|
None
|
|
} else {
|
|
Some(table_name)
|
|
},
|
|
column: physical_name,
|
|
});
|
|
}
|
|
}
|
|
|
|
Ok(targets)
|
|
}
|
|
|
|
async fn table_physical_to_display_map(
|
|
pool: &PgPool,
|
|
profile_name: &str,
|
|
table_name: &str,
|
|
) -> Result<HashMap<String, String>, Status> {
|
|
let rows = sqlx::query(
|
|
r#"
|
|
SELECT tdc.physical_name, tdc.display_name
|
|
FROM schemas s
|
|
JOIN table_definitions td ON td.schema_id = s.id
|
|
JOIN table_definition_columns tdc ON tdc.table_definition_id = td.id
|
|
WHERE s.name = $1
|
|
AND td.table_name = $2
|
|
"#,
|
|
)
|
|
.bind(profile_name)
|
|
.bind(table_name)
|
|
.fetch_all(pool)
|
|
.await
|
|
.map_err(|e| Status::internal(format!("Column mapping lookup failed: {}", e)))?;
|
|
|
|
let mut mapping = HashMap::with_capacity(rows.len());
|
|
for row in rows {
|
|
let physical_name: String = row
|
|
.try_get("physical_name")
|
|
.map_err(|e| Status::internal(format!("Column mapping physical read failed: {}", e)))?;
|
|
let display_name: String = row
|
|
.try_get("display_name")
|
|
.map_err(|e| Status::internal(format!("Column mapping display read failed: {}", e)))?;
|
|
mapping.insert(physical_name, display_name);
|
|
}
|
|
Ok(mapping)
|
|
}
|
|
|
|
async fn table_row_display_column(
|
|
pool: &PgPool,
|
|
profile_name: &str,
|
|
table_name: &str,
|
|
) -> Result<String, Status> {
|
|
sqlx::query_scalar(
|
|
r#"
|
|
SELECT td.row_display_column
|
|
FROM schemas s
|
|
JOIN table_definitions td ON td.schema_id = s.id
|
|
WHERE s.name = $1 AND td.table_name = $2
|
|
"#,
|
|
)
|
|
.bind(profile_name)
|
|
.bind(table_name)
|
|
.fetch_one(pool)
|
|
.await
|
|
.map_err(|e| Status::internal(format!("Row display column lookup failed: {}", e)))
|
|
}
|
|
|
|
fn remap_json_to_display_names(
|
|
value: serde_json::Value,
|
|
physical_to_display: &HashMap<String, String>,
|
|
) -> serde_json::Value {
|
|
match value {
|
|
serde_json::Value::Object(object) => {
|
|
let mut remapped = serde_json::Map::with_capacity(object.len());
|
|
for (key, value) in object {
|
|
let final_key = physical_to_display.get(&key).cloned().unwrap_or(key);
|
|
remapped.insert(final_key, value);
|
|
}
|
|
serde_json::Value::Object(remapped)
|
|
}
|
|
other => other,
|
|
}
|
|
}
|
|
|
|
fn row_display_value(value: &serde_json::Value, column: &str) -> String {
|
|
match value.get(column) {
|
|
Some(serde_json::Value::String(value)) => value.clone(),
|
|
Some(serde_json::Value::Number(value)) => value.to_string(),
|
|
Some(serde_json::Value::Bool(value)) => value.to_string(),
|
|
_ => String::new(),
|
|
}
|
|
}
|
|
|
|
enum ResolvedOrderColumn {
|
|
Position,
|
|
Column(String),
|
|
}
|
|
|
|
fn order_direction_sql(direction: SearchOrderDirection) -> &'static str {
|
|
match direction {
|
|
SearchOrderDirection::Asc => "ASC",
|
|
SearchOrderDirection::Desc => "DESC",
|
|
SearchOrderDirection::Unspecified => unreachable!("order direction is normalized"),
|
|
}
|
|
}
|
|
|
|
async fn resolve_order_column(
|
|
pool: &PgPool,
|
|
profile_name: &str,
|
|
table_name: &str,
|
|
requested_column: &str,
|
|
) -> Result<ResolvedOrderColumn, Status> {
|
|
if requested_column.eq_ignore_ascii_case("position") {
|
|
return Ok(ResolvedOrderColumn::Position);
|
|
}
|
|
|
|
let requested_column = if requested_column.eq_ignore_ascii_case("row_display_column") {
|
|
table_row_display_column(pool, profile_name, table_name).await?
|
|
} else {
|
|
requested_column.to_string()
|
|
};
|
|
let physical_to_display =
|
|
table_physical_to_display_map(pool, profile_name, table_name).await?;
|
|
let physical_column = physical_to_display
|
|
.iter()
|
|
.find(|(physical, display)| {
|
|
physical.eq_ignore_ascii_case(&requested_column)
|
|
|| display.eq_ignore_ascii_case(&requested_column)
|
|
})
|
|
.map(|(physical, _)| physical.clone())
|
|
.unwrap_or_else(|| requested_column.clone());
|
|
|
|
let physical_column = sqlx::query_scalar::<_, String>(
|
|
r#"
|
|
SELECT column_name
|
|
FROM information_schema.columns
|
|
WHERE table_schema = $1 AND table_name = $2 AND LOWER(column_name) = LOWER($3)
|
|
"#,
|
|
)
|
|
.bind(profile_name)
|
|
.bind(table_name)
|
|
.bind(&physical_column)
|
|
.fetch_optional(pool)
|
|
.await
|
|
.map_err(|e| Status::internal(format!("Order column lookup failed: {}", e)))?;
|
|
let Some(physical_column) = physical_column else {
|
|
return Err(Status::invalid_argument(format!(
|
|
"Order column '{}' was not found in table '{}.{}'",
|
|
requested_column, profile_name, table_name
|
|
)));
|
|
};
|
|
|
|
Ok(ResolvedOrderColumn::Column(physical_column))
|
|
}
|
|
|
|
fn order_clause(column: &ResolvedOrderColumn, direction: SearchOrderDirection) -> String {
|
|
let direction = order_direction_sql(direction);
|
|
match column {
|
|
ResolvedOrderColumn::Position => {
|
|
format!("picker_position {}", direction)
|
|
}
|
|
ResolvedOrderColumn::Column(column) => {
|
|
let quoted_column = format!("\"{}\"", column.replace('"', "\"\""));
|
|
format!("{} {} NULLS LAST, id ASC", quoted_column, direction)
|
|
}
|
|
}
|
|
}
|
|
|
|
fn ranked_order_clause(
|
|
column: &ResolvedOrderColumn,
|
|
direction: SearchOrderDirection,
|
|
) -> String {
|
|
let direction = order_direction_sql(direction);
|
|
match column {
|
|
ResolvedOrderColumn::Position => {
|
|
format!("candidate.candidate_group DESC, positioned.picker_position {}", direction)
|
|
}
|
|
ResolvedOrderColumn::Column(column) => {
|
|
let quoted_column = format!("\"{}\"", column.replace('"', "\"\""));
|
|
format!(
|
|
"candidate.candidate_group DESC, positioned.{} {} NULLS LAST, positioned.id ASC",
|
|
quoted_column, direction
|
|
)
|
|
}
|
|
}
|
|
}
|
|
|
|
fn relevance_group(score: f32, best_score: f32) -> i32 {
|
|
if best_score <= 0.0 || !best_score.is_finite() || !score.is_finite() {
|
|
return 0;
|
|
}
|
|
|
|
((score.max(0.0) / best_score / SEARCH_SCORE_GROUP_WIDTH).floor() as i32).clamp(0, 3)
|
|
}
|
|
|
|
async fn fetch_ordered_rows(
|
|
pool: &PgPool,
|
|
profile_name: &str,
|
|
table_name: &str,
|
|
order: Option<&NormalizedSearchOrder>,
|
|
limit: usize,
|
|
offset: usize,
|
|
) -> Result<Vec<Hit>, Status> {
|
|
let physical_to_display = table_physical_to_display_map(pool, profile_name, table_name).await?;
|
|
let display_column = table_row_display_column(pool, profile_name, table_name).await?;
|
|
let (resolved_order, direction) = match order {
|
|
Some(order) => (
|
|
resolve_order_column(pool, profile_name, table_name, &order.column).await?,
|
|
order.direction,
|
|
),
|
|
None => (ResolvedOrderColumn::Position, SearchOrderDirection::Desc),
|
|
};
|
|
let sql = format!(
|
|
"WITH positioned AS (\
|
|
SELECT t.*, ROW_NUMBER() OVER (ORDER BY id ASC) AS picker_position \
|
|
FROM {} t WHERE deleted = FALSE\
|
|
) \
|
|
SELECT id, to_jsonb(positioned) - 'picker_position' AS data, picker_position \
|
|
FROM positioned ORDER BY {} LIMIT $1 OFFSET $2",
|
|
qualify_profile_table(profile_name, table_name),
|
|
order_clause(&resolved_order, direction),
|
|
);
|
|
|
|
let rows = sqlx::query(AssertSqlSafe(sql))
|
|
.bind(limit as i64)
|
|
.bind(offset as i64)
|
|
.fetch_all(pool)
|
|
.await
|
|
.map_err(|e| Status::internal(format!("DB query for default results failed: {}", e)))?;
|
|
|
|
Ok(rows
|
|
.into_iter()
|
|
.map(|row| {
|
|
let id: i64 = row.try_get("id").unwrap_or_default();
|
|
let json_data: serde_json::Value = row.try_get("data").unwrap_or_default();
|
|
let position: i64 = row.try_get("picker_position").unwrap_or_default();
|
|
let json_data = remap_json_to_display_names(json_data, &physical_to_display);
|
|
let row_display_value = row_display_value(&json_data, &display_column);
|
|
Hit {
|
|
id,
|
|
score: 0.0,
|
|
content_json: json_data.to_string(),
|
|
table_name: table_name.to_string(),
|
|
row_display_value,
|
|
row_display_column: display_column.clone(),
|
|
position: u64::try_from(position).ok(),
|
|
}
|
|
})
|
|
.collect())
|
|
}
|
|
|
|
async fn run_search(
|
|
pool: &PgPool,
|
|
profile: &ProfileIndex,
|
|
profile_name: &str,
|
|
table_filter: Option<&str>,
|
|
free_query: &str,
|
|
must: &[SearchConstraint],
|
|
limit: usize,
|
|
offset: usize,
|
|
order: Option<&NormalizedSearchOrder>,
|
|
) -> Result<Vec<Hit>, Status> {
|
|
let master_query = build_master_query(
|
|
&profile.index,
|
|
&profile.fields,
|
|
free_query,
|
|
must,
|
|
table_filter,
|
|
)?;
|
|
|
|
let searcher = profile.reader.searcher();
|
|
let window_limit = if order.is_some() {
|
|
searcher.num_docs() as usize
|
|
} else {
|
|
offset.saturating_add(limit)
|
|
};
|
|
if window_limit == 0 {
|
|
return Ok(Vec::new());
|
|
}
|
|
let top_docs = searcher
|
|
.search(&*master_query, &TopDocs::with_limit(window_limit).order_by_score())
|
|
.map_err(|e| Status::internal(format!("Search failed: {}", e)))?;
|
|
|
|
if top_docs.is_empty() {
|
|
return Ok(vec![]);
|
|
}
|
|
|
|
let best_score = top_docs.first().map(|(score, _)| *score).unwrap_or_default();
|
|
let score_floor = if best_score > 0.0 {
|
|
best_score * SEARCH_SCORE_RELATIVE_FLOOR
|
|
} else {
|
|
0.0
|
|
};
|
|
let eligible_docs = top_docs
|
|
.into_iter()
|
|
.filter(|(score, _)| *score >= score_floor)
|
|
.collect::<Vec<_>>();
|
|
let page_docs = if order.is_some() {
|
|
eligible_docs
|
|
} else {
|
|
eligible_docs
|
|
.into_iter()
|
|
.skip(offset)
|
|
.take(limit)
|
|
.collect::<Vec<_>>()
|
|
};
|
|
|
|
let mut candidates: Vec<(f32, i64, String)> = Vec::with_capacity(page_docs.len());
|
|
for (score, doc_address) in page_docs {
|
|
let doc: TantivyDocument = searcher
|
|
.doc(doc_address)
|
|
.map_err(|e| Status::internal(format!("Failed to retrieve document: {}", e)))?;
|
|
let Some(pg_id) = doc
|
|
.get_first(profile.fields.pg_id)
|
|
.and_then(|value| value.as_u64())
|
|
else {
|
|
continue;
|
|
};
|
|
let Some(table_name) = doc
|
|
.get_first(profile.fields.table_name)
|
|
.and_then(|value| value.as_str())
|
|
else {
|
|
continue;
|
|
};
|
|
candidates.push((score, pg_id as i64, table_name.to_string()));
|
|
}
|
|
|
|
if candidates.is_empty() {
|
|
return Ok(vec![]);
|
|
}
|
|
|
|
if let Some(order) = order {
|
|
let table_name = table_filter.expect("ordered searches require a normalized table filter");
|
|
return fetch_ordered_candidate_rows(
|
|
pool,
|
|
profile_name,
|
|
table_name,
|
|
&candidates,
|
|
order,
|
|
limit,
|
|
offset,
|
|
)
|
|
.await;
|
|
}
|
|
|
|
let mut ids_by_table: HashMap<String, Vec<i64>> = HashMap::new();
|
|
for (_, pg_id, table_name) in &candidates {
|
|
ids_by_table
|
|
.entry(table_name.clone())
|
|
.or_default()
|
|
.push(*pg_id);
|
|
}
|
|
|
|
let mut content_map: HashMap<(String, i64), (String, String, String)> = HashMap::new();
|
|
for (table_name, pg_ids) in ids_by_table {
|
|
validate_identifier(&table_name, "table_name")?;
|
|
let physical_to_display =
|
|
table_physical_to_display_map(pool, profile_name, &table_name).await?;
|
|
let display_column = table_row_display_column(pool, profile_name, &table_name).await?;
|
|
let sql = format!(
|
|
"SELECT id, to_jsonb(t) AS data FROM {} t WHERE deleted = FALSE AND id = ANY($1)",
|
|
qualify_profile_table(profile_name, &table_name)
|
|
);
|
|
|
|
let rows = sqlx::query(AssertSqlSafe(sql))
|
|
.bind(&pg_ids)
|
|
.fetch_all(pool)
|
|
.await
|
|
.map_err(|e| Status::internal(format!("Database query failed: {}", e)))?;
|
|
|
|
for row in rows {
|
|
let id: i64 = row.try_get("id").unwrap_or_default();
|
|
let json_data: serde_json::Value = row.try_get("data").unwrap_or_default();
|
|
let json_data = remap_json_to_display_names(json_data, &physical_to_display);
|
|
let display_value = row_display_value(&json_data, &display_column);
|
|
content_map.insert(
|
|
(table_name.clone(), id),
|
|
(json_data.to_string(), display_value, display_column.clone()),
|
|
);
|
|
}
|
|
}
|
|
|
|
Ok(candidates
|
|
.into_iter()
|
|
.filter_map(|(score, pg_id, table_name)| {
|
|
content_map
|
|
.get(&(table_name.clone(), pg_id))
|
|
.map(|(content_json, row_display_value, row_display_column)| Hit {
|
|
id: pg_id,
|
|
score,
|
|
content_json: content_json.clone(),
|
|
table_name,
|
|
row_display_value: row_display_value.clone(),
|
|
row_display_column: row_display_column.clone(),
|
|
position: None,
|
|
})
|
|
})
|
|
.collect())
|
|
}
|
|
|
|
async fn fetch_ordered_candidate_rows(
|
|
pool: &PgPool,
|
|
profile_name: &str,
|
|
table_name: &str,
|
|
candidates: &[(f32, i64, String)],
|
|
order: &NormalizedSearchOrder,
|
|
limit: usize,
|
|
offset: usize,
|
|
) -> Result<Vec<Hit>, Status> {
|
|
let physical_to_display = table_physical_to_display_map(pool, profile_name, table_name).await?;
|
|
let display_column = table_row_display_column(pool, profile_name, table_name).await?;
|
|
let resolved_order =
|
|
resolve_order_column(pool, profile_name, table_name, &order.column).await?;
|
|
let sql = format!(
|
|
"WITH positioned AS (\
|
|
SELECT t.*, ROW_NUMBER() OVER (ORDER BY id ASC) AS picker_position \
|
|
FROM {} t WHERE deleted = FALSE\
|
|
) \
|
|
SELECT positioned.id, to_jsonb(positioned) - 'picker_position' AS data, \
|
|
picker_position, candidate_score \
|
|
FROM positioned \
|
|
JOIN UNNEST($1::BIGINT[], $2::REAL[], $3::INTEGER[]) \
|
|
AS candidate(candidate_id, candidate_score, candidate_group) \
|
|
ON candidate_id = positioned.id \
|
|
ORDER BY {} LIMIT $4 OFFSET $5",
|
|
qualify_profile_table(profile_name, table_name),
|
|
ranked_order_clause(&resolved_order, order.direction),
|
|
);
|
|
let ids = candidates
|
|
.iter()
|
|
.map(|(_, id, _)| *id)
|
|
.collect::<Vec<_>>();
|
|
let scores = candidates
|
|
.iter()
|
|
.map(|(score, _, _)| *score)
|
|
.collect::<Vec<_>>();
|
|
let best_score = scores
|
|
.iter()
|
|
.copied()
|
|
.max_by(f32::total_cmp)
|
|
.unwrap_or_default();
|
|
let groups = scores
|
|
.iter()
|
|
.map(|score| relevance_group(*score, best_score))
|
|
.collect::<Vec<_>>();
|
|
let rows = sqlx::query(AssertSqlSafe(sql))
|
|
.bind(&ids)
|
|
.bind(&scores)
|
|
.bind(&groups)
|
|
.bind(limit as i64)
|
|
.bind(offset as i64)
|
|
.fetch_all(pool)
|
|
.await
|
|
.map_err(|e| Status::internal(format!("Ordered search query failed: {}", e)))?;
|
|
|
|
Ok(rows
|
|
.into_iter()
|
|
.map(|row| {
|
|
let id: i64 = row.try_get("id").unwrap_or_default();
|
|
let json_data: serde_json::Value = row.try_get("data").unwrap_or_default();
|
|
let position: i64 = row.try_get("picker_position").unwrap_or_default();
|
|
let score: f32 = row.try_get("candidate_score").unwrap_or_default();
|
|
let json_data = remap_json_to_display_names(json_data, &physical_to_display);
|
|
let display_value = row_display_value(&json_data, &display_column);
|
|
Hit {
|
|
id,
|
|
score,
|
|
content_json: json_data.to_string(),
|
|
table_name: table_name.to_string(),
|
|
row_display_value: display_value,
|
|
row_display_column: display_column.clone(),
|
|
position: u64::try_from(position).ok(),
|
|
}
|
|
})
|
|
.collect())
|
|
}
|
|
|
|
#[tonic::async_trait]
|
|
impl Searcher for SearcherService {
|
|
async fn search(
|
|
&self,
|
|
request: Request<SearchRequest>,
|
|
) -> Result<Response<SearchResponse>, Status> {
|
|
self.run_rpc(request).await
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use common::proto::komp_ac::search::{SearchOrder, SearchOrderDirection};
|
|
|
|
fn request(order: Option<SearchOrder>, table_name: Option<&str>) -> SearchRequest {
|
|
SearchRequest {
|
|
profile_name: "finance".to_string(),
|
|
table_name: table_name.map(str::to_string),
|
|
free_query: String::new(),
|
|
must: Vec::new(),
|
|
limit: Some(20),
|
|
offset: Some(40),
|
|
order,
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn normalizes_typed_column_order() {
|
|
let normalized = normalize_request(request(
|
|
Some(SearchOrder {
|
|
column: "created_at".to_string(),
|
|
direction: SearchOrderDirection::Desc as i32,
|
|
}),
|
|
Some("invoice"),
|
|
))
|
|
.unwrap();
|
|
let order = normalized.order.unwrap();
|
|
|
|
assert_eq!(order.column, "created_at");
|
|
assert_eq!(order.direction, SearchOrderDirection::Desc);
|
|
}
|
|
|
|
#[test]
|
|
fn ordered_request_requires_a_table() {
|
|
let result = normalize_request(request(
|
|
Some(SearchOrder {
|
|
column: "position".to_string(),
|
|
direction: SearchOrderDirection::Desc as i32,
|
|
}),
|
|
None,
|
|
));
|
|
|
|
assert!(result.is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn column_order_has_stable_id_tie_breaker() {
|
|
assert_eq!(
|
|
order_clause(
|
|
&ResolvedOrderColumn::Column("1".to_string()),
|
|
SearchOrderDirection::Asc,
|
|
),
|
|
"\"1\" ASC NULLS LAST, id ASC"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn ranked_order_preserves_relevance_before_column_order() {
|
|
assert_eq!(
|
|
ranked_order_clause(
|
|
&ResolvedOrderColumn::Column("created_at".to_string()),
|
|
SearchOrderDirection::Desc,
|
|
),
|
|
"candidate.candidate_group DESC, positioned.\"created_at\" DESC NULLS LAST, positioned.id ASC"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn relevance_groups_nearby_high_scores_and_separates_low_scores() {
|
|
let best_score = 8.7;
|
|
|
|
assert_eq!(relevance_group(8.7, best_score), 3);
|
|
assert_eq!(relevance_group(8.6, best_score), 3);
|
|
assert_eq!(relevance_group(8.4, best_score), 3);
|
|
assert_eq!(relevance_group(7.0, best_score), 3);
|
|
assert_eq!(relevance_group(3.0, best_score), 1);
|
|
}
|
|
}
|