From 77deb033f7de02239dab33c1502db02c89b343ce Mon Sep 17 00:00:00 2001 From: aecsocket <43144841+aecsocket@users.noreply.github.com> Date: Wed, 15 Jul 2026 13:49:33 +0100 Subject: [PATCH] wip: projects/versions collections --- apps/labrinth/src/queue/server_ping.rs | 20 +- .../routes/internal/moderation/tech_review.rs | 1 - apps/labrinth/src/routes/v3/organizations.rs | 3 - .../src/routes/v3/project_creation.rs | 2 - .../src/routes/v3/project_creation/new.rs | 1 - apps/labrinth/src/routes/v3/projects.rs | 82 +- apps/labrinth/src/routes/v3/versions.rs | 13 +- apps/labrinth/src/search/backend/common.rs | 17 - apps/labrinth/src/search/backend/mod.rs | 4 +- .../src/search/backend/typesense/mod.rs | 868 +++++++++++++----- apps/labrinth/src/search/incremental.rs | 46 +- .../src/search/incremental/consume.rs | 45 +- apps/labrinth/src/search/indexing.rs | 329 ++++--- apps/labrinth/src/search/mod.rs | 37 +- scripts/convert-typesense-project-docs.py | 443 +++++++++ 15 files changed, 1446 insertions(+), 465 deletions(-) create mode 100755 scripts/convert-typesense-project-docs.py diff --git a/apps/labrinth/src/queue/server_ping.rs b/apps/labrinth/src/queue/server_ping.rs index 95359701d2..5b6305961b 100644 --- a/apps/labrinth/src/queue/server_ping.rs +++ b/apps/labrinth/src/queue/server_ping.rs @@ -1,9 +1,9 @@ use crate::database::DBProject; -use crate::database::models::{DBProjectId, DBVersionId}; +use crate::database::models::DBProjectId; use crate::database::redis::RedisPool; use crate::env::ENV; use crate::models::exp; -use crate::models::ids::{ProjectId, VersionId}; +use crate::models::ids::ProjectId; use crate::models::projects::ProjectStatus; use crate::search::incremental::IncrementalSearchQueue; use crate::{database::PgPool, util::error::Context}; @@ -175,26 +175,14 @@ impl ServerPingQueue { } if updated_project { - let version_ids = sqlx::query_scalar!( - "SELECT id FROM versions WHERE mod_id = $1", - DBProjectId::from(*project_id) as DBProjectId, - ) - .fetch_all(&self.db) - .await - .wrap_err("failed to fetch project version IDs")? - .into_iter() - .map(|version_id| VersionId::from(DBVersionId(version_id))) - .collect::>(); - let clear_cache = DBProject::clear_cache( (*project_id).into(), None, None, &self.redis, ); - let queue_search = self - .incremental_search_queue - .push(*project_id, version_ids); + let queue_search = + self.incremental_search_queue.push(*project_id); let (clear_cache_result, _) = join(clear_cache, queue_search).await; diff --git a/apps/labrinth/src/routes/internal/moderation/tech_review.rs b/apps/labrinth/src/routes/internal/moderation/tech_review.rs index 0e1df4ec89..745509e834 100644 --- a/apps/labrinth/src/routes/internal/moderation/tech_review.rs +++ b/apps/labrinth/src/routes/internal/moderation/tech_review.rs @@ -1161,7 +1161,6 @@ pub async fn submit_report( if verdict == DelphiVerdict::Unsafe { crate::routes::v3::projects::clear_project_cache_and_queue_search( - &pool, &redis, &search_state, project_id, diff --git a/apps/labrinth/src/routes/v3/organizations.rs b/apps/labrinth/src/routes/v3/organizations.rs index 971e1fe191..8b71d75912 100644 --- a/apps/labrinth/src/routes/v3/organizations.rs +++ b/apps/labrinth/src/routes/v3/organizations.rs @@ -810,7 +810,6 @@ pub async fn organization_delete( for project_id in organization_project_ids { super::projects::clear_project_cache_and_queue_search( - &pool, &redis, &search_state, project_id, @@ -980,7 +979,6 @@ pub async fn organization_projects_add( ) .await?; super::projects::clear_project_cache_and_queue_search( - &pool, &redis, &search_state, project_item.inner.id, @@ -1173,7 +1171,6 @@ pub async fn organization_projects_remove( ) .await?; super::projects::clear_project_cache_and_queue_search( - &pool, &redis, &search_state, project_item.inner.id, diff --git a/apps/labrinth/src/routes/v3/project_creation.rs b/apps/labrinth/src/routes/v3/project_creation.rs index 137ff5809f..3640c34014 100644 --- a/apps/labrinth/src/routes/v3/project_creation.rs +++ b/apps/labrinth/src/routes/v3/project_creation.rs @@ -358,7 +358,6 @@ pub async fn project_create_internal( } else { transaction.commit().await?; super::projects::clear_project_cache_and_queue_search( - &client, &redis, &search_state, project_id.into(), @@ -425,7 +424,6 @@ pub async fn project_create_with_id( } else { transaction.commit().await?; super::projects::clear_project_cache_and_queue_search( - &client, &redis, &search_state, project_id.into(), diff --git a/apps/labrinth/src/routes/v3/project_creation/new.rs b/apps/labrinth/src/routes/v3/project_creation/new.rs index 51f5a1f3e8..968f2ba308 100644 --- a/apps/labrinth/src/routes/v3/project_creation/new.rs +++ b/apps/labrinth/src/routes/v3/project_creation/new.rs @@ -345,7 +345,6 @@ pub async fn create( .wrap_internal_err("failed to commit transaction")?; super::super::projects::clear_project_cache_and_queue_search( - &db, &redis, &search_state, project_id.into(), diff --git a/apps/labrinth/src/routes/v3/projects.rs b/apps/labrinth/src/routes/v3/projects.rs index df93ab527a..43ced679ee 100644 --- a/apps/labrinth/src/routes/v3/projects.rs +++ b/apps/labrinth/src/routes/v3/projects.rs @@ -72,35 +72,42 @@ pub fn project_config(cfg: &mut actix_web::web::ServiceConfig) { } pub async fn clear_project_cache_and_queue_search( - pool: &PgPool, redis: &RedisPool, search_state: &SearchState, project_id: db_ids::DBProjectId, slug: Option, clear_dependencies: Option, ) -> Result<(), ApiError> { - let version_ids = sqlx::query_scalar!( - "SELECT id FROM versions WHERE mod_id = $1", - project_id as db_ids::DBProjectId, - ) - .fetch_all(pool) - .await - .wrap_internal_err("failed to fetch project version IDs")? - .into_iter() - .map(|version_id| VersionId::from(db_ids::DBVersionId(version_id))) - .collect::>(); - - clear_project_cache_and_queue_search_versions( + clear_project_cache_and_queue_search_inner( redis, search_state, project_id, slug, clear_dependencies, - version_ids, ) .await } +pub async fn clear_project_cache_and_queue_search_inner( + redis: &RedisPool, + search_state: &SearchState, + project_id: db_ids::DBProjectId, + slug: Option, + clear_dependencies: Option, +) -> Result<(), ApiError> { + db_models::DBProject::clear_cache( + project_id, + slug, + clear_dependencies, + redis, + ) + .await?; + + search_state.queue.push(project_id.into()).await; + + Ok(()) +} + pub async fn clear_project_cache_and_queue_search_versions( redis: &RedisPool, search_state: &SearchState, @@ -119,7 +126,7 @@ pub async fn clear_project_cache_and_queue_search_versions( search_state .queue - .push(project_id.into(), version_ids) + .push_versions(project_id.into(), version_ids) .await; Ok(()) @@ -1133,6 +1140,9 @@ pub async fn project_edit_internal( Ok(()) } + let reindex_version_project_types = + new_project.minecraft_java_server.is_some(); + update( &mut transaction, id, @@ -1202,15 +1212,26 @@ pub async fn project_edit_internal( transaction.commit().await?; - clear_project_cache_and_queue_search( - &pool, - &redis, - &search_state, - project_item.inner.id, - project_item.inner.slug, - None, - ) - .await?; + if reindex_version_project_types { + clear_project_cache_and_queue_search_versions( + &redis, + &search_state, + project_item.inner.id, + project_item.inner.slug, + None, + project_item.versions.iter().copied().map(VersionId::from), + ) + .await?; + } else { + clear_project_cache_and_queue_search( + &redis, + &search_state, + project_item.inner.id, + project_item.inner.slug, + None, + ) + .await?; + } // Remove no longer searchable projects from search index if let (true, Some(false)) = ( @@ -1764,7 +1785,6 @@ pub async fn projects_edit( for (project_id, slug) in changed_projects { clear_project_cache_and_queue_search( - &pool, &redis, &search_state, project_id, @@ -1996,7 +2016,6 @@ pub async fn project_icon_edit_internal( transaction.commit().await?; clear_project_cache_and_queue_search( - &pool, &redis, &search_state, project_item.inner.id, @@ -2116,7 +2135,6 @@ pub async fn delete_project_icon_internal( transaction.commit().await?; clear_project_cache_and_queue_search( - &pool, &redis, &search_state, project_item.inner.id, @@ -2317,7 +2335,6 @@ pub async fn add_gallery_item_internal( transaction.commit().await?; clear_project_cache_and_queue_search( - &pool, &redis, &search_state, project_item.inner.id, @@ -2536,7 +2553,6 @@ pub async fn edit_gallery_item_internal( transaction.commit().await?; clear_project_cache_and_queue_search( - &pool, &redis, &search_state, project_item.inner.id, @@ -2685,7 +2701,6 @@ pub async fn delete_gallery_item_internal( transaction.commit().await?; clear_project_cache_and_queue_search( - &pool, &redis, &search_state, project_item.inner.id, @@ -2834,6 +2849,13 @@ pub async fn project_delete_internal( &redis, ) .await?; + search_state + .backend + .remove_project_version_documents(&[project.inner.id.into()]) + .await + .wrap_internal_err( + "failed to remove project versions from search index", + )?; search_state .backend .remove_project_documents(&[project.inner.id.into()]) diff --git a/apps/labrinth/src/routes/v3/versions.rs b/apps/labrinth/src/routes/v3/versions.rs index 3cf7d80c01..f912956fd1 100644 --- a/apps/labrinth/src/routes/v3/versions.rs +++ b/apps/labrinth/src/routes/v3/versions.rs @@ -27,6 +27,7 @@ use crate::models::teams::ProjectPermissions; use crate::queue::file_scan::get_files_missing_attribution; use crate::queue::session::AuthQueue; use crate::routes::internal::delphi; +use crate::search::incremental::consume::reindex_project; use crate::search::{SearchBackend, SearchState}; use crate::util::error::Context; use crate::util::img; @@ -1257,9 +1258,17 @@ pub async fn version_delete( ) .await?; search_backend - .remove_documents(&[version.inner.id.into()]) + .remove_version_documents(&[version.inner.id.into()]) .await - .wrap_internal_err("failed to remove documents")?; + .wrap_internal_err("failed to remove version search document")?; + reindex_project( + &pool, + &redis, + search_backend.as_ref(), + version.inner.project_id.into(), + ) + .await + .wrap_internal_err("failed to reindex project")?; if result.is_some() { Ok(HttpResponse::NoContent().body("")) } else { diff --git a/apps/labrinth/src/search/backend/common.rs b/apps/labrinth/src/search/backend/common.rs index 5192a83dc3..9589275186 100644 --- a/apps/labrinth/src/search/backend/common.rs +++ b/apps/labrinth/src/search/backend/common.rs @@ -50,14 +50,7 @@ pub enum SearchIndex { MinecraftJavaServerPlayersOnline, } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum SearchIndexName { - Projects, - ProjectsFiltered, -} - pub struct SearchSort { - pub index_name: SearchIndexName, pub index: SearchIndex, } @@ -65,16 +58,12 @@ pub fn parse_search_index( index: &str, new_filters: Option<&str>, ) -> Result { - let projects_name = SearchIndexName::Projects; - let projects_filtered_name = SearchIndexName::ProjectsFiltered; - // TODO: this is a dumb hack, the frontend should pass the project type it's filtering directly let is_server = new_filters .is_some_and(|f| f.contains("project_types = minecraft_java_server")); Ok(match index { "relevance" => SearchSort { - index_name: projects_name, index: if is_server { SearchIndex::MinecraftJavaServerVerifiedPlays2w } else { @@ -82,27 +71,21 @@ pub fn parse_search_index( }, }, "downloads" => SearchSort { - index_name: projects_filtered_name, index: SearchIndex::Downloads, }, "follows" => SearchSort { - index_name: projects_name, index: SearchIndex::Follows, }, "updated" | "date_modified" => SearchSort { - index_name: projects_name, index: SearchIndex::Updated, }, "newest" | "date_created" => SearchSort { - index_name: projects_name, index: SearchIndex::Newest, }, "minecraft_java_server.verified_plays_2w" => SearchSort { - index_name: projects_name, index: SearchIndex::MinecraftJavaServerVerifiedPlays2w, }, "minecraft_java_server.ping.data.players_online" => SearchSort { - index_name: projects_name, index: SearchIndex::MinecraftJavaServerPlayersOnline, }, i => return Err(ApiError::Request(eyre!("invalid index '{i}'"))), diff --git a/apps/labrinth/src/search/backend/mod.rs b/apps/labrinth/src/search/backend/mod.rs index 307cde1040..544cf2a557 100644 --- a/apps/labrinth/src/search/backend/mod.rs +++ b/apps/labrinth/src/search/backend/mod.rs @@ -2,7 +2,7 @@ mod common; pub mod typesense; pub use common::{ - ParsedSearchRequest, SearchIndex, SearchIndexName, SearchSort, - combined_search_filters, parse_search_index, parse_search_request, + ParsedSearchRequest, SearchIndex, SearchSort, combined_search_filters, + parse_search_index, parse_search_request, }; pub use typesense::{Typesense, TypesenseConfig}; diff --git a/apps/labrinth/src/search/backend/typesense/mod.rs b/apps/labrinth/src/search/backend/typesense/mod.rs index f34c730e26..8033006d60 100644 --- a/apps/labrinth/src/search/backend/typesense/mod.rs +++ b/apps/labrinth/src/search/backend/typesense/mod.rs @@ -3,11 +3,12 @@ use std::sync::LazyLock; use ariadne::ids::base62_impl::to_base62; use async_trait::async_trait; use eyre::{Result, eyre}; +use itertools::Itertools; use regex::Regex; use reqwest::Method; use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; -use tracing::{debug, info, warn}; +use tracing::{debug, info}; use crate::database::PgPool; use crate::database::redis::RedisPool; @@ -15,13 +16,13 @@ use crate::env::ENV; use crate::models::ids::{ProjectId, VersionId}; use crate::routes::ApiError; use crate::search::backend::{ - SearchIndex, SearchIndexName, combined_search_filters, parse_search_index, + SearchIndex, combined_search_filters, parse_search_index, parse_search_request, }; use crate::search::indexing::index_local; use crate::search::{ ResultSearchProject, SearchBackend, SearchField, SearchRequest, - SearchResults, TasksCancelFilter, UploadSearchProject, + SearchResults, TasksCancelFilter, UploadSearchProject, UploadSearchVersion, }; use crate::util::error::Context; @@ -239,6 +240,22 @@ impl TypesenseClient { Ok(()) } + async fn delete_alias_if_exists(&self, alias: &str) -> Result<()> { + let resp = self + .request(Method::DELETE, &format!("/aliases/{alias}")) + .send() + .await + .wrap_err("failed to DELETE Typesense alias")?; + if resp.status() == reqwest::StatusCode::NOT_FOUND { + return Ok(()); + } + if !resp.status().is_success() { + let body = resp.json::().await.unwrap_or_default(); + return Err(eyre!("failed to delete alias `{alias}`: {body}")); + } + Ok(()) + } + async fn collection_exists(&self, name: &str) -> Result { let resp = self .request(Method::GET, &format!("/collections/{name}")) @@ -304,21 +321,35 @@ impl TypesenseClient { )); } // Typesense always returns HTTP 200; individual lines signal per-doc success. - let error_count = body + let failures = body .lines() - .filter(|l| !l.trim().is_empty()) - .filter(|l| { - serde_json::from_str::(l) - .ok() - .and_then(|v| v["success"].as_bool()) - .map(|ok| !ok) - .unwrap_or(false) + .filter(|line| !line.trim().is_empty()) + .filter_map(|line| match serde_json::from_str::(line) { + Ok(result) if result["success"].as_bool() == Some(true) => None, + Ok(result) => Some( + result["error"] + .as_str() + .unwrap_or( + "Typesense returned an unsuccessful import result", + ) + .to_string(), + ), + Err(err) => Some(format!( + "failed to parse Typesense import result: {err}" + )), }) - .count(); - if error_count > 0 { - warn!( - "{error_count} document(s) failed to import into `{collection}`" - ); + .collect::>(); + if !failures.is_empty() { + let failure_count = failures.len(); + let errors = failures + .into_iter() + .unique() + .take(10) + .collect::>() + .join("; "); + return Err(eyre!( + "{failure_count} document(s) failed to import into `{collection}`: {errors}" + )); } Ok(()) } @@ -368,6 +399,18 @@ pub struct TypesenseFieldSpec { } impl SearchField { + const fn is_version_field(self) -> bool { + matches!( + self, + Self::Categories + | Self::ProjectTypes + | Self::Environment + | Self::GameVersions + | Self::ClientSide + | Self::ServerSide + ) + } + pub const fn typesense_spec(self) -> TypesenseFieldSpec { match self { SearchField::Categories => TypesenseFieldSpec { @@ -569,7 +612,7 @@ impl Typesense { Self { config, client } } - fn collection_schema(name: &str) -> Value { + fn project_collection_schema(name: &str) -> Value { let mut fields = vec![ json!({"name": "summary", "type": "string", "facet": false}), json!({"name": "slug", "type": "string", "facet": false}), @@ -585,6 +628,7 @@ impl Typesense { json!({"name": "minecraft_java_server.is_online", "type": "bool", "sort": true, "optional": true}), json!({"name": "minecraft_java_server.ping.data.players_online", "type": "int32", "sort": true, "optional": true}), json!({"name": "dependencies", "type": "object[]", "optional": true}), + json!({"name": "project_categories", "type": "string[]", "facet": true, "optional": true}), ]; fields.extend(TYPESENSE_SEARCH_FIELDS.iter().cloned()); @@ -596,6 +640,50 @@ impl Typesense { }) } + fn version_collection_schema( + name: &str, + projects_collection: &str, + ) -> Value { + use strum::IntoEnumIterator; + + let mut fields = SearchField::iter() + .filter(|field| field.is_version_field()) + .map(|field| { + let spec = field.typesense_spec(); + json!({ + "name": spec.path, + "type": spec.ty, + "facet": spec.facet, + "optional": spec.optional, + "token_separators": spec.token_separators, + }) + }) + .collect::>(); + fields.extend([ + json!({ + "name": "version_id", + "type": "string", + }), + json!({ + "name": "project_id", + "type": "string", + "reference": format!("{projects_collection}.id"), + "async_reference": true, + "cascade_delete": false, + }), + json!({ + "name": "version_published_timestamp", + "type": "int64", + "sort": true, + }), + ]); + + json!({ + "name": name, + "fields": fields, + }) + } + fn text_match_sort_field(request_config: &RequestConfig) -> String { match request_config.bucketing { Bucketing::Buckets(count) => { @@ -675,12 +763,7 @@ impl Typesense { request_config: &RequestConfig, ) -> Result<(String, String), ApiError> { let sort = parse_search_index(index, new_filters)?; - let alias = match sort.index_name { - SearchIndexName::Projects => self.config.get_alias_name("projects"), - SearchIndexName::ProjectsFiltered => { - self.config.get_alias_name("projects_filtered") - } - }; + let alias = self.config.get_alias_name("projects"); Ok((alias, self.get_sort_fields(sort.index, request_config))) } @@ -689,7 +772,10 @@ impl Typesense { /// Handles the new-style filter string, legacy facets JSON, and the legacy /// `filters`/`version` fields, translating each from Meilisearch filter /// syntax to Typesense filter syntax. - fn build_filter(info: &SearchRequest) -> Result, ApiError> { + fn build_filter( + info: &SearchRequest, + versions_collection: &str, + ) -> Result, ApiError> { let facet_part = if let Some(facets_json) = info.facets.as_deref() { Some( facets_to_typesense(facets_json) @@ -710,29 +796,61 @@ impl Typesense { let filter_part = new_filters_part.or(legacy_part); - Ok(match (facet_part, filter_part) { + let filter = match (facet_part, filter_part) { (Some(f), Some(l)) if !f.is_empty() && !l.is_empty() => { Some(format!("({f}) && ({l})")) } (Some(f), _) if !f.is_empty() => Some(f), (_, Some(l)) if !l.is_empty() => Some(l), _ => None, - }) + }; + + filter + .map(|filter| { + rewrite_filter_for_join(&filter, versions_collection) + .wrap_request_err("failed to rewrite search filter") + }) + .transpose() } - /// Ensures the alias and its backing collection both exist, creating them - /// when necessary so reads succeed before the first full index run. - async fn ensure_collection(&self, alias: &str) -> Result<()> { - if self.client.get_alias(alias).await?.is_some() { - return Ok(()); - } - let name = self.config.get_next_collection_name(alias, false); - if !self.client.collection_exists(&name).await? { + async fn ensure_collections(&self) -> Result<()> { + let projects_alias = self.config.get_alias_name("projects"); + let projects_collection = if let Some(collection) = + self.client.get_alias(&projects_alias).await? + { + collection + } else { + let collection = + self.config.get_next_collection_name(&projects_alias, false); + if !self.client.collection_exists(&collection).await? { + self.client + .create_collection(&Self::project_collection_schema( + &collection, + )) + .await?; + } self.client - .create_collection(&Self::collection_schema(&name)) + .upsert_alias(&projects_alias, &collection) + .await?; + collection + }; + + let versions_alias = self.config.get_alias_name("versions"); + if self.client.get_alias(&versions_alias).await?.is_none() { + let collection = + self.config.get_next_collection_name(&versions_alias, false); + if !self.client.collection_exists(&collection).await? { + self.client + .create_collection(&Self::version_collection_schema( + &collection, + &projects_collection, + )) + .await?; + } + self.client + .upsert_alias(&versions_alias, &collection) .await?; } - self.client.upsert_alias(alias, &name).await?; Ok(()) } @@ -780,31 +898,83 @@ impl Typesense { Ok(()) } - async fn existing_write_collections(&self) -> Result> { + async fn import_version_document_batches( + &self, + collections: &[String], + documents: &[UploadSearchVersion], + ) -> Result<()> { + let batch_size = self.config.import_batch_size.max(1); + + for batch in documents.chunks(batch_size) { + let jsonl = version_documents_to_jsonl(batch)?; + + for collection in collections { + info!( + collection, + document_count = batch.len(), + content_length_bytes = jsonl.len(), + "sending Typesense version document import" + ); + self.client + .import_documents(collection, jsonl.clone()) + .await?; + } + } + + Ok(()) + } + + async fn existing_write_collections( + &self, + alias: &str, + ) -> Result> { let mut collections = Vec::new(); - for alias in [ - self.config.get_alias_name("projects"), - self.config.get_alias_name("projects_filtered"), - ] { - let live = self.client.get_alias(&alias).await?; - let shadow_alt = self.config.get_next_collection_name(&alias, true); - let shadow_current = - self.config.get_next_collection_name(&alias, false); + let live = self.client.get_alias(alias).await?; + let shadow_alt = self.config.get_next_collection_name(alias, true); + let shadow_current = self.config.get_next_collection_name(alias, false); - for collection in - live.into_iter().chain([shadow_alt, shadow_current]) + for collection in live.into_iter().chain([shadow_alt, shadow_current]) { + if !collections.contains(&collection) + && self.client.collection_exists(&collection).await? { - if !collections.contains(&collection) - && self.client.collection_exists(&collection).await? - { - collections.push(collection); - } + collections.push(collection); } } Ok(collections) } + + async fn delete_from_write_collections( + &self, + alias: &str, + filter: &str, + ) -> Result<()> { + for collection in self.existing_write_collections(alias).await? { + self.delete_documents_by_filter_if_exists(&collection, filter) + .await?; + } + Ok(()) + } + + async fn delete_legacy_filtered_collections(&self) -> Result<()> { + let alias = self.config.get_alias_name("projects_filtered"); + let live = self.client.get_alias(&alias).await?; + let shadow_alt = self.config.get_next_collection_name(&alias, true); + let shadow_current = + self.config.get_next_collection_name(&alias, false); + + self.client.delete_alias_if_exists(&alias).await?; + for collection in live + .into_iter() + .chain([shadow_alt, shadow_current]) + .unique() + { + self.client.delete_collection_if_exists(&collection).await?; + } + + Ok(()) + } } #[async_trait] @@ -819,7 +989,8 @@ impl SearchBackend for Typesense { info.new_filters.as_deref(), &info.typesense_config, )?; - let filter_by = Self::build_filter(info)?; + let versions_alias = self.config.get_alias_name("versions"); + let filter_by = Self::build_filter(info, &versions_alias)?; let q = if parsed.query.is_empty() { "*" @@ -857,8 +1028,6 @@ impl SearchBackend for Typesense { ("sort_by", sort_by.to_string()), ("page", parsed.page.to_string()), ("per_page", parsed.hits_per_page.to_string()), - ("group_by", "project_id".to_string()), - ("group_limit", "1".to_string()), ( "max_candidates", info.typesense_config.max_candidates.to_string(), @@ -874,6 +1043,14 @@ impl SearchBackend for Typesense { } if let Some(filter) = &filter_by { params.push(("filter_by", filter.clone())); + if filter.contains(&format!("${versions_alias}(")) { + params.push(( + "include_fields", + format!( + "${versions_alias}(version_id, sort_by: version_published_timestamp:desc, limit:1, strategy: nest_array) as matching_versions" + ), + )); + } } let resp = self @@ -916,16 +1093,33 @@ impl SearchBackend for Typesense { let total_hits = body["found"].as_u64().unwrap_or(0) as usize; - let hits = body["grouped_hits"] + let hits = body["hits"] .as_array() .cloned() .unwrap_or_default() .into_iter() - .filter_map(|group| { - let hit = group["hits"].as_array()?.first()?.clone(); + .filter_map(|hit| { let mut doc = hit.get("document")?.clone(); if let Some(obj) = doc.as_object_mut() { obj.remove("id"); + let matching_version_id = + obj.remove("matching_versions").and_then(|versions| { + versions + .as_array() + .and_then(|versions| versions.first()) + .or_else(|| { + versions.as_object().map(|_| &versions) + }) + .and_then(|version| version.get("version_id")) + .and_then(Value::as_str) + .map(ToString::to_string) + }); + if let Some(version_id) = matching_version_id { + obj.insert( + "version_id".to_string(), + Value::String(version_id), + ); + } } let metadata = info.show_metadata.then(|| { @@ -967,54 +1161,55 @@ impl SearchBackend for Typesense { info!("starting project indexing"); let projects_alias = self.config.get_alias_name("projects"); - let filtered_alias = self.config.get_alias_name("projects_filtered"); + let versions_alias = self.config.get_alias_name("versions"); - // Guarantee current aliases exist so reads keep working during re-index. - self.ensure_collection(&projects_alias).await?; - self.ensure_collection(&filtered_alias).await?; + self.ensure_collections().await?; - // Toggle the shadow collection name between __current and __alt. let projects_current = self.client.get_alias(&projects_alias).await?; - let filtered_current = self.client.get_alias(&filtered_alias).await?; + let versions_current = self.client.get_alias(&versions_alias).await?; let projects_use_alt = !projects_current .as_deref() .is_some_and(|n| n.ends_with("__alt")); - let filtered_use_alt = !filtered_current + let versions_use_alt = !versions_current .as_deref() .is_some_and(|n| n.ends_with("__alt")); let projects_next = self .config .get_next_collection_name(&projects_alias, projects_use_alt); - let filtered_next = self + let versions_next = self .config - .get_next_collection_name(&filtered_alias, filtered_use_alt); + .get_next_collection_name(&versions_alias, versions_use_alt); - info!("shadow collections `{projects_next}` and `{filtered_next}`"); + info!("shadow collections `{projects_next}` and `{versions_next}`"); + self.client + .delete_collection_if_exists(&versions_next) + .await?; self.client .delete_collection_if_exists(&projects_next) .await?; self.client - .delete_collection_if_exists(&filtered_next) + .create_collection(&Self::project_collection_schema(&projects_next)) .await?; self.client - .create_collection(&Self::collection_schema(&projects_next)) - .await?; - self.client - .create_collection(&Self::collection_schema(&filtered_next)) + .create_collection(&Self::version_collection_schema( + &versions_next, + &projects_next, + )) .await?; let mut cursor = 0_i64; let mut chunk_idx = 0_usize; - let mut total = 0_usize; + let mut total_projects = 0_usize; + let mut total_versions = 0_usize; loop { info!("fetching index chunk {chunk_idx}"); chunk_idx += 1; - let (uploads, next_cursor) = index_local( + let (documents, next_cursor) = index_local( &ro_pool, &redis, cursor, @@ -1023,19 +1218,25 @@ impl SearchBackend for Typesense { .await .wrap_err("failed to fetch projects from local DB")?; - if uploads.is_empty() { + if documents.projects.is_empty() { info!( - "no more documents; indexed {total} in {chunk_idx} chunks" + "no more documents; indexed {total_projects} projects and {total_versions} versions in {chunk_idx} chunks" ); break; } - total += uploads.len(); + total_projects += documents.projects.len(); + total_versions += documents.versions.len(); cursor = next_cursor; self.import_document_batches( - &[projects_next.clone(), filtered_next.clone()], - &uploads, + std::slice::from_ref(&projects_next), + &documents.projects, + ) + .await?; + self.import_version_document_batches( + std::slice::from_ref(&versions_next), + &documents.versions, ) .await?; } @@ -1045,16 +1246,18 @@ impl SearchBackend for Typesense { .upsert_alias(&projects_alias, &projects_next) .await?; self.client - .upsert_alias(&filtered_alias, &filtered_next) + .upsert_alias(&versions_alias, &versions_next) .await?; info!("cleaning up old collections"); + if let Some(old) = versions_current { + self.client.delete_collection_if_exists(&old).await?; + } if let Some(old) = projects_current { self.client.delete_collection_if_exists(&old).await?; } - if let Some(old) = filtered_current { - self.client.delete_collection_if_exists(&old).await?; - } + + self.delete_legacy_filtered_collections().await?; info!("indexing complete"); Ok(()) @@ -1068,7 +1271,8 @@ impl SearchBackend for Typesense { return Ok(()); } - let collections = self.existing_write_collections().await?; + let alias = self.config.get_alias_name("projects"); + let collections = self.existing_write_collections(&alias).await?; debug!( ?collections, num_documents = documents.len(), @@ -1081,6 +1285,28 @@ impl SearchBackend for Typesense { Ok(()) } + async fn index_version_documents( + &self, + documents: &[UploadSearchVersion], + ) -> eyre::Result<()> { + if documents.is_empty() { + return Ok(()); + } + + let alias = self.config.get_alias_name("versions"); + let collections = self.existing_write_collections(&alias).await?; + debug!( + ?collections, + num_documents = documents.len(), + "Inserting version documents into collections", + ); + self.import_version_document_batches(&collections, documents) + .await?; + + debug!("Done importing version documents"); + Ok(()) + } + async fn remove_project_documents( &self, ids: &[ProjectId], @@ -1096,139 +1322,40 @@ impl SearchBackend for Typesense { .join(", "); let filter = format!("project_id:[{id_list}]"); - for alias in [ - self.config.get_alias_name("projects"), - self.config.get_alias_name("projects_filtered"), - ] { - debug!("Performing removal on alias {alias:?}"); - - let live = self.client.get_alias(&alias).await?; - debug!("Got live alias {live:?}"); - - let shadow_alt = self.config.get_next_collection_name(&alias, true); - debug!("Got shadow alt {shadow_alt:?}"); - - let shadow_current = - self.config.get_next_collection_name(&alias, false); - debug!("Got shadow current {shadow_current:?}"); - - let delete_live = async { - if let Some(collection) = live.as_deref() { - debug!("Working on collection {collection:?}"); - debug!( - filter_len = filter.len(), - "Collection exists, deleting by filter" - ); - self.delete_documents_by_filter_if_exists( - collection, &filter, - ) - .await?; - } - - Ok::<(), eyre::Report>(()) - }; - let delete_shadow_alt = async { - if live.as_deref() != Some(shadow_alt.as_str()) { - debug!("Working on collection {shadow_alt:?}"); - self.delete_documents_by_filter_if_exists( - &shadow_alt, - &filter, - ) - .await?; - } - - Ok::<(), eyre::Report>(()) - }; - let delete_shadow_current = async { - if live.as_deref() != Some(shadow_current.as_str()) { - debug!("Working on collection {shadow_current:?}"); - self.delete_documents_by_filter_if_exists( - &shadow_current, - &filter, - ) - .await?; - } - - Ok::<(), eyre::Report>(()) - }; - let (live_result, shadow_alt_result, shadow_current_result) = tokio::join!( - delete_live, - delete_shadow_alt, - delete_shadow_current - ); - live_result?; - shadow_alt_result?; - shadow_current_result?; - } + let projects_alias = self.config.get_alias_name("projects"); + self.delete_from_write_collections(&projects_alias, &filter) + .await?; debug!("Done"); Ok(()) } - async fn remove_documents(&self, ids: &[VersionId]) -> eyre::Result<()> { + async fn remove_project_version_documents( + &self, + ids: &[ProjectId], + ) -> eyre::Result<()> { if ids.is_empty() { return Ok(()); } - let id_list = ids - .iter() - .map(|id| to_base62(id.0)) - .collect::>() - .join(", "); - let filter = format!("id:[{id_list}]"); + let id_list = ids.iter().map(ToString::to_string).join(", "); + let filter = format!("project_id:[{id_list}]"); + let alias = self.config.get_alias_name("versions"); + self.delete_from_write_collections(&alias, &filter).await + } - for alias in [ - self.config.get_alias_name("projects"), - self.config.get_alias_name("projects_filtered"), - ] { - // Delete from both the live collection and any shadow collections. - let live = self.client.get_alias(&alias).await?; - let shadow_alt = self.config.get_next_collection_name(&alias, true); - let shadow_current = - self.config.get_next_collection_name(&alias, false); - - let delete_live = async { - if let Some(collection) = live.as_deref() { - self.delete_documents_by_filter_if_exists( - collection, &filter, - ) - .await?; - } - - Ok::<(), eyre::Report>(()) - }; - let delete_shadow_alt = async { - if live.as_deref() != Some(shadow_alt.as_str()) { - self.delete_documents_by_filter_if_exists( - &shadow_alt, - &filter, - ) - .await?; - } - - Ok::<(), eyre::Report>(()) - }; - let delete_shadow_current = async { - if live.as_deref() != Some(shadow_current.as_str()) { - self.delete_documents_by_filter_if_exists( - &shadow_current, - &filter, - ) - .await?; - } - - Ok::<(), eyre::Report>(()) - }; - let (live_result, shadow_alt_result, shadow_current_result) = tokio::join!( - delete_live, - delete_shadow_alt, - delete_shadow_current - ); - live_result?; - shadow_alt_result?; - shadow_current_result?; + async fn remove_version_documents( + &self, + ids: &[VersionId], + ) -> eyre::Result<()> { + if ids.is_empty() { + return Ok(()); } - Ok(()) + + let id_list = ids.iter().map(ToString::to_string).join(", "); + let filter = format!("id:[{id_list}]"); + let alias = self.config.get_alias_name("versions"); + self.delete_from_write_collections(&alias, &filter).await } async fn tasks(&self) -> eyre::Result { @@ -1246,7 +1373,7 @@ impl SearchBackend for Typesense { /// Serialises a batch of [`UploadSearchProject`]s to a JSONL string suitable /// for the Typesense bulk-import endpoint. Each document gets an `id` field -/// equal to `version_id` so Typesense can use it as the primary key. +/// equal to `project_id` so Typesense can use it as the primary key. fn documents_to_jsonl(uploads: &[UploadSearchProject]) -> Result { let mut out = String::new(); for upload in uploads { @@ -1254,7 +1381,7 @@ fn documents_to_jsonl(uploads: &[UploadSearchProject]) -> Result { .wrap_err("failed to serialise UploadSearchProject")?; if let Some(obj) = doc.as_object_mut() { let id = obj - .get("version_id") + .get("project_id") .and_then(Value::as_str) .unwrap_or_default() .to_string(); @@ -1278,6 +1405,271 @@ fn documents_to_jsonl(uploads: &[UploadSearchProject]) -> Result { Ok(out) } +fn version_documents_to_jsonl( + uploads: &[UploadSearchVersion], +) -> Result { + let mut out = String::new(); + for upload in uploads { + let mut document = serde_json::to_value(upload) + .wrap_err("failed to serialise UploadSearchVersion")?; + if let Some(object) = document.as_object_mut() { + object.insert( + "id".to_string(), + Value::String(upload.version_id.clone()), + ); + } + out.push_str(&serde_json::to_string(&document)?); + out.push('\n'); + } + Ok(out) +} + +#[derive(Clone, Default)] +struct JoinedFilterClause { + project: Vec, + version: Vec, +} + +fn rewrite_filter_for_join( + filter: &str, + versions_collection: &str, +) -> Result { + const MAX_CLAUSES: usize = 256; + + fn parse(expression: &str) -> Result> { + let expression = trim_outer_parentheses(expression.trim()); + + let or_parts = split_top_level(expression, "||"); + if or_parts.len() > 1 { + let mut clauses = Vec::new(); + for part in or_parts { + clauses.extend(parse(part)?); + if clauses.len() > MAX_CLAUSES { + return Err(eyre!( + "search filter has too many boolean clauses" + )); + } + } + return Ok(clauses); + } + + let and_parts = split_top_level(expression, "&&"); + if and_parts.len() > 1 { + let mut clauses = vec![JoinedFilterClause::default()]; + for part in and_parts { + let right = parse(part)?; + if clauses.len().saturating_mul(right.len()) > MAX_CLAUSES { + return Err(eyre!( + "search filter has too many boolean clauses" + )); + } + clauses = clauses + .into_iter() + .cartesian_product(right) + .map(|(mut left, right)| { + left.project.extend(right.project); + left.version.extend(right.version); + left + }) + .collect(); + } + return Ok(clauses); + } + + let field = filter_field(expression).ok_or_else(|| { + eyre!("could not determine filter field in `{expression}`") + })?; + let mut clause = JoinedFilterClause::default(); + if field == "categories" { + let project_expression = + expression.replacen("categories", "project_categories", 1); + if is_negative_filter(expression) { + clause.project.push(project_expression); + clause.version.push(expression.to_string()); + Ok(vec![clause]) + } else { + Ok(vec![ + JoinedFilterClause { + project: vec![project_expression], + version: Vec::new(), + }, + JoinedFilterClause { + project: Vec::new(), + version: vec![expression.to_string()], + }, + ]) + } + } else { + if is_version_filter_field(field) { + clause.version.push(expression.to_string()); + } else { + clause.project.push(expression.to_string()); + } + Ok(vec![clause]) + } + } + + let clauses = parse(filter)?; + Ok(clauses + .into_iter() + .map(|clause| { + let mut parts = clause.project; + if !clause.version.is_empty() { + parts.push(format!( + "${versions_collection}({})", + clause.version.join(" && ") + )); + } + if parts.len() == 1 { + parts.pop().unwrap_or_default() + } else { + format!("({})", parts.join(" && ")) + } + }) + .join(" || ")) +} + +fn is_version_filter_field(field: &str) -> bool { + matches!( + field, + "categories" + | "project_types" + | "environment" + | "game_versions" + | "client_side" + | "server_side" + ) +} + +fn is_negative_filter(expression: &str) -> bool { + expression + .split_once(':') + .is_some_and(|(_, value)| value.trim_start().starts_with("!=")) +} + +fn filter_field(expression: &str) -> Option<&str> { + let operator = expression.find(':')?; + let field = expression[..operator].trim(); + (!field.is_empty() + && field.chars().all(|character| { + character.is_ascii_alphanumeric() || "_.".contains(character) + })) + .then_some(field) +} + +fn trim_outer_parentheses(mut expression: &str) -> &str { + while expression.starts_with('(') + && expression.ends_with(')') + && matching_outer_parentheses(expression) + { + expression = expression[1..expression.len() - 1].trim(); + } + expression +} + +fn matching_outer_parentheses(expression: &str) -> bool { + let mut depth = 0; + let mut quote = None; + let mut escaped = false; + + for (index, character) in expression.char_indices() { + if escaped { + escaped = false; + continue; + } + if character == '\\' { + escaped = true; + continue; + } + if let Some(active_quote) = quote { + if character == active_quote { + quote = None; + } + continue; + } + if matches!(character, '\'' | '"' | '`') { + quote = Some(character); + continue; + } + match character { + '(' => depth += 1, + ')' => { + depth -= 1; + if depth == 0 && index + character.len_utf8() < expression.len() + { + return false; + } + } + _ => {} + } + } + + depth == 0 +} + +fn split_top_level<'a>(expression: &'a str, operator: &str) -> Vec<&'a str> { + let mut parts = Vec::new(); + let mut start = 0; + let mut parentheses = 0; + let mut brackets = 0; + let mut quote = None; + let mut escaped = false; + let bytes = expression.as_bytes(); + let mut index = 0; + + while index < bytes.len() { + let character = expression[index..].chars().next().unwrap_or_default(); + let width = character.len_utf8(); + if escaped { + escaped = false; + index += width; + continue; + } + if character == '\\' { + escaped = true; + index += width; + continue; + } + if let Some(active_quote) = quote { + if character == active_quote { + quote = None; + } + index += width; + continue; + } + if matches!(character, '\'' | '"' | '`') { + quote = Some(character); + index += width; + continue; + } + match character { + '(' => parentheses += 1, + ')' => parentheses -= 1, + '[' => brackets += 1, + ']' => brackets -= 1, + _ => {} + } + + if parentheses == 0 + && brackets == 0 + && expression[index..].starts_with(operator) + { + parts.push(expression[start..index].trim()); + index += operator.len(); + start = index; + continue; + } + index += width; + } + + if parts.is_empty() { + vec![expression] + } else { + parts.push(expression[start..].trim()); + parts + } +} + /// Translates a Meilisearch filter expression into Typesense `filter_by` /// syntax. /// @@ -1410,3 +1802,61 @@ fn condition_to_typesense_filter(cond: &str) -> String { } cond.to_string() } + +#[cfg(test)] +mod tests { + use super::rewrite_filter_for_join; + + #[test] + fn project_filters_do_not_join_versions() { + assert_eq!( + rewrite_filter_for_join("license:= MIT", "versions").unwrap(), + "license:= MIT" + ); + } + + #[test] + fn correlated_version_filters_share_one_join() { + assert_eq!( + rewrite_filter_for_join( + "categories:= fabric && game_versions:= 1.21", + "versions", + ) + .unwrap(), + "(project_categories:= fabric && $versions(game_versions:= 1.21)) || $versions(categories:= fabric && game_versions:= 1.21)" + ); + } + + #[test] + fn project_and_version_filters_are_partitioned() { + assert_eq!( + rewrite_filter_for_join( + "license:= MIT && categories:= fabric", + "versions", + ) + .unwrap(), + "(license:= MIT && project_categories:= fabric) || (license:= MIT && $versions(categories:= fabric))" + ); + } + + #[test] + fn mixed_boolean_filters_preserve_version_correlation() { + assert_eq!( + rewrite_filter_for_join( + "(license:= MIT || categories:= fabric) && game_versions:= 1.21", + "versions", + ) + .unwrap(), + "(license:= MIT && $versions(game_versions:= 1.21)) || (project_categories:= fabric && $versions(game_versions:= 1.21)) || $versions(categories:= fabric && game_versions:= 1.21)" + ); + } + + #[test] + fn negative_categories_require_project_and_version_exclusion() { + assert_eq!( + rewrite_filter_for_join("categories:!= fabric", "versions") + .unwrap(), + "(project_categories:!= fabric && $versions(categories:!= fabric))" + ); + } +} diff --git a/apps/labrinth/src/search/incremental.rs b/apps/labrinth/src/search/incremental.rs index 96e8b7e096..a418e7638c 100644 --- a/apps/labrinth/src/search/incremental.rs +++ b/apps/labrinth/src/search/incremental.rs @@ -36,7 +36,11 @@ impl IncrementalSearchQueue { } } - pub async fn push( + pub async fn push(&self, project_id: ProjectId) { + self.operations.lock().await.push_project_change(project_id); + } + + pub async fn push_versions( &self, project_id: ProjectId, version_ids: impl IntoIterator, @@ -44,7 +48,7 @@ impl IncrementalSearchQueue { self.operations .lock() .await - .push_project_change(project_id, version_ids); + .push_version_change(project_id, version_ids); } pub async fn push_project_removal(&self, project_id: ProjectId) { @@ -123,22 +127,27 @@ impl PendingSearchIndexOperations { && self.removed_project_ids.is_empty() } - fn push_project_change( + fn push_project_change(&mut self, project_id: ProjectId) { + if !self.removed_project_ids.contains(&project_id) { + self.changed_project_ids.insert(project_id); + } + } + + fn push_version_change( &mut self, project_id: ProjectId, version_ids: impl IntoIterator, ) { - if !self.removed_project_ids.contains(&project_id) { - let version_ids = version_ids.into_iter().collect::>(); - if version_ids.is_empty() { - self.changed_project_versions.remove(&project_id); - self.changed_project_ids.insert(project_id); - } else if !self.changed_project_ids.contains(&project_id) { - self.changed_project_versions - .entry(project_id) - .or_default() - .extend(version_ids); - } + if self.removed_project_ids.contains(&project_id) { + return; + } + + let version_ids = version_ids.into_iter().collect::>(); + if !version_ids.is_empty() { + self.changed_project_versions + .entry(project_id) + .or_default() + .extend(version_ids); } } @@ -151,16 +160,12 @@ impl PendingSearchIndexOperations { fn push_event(&mut self, event: SearchProjectIndexQueueEventData) { match event { SearchProjectIndexQueueEventData::Change { project_id } => { - self.push_project_change(project_id, []) + self.push_project_change(project_id) } SearchProjectIndexQueueEventData::VersionChange { project_id, version_ids, - } => { - if !version_ids.is_empty() { - self.push_project_change(project_id, version_ids) - } - } + } => self.push_version_change(project_id, version_ids), SearchProjectIndexQueueEventData::Removal { project_id } => { self.push_project_removal(project_id) } @@ -188,7 +193,6 @@ impl PendingSearchIndexOperations { } }, )); - events } } diff --git a/apps/labrinth/src/search/incremental/consume.rs b/apps/labrinth/src/search/incremental/consume.rs index 2cfd315f24..94f5274169 100644 --- a/apps/labrinth/src/search/incremental/consume.rs +++ b/apps/labrinth/src/search/incremental/consume.rs @@ -191,7 +191,9 @@ async fn consume_batch( .retain(|project_id| !project_ids_to_remove.contains(project_id)); project_ids_with_version_changes .retain(|project_id| !project_ids_to_remove.contains(project_id)); - + project_ids_to_change.retain(|project_id| { + !project_ids_with_version_changes.contains(project_id) + }); let project_ids_to_change = project_ids_to_change.into_iter().collect::>(); let project_ids_with_version_changes = project_ids_with_version_changes @@ -204,7 +206,7 @@ async fn consume_batch( info!( kafka.message_count = messages_to_commit.len(), - "Read all Kafka messages in {:.2?}, found {} projects to change, {} projects with version changes, {} versions to change, and {} projects to remove", + "Read all Kafka messages in {:.2?}, found {} projects to change, {} projects with {} version changes, and {} projects to remove", start.elapsed(), project_ids_to_change.len(), project_ids_with_version_changes.len(), @@ -219,6 +221,10 @@ async fn consume_batch( project_count = project_ids_to_remove.len(), "Removing project documents" ); + search_backend + .remove_project_version_documents(&project_ids_to_remove) + .await + .wrap_err("failed to remove project version documents")?; search_backend .remove_project_documents(&project_ids_to_remove) .await @@ -232,12 +238,8 @@ async fn consume_batch( if !version_ids_to_change.is_empty() { let operation_start = Instant::now(); - info!( - version_count = version_ids_to_change.len(), - "Removing changed version documents", - ); search_backend - .remove_documents(&version_ids_to_change) + .remove_version_documents(&version_ids_to_change) .await .wrap_err("failed to remove changed version documents")?; info!( @@ -249,12 +251,7 @@ async fn consume_batch( if !project_ids_with_version_changes.is_empty() { let operation_start = Instant::now(); - info!( - project_count = project_ids_with_version_changes.len(), - version_count = version_ids_to_change.len(), - "Indexing changed project versions" - ); - index_changed_project_versions( + reindex_changed_project_versions( ro_pool, redis_pool, search_backend, @@ -262,11 +259,11 @@ async fn consume_batch( &version_ids_to_change, ) .await - .wrap_err("failed to index changed project version batch")?; + .wrap_err("failed to reindex changed project versions")?; info!( project_count = project_ids_with_version_changes.len(), version_count = version_ids_to_change.len(), - "Indexed changed project versions in {:.2?}", + "Reindexed changed project versions in {:.2?}", operation_start.elapsed() ); } @@ -275,19 +272,19 @@ async fn consume_batch( let operation_start = Instant::now(); info!( project_count = project_ids_to_change.len(), - "Indexing changed projects" + "Reindexing changed projects" ); - index_changed_projects( + reindex_projects( ro_pool, redis_pool, search_backend, &project_ids_to_change, ) .await - .wrap_err("failed to index changed project batch")?; + .wrap_err("failed to reindex changed project batch")?; info!( project_count = project_ids_to_change.len(), - "Indexed changed projects in {:.2?}", + "Reindexed changed projects in {:.2?}", operation_start.elapsed() ); } @@ -356,7 +353,7 @@ async fn index_changed_projects( Ok(()) } -async fn index_changed_project_versions( +async fn reindex_changed_project_versions( ro_pool: &PgPool, redis_pool: &RedisPool, search_backend: &dyn SearchBackend, @@ -383,9 +380,11 @@ async fn index_changed_project_versions( ) })?; - info!("Fetched all project version documents, indexing into backend"); - - search_backend.index_documents(&documents).await?; + search_backend.remove_project_documents(project_ids).await?; + search_backend.index_documents(&documents.projects).await?; + search_backend + .index_version_documents(&documents.versions) + .await?; Ok(()) } diff --git a/apps/labrinth/src/search/indexing.rs b/apps/labrinth/src/search/indexing.rs index 503e34e002..792aa78eea 100644 --- a/apps/labrinth/src/search/indexing.rs +++ b/apps/labrinth/src/search/indexing.rs @@ -24,7 +24,10 @@ use crate::models::ids::{ProjectId, VersionId}; use crate::models::projects::{DependencyType, from_duplicate_version_fields}; use crate::models::v2::projects::LegacyProject; use crate::routes::v2_reroute; -use crate::search::{SearchProjectDependency, UploadSearchProject}; +use crate::search::{ + SearchDocumentBatch, SearchProjectDependency, UploadSearchProject, + UploadSearchVersion, +}; use crate::util::error::Context; struct PartialProject { @@ -68,7 +71,7 @@ pub async fn index_local( redis: &RedisPool, cursor: i64, limit: i64, -) -> eyre::Result<(Vec, i64)> { +) -> eyre::Result<(SearchDocumentBatch, i64)> { info!("Indexing local projects!"); let searchable_statuses = searchable_statuses(); @@ -111,12 +114,11 @@ pub async fn index_local( let project_ids = db_projects.iter().map(|x| x.id.0).collect::>(); let Some(largest) = project_ids.iter().max() else { - return Ok((vec![], i64::MAX)); + return Ok((SearchDocumentBatch::default(), i64::MAX)); }; - let uploads = - build_search_documents(pool, redis, db_projects, None).await?; - Ok((uploads, *largest)) + let documents = build_search_documents(pool, redis, db_projects).await?; + Ok((documents, *largest)) } pub async fn index_project_documents( @@ -164,7 +166,9 @@ pub async fn index_project_documents( info!("Fetched partial projects"); - build_search_documents(pool, redis, db_projects, None).await + Ok(build_search_documents(pool, redis, db_projects) + .await? + .projects) } pub async fn index_project_version_documents( @@ -172,16 +176,33 @@ pub async fn index_project_version_documents( redis: &RedisPool, project_ids: &[ProjectId], version_ids: &[VersionId], -) -> eyre::Result> { +) -> eyre::Result { + let projects = + index_project_document_batch(pool, redis, project_ids).await?; + let version_ids = version_ids + .iter() + .map(ToString::to_string) + .collect::>(); + Ok(SearchDocumentBatch { + projects: projects.projects, + versions: projects + .versions + .into_iter() + .filter(|version| version_ids.contains(&version.version_id)) + .collect(), + }) +} + +async fn index_project_document_batch( + pool: &PgPool, + redis: &RedisPool, + project_ids: &[ProjectId], +) -> eyre::Result { let searchable_statuses = searchable_statuses(); let project_ids = project_ids .iter() .map(|project_id| DBProjectId::from(*project_id).0) .collect::>(); - let version_ids = version_ids - .iter() - .map(|version_id| DBVersionId::from(*version_id)) - .collect::>(); let db_projects = sqlx::query!( r#" @@ -215,15 +236,14 @@ pub async fn index_project_version_documents( .await .wrap_err("failed to fetch project")?; - build_search_documents(pool, redis, db_projects, Some(&version_ids)).await + build_search_documents(pool, redis, db_projects).await } async fn build_search_documents( pool: &PgPool, redis: &RedisPool, db_projects: Vec, - version_ids_to_index: Option<&HashSet>, -) -> eyre::Result> { +) -> eyre::Result { let searchable_statuses = searchable_statuses(); let project_ids = db_projects.iter().map(|x| x.id.0).collect::>(); let project_components = db_projects @@ -391,7 +411,7 @@ async fn build_search_documents( .await?; info!("Getting all loader fields!"); - let loader_fields: Vec = sqlx::query!( + let loader_field_definitions: Vec = sqlx::query!( " SELECT DISTINCT id, field, field_type, enum_type, min_val, max_val, optional FROM loader_fields lf @@ -409,7 +429,8 @@ async fn build_search_documents( }) .try_collect() .await?; - let loader_fields: Vec<&QueryLoaderField> = loader_fields.iter().collect(); + let loader_field_definitions: Vec<&QueryLoaderField> = + loader_field_definitions.iter().collect(); info!("Getting all loader field enum values!"); @@ -434,7 +455,8 @@ async fn build_search_documents( .await?; info!("Indexing loaders, project types!"); - let mut uploads = Vec::new(); + let mut project_uploads = Vec::new(); + let mut version_uploads = Vec::new(); let total_len = db_projects.len(); let mut count = 0; @@ -533,21 +555,34 @@ async fn build_search_documents( .collect::>(); if let Some(versions) = versions.remove(&project.id) { - // Aggregated project loader fields + let Some(latest_version) = versions.iter().max_by(|a, b| { + a.date_published + .cmp(&b.date_published) + .then_with(|| a.id.0.cmp(&b.id.0)) + }) else { + continue; + }; + let project_version_fields = versions .iter() .flat_map(|x| x.version_fields.clone()) .collect::>(); let aggregated_version_fields = VersionField::from_query_json( project_version_fields, - &loader_fields, + &loader_field_definitions, &loader_field_enum_values, true, ); - let project_loader_fields = + let unvectorized_loader_fields = aggregated_version_fields + .iter() + .map(|field| { + (field.field_name.clone(), field.value.serialize_internal()) + }) + .collect(); + let mut loader_fields = from_duplicate_version_fields(aggregated_version_fields); + let project_loader_fields = loader_fields.clone(); - // aggregated project loaders let mut project_loaders = versions .iter() .flat_map(|x| x.loaders.clone()) @@ -555,162 +590,184 @@ async fn build_search_documents( project_loaders.sort(); project_loaders.dedup(); - // all valid project types across every version of the project, so that - // filters can exclude projects that have *any* version of a given - // project type (unlike the version-specific `project_types` field). - let mut all_project_types = versions + let mut project_types = versions .iter() .flat_map(|x| x.project_types.clone()) .collect::>(); - all_project_types.sort(); - all_project_types.dedup(); + project_types.sort(); + project_types.dedup(); exp::compat::correct_project_types( &project.components, - &mut all_project_types, + &mut project_types, ); - for version in versions { - if let Some(version_ids_to_index) = version_ids_to_index - && !version_ids_to_index.contains(&version.id) - { - continue; - } - + let project_id = ProjectId::from(project.id).to_string(); + version_uploads.extend(versions.iter().map(|version| { let version_fields = VersionField::from_query_json( - version.version_fields, - &loader_fields, + version.version_fields.clone(), + &loader_field_definitions, &loader_field_enum_values, false, ); let unvectorized_loader_fields = version_fields .iter() - .map(|vf| { - (vf.field_name.clone(), vf.value.serialize_internal()) + .map(|field| { + ( + field.field_name.clone(), + field.value.serialize_internal(), + ) }) .collect(); - let mut loader_fields = - from_duplicate_version_fields(version_fields); - let mut project_types = version.project_types; - + let mut fields = from_duplicate_version_fields(version_fields); + let mut version_project_types = version.project_types.clone(); exp::compat::correct_project_types( &project.components, - &mut project_types, + &mut version_project_types, ); - let mut version_loaders = version.loaders; - - // Uses version loaders, not project loaders. - let mut categories = categories.clone(); - categories.append(&mut version_loaders.clone()); - - let display_categories = display_categories.clone(); - categories.append(&mut version_loaders); - - // SPECIAL BEHAVIOUR - // Todo: revisit. - // For consistency with v2 searching, we consider the loader field 'mrpack_loaders' to be a category. - // These were previously considered the loader, and in v2, the loader is a category for searching. - // So to avoid breakage or awkward conversions, we just consider those loader_fields to be categories. - // The loaders are kept in loader_fields as well, so that no information is lost on retrieval. - let mrpack_loaders = loader_fields + let mut version_categories = version.loaders.clone(); + let mrpack_loaders = fields .get("mrpack_loaders") - .cloned() - .map(|x| { - x.into_iter() - .filter_map(|x| x.as_str().map(String::from)) - .collect::>() - }) - .unwrap_or_default(); - categories.extend(mrpack_loaders); - if loader_fields.contains_key("mrpack_loaders") { - categories.retain(|x| *x != "mrpack"); + .into_iter() + .flatten() + .filter_map(|value| value.as_str().map(String::from)) + .collect::>(); + version_categories.extend(mrpack_loaders); + if fields.contains_key("mrpack_loaders") { + version_categories.retain(|category| category != "mrpack"); } + version_categories.sort(); + version_categories.dedup(); - // SPECIAL BEHAVIOUR: - // For consistency with v2 searching, we manually input the - // client_side and server_side fields from the loader fields into - // separate loader fields. - // 'client_side' and 'server_side' remain supported by meilisearch even though they are no longer v3 fields. let (_, v2_og_project_type) = - LegacyProject::get_project_type(&project_types); + LegacyProject::get_project_type(&version_project_types); let (client_side, server_side) = v2_reroute::convert_v3_side_types_to_v2_side_types( &unvectorized_loader_fields, Some(&v2_og_project_type), ); - if let Ok(client_side) = serde_json::to_value(client_side) { - loader_fields - .insert("client_side".to_string(), vec![client_side]); + fields.insert("client_side".to_string(), vec![client_side]); } if let Ok(server_side) = serde_json::to_value(server_side) { - loader_fields - .insert("server_side".to_string(), vec![server_side]); + fields.insert("server_side".to_string(), vec![server_side]); } - - let components = project - .components - .clone() - .into_query( - ProjectId::from(project.id), - &project_query_context, + fields.retain(|field, _| { + matches!( + field.as_str(), + "environment" + | "game_versions" + | "client_side" + | "server_side" ) - .wrap_err("failed to populate query components")?; + }); - let usp = UploadSearchProject { - version_id: crate::models::ids::VersionId::from(version.id) - .to_string(), - project_id: crate::models::ids::ProjectId::from(project.id) - .to_string(), - name: project.name.clone(), - indexed_name: normalize_for_search(&project.name), - summary: project.summary.clone(), - categories: categories.clone(), - display_categories: display_categories.clone(), - follows: project.follows, - downloads: project.downloads, - log_downloads: (project.downloads.max(1) as f64).ln(), - icon_url: project.icon_url.clone(), - author: username.clone(), - author_id: ariadne::ids::UserId::from(user_id).to_string(), - organization: org_name.clone(), - organization_id: org_id.map(|e| { - crate::models::ids::OrganizationId::from(e).to_string() - }), - indexed_author: normalize_for_search(&username), - date_created: project.approved, - created_timestamp: project.approved.timestamp(), - date_modified: project.updated, - modified_timestamp: project.updated.timestamp(), + UploadSearchVersion { + version_id: VersionId::from(version.id).to_string(), + project_id: project_id.clone(), + categories: version_categories, + project_types: version_project_types, version_published_timestamp: version .date_published .timestamp(), - license: license.clone(), - slug: project.slug.clone(), - // TODO - project_types, - all_project_types: all_project_types.clone(), - gallery: gallery.clone(), - featured_gallery: featured_gallery.clone(), - open_source, - color: project.color.map(|x| x as u32), - dependency_project_ids: dependency_project_ids.clone(), - compatible_dependency_project_ids: - compatible_dependency_project_ids.clone(), - dependencies: dependencies.clone(), - loader_fields, - project_loader_fields: project_loader_fields.clone(), - // 'loaders' is aggregate of all versions' loaders - loaders: project_loaders.clone(), - components, - }; + loader_fields: fields, + } + })); - uploads.push(usp); + let project_categories = categories.clone(); + let mut categories = categories; + categories.extend(project_loaders.iter().cloned()); + + let mrpack_loaders = loader_fields + .get("mrpack_loaders") + .into_iter() + .flatten() + .filter_map(|value| value.as_str().map(String::from)) + .collect::>(); + categories.extend(mrpack_loaders); + if loader_fields.contains_key("mrpack_loaders") { + categories.retain(|category| category != "mrpack"); } + categories.sort(); + categories.dedup(); + + let (_, v2_og_project_type) = + LegacyProject::get_project_type(&project_types); + let (client_side, server_side) = + v2_reroute::convert_v3_side_types_to_v2_side_types( + &unvectorized_loader_fields, + Some(&v2_og_project_type), + ); + + if let Ok(client_side) = serde_json::to_value(client_side) { + loader_fields + .insert("client_side".to_string(), vec![client_side]); + } + if let Ok(server_side) = serde_json::to_value(server_side) { + loader_fields + .insert("server_side".to_string(), vec![server_side]); + } + + let components = project + .components + .clone() + .into_query(ProjectId::from(project.id), &project_query_context) + .wrap_err("failed to populate query components")?; + let indexed_name = normalize_for_search(&project.name); + + project_uploads.push(UploadSearchProject { + version_id: crate::models::ids::VersionId::from( + latest_version.id, + ) + .to_string(), + project_id, + name: project.name, + indexed_name, + summary: project.summary, + categories, + project_categories, + display_categories, + follows: project.follows, + downloads: project.downloads, + log_downloads: (project.downloads.max(1) as f64).ln(), + icon_url: project.icon_url, + author: username.clone(), + author_id: ariadne::ids::UserId::from(user_id).to_string(), + organization: org_name, + organization_id: org_id.map(|id| { + crate::models::ids::OrganizationId::from(id).to_string() + }), + indexed_author: normalize_for_search(&username), + date_created: project.approved, + created_timestamp: project.approved.timestamp(), + date_modified: project.updated, + modified_timestamp: project.updated.timestamp(), + version_published_timestamp: latest_version + .date_published + .timestamp(), + license, + slug: project.slug, + project_types: project_types.clone(), + all_project_types: project_types, + gallery, + featured_gallery, + open_source, + color: project.color.map(|x| x as u32), + dependency_project_ids, + compatible_dependency_project_ids, + dependencies, + project_loader_fields, + loader_fields, + loaders: project_loaders, + components, + }); } } - Ok(uploads) + Ok(SearchDocumentBatch { + projects: project_uploads, + versions: version_uploads, + }) } struct PartialVersion { diff --git a/apps/labrinth/src/search/mod.rs b/apps/labrinth/src/search/mod.rs index 7b1db63299..0b53f387ea 100644 --- a/apps/labrinth/src/search/mod.rs +++ b/apps/labrinth/src/search/mod.rs @@ -116,12 +116,25 @@ pub trait SearchBackend: Send + Sync { documents: &[UploadSearchProject], ) -> eyre::Result<()>; + async fn index_version_documents( + &self, + documents: &[UploadSearchVersion], + ) -> eyre::Result<()>; + async fn remove_project_documents( &self, ids: &[ProjectId], ) -> eyre::Result<()>; - async fn remove_documents(&self, ids: &[VersionId]) -> eyre::Result<()>; + async fn remove_project_version_documents( + &self, + ids: &[ProjectId], + ) -> eyre::Result<()>; + + async fn remove_version_documents( + &self, + ids: &[VersionId], + ) -> eyre::Result<()>; async fn tasks(&self) -> eyre::Result; @@ -238,6 +251,7 @@ impl FromStr for SearchBackendKind { /// serialized as `null`. #[derive(Serialize, Deserialize, Debug, Clone)] pub struct UploadSearchProject { + /// ID of the most recently published version. pub version_id: String, pub project_id: String, // @@ -256,6 +270,7 @@ pub struct UploadSearchProject { pub indexed_name: String, pub summary: String, pub categories: Vec, + pub project_categories: Vec, pub display_categories: Vec, pub follows: i32, pub downloads: i32, @@ -274,7 +289,7 @@ pub struct UploadSearchProject { pub date_modified: DateTime, /// Unix timestamp of the last major modification pub modified_timestamp: i64, - /// Unix timestamp of the publication date of the version + /// Unix timestamp of the most recently published version. pub version_published_timestamp: i64, pub open_source: bool, #[serde(skip_serializing_if = "Option::is_none")] @@ -296,6 +311,23 @@ pub struct UploadSearchProject { pub loader_fields: HashMap>, } +#[derive(Serialize, Deserialize, Debug, Clone)] +pub struct UploadSearchVersion { + pub version_id: String, + pub project_id: String, + pub categories: Vec, + pub project_types: Vec, + pub version_published_timestamp: i64, + #[serde(flatten)] + pub loader_fields: HashMap>, +} + +#[derive(Debug, Default)] +pub struct SearchDocumentBatch { + pub projects: Vec, + pub versions: Vec, +} + /// Nullable fields in Typesense-bound documents should use /// `skip_serializing_if = "Option::is_none"` so they are omitted instead of /// serialized as `null`. @@ -320,6 +352,7 @@ pub struct SearchResults { #[derive(Serialize, Deserialize, Debug, Clone, ToSchema)] pub struct ResultSearchProject { + /// ID of the most recently published version. pub version_id: String, pub project_id: String, pub project_types: Vec, diff --git a/scripts/convert-typesense-project-docs.py b/scripts/convert-typesense-project-docs.py new file mode 100755 index 0000000000..5cfd594028 --- /dev/null +++ b/scripts/convert-typesense-project-docs.py @@ -0,0 +1,443 @@ +#!/usr/bin/env python3 +"""Split legacy Typesense JSONL into project and version documents.""" + +import argparse +import json +import os +import shutil +import sys +import tempfile +import time +import zlib +from collections import OrderedDict + + +VERSION_FILTER_PATHS = ( + "project_types", + "environment", + "game_versions", + "client_side", + "server_side", +) + +BASE62_DIGITS = { + character: index + for index, character in enumerate( + "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz" + ) +} + + +class ConversionError(Exception): + pass + + +def parse_args(): + parser = argparse.ArgumentParser( + description=( + "Convert a legacy Typesense JSONL export from one document per " + "version into separate project and version collections." + ) + ) + parser.add_argument("input", help="Legacy Typesense JSONL export") + parser.add_argument("projects_output", help="Destination project-document JSONL") + parser.add_argument("versions_output", help="Destination version-document JSONL") + parser.add_argument( + "--shards", + type=int, + default=512, + help="Temporary hash shards used to bound memory usage (default: 512)", + ) + parser.add_argument( + "--max-open-shards", + type=int, + default=64, + help="Maximum temporary shard files held open at once (default: 64)", + ) + parser.add_argument( + "--progress-interval", + type=float, + default=10.0, + help="Seconds between progress reports (default: 10)", + ) + parser.add_argument( + "--keep-temporary", + action="store_true", + help="Keep temporary shard files after completion or failure", + ) + return parser.parse_args() + + +def base62_value(value): + result = 0 + try: + for character in value: + result = result * 62 + BASE62_DIGITS[character] + except KeyError: + return -1 + return result + + +def get_path(document, path): + value = document + for segment in path.split("."): + if not isinstance(value, dict) or segment not in value: + return None, False + value = value[segment] + return value, True + + +def set_path(document, path, value): + segments = path.split(".") + current = document + for segment in segments[:-1]: + child = current.get(segment) + if not isinstance(child, dict): + child = {} + current[segment] = child + current = child + current[segments[-1]] = value + + +def unique_sorted(values): + return sorted(set(values)) + + +def extend_values(target, value): + if isinstance(value, list): + target.extend(value) + elif value is not None: + target.append(value) + + +def version_filter_document(document): + version_id = document.get("version_id") or document.get("id") + result = { + "id": str(version_id), + "version_id": str(version_id), + "project_id": str(document["project_id"]), + "version_published_timestamp": document.get( + "version_published_timestamp", -1 + ), + } + loaders = list(document.get("loaders") or []) + mrpack_loaders = list(document.get("mrpack_loaders") or []) + loaders.extend(mrpack_loaders) + if mrpack_loaders: + loaders = [loader for loader in loaders if loader != "mrpack"] + result["categories"] = unique_sorted(loaders) + for path in VERSION_FILTER_PATHS: + value, exists = get_path(document, path) + if exists and value is not None: + set_path(result, path, value) + return result + + +class ProjectAccumulator: + def __init__(self, document): + self.latest_document = document + self.latest_key = self.version_key(document) + self.versions = {} + self.categories = [] + self.version_categories = [] + self.loaders = [] + self.project_types = [] + self.client_side = [] + self.server_side = [] + self.add(document) + + @staticmethod + def version_key(document): + return ( + document.get("version_published_timestamp", -1), + base62_value(str(document.get("version_id", ""))), + ) + + def add(self, document): + project_id = document.get("project_id") + if project_id != self.latest_document.get("project_id"): + raise ConversionError("attempted to combine different projects") + + version_id = document.get("version_id") or document.get("id") + if not version_id: + raise ConversionError(f"project `{project_id}` has a version without an ID") + + key = self.version_key(document) + if key > self.latest_key: + self.latest_document = document + self.latest_key = key + + version_document = version_filter_document(document) + self.versions[str(version_id)] = (key, version_document) + extend_values(self.categories, document.get("categories")) + extend_values(self.version_categories, version_document.get("categories")) + extend_values(self.loaders, document.get("loaders")) + extend_values(self.project_types, document.get("project_types")) + extend_values(self.project_types, document.get("all_project_types")) + extend_values(self.client_side, document.get("client_side")) + extend_values(self.server_side, document.get("server_side")) + + def finish(self): + result = dict(self.latest_document) + project_id = str(result["project_id"]) + all_project_types = unique_sorted(self.project_types) + + result["id"] = project_id + result["version_id"] = str( + self.latest_document.get("version_id") + or self.latest_document.get("id") + ) + result["categories"] = unique_sorted(self.categories) + result["project_categories"] = unique_sorted( + set(self.categories) - set(self.version_categories) + ) + result["loaders"] = unique_sorted(self.loaders) + result["project_types"] = all_project_types + result["all_project_types"] = all_project_types + + project_loader_fields = result.get("project_loader_fields") + if not isinstance(project_loader_fields, dict): + project_loader_fields = {} + result["project_loader_fields"] = project_loader_fields + for field, value in project_loader_fields.items(): + result[field] = value + + if self.client_side: + result["client_side"] = unique_sorted(self.client_side) + if self.server_side: + result["server_side"] = unique_sorted(self.server_side) + + versions = [ + version + for _, version in sorted( + self.versions.values(), key=lambda item: item[0] + ) + ] + result.pop("versions", None) + return result, versions + + +class ShardWriter: + def __init__(self, directory, shard_count, max_open): + self.directory = directory + self.shard_count = shard_count + self.max_open = max_open + self.handles = OrderedDict() + + def path(self, shard): + return os.path.join(self.directory, f"shard-{shard:04d}.jsonl") + + def write(self, project_id, line): + shard = zlib.crc32(project_id.encode("utf-8")) % self.shard_count + handle = self.handles.pop(shard, None) + if handle is None: + if len(self.handles) >= self.max_open: + _, oldest = self.handles.popitem(last=False) + oldest.close() + handle = open(self.path(shard), "ab") + self.handles[shard] = handle + handle.write(line) + + def close(self): + for handle in self.handles.values(): + handle.close() + self.handles.clear() + + +def shard_input(args, temporary_directory): + writer = ShardWriter( + temporary_directory, + args.shards, + args.max_open_shards, + ) + input_size = os.path.getsize(args.input) + bytes_read = 0 + document_count = 0 + started_at = time.monotonic() + last_report = started_at + + try: + with open(args.input, "rb") as input_file: + for line_number, line in enumerate(input_file, start=1): + bytes_read += len(line) + if not line.strip(): + continue + try: + document = json.loads(line) + except json.JSONDecodeError as error: + raise ConversionError( + f"invalid JSON on input line {line_number}: {error}" + ) from error + project_id = document.get("project_id") + if not project_id: + raise ConversionError( + f"input line {line_number} has no `project_id`" + ) + writer.write(str(project_id), line) + document_count += 1 + + now = time.monotonic() + if now - last_report >= args.progress_interval: + percent = bytes_read / input_size * 100 if input_size else 100 + print( + f"sharding: {percent:.1f}% ({document_count:,} documents)", + flush=True, + ) + last_report = now + finally: + writer.close() + + print( + f"sharding complete: {document_count:,} version documents in " + f"{time.monotonic() - started_at:.1f}s", + flush=True, + ) + return document_count + + +def convert_shards( + args, + temporary_directory, + partial_projects_output, + partial_versions_output, +): + project_count = 0 + version_count = 0 + started_at = time.monotonic() + last_report = started_at + + with ( + open(partial_projects_output, "w", encoding="utf-8") as projects_file, + open(partial_versions_output, "w", encoding="utf-8") as versions_file, + ): + for shard in range(args.shards): + path = os.path.join(temporary_directory, f"shard-{shard:04d}.jsonl") + if not os.path.exists(path): + continue + + projects = {} + with open(path, "rb") as shard_file: + for line_number, line in enumerate(shard_file, start=1): + try: + document = json.loads(line) + except json.JSONDecodeError as error: + raise ConversionError( + f"invalid JSON in shard {shard}, line {line_number}: {error}" + ) from error + project_id = str(document["project_id"]) + if project_id in projects: + projects[project_id].add(document) + else: + projects[project_id] = ProjectAccumulator(document) + version_count += 1 + + for project_id in sorted(projects): + project, versions = projects[project_id].finish() + json.dump( + project, + projects_file, + separators=(",", ":"), + ensure_ascii=False, + ) + projects_file.write("\n") + for version in versions: + json.dump( + version, + versions_file, + separators=(",", ":"), + ensure_ascii=False, + ) + versions_file.write("\n") + project_count += 1 + + os.remove(path) + now = time.monotonic() + if now - last_report >= args.progress_interval: + print( + f"converting: shard {shard + 1}/{args.shards}, " + f"{project_count:,} projects", + flush=True, + ) + last_report = now + + print( + f"conversion complete: {version_count:,} versions into " + f"{project_count:,} projects in {time.monotonic() - started_at:.1f}s", + flush=True, + ) + return project_count, version_count + + +def main(): + args = parse_args() + if args.shards <= 0: + raise ConversionError("--shards must be greater than zero") + if args.max_open_shards <= 0: + raise ConversionError("--max-open-shards must be greater than zero") + if args.progress_interval <= 0: + raise ConversionError("--progress-interval must be greater than zero") + if not os.path.isfile(args.input): + raise ConversionError(f"input file does not exist: {args.input}") + paths = { + os.path.abspath(args.input), + os.path.abspath(args.projects_output), + os.path.abspath(args.versions_output), + } + if len(paths) != 3: + raise ConversionError("input and output paths must all differ") + + projects_directory = os.path.dirname(os.path.abspath(args.projects_output)) + versions_directory = os.path.dirname(os.path.abspath(args.versions_output)) + os.makedirs(projects_directory, exist_ok=True) + os.makedirs(versions_directory, exist_ok=True) + temporary_directory = tempfile.mkdtemp( + prefix="typesense-project-convert-", + dir=projects_directory, + ) + partial_projects_output = f"{args.projects_output}.partial" + partial_versions_output = f"{args.versions_output}.partial" + + try: + expected_versions = shard_input(args, temporary_directory) + project_count, version_count = convert_shards( + args, + temporary_directory, + partial_projects_output, + partial_versions_output, + ) + if version_count != expected_versions: + raise ConversionError( + f"sharded {expected_versions} versions but converted {version_count}" + ) + os.replace(partial_projects_output, args.projects_output) + os.replace(partial_versions_output, args.versions_output) + print( + f"wrote {project_count:,} project documents to " + f"{args.projects_output} and {version_count:,} version documents to " + f"{args.versions_output}", + flush=True, + ) + except Exception: + for partial_output in ( + partial_projects_output, + partial_versions_output, + ): + if os.path.exists(partial_output): + os.remove(partial_output) + raise + finally: + if args.keep_temporary: + print(f"temporary shards kept at {temporary_directory}", file=sys.stderr) + else: + shutil.rmtree(temporary_directory, ignore_errors=True) + + +if __name__ == "__main__": + try: + main() + except KeyboardInterrupt: + print("interrupted", file=sys.stderr) + sys.exit(130) + except (ConversionError, OSError) as error: + print(error, file=sys.stderr) + sys.exit(1)