1use crate::dialect::{PlatformDialect, PlatformRegistry};
10use crate::event::{EventRegistry, EventSink, RazeriEvent};
11use crate::language::{LanguageRegistry, LanguageRendererFactory};
12use crate::llm::{LlmClientFactory, LlmRegistry};
13use crate::model::{LlmModelAdapter, ModelRegistry};
14use crate::registry::PluginRegistry;
15use crate::storage::{StorageProviderFactory, StorageRegistry};
16use chrono::Utc;
17use razeri_plugin_contract::StoragePlatform;
18use std::sync::Arc;
19
20pub struct PluginRegistryBuilder {
32 dialects: PlatformRegistry,
33 storage: StorageRegistry,
34 llm: LlmRegistry,
35 models: ModelRegistry,
36 language: LanguageRegistry,
37 events: EventRegistry,
38}
39
40impl PluginRegistryBuilder {
41 #[must_use]
43 pub fn new() -> Self {
44 Self {
45 dialects: PlatformRegistry::new(),
46 storage: StorageRegistry::new(),
47 llm: LlmRegistry::new(),
48 models: ModelRegistry::new(),
49 language: LanguageRegistry::new(),
50 events: EventRegistry::new(),
51 }
52 }
53
54 #[must_use]
58 pub fn with_dialect(mut self, name: &str, dialect: Arc<dyn PlatformDialect>) -> Self {
59 self.dialects.register(name.to_string(), dialect);
60 self
61 }
62
63 pub fn register_dialect(&mut self, dialect: Arc<dyn PlatformDialect>) {
69 let name = dialect.info().name.to_string();
70 self.dialects.register(name, dialect);
71 }
72
73 #[must_use]
77 pub fn with_storage(
78 mut self,
79 platform: StoragePlatform,
80 factory: Arc<dyn StorageProviderFactory>,
81 ) -> Self {
82 self.storage.register_known(platform, factory);
83 self
84 }
85
86 #[must_use]
88 pub fn with_storage_extension(
89 mut self,
90 name: String,
91 factory: Arc<dyn StorageProviderFactory>,
92 ) -> Self {
93 self.storage.register_extension(name, factory);
94 self
95 }
96
97 #[must_use]
101 pub fn with_llm(mut self, name: String, factory: Arc<dyn LlmClientFactory>) -> Self {
102 self.llm.register(name, factory);
103 self
104 }
105
106 #[must_use]
110 pub fn with_model(mut self, adapter: Arc<dyn LlmModelAdapter>) -> Self {
111 self.models.register(adapter);
112 self
113 }
114
115 #[must_use]
119 pub fn with_language(
120 mut self,
121 language: String,
122 factory: Arc<dyn LanguageRendererFactory>,
123 ) -> Self {
124 self.language.register(language, factory);
125 self
126 }
127
128 #[must_use]
132 pub fn with_event_sink(mut self, sink: Arc<dyn EventSink>) -> Self {
133 self.events.register(sink);
134 self
135 }
136
137 #[must_use]
141 pub fn with_auto_discovered(mut self) -> Self {
142 for reg in inventory::iter::<DialectRegistration> {
143 let dialect = (reg.factory)();
144 let name = dialect.info().name.to_string();
145 self.dialects.register(name, dialect);
146 self.record_plugin_loaded(reg.name, "dialect", "inventory");
147 }
148 for reg in inventory::iter::<StorageRegistration> {
149 let factory = (reg.factory)();
150 if let Some(platform) = factory.known_platform() {
151 self.storage.register_known(platform, factory);
152 } else {
153 self.storage.register_extension(reg.name.to_string(), factory);
154 }
155 self.record_plugin_loaded(reg.name, "storage", "inventory");
156 }
157 for reg in inventory::iter::<LlmRegistration> {
158 let factory = (reg.factory)();
159 self.llm.register(reg.name.to_string(), factory);
160 self.record_plugin_loaded(reg.name, "llm", "inventory");
161 }
162 for reg in inventory::iter::<ModelRegistration> {
163 let adapter = (reg.factory)();
164 self.models.register(adapter);
165 self.record_plugin_loaded(reg.model_id, "model", "inventory");
166 }
167 for reg in inventory::iter::<LanguageRegistration> {
168 let factory = (reg.factory)();
169 self.language.register(reg.language.to_string(), factory);
170 self.record_plugin_loaded(reg.language, "language", "inventory");
171 }
172 for reg in inventory::iter::<EventSinkRegistration> {
173 let sink = (reg.factory)();
174 self.events.register(sink);
175 self.record_plugin_loaded(reg.name, "event_sink", "inventory");
176 }
177 for reg in inventory::iter::<PluginRegistration> {
179 (reg.register)(&mut self);
180 }
181 self
182 }
183
184 fn record_plugin_loaded(&self, name: &str, plugin_type: &str, source: &str) {
186 if let Ok(event) = RazeriEvent::plugin_loaded(
187 format!("plugin_loaded:{plugin_type}:{name}"),
188 Utc::now(),
189 name,
190 plugin_type,
191 source,
192 ) {
193 self.events.record(event);
194 }
195 }
196
197 #[cfg(all(feature = "dynamic", not(target_arch = "wasm32")))]
208 #[must_use]
209 pub fn with_dynamic_plugins(mut self, plugin_dir: &std::path::Path) -> Self {
210 let results = crate::dynamic::load_plugins_from_dir(plugin_dir, &mut self);
211 for result in &results {
212 match result {
213 crate::dynamic::PluginLoadResult::Loaded { name, path } => {
214 tracing::info!("Loaded dynamic plugin '{name}' from {}", path.display());
215 self.record_plugin_loaded(name, "dynamic", &path.display().to_string());
216 }
217 crate::dynamic::PluginLoadResult::Refused { path, refusal } => {
218 tracing::warn!("Plugin {} was not loaded: {refusal}", path.display());
219 }
220 crate::dynamic::PluginLoadResult::Failed { path, error } => {
221 tracing::warn!("Failed to load plugin from {}: {error}", path.display());
222 }
223 }
224 }
225 self
226 }
227
228 pub fn build(self) -> PluginRegistry {
232 PluginRegistry {
233 dialects: self.dialects,
234 storage: self.storage,
235 llm: self.llm,
236 models: self.models,
237 language: self.language,
238 events: self.events,
239 }
240 }
241}
242
243impl Default for PluginRegistryBuilder {
244 fn default() -> Self {
248 Self::new()
249 }
250}
251
252pub struct DialectRegistration {
258 pub name: &'static str,
260 pub factory: fn() -> Arc<dyn PlatformDialect>,
262}
263
264impl DialectRegistration {
265 pub const fn new(name: &'static str, factory: fn() -> Arc<dyn PlatformDialect>) -> Self {
267 Self { name, factory }
268 }
269}
270
271inventory::collect!(DialectRegistration);
272
273pub struct PluginRegistration {
279 pub register: fn(&mut PluginRegistryBuilder),
281}
282
283impl PluginRegistration {
284 pub const fn new(register: fn(&mut PluginRegistryBuilder)) -> Self {
286 Self { register }
287 }
288}
289
290inventory::collect!(PluginRegistration);
291
292pub struct StorageRegistration {
294 pub name: &'static str,
296 pub factory: fn() -> Arc<dyn StorageProviderFactory>,
298}
299
300impl StorageRegistration {
301 pub const fn new(name: &'static str, factory: fn() -> Arc<dyn StorageProviderFactory>) -> Self {
303 Self { name, factory }
304 }
305}
306
307inventory::collect!(StorageRegistration);
308
309pub struct LlmRegistration {
311 pub name: &'static str,
313 pub factory: fn() -> Arc<dyn LlmClientFactory>,
315}
316
317impl LlmRegistration {
318 pub const fn new(name: &'static str, factory: fn() -> Arc<dyn LlmClientFactory>) -> Self {
320 Self { name, factory }
321 }
322}
323
324inventory::collect!(LlmRegistration);
325
326pub struct ModelRegistration {
328 pub model_id: &'static str,
330 pub factory: fn() -> Arc<dyn LlmModelAdapter>,
332}
333
334impl ModelRegistration {
335 pub const fn new(model_id: &'static str, factory: fn() -> Arc<dyn LlmModelAdapter>) -> Self {
337 Self { model_id, factory }
338 }
339}
340
341inventory::collect!(ModelRegistration);
342
343pub struct LanguageRegistration {
345 pub language: &'static str,
347 pub factory: fn() -> Arc<dyn LanguageRendererFactory>,
349}
350
351impl LanguageRegistration {
352 pub const fn new(
354 language: &'static str,
355 factory: fn() -> Arc<dyn LanguageRendererFactory>,
356 ) -> Self {
357 Self { language, factory }
358 }
359}
360
361inventory::collect!(LanguageRegistration);
362
363pub struct EventSinkRegistration {
365 pub name: &'static str,
367 pub factory: fn() -> Arc<dyn EventSink>,
369}
370
371impl EventSinkRegistration {
372 pub const fn new(name: &'static str, factory: fn() -> Arc<dyn EventSink>) -> Self {
374 Self { name, factory }
375 }
376}
377
378inventory::collect!(EventSinkRegistration);
379
380#[cfg(test)]
381mod tests {
382 use super::*;
383 use crate::dialect::{DialectInfo, PlatformDialect};
384 use crate::error::{EventSinkError, LanguageError, LlmError};
385 use crate::event::RazeriEvent;
386 use crate::llm::{ChatResponse, LlmClient, Message, ModelInfo, ToolDefinition};
387 use crate::model::ModelCapabilities;
388 use async_trait::async_trait;
389 use razeri_plugin_contract::DataType;
390 use razeri_plugin_contract::DialectFamily;
391 use razeri_plugin_contract::{
392 DdlRenderer, DialectCapabilities, DmlRenderer, ExpressionRenderer, FormulaRenderer,
393 ProcedureRenderer, ReplicationRenderer, SqlRenderer, StorageRenderer, TaskRenderer,
394 };
395 use std::collections::HashMap;
396
397 struct TestDialect;
400
401 impl SqlRenderer for TestDialect {
402 fn render_data_type(&self, _dt: &DataType) -> String {
407 "TEXT".to_string()
408 }
409 fn quote_identifier(&self, name: &str) -> String {
414 format!("\"{name}\"")
415 }
416 fn qualify_name(&self, _schema: Option<&str>, name: &str) -> String {
422 name.to_string()
423 }
424 fn render_current_timestamp(&self) -> String {
428 "NOW()".to_string()
429 }
430 }
431
432 impl ExpressionRenderer for TestDialect {}
433 impl DdlRenderer for TestDialect {}
434 impl DmlRenderer for TestDialect {}
435 impl ProcedureRenderer for TestDialect {}
436 impl FormulaRenderer for TestDialect {}
437 impl TaskRenderer for TestDialect {}
438 impl ReplicationRenderer for TestDialect {}
439 impl StorageRenderer for TestDialect {}
440
441 impl DialectCapabilities for TestDialect {
442 fn merge_shape(&self) -> razeri_plugin_contract::MergeShape {
446 razeri_plugin_contract::MergeShape::Merge(
447 razeri_plugin_contract::MergeRestrictions::none(),
448 )
449 }
450 fn supports_identity(&self) -> bool {
454 true
455 }
456 fn supports_sequences(&self) -> bool {
460 true
461 }
462 fn supports_stored_procedures(&self) -> bool {
466 true
467 }
468 }
469
470 impl PlatformDialect for TestDialect {
471 fn info(&self) -> DialectInfo {
475 DialectInfo {
476 name: "test",
477 display_name: "Test DB",
478 aliases: &[],
479 dialect_family: DialectFamily::PostgreSql,
480 family: "PostgreSQL",
481 is_cloud_native: false,
482 is_olap: false,
483 is_postgres_compatible: false,
484 versions: &[],
485 deploy_script: None,
486 }
487 }
488 }
489
490 struct TestStorageFactory;
491 #[async_trait]
492 impl StorageProviderFactory for TestStorageFactory {
493 fn name(&self) -> &'static str {
497 "test-storage"
498 }
499 fn display_name(&self) -> &'static str {
503 "Test Storage"
504 }
505 fn known_platform(&self) -> Option<StoragePlatform> {
509 None
510 }
511 async fn create(
512 &self,
513 _config: &HashMap<String, String>,
514 ) -> Result<Box<dyn crate::storage::StorageOps>, crate::error::StorageError> {
515 unimplemented!("test only")
516 }
517 fn validate_config(&self, _config: &HashMap<String, String>) -> Vec<String> {
522 vec![]
523 }
524 fn required_config_keys(&self) -> Vec<&str> {
528 vec![]
529 }
530 }
531
532 struct TestLlmFactory;
533 impl LlmClientFactory for TestLlmFactory {
534 fn name(&self) -> &'static str {
538 "test-llm"
539 }
540 fn display_name(&self) -> &'static str {
544 "Test LLM"
545 }
546 fn create(
551 &self,
552 _config: &HashMap<String, String>,
553 ) -> Result<Arc<dyn LlmClient>, LlmError> {
554 Ok(Arc::new(TestLlmClient))
555 }
556 fn validate_config(&self, _config: &HashMap<String, String>) -> Vec<String> {
561 vec![]
562 }
563 fn required_config_keys(&self) -> Vec<&str> {
567 vec![]
568 }
569 }
570
571 struct TestLlmClient;
572 #[async_trait]
573 impl LlmClient for TestLlmClient {
574 async fn chat(
575 &self,
576 _messages: &[Message],
577 _model: Option<&str>,
578 _tools: Option<Vec<ToolDefinition>>,
579 ) -> Result<ChatResponse, LlmError> {
580 Ok(ChatResponse { content: "ok".into(), model: "test".into(), usage: None })
581 }
582 async fn chat_json(
583 &self,
584 _messages: &[Message],
585 _model: Option<&str>,
586 _tools: Option<Vec<ToolDefinition>>,
587 ) -> Result<serde_json::Value, LlmError> {
588 Ok(serde_json::json!({}))
589 }
590 async fn list_models(&self) -> Result<Vec<ModelInfo>, LlmError> {
591 Ok(vec![])
592 }
593 async fn health_check(&self) -> Result<bool, LlmError> {
594 Ok(true)
595 }
596 }
597
598 struct TestModelAdapter;
599 impl LlmModelAdapter for TestModelAdapter {
600 fn model_id(&self) -> &'static str {
604 "test-model"
605 }
606 fn provider_name(&self) -> &'static str {
610 "test-llm"
611 }
612 fn display_name(&self) -> &'static str {
616 "Test Model"
617 }
618 fn capabilities(&self) -> ModelCapabilities {
622 ModelCapabilities {
623 max_context_tokens: 4096,
624 supports_json_mode: true,
625 supports_tools: false,
626 supports_streaming: false,
627 }
628 }
629 }
630
631 struct TestLanguageFactory;
632 impl LanguageRendererFactory for TestLanguageFactory {
633 fn language(&self) -> &'static str {
637 "test-lang"
638 }
639 fn display_name(&self) -> &'static str {
643 "Test Language"
644 }
645 fn create(
650 &self,
651 _config: &HashMap<String, String>,
652 ) -> Result<Arc<dyn crate::language::LanguageRenderer>, LanguageError> {
653 unimplemented!("test only")
654 }
655 }
656
657 struct TestEventSink;
658 #[async_trait]
659 impl crate::event::EventSink for TestEventSink {
660 fn name(&self) -> &'static str {
664 "test-sink"
665 }
666 async fn on_event(&self, _event: &RazeriEvent) -> Result<(), EventSinkError> {
667 Ok(())
668 }
669 async fn flush(&self) -> Result<(), EventSinkError> {
670 Ok(())
671 }
672 }
673
674 #[test]
680 fn builder_creates_registry_with_all_sub_registries() {
681 let registry = PluginRegistryBuilder::new()
682 .with_dialect("postgres", Arc::new(TestDialect))
683 .with_storage_extension("test".into(), Arc::new(TestStorageFactory))
684 .with_llm("test-llm".into(), Arc::new(TestLlmFactory))
685 .with_model(Arc::new(TestModelAdapter))
686 .with_language("test-lang".into(), Arc::new(TestLanguageFactory))
687 .with_event_sink(Arc::new(TestEventSink))
688 .build();
689
690 assert!(registry.dialects.get_by_name("postgres").is_some());
691 assert!(registry.llm.get("test-llm").is_some());
692 assert!(registry.models.get("test-llm", "test-model").is_some());
693 assert!(registry.language.get("test-lang").is_some());
694 assert!(!registry.events.is_empty());
695 }
696
697 #[test]
701 fn builder_chaining_is_ergonomic() {
702 let _registry = PluginRegistryBuilder::new()
704 .with_dialect("postgres", Arc::new(TestDialect))
705 .with_dialect("snowflake", Arc::new(TestDialect))
706 .with_llm("a".into(), Arc::new(TestLlmFactory))
707 .with_llm("b".into(), Arc::new(TestLlmFactory))
708 .with_event_sink(Arc::new(TestEventSink))
709 .with_event_sink(Arc::new(TestEventSink))
710 .build();
711 }
712
713 #[test]
717 fn builder_with_mixed_known_and_extension() {
718 let registry = PluginRegistryBuilder::new()
719 .with_dialect("postgres", Arc::new(TestDialect))
720 .with_dialect("teradata", Arc::new(TestDialect))
721 .with_storage(StoragePlatform::S3, Arc::new(TestStorageFactory))
722 .with_storage_extension("minio".into(), Arc::new(TestStorageFactory))
723 .build();
724
725 assert!(registry.dialects.get_by_name("postgres").is_some());
726 assert!(registry.dialects.get_by_name("teradata").is_some());
727 }
728
729 #[test]
734 fn builder_with_auto_discovered_finds_inventory_entries() {
735 let registry = PluginRegistryBuilder::new().with_auto_discovered().build();
739 let mut names = registry.dialects.registered_names();
740 names.sort();
741 assert_eq!(names, vec!["linked_by_export".to_owned(), "outsidedb".to_owned()]);
742 }
743
744 #[test]
748 fn manual_registration_records_plugin_loaded_events() {
749 let builder = PluginRegistryBuilder::new()
750 .with_dialect("postgres", Arc::new(TestDialect))
751 .with_llm("test-llm".into(), Arc::new(TestLlmFactory));
752
753 builder.record_plugin_loaded("TestDialect", "dialect", "manual");
757 builder.record_plugin_loaded("TestLlm", "llm", "manual");
758
759 let registry = builder.build();
760 let events = registry.events.events();
761 assert_eq!(events.len(), 2);
762
763 assert_eq!(events[0].category(), crate::event::EventCategory::Plugin);
764 assert_eq!(events[0].summary().title, "Plugin loaded");
765 assert!(events[0].summary().detail.contains("dialect"));
766
767 assert_eq!(events[1].category(), crate::event::EventCategory::Plugin);
768 assert!(events[1].summary().detail.contains("llm"));
769 }
770
771 #[test]
778 fn auto_discovery_registers_no_sink_this_crate_ships() {
779 let registry = PluginRegistryBuilder::new().with_auto_discovered().build();
780 assert_eq!(registry.events.sink_names(), Vec::<String>::new());
781 let sink_events = registry
782 .events
783 .events()
784 .into_iter()
785 .filter(|e| e.summary().detail.contains("event_sink"))
786 .count();
787 assert_eq!(sink_events, 0, "no sink was loaded, so none is reported loaded");
788 }
789
790 struct Linked;
792
793 impl SqlRenderer for Linked {
794 fn render_data_type(&self, _dt: &DataType) -> String {
795 "TEXT".to_string()
796 }
797 fn quote_identifier(&self, name: &str) -> String {
798 format!("\"{name}\"")
799 }
800 fn qualify_name(&self, _schema: Option<&str>, name: &str) -> String {
801 name.to_string()
802 }
803 fn render_current_timestamp(&self) -> String {
804 "NOW()".to_string()
805 }
806 }
807 impl ExpressionRenderer for Linked {}
808 impl DdlRenderer for Linked {}
809 impl DmlRenderer for Linked {}
810 impl ProcedureRenderer for Linked {}
811 impl FormulaRenderer for Linked {}
812 impl TaskRenderer for Linked {}
813 impl ReplicationRenderer for Linked {}
814 impl StorageRenderer for Linked {}
815 impl DialectCapabilities for Linked {}
816 impl PlatformDialect for Linked {
817 fn info(&self) -> DialectInfo {
818 DialectInfo {
819 name: "linked_by_export",
820 display_name: "Linked by export",
821 aliases: &[],
822 dialect_family: DialectFamily::PostgreSql,
823 family: "PostgreSQL",
824 is_cloud_native: false,
825 is_olap: false,
826 is_postgres_compatible: false,
827 versions: &[],
828 deploy_script: None,
829 }
830 }
831 }
832
833 fn register_linked(builder: &mut PluginRegistryBuilder) {
834 builder.register_dialect(Arc::new(Linked));
835 }
836
837 crate::export_plugin!(register_linked);
838
839 #[test]
843 fn a_plugin_exported_with_the_macro_and_linked_in_is_discovered() {
844 let registry = PluginRegistryBuilder::new().with_auto_discovered().build();
845 assert!(
846 registry.dialects.get_by_name("linked_by_export").is_some(),
847 "the builder did not find the linked plugin: {:?}",
848 registry.dialects.registered_names()
849 );
850 assert!(
851 PlatformRegistry::with_auto_discovered().get_by_name("linked_by_export").is_some(),
852 "the dialect registry did not find the linked plugin"
853 );
854 }
855
856 #[test]
858 fn the_macro_exports_this_build() {
859 let exported = unsafe { crate::dynamic::read_declaration(razeri_plugin_build()) };
861 assert_eq!(exported, Ok(crate::dynamic::BuildDeclaration::this_build()));
862 }
863}