diff --git a/python/cocoindex/auth_registry.py b/python/cocoindex/auth_registry.py index 87c202fa..f37f82f1 100644 --- a/python/cocoindex/auth_registry.py +++ b/python/cocoindex/auth_registry.py @@ -4,10 +4,9 @@ from dataclasses import dataclass from typing import Generic, TypeVar -import threading from . import _engine # type: ignore -from .convert import dump_engine_object +from .convert import dump_engine_object, load_engine_object T = TypeVar("T") @@ -38,3 +37,8 @@ def add_auth_entry(key: str, value: T) -> AuthEntryReference[T]: def ref_auth_entry(key: str) -> AuthEntryReference[T]: """Reference an auth entry by its key.""" return AuthEntryReference(key) + + +def get_auth_entry(cls: type[T], ref: TransientAuthEntryReference[T]) -> T: + """Get an auth entry by its key.""" + return load_engine_object(cls, _engine.get_auth_entry(ref.key)) diff --git a/src/base/spec.rs b/src/base/spec.rs index ed55d345..3c01ae2e 100644 --- a/src/base/spec.rs +++ b/src/base/spec.rs @@ -516,6 +516,14 @@ pub struct TransientFlowSpec { pub output_value: ValueMapping, } +impl AuthEntryReference { + pub fn new(key: String) -> Self { + Self { + key, + _phantom: std::marker::PhantomData, + } + } +} pub struct AuthEntryReference { pub key: String, _phantom: std::marker::PhantomData, @@ -535,10 +543,7 @@ impl fmt::Display for AuthEntryReference { impl Clone for AuthEntryReference { fn clone(&self) -> Self { - Self { - key: self.key.clone(), - _phantom: std::marker::PhantomData, - } + Self::new(self.key.clone()) } } @@ -562,10 +567,7 @@ impl<'de, T> Deserialize<'de> for AuthEntryReference { D: serde::Deserializer<'de>, { let untyped_ref = UntypedAuthEntryReference::::deserialize(deserializer)?; - Ok(AuthEntryReference { - key: untyped_ref.key, - _phantom: std::marker::PhantomData, - }) + Ok(AuthEntryReference::new(untyped_ref.key)) } } diff --git a/src/py/mod.rs b/src/py/mod.rs index 85821bd2..2a9ac71a 100644 --- a/src/py/mod.rs +++ b/src/py/mod.rs @@ -2,7 +2,7 @@ use crate::execution::evaluator::evaluate_transient_flow; use crate::prelude::*; use crate::base::schema::{FieldSchema, ValueType}; -use crate::base::spec::{NamedSpec, OutputMode, ReactiveOpSpec, SpecFormatter}; +use crate::base::spec::{AuthEntryReference, NamedSpec, OutputMode, ReactiveOpSpec, SpecFormatter}; use crate::lib_context::{ QueryHandlerContext, clear_lib_context, get_auth_registry, init_lib_context, }; @@ -630,6 +630,13 @@ fn add_transient_auth_entry(value: Pythonized) -> PyResult PyResult> { + let auth_ref = AuthEntryReference::new(key); + let json_value: serde_json::Value = get_auth_registry().get(&auth_ref).into_py_result()?; + Ok(Pythonized(json_value)) +} + #[pyfunction] fn get_app_namespace(py: Python<'_>) -> PyResult { let app_namespace = py @@ -671,6 +678,7 @@ fn cocoindex_engine(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(remove_flow_context, m)?)?; m.add_function(wrap_pyfunction!(add_auth_entry, m)?)?; m.add_function(wrap_pyfunction!(add_transient_auth_entry, m)?)?; + m.add_function(wrap_pyfunction!(get_auth_entry, m)?)?; m.add_function(wrap_pyfunction!(get_app_namespace, m)?)?; m.add_class::()?;