Skip to main content

recoco_core/ops/
registration.rs

1// ReCoco is a Rust-only fork of CocoIndex, by [CocoIndex](https://CocoIndex)
2// Original code from CocoIndex is copyrighted by CocoIndex
3// SPDX-FileCopyrightText: 2025-2026 CocoIndex (upstream)
4// SPDX-FileContributor: CocoIndex Contributors
5//
6// All modifications from the upstream for ReCoco are copyrighted by Knitli Inc.
7// SPDX-FileCopyrightText: 2026 Knitli Inc. (ReCoco)
8// SPDX-FileContributor: Adam Poulemanos <adam@knit.li>
9//
10// Both the upstream CocoIndex code and the ReCoco modifications are licensed under the Apache-2.0 License.
11// SPDX-License-Identifier: Apache-2.0
12
13#[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}