haste_worker/search_indexing/
mod.rs

1use crate::{
2    indexing_lock::{IndexLockProvider, postgres::TenantLockIndex},
3    traits::Worker,
4};
5use haste_fhir_model::r4::generated::resources::ResourceTypeError;
6use haste_fhir_operation_error::{OperationOutcomeError, derive::OperationOutcomeError};
7use haste_fhir_search::{
8    IndexResource, SearchEngine,
9    elastic_search::{
10        ElasticSearchEngine, create_es_client,
11        search_parameter_resolver::ElasticSearchParameterResolver,
12    },
13};
14use haste_fhirpath::FHIRPathError;
15use haste_jwt::{TenantId, VersionId};
16use haste_repository::{
17    fhir::FHIRRepository, pg::PGConnection, sequence::ResourceSequential,
18    types::SupportedFHIRVersions,
19};
20use serde::{Deserialize, Serialize};
21use sqlx::{Acquire, query_as, types::time::OffsetDateTime};
22use std::sync::Arc;
23use tokio::{sync::Mutex, task::JoinHandle};
24
25#[derive(OperationOutcomeError, Debug)]
26pub enum IndexingWorkerError {
27    #[fatal(code = "exception", diagnostic = "Database error: '{arg0}'")]
28    DatabaseConnectionError(#[from] sqlx::Error),
29    #[fatal(code = "exception", diagnostic = "Lock error: '{arg0}'")]
30    OperationError(#[from] OperationOutcomeError),
31    #[fatal(code = "exception", diagnostic = "Elasticsearch error: '{arg0}'")]
32    ElasticsearchError(#[from] elasticsearch::Error),
33    #[fatal(code = "exception", diagnostic = "FHIRPath error: '{arg0}'")]
34    FHIRPathError(#[from] FHIRPathError),
35    #[fatal(
36        code = "exception",
37        diagnostic = "Missing search parameters for resource: '{arg0}'"
38    )]
39    MissingSearchParameters(String),
40    #[fatal(
41        code = "exception",
42        diagnostic = "Fatal error occurred during indexing"
43    )]
44    Fatal,
45    #[fatal(
46        code = "exception",
47        diagnostic = "Artifact error: Invalid resource type '{arg0}'"
48    )]
49    ResourceTypeError(#[from] ResourceTypeError),
50}
51
52struct TenantReturn {
53    id: TenantId,
54    created_at: OffsetDateTime,
55}
56
57async fn get_tenants(
58    repo: &PGConnection,
59    cursor: &OffsetDateTime,
60    count: usize,
61) -> Result<Vec<TenantReturn>, OperationOutcomeError> {
62    match repo {
63        PGConnection::Pool(pool, _) => {
64            let mut connection = pool.acquire().await.map_err(IndexingWorkerError::from)?;
65            let conn = connection
66                .acquire()
67                .await
68                .map_err(IndexingWorkerError::from)?;
69            let result = query_as!(
70                TenantReturn,
71                r#"SELECT id as "id: TenantId", created_at FROM tenants WHERE created_at > $1 ORDER BY created_at DESC LIMIT $2"#,
72                cursor,
73                count as i64
74            )
75            .fetch_all(&mut *conn)
76            .await
77            .map_err(IndexingWorkerError::from)?;
78
79            Ok(result)
80        }
81        PGConnection::Transaction(tx, _) => {
82            let mut connection = tx.lock().await;
83            let conn = connection
84                .acquire()
85                .await
86                .map_err(IndexingWorkerError::from)?;
87            let result = query_as!(
88                TenantReturn,
89                r#"SELECT id as "id: TenantId", created_at FROM tenants WHERE created_at > $1 ORDER BY created_at DESC LIMIT $2"#,
90                cursor,
91                count as i64
92            )
93            .fetch_all(&mut *conn)
94            .await
95            .map_err(IndexingWorkerError::from)?;
96
97            Ok(result)
98        }
99    }
100}
101
102static TOTAL_INDEXED: std::sync::LazyLock<Mutex<usize>> =
103    std::sync::LazyLock::new(|| Mutex::new(0));
104
105async fn index_tenant_next_sequence<
106    Repo: ResourceSequential + IndexLockProvider<TenantId, TenantLockIndex>,
107    Engine: SearchEngine,
108>(
109    max_concurrent_limit: u64,
110    search_client: Arc<Engine>,
111    repo: &Repo,
112    tenant_id: &TenantId,
113) -> Result<(), IndexingWorkerError> {
114    let start = std::time::Instant::now();
115    let tenant_locks = repo.get_available_locks(vec![tenant_id]).await?;
116
117    if tenant_locks.is_empty() {
118        tracing::info!(
119            "No available locks for tenant '{}', skipping indexing.",
120            tenant_id
121        );
122        return Ok(());
123    }
124
125    tracing::info!(
126        "Acquired lock for tenant '{}', starting indexing from sequence {}.",
127        tenant_id,
128        tenant_locks[0].index_sequence_position
129    );
130
131    let resources = repo
132        .get_sequence(
133            tenant_id,
134            tenant_locks[0].index_sequence_position as u64,
135            Some(max_concurrent_limit),
136        )
137        .await?;
138
139    let resources_total = resources.len();
140    let start_sequence = resources.first().map(|r| r.sequence);
141    let last_value = resources.last().cloned();
142
143    // Perform indexing if there are resources to index.
144    if !resources.is_empty() {
145        let result = search_client
146            .index(
147                SupportedFHIRVersions::R4,
148                resources
149                    .into_iter()
150                    .map(|r| IndexResource {
151                        tenant: r.tenant,
152                        id: r.id,
153                        version_id: VersionId::new(r.version_id),
154                        project: r.project,
155                        fhir_method: r.fhir_method,
156                        resource_type: r.resource_type,
157                        resource: r.resource.0,
158                    })
159                    .collect(),
160            )
161            .await?;
162
163        if result.0 != resources_total {
164            tracing::error!(
165                "Indexed resource count '{}' does not match retrieved resource count '{}'",
166                result.0,
167                resources_total
168            );
169            return Err(IndexingWorkerError::Fatal);
170        }
171
172        if let Some(resource) = last_value {
173            let diff = (resource.sequence + 1) - start_sequence.unwrap_or(0);
174            let total = resources_total;
175
176            if total != diff as usize {
177                tracing::event!(
178                    tracing::Level::INFO,
179                    // safe_seq = resource.max_safe_seq.unwrap_or(0),
180                    first_seq = start_sequence.unwrap_or(0),
181                    last_seq = resource.sequence,
182                    total = resources_total,
183                    diff = (resource.sequence + 1) - start_sequence.unwrap_or(0)
184                );
185            }
186
187            tracing::trace!(
188                "Updating lock for tenant '{}' to sequence position {}.",
189                tenant_id,
190                resource.sequence
191            );
192
193            repo.update_lock(
194                &tenant_id,
195                TenantLockIndex {
196                    id: tenant_id.clone(),
197                    index_sequence_position: resource.sequence as i64,
198                },
199            )
200            .await?;
201
202            let elapsed = start.elapsed();
203            tracing::trace!(
204                "Indexed {} resources for tenant '{}' in {:.2?} (up to sequence {})",
205                result.0,
206                tenant_id.as_ref(),
207                elapsed,
208                resource.sequence
209            );
210        }
211
212        *(TOTAL_INDEXED.lock().await) += result.0;
213    }
214
215    Ok(())
216}
217
218async fn index_for_tenant<
219    Search: SearchEngine,
220    Repository: FHIRRepository + ResourceSequential + IndexLockProvider<TenantId, TenantLockIndex>,
221>(
222    max_concurrent_limit: u64,
223    repo: Arc<Repository>,
224    search_client: Arc<Search>,
225    tenant_id: &TenantId,
226) -> Result<(), IndexingWorkerError> {
227    let tx = repo
228        .transaction(false)
229        .await
230        .map_err(IndexingWorkerError::from)?;
231    let res =
232        index_tenant_next_sequence(max_concurrent_limit, search_client, &tx, &tenant_id).await;
233
234    match res {
235        Ok(res) => {
236            tx.commit().await?;
237            Ok(res)
238        }
239        Err(e) => {
240            tx.rollback().await?;
241            Err(e)
242        }
243    }
244}
245
246pub enum IndexingWorkerEnvironmentVariables {
247    DatabaseURL,
248    ElasticSearchURL,
249    ElasticSearchUsername,
250    ElasticSearchPassword,
251}
252
253impl From<IndexingWorkerEnvironmentVariables> for String {
254    fn from(value: IndexingWorkerEnvironmentVariables) -> Self {
255        match value {
256            IndexingWorkerEnvironmentVariables::DatabaseURL => "DATABASE_URL".to_string(),
257            IndexingWorkerEnvironmentVariables::ElasticSearchURL => "ELASTICSEARCH_URL".to_string(),
258            IndexingWorkerEnvironmentVariables::ElasticSearchUsername => {
259                "ELASTICSEARCH_USERNAME".to_string()
260            }
261            IndexingWorkerEnvironmentVariables::ElasticSearchPassword => {
262                "ELASTICSEARCH_PASSWORD".to_string()
263            }
264        }
265    }
266}
267
268pub struct IndexingWorker {
269    max_concurrent_limit: Option<u64>,
270    running: Arc<tokio::sync::Mutex<bool>>,
271    repo: Arc<PGConnection>,
272    search_engine: Arc<ElasticSearchEngine<ElasticSearchParameterResolver<PGConnection>>>,
273}
274
275#[derive(Clone, Deserialize, Serialize)]
276#[serde(default)]
277pub struct WorkerEnvironment {
278    pub max_concurrent_limit: Option<u64>,
279    pub repo: RepoConfig,
280    pub search: SearchConfig,
281}
282
283// Repo backend where the FHIR server stores its data/resources.
284#[derive(Clone, Deserialize, Serialize)]
285#[serde(tag = "backend", rename_all = "snake_case")]
286pub enum RepoConfig {
287    Postgres(PostgresConfig),
288}
289
290#[derive(Clone, Deserialize, Serialize)]
291pub struct PostgresConfig {
292    pub database_url: String,
293    pub max_connections: u32,
294}
295
296#[derive(Clone, Deserialize, Serialize)]
297pub struct ElasticsearchConfig {
298    pub url: String,
299    pub username: String,
300    pub password: String,
301}
302
303// Search backend where the FHIR server stores its search indices.
304#[derive(Clone, Deserialize, Serialize)]
305#[serde(tag = "backend", rename_all = "snake_case")]
306pub enum SearchConfig {
307    Elasticsearch(ElasticsearchConfig),
308}
309
310impl Default for WorkerEnvironment {
311    fn default() -> Self {
312        Self {
313            max_concurrent_limit: Some(1000),
314            repo: RepoConfig::default(),
315            search: SearchConfig::default(),
316        }
317    }
318}
319
320impl Default for RepoConfig {
321    fn default() -> Self {
322        RepoConfig::Postgres(PostgresConfig::default())
323    }
324}
325impl Default for PostgresConfig {
326    fn default() -> Self {
327        Self {
328            database_url: "postgresql://postgres:postgres@localhost:5432/haste_health".into(),
329            max_connections: 10,
330        }
331    }
332}
333impl Default for SearchConfig {
334    fn default() -> Self {
335        SearchConfig::Elasticsearch(ElasticsearchConfig::default())
336    }
337}
338impl Default for ElasticsearchConfig {
339    fn default() -> Self {
340        Self {
341            url: "http://localhost:9200".into(),
342            username: "elastic".into(),
343            password: "elastic".into(),
344        }
345    }
346}
347
348async fn create_repo(config: &RepoConfig) -> Result<Arc<PGConnection>, OperationOutcomeError> {
349    match config {
350        RepoConfig::Postgres(pg_config) => {
351            let pool = sqlx::PgPool::connect(&pg_config.database_url)
352                .await
353                .map_err(IndexingWorkerError::from)?;
354            Ok(Arc::new(PGConnection::pool(pool)))
355        }
356    }
357}
358
359async fn create_search_engine(
360    config: &SearchConfig,
361    repo: Arc<PGConnection>,
362) -> Result<
363    Arc<ElasticSearchEngine<ElasticSearchParameterResolver<PGConnection>>>,
364    OperationOutcomeError,
365> {
366    match config {
367        SearchConfig::Elasticsearch(elasticsearch_config) => {
368            let es_client = create_es_client(
369                &elasticsearch_config.url,
370                elasticsearch_config.username.clone(),
371                elasticsearch_config.password.clone(),
372            )?;
373            let search_engine = Arc::new(ElasticSearchEngine::new(
374                Arc::new(ElasticSearchParameterResolver::new(
375                    es_client.clone(),
376                    repo.clone(),
377                )),
378                Arc::new(haste_fhirpath::FPEngine::new()),
379                es_client,
380            ));
381
382            Ok(search_engine)
383        }
384    }
385}
386
387impl IndexingWorker {
388    pub async fn new(config: Arc<WorkerEnvironment>) -> Result<Self, OperationOutcomeError> {
389        let repo = create_repo(&config.repo).await?;
390        let search_engine = create_search_engine(&config.search, repo.clone()).await?;
391
392        let mut attempts = 0;
393        while !search_engine.is_connected().await.is_ok() && attempts < 5 {
394            tracing::error!("Elasticsearch is not connected, retrying in 5 seconds...");
395            tokio::time::sleep(std::time::Duration::from_secs(5)).await;
396            attempts += 1;
397        }
398
399        if !search_engine.is_connected().await.is_ok() {
400            return Err(OperationOutcomeError::fatal(
401                haste_fhir_model::r4::generated::terminology::IssueType::Exception(None),
402                "Elasticsearch is not connected after 5 attempts".to_string(),
403            ));
404        }
405
406        Ok(Self {
407            max_concurrent_limit: config.max_concurrent_limit,
408            running: Arc::new(tokio::sync::Mutex::new(true)),
409            repo,
410            search_engine,
411        })
412    }
413}
414
415impl Worker for IndexingWorker {
416    async fn run(&self) -> Result<JoinHandle<()>, OperationOutcomeError> {
417        let mut cursor = OffsetDateTime::UNIX_EPOCH;
418        let tenants_limit: usize = 100;
419
420        tracing::info!("Starting indexing worker...");
421
422        let mut k = *TOTAL_INDEXED.lock().await;
423
424        let repo = self.repo.clone();
425        let search_engine: Arc<ElasticSearchEngine<ElasticSearchParameterResolver<PGConnection>>> =
426            self.search_engine.clone();
427        let running = self.running.clone();
428        let max_concurrent_limit = self.max_concurrent_limit.unwrap_or(1000);
429
430        let spawned = tokio::spawn(async move {
431            while *running.lock().await {
432                let tenants_to_check = get_tenants(repo.as_ref(), &cursor, tenants_limit).await;
433                if let Ok(tenants_to_check) = tenants_to_check {
434                    if tenants_to_check.is_empty() || tenants_to_check.len() < tenants_limit {
435                        cursor = OffsetDateTime::UNIX_EPOCH; // Reset cursor if no tenants found
436                    } else {
437                        cursor = tenants_to_check[0].created_at;
438                    }
439
440                    for tenant in tenants_to_check {
441                        tracing::info!("Indexing tenant: '{}'", &tenant.id);
442
443                        let result = index_for_tenant(
444                            max_concurrent_limit,
445                            repo.clone(),
446                            search_engine.clone(),
447                            &tenant.id,
448                        )
449                        .await;
450
451                        if let Err(_error) = result {
452                            tracing::error!(
453                                "Failed to index tenant: '{}' cause: '{:?}'",
454                                &tenant.id,
455                                _error
456                            );
457                        }
458                    }
459                } else if let Err(error) = tenants_to_check {
460                    tracing::error!("Failed to retrieve tenants: {:?}", error);
461                }
462
463                if k != *TOTAL_INDEXED.lock().await {
464                    k = *TOTAL_INDEXED.lock().await;
465                    tracing::info!("TOTAL INDEXED SO FAR: {}", k);
466                }
467            }
468        });
469
470        Ok(spawned)
471    }
472
473    async fn stop(&mut self) -> Result<(), OperationOutcomeError> {
474        let mut running = self.running.lock().await;
475        *running = false;
476        Ok(())
477    }
478}
haste_worker/search_indexing/mod.rs

haste_worker/search_indexing/
mod.rs
