refactor: introduce PersistenceContext to encapsulate database pool and setup states

Author: vumichien

src/builder/flow_builder.rs

@@ -259,7 +259,8 @@ impl FlowBuilder {
     pub fn new(name: &str) -> PyResult<Self> {
         let lib_context = get_lib_context().into_py_result()?;
         let existing_flow_ss = lib_context
-            .all_setup_states
+            .require_all_setup_states()
+            .into_py_result()?
             .read()
             .unwrap()
             .flows

src/lib_context.rs

@@ -77,11 +77,15 @@ impl DbPools {
     }
 }
 
+pub struct PersistenceContext {
+    pub builtin_db_pool: PgPool,
+    pub all_setup_states: RwLock<setup::AllSetupState<setup::ExistingMode>>,
+}
+
 pub struct LibContext {
     pub db_pools: DbPools,
-    pub builtin_db_pool: Option<PgPool>,
+    pub persistence_ctx: Option<PersistenceContext>,
     pub flows: Mutex<BTreeMap<String, Arc<FlowContext>>>,
-    pub all_setup_states: RwLock<setup::AllSetupState<setup::ExistingMode>>,
 }
 
 impl LibContext {
@@ -100,8 +104,16 @@ impl LibContext {
     }
 
     pub fn require_builtin_db_pool(&self) -> Result<&PgPool> {
-        self.builtin_db_pool
+        self.persistence_ctx
+            .as_ref()
+            .map(|ctx| &ctx.builtin_db_pool)
+            .ok_or_else(|| anyhow!("Database is required for this operation. Please set COCOINDEX_DATABASE_URL environment variable and call cocoindex.init() with database settings."))
+    }
+
+    pub fn require_all_setup_states(&self) -> Result<&RwLock<setup::AllSetupState<setup::ExistingMode>>> {
+        self.persistence_ctx
             .as_ref()
+            .map(|ctx| &ctx.all_setup_states)
             .ok_or_else(|| anyhow!("Database is required for this operation. Please set COCOINDEX_DATABASE_URL environment variable and call cocoindex.init() with database settings."))
     }
 }
@@ -123,22 +135,24 @@ pub fn create_lib_context(settings: settings::Settings) -> Result<LibContext> {
     });
 
     let db_pools = DbPools::default();
-    let (pool, all_setup_states) = if let Some(database_spec) = &settings.database {
+    let persistence_ctx = if let Some(database_spec) = &settings.database {
         let (pool, all_setup_states) = get_runtime().block_on(async {
             let pool = db_pools.get_pool(database_spec).await?;
             let existing_ss = setup::get_existing_setup_state(&pool).await?;
-            anyhow::Ok((Some(pool), existing_ss))
+            anyhow::Ok((pool, existing_ss))
         })?;
-        (pool, all_setup_states)
+        Some(PersistenceContext {
+            builtin_db_pool: pool,
+            all_setup_states: RwLock::new(all_setup_states),
+        })
     } else {
-        // No database configured - create empty setup states
-        (None, setup::AllSetupState::default())
+        // No database configured
+        None
     };
 
     Ok(LibContext {
         db_pools,
-        builtin_db_pool: pool,
-        all_setup_states: RwLock::new(all_setup_states),
+        persistence_ctx,
         flows: Mutex::new(BTreeMap::new()),
     })
 }
@@ -185,4 +199,35 @@ mod tests {
         assert!(settings.database.is_none());
         assert_eq!(settings.app_namespace, "test");
     }
+
+    #[test]
+    fn test_lib_context_without_database() {
+        let settings = settings::Settings {
+            database: None,
+            app_namespace: "test".to_string(),
+        };
+
+        let lib_context = create_lib_context(settings).unwrap();
+        assert!(lib_context.persistence_ctx.is_none());
+        assert!(lib_context.require_builtin_db_pool().is_err());
+        assert!(lib_context.require_all_setup_states().is_err());
+    }
+
+    #[test]
+    fn test_persistence_context_type_safety() {
+        // This test ensures that PersistenceContext groups related fields together
+        let settings = settings::Settings {
+            database: Some(settings::DatabaseConnectionSpec {
+                url: "postgresql://test".to_string(),
+                user: None,
+                password: None,
+            }),
+            app_namespace: "test".to_string(),
+        };
+
+        // This would fail at runtime due to invalid connection, but we're testing the structure
+        let result = create_lib_context(settings);
+        // We expect this to fail due to invalid connection, but the structure should be correct
+        assert!(result.is_err());
+    }
 }

src/py/mod.rs

@@ -394,7 +394,7 @@ impl SetupStatus {
 fn sync_setup(py: Python<'_>) -> PyResult<SetupStatus> {
     let lib_context = get_lib_context().into_py_result()?;
     let flows = lib_context.flows.lock().unwrap();
-    let all_setup_states = lib_context.all_setup_states.read().unwrap();
+    let all_setup_states = lib_context.require_all_setup_states().into_py_result()?.read().unwrap();
     py.allow_threads(|| {
         get_runtime()
             .block_on(async {
@@ -408,7 +408,7 @@ fn sync_setup(py: Python<'_>) -> PyResult<SetupStatus> {
 #[pyfunction]
 fn drop_setup(py: Python<'_>, flow_names: Vec<String>) -> PyResult<SetupStatus> {
     let lib_context = get_lib_context().into_py_result()?;
-    let all_setup_states = lib_context.all_setup_states.read().unwrap();
+    let all_setup_states = lib_context.require_all_setup_states().into_py_result()?.read().unwrap();
     py.allow_threads(|| {
         get_runtime()
             .block_on(async {
@@ -422,7 +422,7 @@ fn drop_setup(py: Python<'_>, flow_names: Vec<String>) -> PyResult<SetupStatus>
 #[pyfunction]
 fn flow_names_with_setup() -> PyResult<Vec<String>> {
     let lib_context = get_lib_context().into_py_result()?;
-    let all_setup_states = lib_context.all_setup_states.read().unwrap();
+    let all_setup_states = lib_context.require_all_setup_states().into_py_result()?.read().unwrap();
     let flow_names = all_setup_states.flows.keys().cloned().collect();
     Ok(flow_names)
 }