haste_indexing_worker/
lib.rs

1use crate::indexing_lock::IndexLockProvider;
2use haste_config::get_config;
3use haste_fhir_model::r4::generated::resources::ResourceTypeError;
4use haste_fhir_operation_error::{OperationOutcomeError, derive::OperationOutcomeError};
5use haste_fhir_search::{IndexResource, SearchEngine, elastic_search::ElasticSearchEngine};
6use haste_fhirpath::FHIRPathError;
7use haste_jwt::TenantId;
8use haste_repository::{fhir::FHIRRepository, types::SupportedFHIRVersions};
9use sqlx::{Pool, Postgres, query_as, types::time::OffsetDateTime};
10use std::sync::Arc;
11use tokio::sync::Mutex;
12
13mod indexing_lock;
14
15#[derive(OperationOutcomeError, Debug)]
16pub enum IndexingWorkerError {
17    #[fatal(code = "exception", diagnostic = "Database error: '{arg0}'")]
18    DatabaseConnectionError(#[from] sqlx::Error),
19    #[fatal(code = "exception", diagnostic = "Lock error: '{arg0}'")]
20    OperationError(#[from] OperationOutcomeError),
21    #[fatal(code = "exception", diagnostic = "Elasticsearch error: '{arg0}'")]
22    ElasticsearchError(#[from] elasticsearch::Error),
23    #[fatal(code = "exception", diagnostic = "FHIRPath error: '{arg0}'")]
24    FHIRPathError(#[from] FHIRPathError),
25    #[fatal(
26        code = "exception",
27        diagnostic = "Missing search parameters for resource: '{arg0}'"
28    )]
29    MissingSearchParameters(String),
30    #[fatal(
31        code = "exception",
32        diagnostic = "Fatal error occurred during indexing"
33    )]
34    Fatal,
35    #[fatal(
36        code = "exception",
37        diagnostic = "Artifact error: Invalid resource type '{arg0}'"
38    )]
39    ResourceTypeError(#[from] ResourceTypeError),
40}
41
42struct TenantReturn {
43    id: TenantId,
44    created_at: OffsetDateTime,
45}
46
47async fn get_tenants(
48    client: &Pool<Postgres>,
49    cursor: &OffsetDateTime,
50    count: usize,
51) -> Result<Vec<TenantReturn>, OperationOutcomeError> {
52    let result = query_as!(
53        TenantReturn,
54        r#"SELECT id as "id: TenantId", created_at FROM tenants WHERE created_at > $1 ORDER BY created_at DESC LIMIT $2"#,
55        cursor,
56        count as i64
57    )
58    .fetch_all(client)
59    .await
60    .map_err(IndexingWorkerError::from)?;
61
62    Ok(result)
63}
64
65static TOTAL_INDEXED: std::sync::LazyLock<Mutex<usize>> =
66    std::sync::LazyLock::new(|| Mutex::new(0));
67
68async fn index_tenant_next_sequence<
69    Repo: FHIRRepository + IndexLockProvider,
70    Engine: SearchEngine,
71>(
72    search_client: Arc<Engine>,
73    tx: &Repo,
74    tenant_id: &TenantId,
75) -> Result<(), IndexingWorkerError> {
76    let start = std::time::Instant::now();
77    let tenant_locks = tx.get_available_locks(vec![tenant_id]).await?;
78
79    if tenant_locks.is_empty() {
80        return Ok(());
81    }
82
83    let resources = tx
84        .get_sequence(
85            tenant_id,
86            tenant_locks[0].index_sequence_position as u64,
87            Some(1000),
88        )
89        .await?;
90
91    // Perform indexing if there are resources to index.
92    if !resources.is_empty() {
93        let result = search_client
94            .index(
95                &SupportedFHIRVersions::R4,
96                &tenant_id,
97                resources
98                    .iter()
99                    .map(|r| IndexResource {
100                        id: &r.id,
101                        version_id: &r.version_id,
102                        project: &r.project,
103                        fhir_method: &r.fhir_method,
104                        resource_type: &r.resource_type,
105                        resource: &r.resource.0,
106                    })
107                    .collect(),
108            )
109            .await?;
110
111        if result.0 != resources.len() {
112            tracing::error!(
113                "Indexed resource count '{}' does not match retrieved resource count '{}'",
114                result.0,
115                resources.len()
116            );
117            return Err(IndexingWorkerError::Fatal);
118        }
119
120        if let Some(resource) = resources.last() {
121            let diff = (resource.sequence + 1) - resources[0].sequence;
122            let total = resources.len();
123
124            if total != diff as usize {
125                tracing::event!(
126                    tracing::Level::INFO,
127                    // safe_seq = resource.max_safe_seq.unwrap_or(0),
128                    first_seq = resources[0].sequence,
129                    last_seq = resource.sequence,
130                    total = resources.len(),
131                    diff = (resource.sequence + 1) - resources[0].sequence
132                );
133            }
134
135            tx.update_lock(tenant_id.as_ref(), resource.sequence as usize)
136                .await?;
137
138            let elapsed = start.elapsed();
139            tracing::info!(
140                "Indexed {} resources for tenant '{}' in {:.2?} (up to sequence {})",
141                result.0,
142                tenant_id.as_ref(),
143                elapsed,
144                resource.sequence
145            );
146        }
147
148        *(TOTAL_INDEXED.lock().await) += result.0;
149    }
150
151    Ok(())
152}
153
154async fn index_for_tenant<Search: SearchEngine, Repository: FHIRRepository + IndexLockProvider>(
155    repo: Arc<Repository>,
156    search_client: Arc<Search>,
157    tenant_id: &TenantId,
158) -> Result<(), IndexingWorkerError> {
159    let search_client = search_client.clone();
160
161    let tx = repo.transaction(false).await.unwrap();
162
163    let res = index_tenant_next_sequence(search_client, &tx, &tenant_id).await;
164
165    match res {
166        Ok(res) => {
167            tx.commit().await?;
168            Ok(res)
169        }
170        Err(e) => {
171            tx.rollback().await?;
172            Err(e)
173        }
174    }
175}
176
177pub enum IndexingWorkerEnvironmentVariables {
178    DatabaseURL,
179    ElasticSearchURL,
180    ElasticSearchUsername,
181    ElasticSearchPassword,
182}
183
184impl From<IndexingWorkerEnvironmentVariables> for String {
185    fn from(value: IndexingWorkerEnvironmentVariables) -> Self {
186        match value {
187            IndexingWorkerEnvironmentVariables::DatabaseURL => "DATABASE_URL".to_string(),
188            IndexingWorkerEnvironmentVariables::ElasticSearchURL => "ELASTICSEARCH_URL".to_string(),
189            IndexingWorkerEnvironmentVariables::ElasticSearchUsername => {
190                "ELASTICSEARCH_USERNAME".to_string()
191            }
192            IndexingWorkerEnvironmentVariables::ElasticSearchPassword => {
193                "ELASTICSEARCH_PASSWORD".to_string()
194            }
195        }
196    }
197}
198
199pub async fn run_worker() -> Result<(), OperationOutcomeError> {
200    let config = get_config::<IndexingWorkerEnvironmentVariables>("environment".into());
201    let subscriber = tracing_subscriber::FmtSubscriber::new();
202    tracing::subscriber::set_global_default(subscriber).unwrap();
203    let fp_engine = Arc::new(haste_fhirpath::FPEngine::new());
204
205    let search_engine = Arc::new(
206        ElasticSearchEngine::new(
207            fp_engine.clone(),
208            &config
209                .get(IndexingWorkerEnvironmentVariables::ElasticSearchURL)
210                .expect(&format!(
211                    "'{}' variable not set",
212                    String::from(IndexingWorkerEnvironmentVariables::ElasticSearchURL)
213                )),
214            config
215                .get(IndexingWorkerEnvironmentVariables::ElasticSearchUsername)
216                .expect(&format!(
217                    "'{}' variable not set",
218                    String::from(IndexingWorkerEnvironmentVariables::ElasticSearchUsername)
219                )),
220            config
221                .get(IndexingWorkerEnvironmentVariables::ElasticSearchPassword)
222                .expect(&format!(
223                    "'{}' variable not set",
224                    String::from(IndexingWorkerEnvironmentVariables::ElasticSearchPassword)
225                )),
226        )
227        .expect("Failed to create Elasticsearch client"),
228    );
229
230    let mut attempts = 0;
231    while !search_engine.is_connected().await.is_ok() && attempts < 5 {
232        tracing::error!("Elasticsearch is not connected, retrying in 5 seconds...");
233        tokio::time::sleep(std::time::Duration::from_secs(5)).await;
234        attempts += 1;
235    }
236
237    let pg_pool = sqlx::PgPool::connect(
238        &config
239            .get(IndexingWorkerEnvironmentVariables::DatabaseURL)
240            .unwrap(),
241    )
242    .await
243    .expect("Failed to connect to the database");
244
245    let repo = Arc::new(haste_repository::pg::PGConnection::pool(pg_pool.clone()));
246    let mut cursor = OffsetDateTime::UNIX_EPOCH;
247    let tenants_limit: usize = 100;
248
249    tracing::info!("Starting indexing worker...");
250
251    let mut k = *TOTAL_INDEXED.lock().await;
252
253    loop {
254        let tenants_to_check = get_tenants(&pg_pool, &cursor, tenants_limit).await;
255        if let Ok(tenants_to_check) = tenants_to_check {
256            if tenants_to_check.is_empty() || tenants_to_check.len() < tenants_limit {
257                cursor = OffsetDateTime::UNIX_EPOCH; // Reset cursor if no tenants found
258            } else {
259                cursor = tenants_to_check[0].created_at;
260            }
261
262            for tenant in tenants_to_check {
263                let result =
264                    index_for_tenant(repo.clone(), search_engine.clone(), &tenant.id).await;
265
266                if let Err(_error) = result {
267                    tracing::error!(
268                        "Failed to index tenant: '{}' cause: '{:?}'",
269                        &tenant.id,
270                        _error
271                    );
272                }
273            }
274        } else if let Err(error) = tenants_to_check {
275            tracing::error!("Failed to retrieve tenants: {:?}", error);
276        }
277
278        if k != *TOTAL_INDEXED.lock().await {
279            k = *TOTAL_INDEXED.lock().await;
280            tracing::info!("TOTAL INDEXED SO FAR: {}", k);
281        }
282    }
283}
haste_indexing_worker/lib.rs

haste_indexing_worker/
lib.rs
