From 98b7db34ac23da91ece7e9dbba7b3bbbf2fc2476 Mon Sep 17 00:00:00 2001 From: Adam Hendel Date: Wed, 20 May 2026 21:44:52 -0500 Subject: [PATCH 1/2] dedupe fns --- Cargo.lock | 1 + core/Cargo.toml | 1 + core/src/cache.rs | 146 ++++++++++++++++++++++++++++++++++++++++++ core/src/lib.rs | 1 + proxy/src/cache.rs | 147 +++---------------------------------------- server/src/cache.rs | 150 +++----------------------------------------- 6 files changed, 167 insertions(+), 279 deletions(-) create mode 100644 core/src/cache.rs diff --git a/Cargo.lock b/Cargo.lock index 0b5198c6..a6783d76 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4244,6 +4244,7 @@ dependencies = [ "tiktoken-rs", "tokio", "tokio-postgres", + "tracing", "url", "utoipa 5.4.0", "uuid", diff --git a/core/Cargo.toml b/core/Cargo.toml index ec5fc71b..0d22a3f7 100644 --- a/core/Cargo.toml +++ b/core/Cargo.toml @@ -26,6 +26,7 @@ serde_json = "1.0" sqlparser = "0.51" sqlx = { version = "0.8", features = ["runtime-tokio-rustls", "postgres", "uuid", "time"] } thiserror = "2.0.12" +tracing = { workspace = true } tiktoken-rs = "0.7.0" tokio = { version = "1.0", features = ["full"] } tokio-postgres = "0.7" diff --git a/core/src/cache.rs b/core/src/cache.rs new file mode 100644 index 00000000..b0b539c4 --- /dev/null +++ b/core/src/cache.rs @@ -0,0 +1,146 @@ +use std::collections::HashMap; +use std::sync::Arc; +use tokio::sync::RwLock; +use tracing::{error, info}; + +use crate::types::VectorizeJob; + +pub async fn setup_job_change_notifications( + pool: &sqlx::PgPool, +) -> Result<(), Box> { + let mut tx = pool.begin().await?; + + let create_notify_function = r#" + CREATE OR REPLACE FUNCTION vectorize.notify_job_change() + RETURNS TRIGGER AS $$ + BEGIN + IF TG_OP = 'DELETE' THEN + PERFORM pg_notify('vectorize_job_changes', + json_build_object( + 'operation', TG_OP, + 'job_name', OLD.job_name + )::text + ); + RETURN OLD; + ELSE + PERFORM pg_notify('vectorize_job_changes', + json_build_object( + 'operation', TG_OP, + 'job_name', NEW.job_name + )::text + ); + RETURN NEW; + END IF; + END; + $$ LANGUAGE plpgsql; + "#; + + sqlx::query("DROP TRIGGER IF EXISTS job_change_trigger ON vectorize.job;") + .execute(&mut *tx) + .await?; + + let create_trigger = r#" + CREATE TRIGGER job_change_trigger + AFTER INSERT OR UPDATE OR DELETE ON vectorize.job + FOR EACH ROW EXECUTE FUNCTION vectorize.notify_job_change(); + "#; + + sqlx::query(create_notify_function) + .execute(&mut *tx) + .await?; + sqlx::query(create_trigger).execute(&mut *tx).await?; + + tx.commit().await?; + info!("Database trigger for job changes setup successfully"); + Ok(()) +} + +pub async fn try_listen_for_changes( + db_pool: &sqlx::PgPool, + job_cache: &Arc>>, +) -> Result<(), Box> { + let mut listener = sqlx::postgres::PgListener::connect_with(db_pool).await?; + listener.listen("vectorize_job_changes").await?; + + info!("Connected and listening for vectorize job changes"); + + loop { + match listener.recv().await { + Ok(notification) => { + info!( + "Received job change notification: {}", + notification.payload() + ); + + if let Ok(payload) = + serde_json::from_str::(notification.payload()) + { + let operation = payload.get("operation").and_then(|v| v.as_str()); + let job_name = payload.get("job_name").and_then(|v| v.as_str()); + info!( + "Job change detected - Operation: {}, Job: {}", + operation.unwrap_or("unknown"), + job_name.unwrap_or("unknown") + ); + } + + if let Err(e) = refresh_job_cache(db_pool, job_cache).await { + error!("Failed to refresh job cache: {e}"); + } else { + info!("Job cache refreshed successfully"); + } + } + Err(e) => { + error!("Error receiving notification: {e}"); + return Err(e.into()); + } + } + } +} + +pub async fn refresh_job_cache( + db_pool: &sqlx::PgPool, + job_cache: &Arc>>, +) -> Result<(), Box> { + let all_jobs: Vec = sqlx::query_as( + "SELECT job_name, src_table, src_schema, src_columns, primary_key, update_time_col, model FROM vectorize.job", + ) + .fetch_all(db_pool) + .await?; + + let jobmap: HashMap = all_jobs + .into_iter() + .map(|mut item| { + let key = std::mem::take(&mut item.job_name); + (key, item) + }) + .collect(); + + { + let mut jobmap_write = job_cache.write().await; + *jobmap_write = jobmap; + info!("Updated job cache with {} jobs", jobmap_write.len()); + } + + Ok(()) +} + +pub async fn load_initial_job_cache( + pool: &sqlx::PgPool, +) -> Result, sqlx::Error> { + let all_jobs: Vec = sqlx::query_as( + "SELECT job_name, src_table, src_schema, src_columns, primary_key, update_time_col, model FROM vectorize.job", + ) + .fetch_all(pool) + .await?; + + let jobmap: HashMap = all_jobs + .into_iter() + .map(|mut item| { + let key = std::mem::take(&mut item.job_name); + (key, item) + }) + .collect(); + + Ok(jobmap) +} diff --git a/core/src/lib.rs b/core/src/lib.rs index b72417b3..3ef529ae 100644 --- a/core/src/lib.rs +++ b/core/src/lib.rs @@ -1,3 +1,4 @@ +pub mod cache; pub mod config; pub mod db; pub mod errors; diff --git a/proxy/src/cache.rs b/proxy/src/cache.rs index 7ee1e1b4..9da5a92f 100644 --- a/proxy/src/cache.rs +++ b/proxy/src/cache.rs @@ -1,59 +1,19 @@ use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; -use tracing::{error, info}; +use tracing::error; use vectorize_core::types::VectorizeJob; use super::protocol::{ProxyConfig, WireProxyError}; -pub async fn setup_job_change_notifications( - pool: &sqlx::PgPool, -) -> Result<(), Box> { - let mut tx = pool.begin().await?; - - let create_notify_function = r#" - CREATE OR REPLACE FUNCTION vectorize.notify_job_change() - RETURNS TRIGGER AS $$ - BEGIN - IF TG_OP = 'DELETE' THEN - PERFORM pg_notify('vectorize_job_changes', - json_build_object( - 'operation', TG_OP, - 'job_name', OLD.job_name - )::text - ); - RETURN OLD; - ELSE - PERFORM pg_notify('vectorize_job_changes', - json_build_object( - 'operation', TG_OP, - 'job_name', NEW.job_name - )::text - ); - RETURN NEW; - END IF; - END; - $$ LANGUAGE plpgsql; - "#; - - sqlx::query("DROP TRIGGER IF EXISTS job_change_trigger ON vectorize.job;") - .execute(&mut *tx) - .await?; - - let create_trigger = r#" - CREATE TRIGGER job_change_trigger - AFTER INSERT OR UPDATE OR DELETE ON vectorize.job - FOR EACH ROW EXECUTE FUNCTION vectorize.notify_job_change(); - "#; +pub use vectorize_core::cache::setup_job_change_notifications; - sqlx::query(create_notify_function) - .execute(&mut *tx) - .await?; - sqlx::query(create_trigger).execute(&mut *tx).await?; - - tx.commit().await?; - info!("Database trigger for job changes setup successfully"); - Ok(()) +pub async fn load_initial_job_cache( + pool: &sqlx::PgPool, +) -> Result, WireProxyError> { + vectorize_core::cache::load_initial_job_cache(pool) + .await + .map_err(WireProxyError::Database) } pub async fn start_cache_sync_listener( @@ -63,7 +23,7 @@ pub async fn start_cache_sync_listener( let max_retry_delay = Duration::from_secs(60); loop { - match try_listen_for_changes(&config).await { + match vectorize_core::cache::try_listen_for_changes(&config.db_pool, &config.jobmap).await { Ok(_) => retry_delay = Duration::from_secs(1), Err(e) => { error!("Cache sync listener error: {e}. Retrying in {retry_delay:?}"); @@ -73,92 +33,3 @@ pub async fn start_cache_sync_listener( } } } - -async fn try_listen_for_changes( - config: &ProxyConfig, -) -> Result<(), Box> { - let mut listener = sqlx::postgres::PgListener::connect_with(&config.db_pool).await?; - listener.listen("vectorize_job_changes").await?; - - info!("Connected and listening for vectorize job changes"); - - loop { - match listener.recv().await { - Ok(notification) => { - info!( - "Received job change notification: {}", - notification.payload() - ); - - if let Ok(payload) = - serde_json::from_str::(notification.payload()) - { - let operation = payload.get("operation").and_then(|v| v.as_str()); - let job_name = payload.get("job_name").and_then(|v| v.as_str()); - info!( - "Job change detected - Operation: {}, Job: {}", - operation.unwrap_or("unknown"), - job_name.unwrap_or("unknown") - ); - } - - if let Err(e) = refresh_job_cache(config).await { - error!("Failed to refresh job cache: {e}"); - } else { - info!("Job cache refreshed successfully"); - } - } - Err(e) => { - error!("Error receiving notification: {e}"); - return Err(e.into()); - } - } - } -} - -pub async fn refresh_job_cache( - config: &ProxyConfig, -) -> Result<(), Box> { - let all_jobs: Vec = sqlx::query_as( - "SELECT job_name, src_table, src_schema, src_columns, primary_key, update_time_col, model FROM vectorize.job", - ) - .fetch_all(&config.db_pool) - .await?; - - let jobmap: HashMap = all_jobs - .into_iter() - .map(|mut item| { - let key = std::mem::take(&mut item.job_name); - (key, item) - }) - .collect(); - - { - let mut jobmap_write = config.jobmap.write().await; - *jobmap_write = jobmap; - info!("Updated job cache with {} jobs", jobmap_write.len()); - } - - Ok(()) -} - -pub async fn load_initial_job_cache( - pool: &sqlx::PgPool, -) -> Result, WireProxyError> { - let all_jobs: Vec = sqlx::query_as( - "SELECT job_name, src_table, src_schema, src_columns, primary_key, update_time_col, model FROM vectorize.job", - ) - .fetch_all(pool) - .await - .map_err(WireProxyError::Database)?; - - let jobmap: HashMap = all_jobs - .into_iter() - .map(|mut item| { - let key = std::mem::take(&mut item.job_name); - (key, item) - }) - .collect(); - - Ok(jobmap) -} diff --git a/server/src/cache.rs b/server/src/cache.rs index 8e42f10e..7a07f1fc 100644 --- a/server/src/cache.rs +++ b/server/src/cache.rs @@ -1,58 +1,17 @@ use std::collections::HashMap; use std::sync::Arc; use tokio::sync::RwLock; -use tracing::{error, info}; +use tracing::error; use vectorize_core::types::VectorizeJob; -/// Cache sync functions for job change notifications -pub async fn setup_job_change_notifications( - pool: &sqlx::PgPool, -) -> Result<(), Box> { - let mut tx = pool.begin().await?; - - let create_notify_function = r#" - CREATE OR REPLACE FUNCTION vectorize.notify_job_change() - RETURNS TRIGGER AS $$ - BEGIN - IF TG_OP = 'DELETE' THEN - PERFORM pg_notify('vectorize_job_changes', - json_build_object( - 'operation', TG_OP, - 'job_name', OLD.job_name - )::text - ); - RETURN OLD; - ELSE - PERFORM pg_notify('vectorize_job_changes', - json_build_object( - 'operation', TG_OP, - 'job_name', NEW.job_name - )::text - ); - RETURN NEW; - END IF; - END; - $$ LANGUAGE plpgsql; - "#; - - sqlx::query("DROP TRIGGER IF EXISTS job_change_trigger ON vectorize.job;") - .execute(&mut *tx) - .await?; - - let create_trigger = r#" - CREATE TRIGGER job_change_trigger - AFTER INSERT OR UPDATE OR DELETE ON vectorize.job - FOR EACH ROW EXECUTE FUNCTION vectorize.notify_job_change(); - "#; +pub use vectorize_core::cache::{refresh_job_cache, setup_job_change_notifications}; - sqlx::query(create_notify_function) - .execute(&mut *tx) - .await?; - sqlx::query(create_trigger).execute(&mut *tx).await?; - - tx.commit().await?; - info!("Database trigger for job changes setup successfully"); - Ok(()) +pub async fn load_initial_job_cache( + pool: &sqlx::PgPool, +) -> Result, crate::app_state::AppStateError> { + vectorize_core::cache::load_initial_job_cache(pool) + .await + .map_err(crate::app_state::AppStateError::Database) } pub async fn start_cache_sync_listener( @@ -63,7 +22,7 @@ pub async fn start_cache_sync_listener( let max_retry_delay = std::time::Duration::from_secs(60); loop { - match try_listen_for_changes(&db_pool, &job_cache).await { + match vectorize_core::cache::try_listen_for_changes(&db_pool, &job_cache).await { Ok(_) => retry_delay = std::time::Duration::from_secs(1), Err(e) => { error!("Cache sync listener error: {e}. Retrying in {retry_delay:?}"); @@ -73,94 +32,3 @@ pub async fn start_cache_sync_listener( } } } - -async fn try_listen_for_changes( - db_pool: &sqlx::PgPool, - job_cache: &Arc>>, -) -> Result<(), Box> { - let mut listener = sqlx::postgres::PgListener::connect_with(db_pool).await?; - listener.listen("vectorize_job_changes").await?; - - info!("Connected and listening for vectorize job changes"); - - loop { - match listener.recv().await { - Ok(notification) => { - info!( - "Received job change notification: {}", - notification.payload() - ); - - if let Ok(payload) = - serde_json::from_str::(notification.payload()) - { - let operation = payload.get("operation").and_then(|v| v.as_str()); - let job_name = payload.get("job_name").and_then(|v| v.as_str()); - info!( - "Job change detected - Operation: {}, Job: {}", - operation.unwrap_or("unknown"), - job_name.unwrap_or("unknown") - ); - } - - if let Err(e) = refresh_job_cache(db_pool, job_cache).await { - error!("Failed to refresh job cache: {e}"); - } else { - info!("Job cache refreshed successfully"); - } - } - Err(e) => { - error!("Error receiving notification: {e}"); - return Err(e.into()); - } - } - } -} - -pub async fn refresh_job_cache( - db_pool: &sqlx::PgPool, - job_cache: &Arc>>, -) -> Result<(), Box> { - let all_jobs: Vec = sqlx::query_as( - "SELECT job_name, src_table, src_schema, src_columns, primary_key, update_time_col, model FROM vectorize.job", - ) - .fetch_all(db_pool) - .await?; - - let jobmap: HashMap = all_jobs - .into_iter() - .map(|mut item| { - let key = std::mem::take(&mut item.job_name); - (key, item) - }) - .collect(); - - { - let mut jobmap_write = job_cache.write().await; - *jobmap_write = jobmap; - info!("Updated job cache with {} jobs", jobmap_write.len()); - } - - Ok(()) -} - -pub async fn load_initial_job_cache( - pool: &sqlx::PgPool, -) -> Result, crate::app_state::AppStateError> { - let all_jobs: Vec = sqlx::query_as( - "SELECT job_name, src_table, src_schema, src_columns, primary_key, update_time_col, model FROM vectorize.job", - ) - .fetch_all(pool) - .await - .map_err(crate::app_state::AppStateError::Database)?; - - let jobmap: HashMap = all_jobs - .into_iter() - .map(|mut item| { - let key = std::mem::take(&mut item.job_name); - (key, item) - }) - .collect(); - - Ok(jobmap) -} From 8824eb63c84c5334fac175deb2876f6d998f06b2 Mon Sep 17 00:00:00 2001 From: Adam Hendel Date: Thu, 28 May 2026 06:25:02 -0500 Subject: [PATCH 2/2] remove unreachable match arm --- proxy/src/cache.rs | 13 ++++++------- server/src/cache.rs | 11 ++++------- 2 files changed, 10 insertions(+), 14 deletions(-) diff --git a/proxy/src/cache.rs b/proxy/src/cache.rs index 9da5a92f..119914ad 100644 --- a/proxy/src/cache.rs +++ b/proxy/src/cache.rs @@ -23,13 +23,12 @@ pub async fn start_cache_sync_listener( let max_retry_delay = Duration::from_secs(60); loop { - match vectorize_core::cache::try_listen_for_changes(&config.db_pool, &config.jobmap).await { - Ok(_) => retry_delay = Duration::from_secs(1), - Err(e) => { - error!("Cache sync listener error: {e}. Retrying in {retry_delay:?}"); - tokio::time::sleep(retry_delay).await; - retry_delay = std::cmp::min(retry_delay * 2, max_retry_delay); - } + if let Err(e) = + vectorize_core::cache::try_listen_for_changes(&config.db_pool, &config.jobmap).await + { + error!("Cache sync listener error: {e}. Retrying in {retry_delay:?}"); + tokio::time::sleep(retry_delay).await; + retry_delay = std::cmp::min(retry_delay * 2, max_retry_delay); } } } diff --git a/server/src/cache.rs b/server/src/cache.rs index 7a07f1fc..171060f3 100644 --- a/server/src/cache.rs +++ b/server/src/cache.rs @@ -22,13 +22,10 @@ pub async fn start_cache_sync_listener( let max_retry_delay = std::time::Duration::from_secs(60); loop { - match vectorize_core::cache::try_listen_for_changes(&db_pool, &job_cache).await { - Ok(_) => retry_delay = std::time::Duration::from_secs(1), - Err(e) => { - error!("Cache sync listener error: {e}. Retrying in {retry_delay:?}"); - tokio::time::sleep(retry_delay).await; - retry_delay = std::cmp::min(retry_delay * 2, max_retry_delay); - } + if let Err(e) = vectorize_core::cache::try_listen_for_changes(&db_pool, &job_cache).await { + error!("Cache sync listener error: {e}. Retrying in {retry_delay:?}"); + tokio::time::sleep(retry_delay).await; + retry_delay = std::cmp::min(retry_delay * 2, max_retry_delay); } } }