Skip to main content

jellyflow_runtime/schema/registry/
mod.rs

1use std::collections::BTreeMap;
2use std::sync::Arc;
3
4use jellyflow_core::core::{CanvasPoint, NodeId, NodeKindKey, PortId};
5
6use super::migration::NodeKindMigrator;
7use super::types::{NodeInstantiation, NodeInstantiationError, NodeKindViewDescriptor, NodeSchema};
8
9mod plans;
10
11/// Registry for node kinds.
12#[derive(Default, Clone)]
13pub struct NodeRegistry {
14    by_kind: BTreeMap<NodeKindKey, NodeSchema>,
15    by_alias: BTreeMap<NodeKindKey, NodeKindKey>,
16    migrators: BTreeMap<NodeKindKey, Arc<dyn NodeKindMigrator>>,
17}
18
19impl std::fmt::Debug for NodeRegistry {
20    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
21        f.debug_struct("NodeRegistry")
22            .field("schema_count", &self.by_kind.len())
23            .field("alias_count", &self.by_alias.len())
24            .field("migrator_count", &self.migrators.len())
25            .finish()
26    }
27}
28
29impl NodeRegistry {
30    /// Creates an empty registry.
31    pub fn new() -> Self {
32        Self::default()
33    }
34
35    /// Registers a schema.
36    ///
37    /// Aliases are mapped to the schema's canonical kind.
38    pub fn register(&mut self, schema: NodeSchema) {
39        for alias in &schema.kind_aliases {
40            self.by_alias.insert(alias.clone(), schema.kind.clone());
41        }
42        self.by_kind.insert(schema.kind.clone(), schema);
43    }
44
45    /// Registers a per-kind data migrator.
46    ///
47    /// The migrator is stored as an in-memory hook (not serialized as part of the schema data).
48    pub fn register_migrator(
49        &mut self,
50        kind: NodeKindKey,
51        migrator: Arc<dyn NodeKindMigrator>,
52    ) -> &mut Self {
53        self.migrators.insert(kind, migrator);
54        self
55    }
56
57    /// Resolves an input kind to a canonical kind (via aliases).
58    pub fn resolve_kind<'a>(&'a self, kind: &'a NodeKindKey) -> &'a NodeKindKey {
59        if self.by_kind.contains_key(kind) {
60            kind
61        } else {
62            self.by_alias.get(kind).unwrap_or(kind)
63        }
64    }
65
66    /// Looks up a schema by canonical kind key.
67    pub fn get(&self, kind: &NodeKindKey) -> Option<&NodeSchema> {
68        self.by_kind.get(kind)
69    }
70
71    /// Iterates all registered schemas in deterministic order (by kind key).
72    pub fn schemas(&self) -> impl Iterator<Item = &NodeSchema> {
73        self.by_kind.values()
74    }
75
76    /// Returns the adapter-facing descriptor for a node kind or alias.
77    pub fn view_descriptor(&self, kind: &NodeKindKey) -> Option<NodeKindViewDescriptor> {
78        let canonical = self.resolve_kind(kind);
79        self.get(canonical).map(NodeKindViewDescriptor::from_schema)
80    }
81
82    /// Returns adapter-facing node-kind descriptors in deterministic order.
83    pub fn view_descriptors(&self) -> Vec<NodeKindViewDescriptor> {
84        self.schemas()
85            .map(NodeKindViewDescriptor::from_schema)
86            .collect()
87    }
88
89    /// Instantiates a node kind or alias with freshly allocated node and port ids.
90    pub fn instantiate_node(
91        &self,
92        kind: &NodeKindKey,
93        pos: CanvasPoint,
94    ) -> Result<NodeInstantiation, NodeInstantiationError> {
95        let canonical = self.resolve_kind(kind);
96        self.get(canonical)
97            .map(|schema| schema.instantiate(pos))
98            .ok_or_else(|| NodeInstantiationError::MissingSchema(kind.clone()))
99    }
100
101    /// Instantiates a node kind or alias with caller-provided ids.
102    pub fn instantiate_node_with_ids(
103        &self,
104        kind: &NodeKindKey,
105        node_id: NodeId,
106        pos: CanvasPoint,
107        port_ids: impl IntoIterator<Item = PortId>,
108    ) -> Result<NodeInstantiation, NodeInstantiationError> {
109        let canonical = self.resolve_kind(kind);
110        self.get(canonical)
111            .ok_or_else(|| NodeInstantiationError::MissingSchema(kind.clone()))?
112            .instantiate_with_ids(node_id, pos, port_ids)
113    }
114}