Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions python/cocoindex/auth_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")

Expand Down Expand Up @@ -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))
18 changes: 10 additions & 8 deletions src/base/spec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -516,6 +516,14 @@ pub struct TransientFlowSpec {
pub output_value: ValueMapping,
}

impl<T> AuthEntryReference<T> {
pub fn new(key: String) -> Self {
Self {
key,
_phantom: std::marker::PhantomData,
}
}
}
pub struct AuthEntryReference<T> {
pub key: String,
_phantom: std::marker::PhantomData<T>,
Expand All @@ -535,10 +543,7 @@ impl<T> fmt::Display for AuthEntryReference<T> {

impl<T> Clone for AuthEntryReference<T> {
fn clone(&self) -> Self {
Self {
key: self.key.clone(),
_phantom: std::marker::PhantomData,
}
Self::new(self.key.clone())
}
}

Expand All @@ -562,10 +567,7 @@ impl<'de, T> Deserialize<'de> for AuthEntryReference<T> {
D: serde::Deserializer<'de>,
{
let untyped_ref = UntypedAuthEntryReference::<String>::deserialize(deserializer)?;
Ok(AuthEntryReference {
key: untyped_ref.key,
_phantom: std::marker::PhantomData,
})
Ok(AuthEntryReference::new(untyped_ref.key))
}
}

Expand Down
10 changes: 9 additions & 1 deletion src/py/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};
Expand Down Expand Up @@ -630,6 +630,13 @@ fn add_transient_auth_entry(value: Pythonized<serde_json::Value>) -> PyResult<St
.into_py_result()
}

#[pyfunction]
fn get_auth_entry(key: String) -> PyResult<Pythonized<serde_json::Value>> {
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<String> {
let app_namespace = py
Expand Down Expand Up @@ -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::<builder::flow_builder::FlowBuilder>()?;
Expand Down
Loading