refactor(init): make lib initiazation logic async

Author: georgeh0

src/builder/analyzer.rs

@@ -995,7 +995,7 @@ pub async fn analyze_flow(
     impl Future<Output = Result<ExecutionPlan>> + Send + use<>,
 )> {
     let analyzer_ctx = AnalyzerContext {
-        lib_ctx: get_lib_context()?,
+        lib_ctx: get_lib_context().await?,
         flow_ctx,
     };
     let root_data_scope = Arc::new(Mutex::new(DataScopeBuilder::new()));
@@ -1109,7 +1109,7 @@ pub async fn analyze_transient_flow<'a>(
 )> {
     let mut root_data_scope = DataScopeBuilder::new();
     let analyzer_ctx = AnalyzerContext {
-        lib_ctx: get_lib_context()?,
+        lib_ctx: get_lib_context().await?,
         flow_ctx,
     };
     let mut input_fields = vec![];

src/builder/flow_builder.rs

@@ -247,8 +247,12 @@ pub struct FlowBuilder {
 #[pymethods]
 impl FlowBuilder {
     #[new]
-    pub fn new(name: &str) -> PyResult<Self> {
-        let lib_context = get_lib_context().into_py_result()?;
+    pub fn new(py: Python<'_>, name: &str) -> PyResult<Self> {
+        let lib_context = py
+            .allow_threads(|| -> anyhow::Result<Arc<LibContext>> {
+                get_runtime().block_on(get_lib_context())
+            })
+            .into_py_result()?;
         let root_op_scope = OpScope::new(
             spec::ROOT_SCOPE_NAME.to_string(),
             None,

src/execution/db_tracking_setup.rs

@@ -252,7 +252,7 @@ impl ResourceSetupChange for TrackingTableSetupChange {
 
 impl TrackingTableSetupChange {
     pub async fn apply_change(&self) -> Result<()> {
-        let lib_context = get_lib_context()?;
+        let lib_context = get_lib_context().await?;
         let pool = lib_context.require_builtin_db_pool()?;
         if let Some(desired) = &self.desired_state {
             for lagacy_name in self.legacy_tracking_table_names.iter() {

src/lib_context.rs

@@ -267,7 +267,7 @@ pub fn get_auth_registry() -> &'static Arc<AuthRegistry> {
 }
 
 static LIB_INIT: OnceLock<()> = OnceLock::new();
-pub fn create_lib_context(settings: settings::Settings) -> Result<LibContext> {
+pub async fn create_lib_context(settings: settings::Settings) -> Result<LibContext> {
     LIB_INIT.get_or_init(|| {
         let _ = env_logger::try_init();
 
@@ -278,11 +278,8 @@ pub fn create_lib_context(settings: settings::Settings) -> Result<LibContext> {
 
     let db_pools = DbPools::default();
     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((pool, existing_ss))
-        })?;
+        let pool = db_pools.get_pool(database_spec).await?;
+        let all_setup_states = setup::get_existing_setup_state(&pool).await?;
         Some(PersistenceContext {
             builtin_db_pool: pool,
             setup_ctx: tokio::sync::RwLock::new(LibSetupContext {
@@ -308,24 +305,26 @@ pub fn create_lib_context(settings: settings::Settings) -> Result<LibContext> {
     })
 }
 
-pub static LIB_CONTEXT: RwLock<Option<Arc<LibContext>>> = RwLock::new(None);
+static GET_SETTINGS_FN: Mutex<Option<Box<dyn Fn() -> settings::Settings + Send + Sync>>> =
+    Mutex::new(None);
+static LIB_CONTEXT: Mutex<Option<Arc<LibContext>>> = Mutex::new(None);
 
-pub(crate) fn init_lib_context(settings: settings::Settings) -> Result<()> {
-    let mut lib_context_locked = LIB_CONTEXT.write().unwrap();
-    *lib_context_locked = Some(Arc::new(create_lib_context(settings)?));
+pub(crate) async fn init_lib_context(settings: settings::Settings) -> Result<()> {
+    let mut lib_context_locked = LIB_CONTEXT.lock().unwrap();
+    *lib_context_locked = Some(Arc::new(create_lib_context(settings).await?));
     Ok(())
 }
 
-pub(crate) fn get_lib_context() -> Result<Arc<LibContext>> {
-    let lib_context_locked = LIB_CONTEXT.read().unwrap();
+pub(crate) async fn get_lib_context() -> Result<Arc<LibContext>> {
+    let lib_context_locked = LIB_CONTEXT.lock().unwrap();
     lib_context_locked
         .as_ref()
         .cloned()
         .ok_or_else(|| anyhow!("CocoIndex library is not initialized or already stopped"))
 }
 
 pub(crate) fn clear_lib_context() {
-    let mut lib_context_locked = LIB_CONTEXT.write().unwrap();
+    let mut lib_context_locked = LIB_CONTEXT.lock().unwrap();
     *lib_context_locked = None;
 }
 
@@ -339,15 +338,17 @@ mod tests {
         assert!(db_pools.pools.lock().unwrap().is_empty());
     }
 
-    #[test]
-    fn test_lib_context_without_database() {
-        let lib_context = create_lib_context(settings::Settings::default()).unwrap();
+    #[tokio::test]
+    async fn test_lib_context_without_database() {
+        let lib_context = create_lib_context(settings::Settings::default())
+            .await
+            .unwrap();
         assert!(lib_context.persistence_ctx.is_none());
         assert!(lib_context.require_builtin_db_pool().is_err());
     }
 
-    #[test]
-    fn test_persistence_context_type_safety() {
+    #[tokio::test]
+    async fn test_persistence_context_type_safety() {
         // This test ensures that PersistenceContext groups related fields together
         let settings = settings::Settings {
             database: Some(settings::DatabaseConnectionSpec {
@@ -361,7 +362,7 @@ mod tests {
         };
 
         // This would fail at runtime due to invalid connection, but we're testing the structure
-        let result = create_lib_context(settings);
+        let result = create_lib_context(settings).await;
         // We expect this to fail due to invalid connection, but the structure should be correct
         assert!(result.is_err());
     }

src/ops/shared/postgres.rs

@@ -10,7 +10,7 @@ pub async fn get_db_pool(
     db_ref: Option<&spec::AuthEntryReference<DatabaseConnectionSpec>>,
     auth_registry: &AuthRegistry,
 ) -> Result<PgPool> {
-    let lib_context = get_lib_context()?;
+    let lib_context = get_lib_context().await?;
     let db_conn_spec = db_ref
         .as_ref()
         .map(|db_ref| auth_registry.get(db_ref))

src/py/mod.rs

@@ -97,19 +97,17 @@ impl<T> AnyhowIntoPyResult<T> for anyhow::Result<T> {
 #[pyfunction]
 fn init(py: Python<'_>, settings: Pythonized<Settings>) -> PyResult<()> {
     py.allow_threads(|| -> anyhow::Result<()> {
-        init_lib_context(settings.into_inner())?;
-        Ok(())
+        get_runtime().block_on(async move { init_lib_context(settings.into_inner()).await })
     })
     .into_py_result()
 }
 
 #[pyfunction]
 fn start_server(py: Python<'_>, settings: Pythonized<ServerSettings>) -> PyResult<()> {
     py.allow_threads(|| -> anyhow::Result<()> {
-        let server = get_runtime().block_on(server::init_server(
-            get_lib_context()?,
-            settings.into_inner(),
-        ))?;
+        let server = get_runtime().block_on(async move {
+            server::init_server(get_lib_context().await?, settings.into_inner()).await
+        })?;
         get_runtime().spawn(server);
         Ok(())
     })
@@ -202,7 +200,7 @@ impl FlowLiveUpdater {
     ) -> PyResult<Bound<'py, PyAny>> {
         let flow = flow.0.clone();
         future_into_py(py, async move {
-            let lib_context = get_lib_context().into_py_result()?;
+            let lib_context = get_lib_context().await.into_py_result()?;
             let live_updater = execution::FlowLiveUpdater::start(
                 flow,
                 lib_context.require_builtin_db_pool().into_py_result()?,
@@ -262,7 +260,7 @@ impl Flow {
             get_runtime()
                 .block_on(async {
                     let exec_plan = self.0.flow.get_execution_plan().await?;
-                    let lib_context = get_lib_context()?;
+                    let lib_context = get_lib_context().await?;
                     let execution_ctx = self.0.use_execution_ctx().await?;
                     execution::dumper::evaluate_and_dump(
                         &exec_plan,
@@ -457,9 +455,9 @@ pub struct SetupChangeBundle(Arc<setup::SetupChangeBundle>);
 #[pymethods]
 impl SetupChangeBundle {
     pub fn describe_async<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
-        let lib_context = get_lib_context().into_py_result()?;
         let bundle = self.0.clone();
         future_into_py(py, async move {
+            let lib_context = get_lib_context().await.into_py_result()?;
             bundle.describe(&lib_context).await.into_py_result()
         })
     }
@@ -469,10 +467,10 @@ impl SetupChangeBundle {
         py: Python<'py>,
         report_to_stdout: bool,
     ) -> PyResult<Bound<'py, PyAny>> {
-        let lib_context = get_lib_context().into_py_result()?;
         let bundle = self.0.clone();
 
         future_into_py(py, async move {
+            let lib_context = get_lib_context().await.into_py_result()?;
             let mut stdout = None;
             let mut sink = None;
             bundle
@@ -493,7 +491,7 @@ impl SetupChangeBundle {
 #[pyfunction]
 fn flow_names_with_setup_async(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
     future_into_py(py, async move {
-        let lib_context = get_lib_context().into_py_result()?;
+        let lib_context = get_lib_context().await.into_py_result()?;
         let setup_ctx = lib_context
             .require_persistence_ctx()
             .into_py_result()?
@@ -524,11 +522,15 @@ fn make_drop_bundle(flow_names: Vec<String>) -> PyResult<SetupChangeBundle> {
 }
 
 #[pyfunction]
-fn remove_flow_context(flow_name: String) {
-    let lib_context_locked = crate::lib_context::LIB_CONTEXT.read().unwrap();
-    if let Some(lib_context) = lib_context_locked.as_ref() {
-        lib_context.remove_flow_context(&flow_name)
-    }
+fn remove_flow_context(py: Python<'_>, flow_name: String) -> PyResult<()> {
+    py.allow_threads(|| -> anyhow::Result<()> {
+        get_runtime().block_on(async move {
+            let lib_context = get_lib_context().await.into_py_result()?;
+            lib_context.remove_flow_context(&flow_name);
+            Ok(())
+        })
+    })
+    .into_py_result()
 }
 
 #[pyfunction]

src/setup/db_metadata.rs

@@ -355,7 +355,7 @@ impl MetadataTableSetup {
         if !self.metadata_table_missing {
             return Ok(());
         }
-        let lib_context = get_lib_context()?;
+        let lib_context = get_lib_context().await?;
         let pool = lib_context.require_builtin_db_pool()?;
         let query_str = format!(
             "CREATE TABLE IF NOT EXISTS {SETUP_METADATA_TABLE_NAME} (