recoco_core/ops/
registration.rs1#[cfg(any(
14 feature = "function-detect-lang",
15 feature = "function-embed",
16 feature = "function-extract-llm",
17 feature = "function-json",
18 feature = "function-split"
19))]
20use super::functions;
21#[cfg(any(
22 feature = "source-azure",
23 feature = "source-gdrive",
24 feature = "source-local-file",
25 feature = "source-postgres",
26 feature = "source-s3"
27))]
28use super::sources;
29#[cfg(any(
30 feature = "target-kuzu",
31 feature = "target-neo4j",
32 feature = "target-postgres",
33 feature = "target-qdrant"
34))]
35use super::targets;
36#[allow(unused_imports)]
37use super::{factory_bases::*, registry::ExecutorFactoryRegistry, sdk::ExecutorFactory};
38use crate::prelude::*;
39use recoco_utils::client_error;
40use std::sync::{LazyLock, RwLock};
41
42fn register_executor_factories(
43 #[allow(unused_variables)] registry: &mut ExecutorFactoryRegistry,
44) -> Result<()> {
45 #[cfg(feature = "target-kuzu")]
46 let reqwest_client = reqwest::Client::new();
47
48 #[cfg(feature = "source-local-file")]
49 sources::local_file::Factory.register(registry)?;
50 #[cfg(feature = "source-gdrive")]
51 sources::google_drive::Factory.register(registry)?;
52 #[cfg(feature = "source-s3")]
53 sources::amazon_s3::Factory.register(registry)?;
54 #[cfg(feature = "source-azure")]
55 sources::azure_blob::Factory.register(registry)?;
56 #[cfg(feature = "source-postgres")]
57 sources::postgres::Factory.register(registry)?;
58
59 #[cfg(feature = "function-detect-lang")]
60 functions::detect_program_lang::register(registry)?;
61 #[cfg(feature = "function-embed")]
62 functions::embed_text::register(registry)?;
63 #[cfg(feature = "function-extract-llm")]
64 functions::extract_by_llm::Factory.register(registry)?;
65 #[cfg(feature = "function-json")]
66 functions::parse_json::Factory.register(registry)?;
67 #[cfg(feature = "function-split")]
68 functions::split_by_separators::register(registry)?;
69 #[cfg(feature = "function-split")]
70 functions::split_recursively::register(registry)?;
71
72 #[cfg(feature = "target-postgres")]
73 targets::postgres::register(registry)?;
74 #[cfg(feature = "target-qdrant")]
75 targets::qdrant::register(registry)?;
76 #[cfg(feature = "target-kuzu")]
77 targets::kuzu::register(registry, reqwest_client)?;
78
79 #[cfg(feature = "target-neo4j")]
80 targets::neo4j::Factory::new().register(registry)?;
81
82 Ok(())
83}
84
85static EXECUTOR_FACTORY_REGISTRY: LazyLock<RwLock<ExecutorFactoryRegistry>> = LazyLock::new(|| {
86 let mut registry = ExecutorFactoryRegistry::new();
87 register_executor_factories(&mut registry).expect("Failed to register executor factories");
88 RwLock::new(registry)
89});
90
91pub fn get_optional_source_factory(
92 kind: &str,
93) -> Option<std::sync::Arc<dyn super::interface::SourceFactory + Send + Sync>> {
94 let registry = EXECUTOR_FACTORY_REGISTRY.read().unwrap();
95 registry.get_source(kind).cloned()
96}
97
98pub fn get_optional_function_factory(
99 kind: &str,
100) -> Option<std::sync::Arc<dyn super::interface::SimpleFunctionFactory + Send + Sync>> {
101 let registry = EXECUTOR_FACTORY_REGISTRY.read().unwrap();
102 registry.get_function(kind).cloned()
103}
104
105pub fn get_optional_target_factory(
106 kind: &str,
107) -> Option<std::sync::Arc<dyn super::interface::TargetFactory + Send + Sync>> {
108 let registry = EXECUTOR_FACTORY_REGISTRY.read().unwrap();
109 registry.get_target(kind).cloned()
110}
111
112pub fn get_optional_attachment_factory(
113 kind: &str,
114) -> Option<std::sync::Arc<dyn super::interface::TargetAttachmentFactory + Send + Sync>> {
115 let registry = EXECUTOR_FACTORY_REGISTRY.read().unwrap();
116 registry.get_target_attachment(kind).cloned()
117}
118
119pub fn get_source_factory(
120 kind: &str,
121) -> Result<std::sync::Arc<dyn super::interface::SourceFactory + Send + Sync>> {
122 get_optional_source_factory(kind)
123 .ok_or_else(|| client_error!("Source factory not found for op kind: {}", kind))
124}
125
126pub fn get_function_factory(
127 kind: &str,
128) -> Result<std::sync::Arc<dyn super::interface::SimpleFunctionFactory + Send + Sync>> {
129 get_optional_function_factory(kind)
130 .ok_or_else(|| client_error!("Function factory not found for op kind: {}", kind))
131}
132
133pub fn get_target_factory(
134 kind: &str,
135) -> Result<std::sync::Arc<dyn super::interface::TargetFactory + Send + Sync>> {
136 get_optional_target_factory(kind)
137 .ok_or_else(|| client_error!("Target factory not found for op kind: {}", kind))
138}
139
140pub fn get_attachment_factory(
141 kind: &str,
142) -> Result<std::sync::Arc<dyn super::interface::TargetAttachmentFactory + Send + Sync>> {
143 get_optional_attachment_factory(kind)
144 .ok_or_else(|| client_error!("Attachment factory not found for op kind: {}", kind))
145}
146
147pub fn register_factory(name: String, factory: ExecutorFactory) -> Result<()> {
148 let mut registry = EXECUTOR_FACTORY_REGISTRY.write().unwrap();
149 registry.register(name, factory)
150}