Make the auth registry prior to LibContext to be configured before main (#289)

Author: badmonster0

src/builder/analyzed_flow.rs

@@ -21,13 +21,12 @@ impl AnalyzedFlow {
         flow_instance: crate::base::spec::FlowInstanceSpec,
         existing_flow_ss: Option<&setup::FlowSetupState<setup::ExistingMode>>,
         registry: &ExecutorFactoryRegistry,
-        auth_registry: &Arc<AuthRegistry>,
     ) -> Result<Self> {
-        let ctx = analyzer::build_flow_instance_context(&flow_instance.name, auth_registry);
+        let ctx = analyzer::build_flow_instance_context(&flow_instance.name);
         let (data_schema, execution_plan_fut, desired_state) =
             analyzer::analyze_flow(&flow_instance, &ctx, existing_flow_ss, registry)?;
         let setup_status_check =
-            setup::check_flow_setup_status(Some(&desired_state), existing_flow_ss, auth_registry)?;
+            setup::check_flow_setup_status(Some(&desired_state), existing_flow_ss)?;
         let execution_plan = if setup_status_check.is_up_to_date() {
             Some(
                 async move {
@@ -73,9 +72,8 @@ impl AnalyzedTransientFlow {
     pub async fn from_transient_flow(
         transient_flow: spec::TransientFlowSpec,
         registry: &ExecutorFactoryRegistry,
-        auth_registry: &Arc<AuthRegistry>,
     ) -> Result<Self> {
-        let ctx = analyzer::build_flow_instance_context(&transient_flow.name, auth_registry);
+        let ctx = analyzer::build_flow_instance_context(&transient_flow.name);
         let (output_type, data_schema, execution_plan_fut) =
             analyzer::analyze_transient_flow(&transient_flow, &ctx, registry)?;
         Ok(Self {

src/builder/analyzer.rs

@@ -4,9 +4,10 @@ use std::{collections::HashMap, future::Future, sync::Arc};
 
 use super::plan::*;
 use crate::execution::db_tracking_setup;
+use crate::lib_context::get_auth_registry;
 use crate::setup::{
-    self, AuthRegistry, DesiredMode, FlowSetupMetadata, FlowSetupState, ResourceIdentifier,
-    SourceSetupState, TargetSetupState, TargetSetupStateCommon,
+    self, DesiredMode, FlowSetupMetadata, FlowSetupState, ResourceIdentifier, SourceSetupState,
+    TargetSetupState, TargetSetupStateCommon,
 };
 use crate::utils::fingerprint::Fingerprinter;
 use crate::{
@@ -1027,13 +1028,10 @@ impl AnalyzerContext<'_> {
     }
 }
 
-pub fn build_flow_instance_context(
-    flow_inst_name: &str,
-    auth_registry: &Arc<AuthRegistry>,
-) -> Arc<FlowInstanceContext> {
+pub fn build_flow_instance_context(flow_inst_name: &str) -> Arc<FlowInstanceContext> {
     Arc::new(FlowInstanceContext {
         flow_instance_name: flow_inst_name.to_string(),
-        auth_registry: auth_registry.clone(),
+        auth_registry: get_auth_registry().clone(),
     })
 }
 

src/builder/flow_builder.rs

@@ -347,7 +347,7 @@ impl FlowBuilder {
             .get(name)
             .cloned();
         let root_data_scope = Arc::new(Mutex::new(DataScopeBuilder::new()));
-        let flow_inst_context = build_flow_instance_context(name, &lib_context.auth_registry);
+        let flow_inst_context = build_flow_instance_context(name);
         let result = Self {
             lib_context,
             flow_inst_context,
@@ -649,7 +649,6 @@ impl FlowBuilder {
                     spec,
                     self.existing_flow_ss.as_ref(),
                     &crate::ops::executor_factory_registry(),
-                    &self.lib_context.auth_registry,
                 ))
             })
             .into_py_result()?;
@@ -690,7 +689,6 @@ impl FlowBuilder {
                 get_runtime().block_on(super::AnalyzedTransientFlow::from_transient_flow(
                     spec,
                     &crate::ops::executor_factory_registry(),
-                    &self.lib_context.auth_registry,
                 ))
             })
             .into_py_result()?;

src/lib_context.rs

@@ -59,11 +59,11 @@ impl FlowContext {
 }
 
 static TOKIO_RUNTIME: LazyLock<Runtime> = LazyLock::new(|| Runtime::new().unwrap());
+static AUTH_REGISTRY: LazyLock<Arc<AuthRegistry>> = LazyLock::new(|| Arc::new(AuthRegistry::new()));
 
 pub struct LibContext {
     pub pool: PgPool,
     pub flows: Mutex<BTreeMap<String, Arc<FlowContext>>>,
-    pub auth_registry: Arc<AuthRegistry>,
     pub all_setup_states: RwLock<setup::AllSetupState<setup::ExistingMode>>,
 }
 
@@ -87,6 +87,10 @@ pub fn get_runtime() -> &'static Runtime {
     &TOKIO_RUNTIME
 }
 
+pub fn get_auth_registry() -> &'static Arc<AuthRegistry> {
+    &AUTH_REGISTRY
+}
+
 static LIB_INIT: OnceLock<()> = OnceLock::new();
 pub fn create_lib_context(settings: settings::Settings) -> Result<LibContext> {
     LIB_INIT.get_or_init(|| {
@@ -104,7 +108,6 @@ pub fn create_lib_context(settings: settings::Settings) -> Result<LibContext> {
         pool,
         all_setup_states: RwLock::new(all_setup_states),
         flows: Mutex::new(BTreeMap::new()),
-        auth_registry: Arc::new(AuthRegistry::new()),
     })
 }
 

src/py/mod.rs

@@ -2,7 +2,7 @@ use crate::prelude::*;
 
 use crate::base::spec::VectorSimilarityMetric;
 use crate::execution::query;
-use crate::lib_context::{clear_lib_context, init_lib_context};
+use crate::lib_context::{clear_lib_context, get_auth_registry, init_lib_context};
 use crate::ops::interface::QueryResults;
 use crate::ops::py_factory::PyOpArgSchema;
 use crate::ops::{interface::ExecutorFactory, py_factory::PyFunctionFactory, register_factory};
@@ -282,17 +282,15 @@ fn sync_setup() -> PyResult<SetupStatusCheck> {
     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 setup_status = setup::sync_setup(&flows, &all_setup_states, &lib_context.auth_registry)
-        .into_py_result()?;
+    let setup_status = setup::sync_setup(&flows, &all_setup_states).into_py_result()?;
     Ok(SetupStatusCheck(setup_status))
 }
 
 #[pyfunction]
 fn drop_setup(flow_names: Vec<String>) -> PyResult<SetupStatusCheck> {
     let lib_context = get_lib_context().into_py_result()?;
     let all_setup_states = lib_context.all_setup_states.read().unwrap();
-    let setup_status = setup::drop_setup(flow_names, &all_setup_states, &lib_context.auth_registry)
-        .into_py_result()?;
+    let setup_status = setup::drop_setup(flow_names, &all_setup_states).into_py_result()?;
     Ok(SetupStatusCheck(setup_status))
 }
 
@@ -323,9 +321,7 @@ fn apply_setup_changes(py: Python<'_>, setup_status: &SetupStatusCheck) -> PyRes
 
 #[pyfunction]
 fn add_auth_entry(key: String, value: Pythonized<serde_json::Value>) -> PyResult<()> {
-    let lib_context = get_lib_context().into_py_result()?;
-    lib_context
-        .auth_registry
+    get_auth_registry()
         .add(key, value.into_inner())
         .into_py_result()?;
     Ok(())

src/setup/driver.rs

@@ -1,4 +1,4 @@
-use crate::prelude::*;
+use crate::{lib_context::get_auth_registry, prelude::*};
 
 use indexmap::IndexMap;
 use serde::de::DeserializeOwned;
@@ -213,7 +213,6 @@ fn group_resource_states<'a>(
 pub fn check_flow_setup_status(
     desired_state: Option<&FlowSetupState<DesiredMode>>,
     existing_state: Option<&FlowSetupState<ExistingMode>>,
-    auth_registry: &Arc<AuthRegistry>,
 ) -> Result<FlowSetupStatusCheck> {
     let metadata_change = diff_state(
         existing_state.map(|e| &e.metadata),
@@ -294,7 +293,7 @@ pub fn check_flow_setup_status(
                 &resource_id.key,
                 target_state,
                 existing_without_setup_by_user,
-                auth_registry,
+                get_auth_registry(),
             )?)
         };
         target_resources.push(ResourceSetupInfo {
@@ -316,18 +315,13 @@ pub fn check_flow_setup_status(
 pub fn sync_setup(
     flows: &BTreeMap<String, Arc<FlowContext>>,
     all_setup_state: &AllSetupState<ExistingMode>,
-    auth_registry: &Arc<AuthRegistry>,
 ) -> Result<AllSetupStatusCheck> {
     let mut flow_status_checks = BTreeMap::new();
     for (flow_name, flow_context) in flows {
         let existing_state = all_setup_state.flows.get(flow_name);
         flow_status_checks.insert(
             flow_name.clone(),
-            check_flow_setup_status(
-                Some(&flow_context.flow.desired_state),
-                existing_state,
-                auth_registry,
-            )?,
+            check_flow_setup_status(Some(&flow_context.flow.desired_state), existing_state)?,
         );
     }
     Ok(AllSetupStatusCheck {
@@ -342,7 +336,6 @@ pub fn sync_setup(
 pub fn drop_setup(
     flow_names: impl IntoIterator<Item = String>,
     all_setup_state: &AllSetupState<ExistingMode>,
-    auth_registry: &Arc<AuthRegistry>,
 ) -> Result<AllSetupStatusCheck> {
     if !all_setup_state.has_metadata_table {
         api_bail!("CocoIndex metadata table is missing.");
@@ -352,7 +345,7 @@ pub fn drop_setup(
         if let Some(existing_state) = all_setup_state.flows.get(&flow_name) {
             flow_status_checks.insert(
                 flow_name,
-                check_flow_setup_status(None, Some(existing_state), auth_registry)?,
+                check_flow_setup_status(None, Some(existing_state))?,
             );
         }
     }