1use std::sync::Arc;
23
24#[cfg(feature = "arrow")]
25use arrow::{datatypes::Schema, record_batch::RecordBatch};
26use dashmap::{DashMap, mapref::entry::Entry};
27use nautilus_core::Params;
28#[cfg(feature = "python")]
29use pyo3::types::PyAnyMethods;
30
31use crate::data::{CustomData, CustomDataTrait, Data, DataType};
32
33pub type JsonDeserializer =
34 Box<dyn Fn(serde_json::Value) -> Result<Arc<dyn CustomDataTrait>, anyhow::Error> + Send + Sync>;
35#[cfg(feature = "arrow")]
36pub type ArrowEncoder =
37 Box<dyn Fn(&[Arc<dyn CustomDataTrait>]) -> Result<RecordBatch, anyhow::Error> + Send + Sync>;
38#[cfg(feature = "arrow")]
39pub type ArrowDecoder = Box<
40 dyn Fn(
41 &std::collections::HashMap<String, String>,
42 RecordBatch,
43 ) -> Result<Vec<Data>, anyhow::Error>
44 + Send
45 + Sync,
46>;
47
48struct Registries {
49 json: DashMap<String, JsonDeserializer>,
50 #[cfg(feature = "arrow")]
51 arrow: DashMap<String, (Arc<Schema>, ArrowEncoder, ArrowDecoder)>,
52}
53
54fn registries() -> &'static Registries {
55 static REGISTRIES: std::sync::OnceLock<Registries> = std::sync::OnceLock::new();
56 REGISTRIES.get_or_init(|| Registries {
57 json: DashMap::new(),
58 #[cfg(feature = "arrow")]
59 arrow: DashMap::new(),
60 })
61}
62
63pub fn register_json_deserializer(
69 type_name: &str,
70 deserializer: JsonDeserializer,
71) -> Result<(), anyhow::Error> {
72 let reg = registries();
73 match reg.json.entry(type_name.to_string()) {
74 Entry::Occupied(_) => {
75 anyhow::bail!("Custom data type \"{type_name}\" is already registered for JSON");
76 }
77 Entry::Vacant(v) => {
78 v.insert(deserializer);
79 Ok(())
80 }
81 }
82}
83
84pub fn ensure_json_deserializer_registered(
91 type_name: &str,
92 deserializer: JsonDeserializer,
93) -> Result<(), anyhow::Error> {
94 let reg = registries();
95 reg.json
96 .entry(type_name.to_string())
97 .or_insert_with(|| deserializer);
98 Ok(())
99}
100
101fn parse_data_type_from_value(value: &serde_json::Value) -> Option<DataType> {
103 let obj = value.get("data_type")?.as_object()?;
104 let type_name = obj.get("type_name")?.as_str()?;
105 let metadata = obj.get("metadata").and_then(|m| {
106 if m.is_null() {
107 None
108 } else {
109 let p: Params = serde_json::from_value(m.clone()).ok()?;
110 if p.is_empty() { None } else { Some(p) }
111 }
112 });
113 let identifier = obj
114 .get("identifier")
115 .and_then(|v| v.as_str())
116 .map(String::from);
117 Some(DataType::new(type_name, metadata, identifier))
118}
119
120fn parse_envelope_payload(value: &serde_json::Value) -> Result<serde_json::Value, anyhow::Error> {
124 let payload = value
125 .get("payload")
126 .cloned()
127 .ok_or_else(|| anyhow::anyhow!("CustomData JSON missing 'payload' field"))?;
128 Ok(payload)
129}
130
131pub fn deserialize_custom_from_json(
137 type_name: &str,
138 value: &serde_json::Value,
139) -> Result<Option<Data>, anyhow::Error> {
140 let reg = registries();
141 let deserializer_ref = match reg.json.get(type_name) {
142 Some(d) => d,
143 None => return Ok(None),
144 };
145 let data_type = parse_data_type_from_value(value);
146 let payload = parse_envelope_payload(value)?;
147 let arc = deserializer_ref.value()(payload)?;
148 let custom = match data_type {
149 Some(dt) => CustomData::new(arc, dt),
150 None => CustomData::from_arc(arc),
151 };
152 Ok(Some(Data::Custom(custom)))
153}
154
155#[cfg(feature = "arrow")]
160pub fn register_arrow(
161 type_name: &str,
162 schema: Arc<Schema>,
163 encoder: ArrowEncoder,
164 decoder: ArrowDecoder,
165) -> Result<(), anyhow::Error> {
166 let reg = registries();
167 match reg.arrow.entry(type_name.to_string()) {
168 Entry::Occupied(_) => {
169 anyhow::bail!("Custom data type \"{type_name}\" is already registered for Arrow");
170 }
171 Entry::Vacant(v) => {
172 v.insert((schema, encoder, decoder));
173 Ok(())
174 }
175 }
176}
177
178#[cfg(feature = "arrow")]
185pub fn ensure_arrow_registered(
186 type_name: &str,
187 schema: Arc<Schema>,
188 encoder: ArrowEncoder,
189 decoder: ArrowDecoder,
190) -> Result<(), anyhow::Error> {
191 let reg = registries();
192 reg.arrow
193 .entry(type_name.to_string())
194 .or_insert_with(|| (schema, encoder, decoder));
195 Ok(())
196}
197
198#[must_use]
200#[cfg(feature = "arrow")]
201pub fn get_arrow_schema(type_name: &str) -> Option<Arc<Schema>> {
202 let reg = registries();
203 reg.arrow
204 .get(type_name)
205 .map(|entry| Arc::clone(&entry.value().0))
206}
207
208#[cfg(feature = "arrow")]
213pub fn encode_custom_to_arrow(
214 type_name: &str,
215 items: &[Arc<dyn CustomDataTrait>],
216) -> Result<Option<RecordBatch>, anyhow::Error> {
217 let reg = registries();
218 let entry = match reg.arrow.get(type_name) {
219 Some(e) => e,
220 None => return Ok(None),
221 };
222 let encoder = &entry.value().1;
223 encoder(items).map(Some)
224}
225
226#[expect(
231 clippy::implicit_hasher,
232 reason = "callers always use the default hasher"
233)]
234#[cfg(feature = "arrow")]
235pub fn decode_custom_from_arrow(
236 type_name: &str,
237 metadata: &std::collections::HashMap<String, String>,
238 record_batch: RecordBatch,
239) -> Result<Option<Vec<Data>>, anyhow::Error> {
240 let reg = registries();
241 let entry = match reg.arrow.get(type_name) {
242 Some(e) => e,
243 None => return Ok(None),
244 };
245 let decoder = &entry.value().2;
246 decoder(metadata, record_batch).map(Some)
247}
248
249#[cfg(feature = "python")]
250pub type PyExtractor = Box<
251 dyn for<'a> Fn(&pyo3::Bound<'a, pyo3::PyAny>) -> Option<Arc<dyn CustomDataTrait>> + Send + Sync,
252>;
253
254#[cfg(feature = "python")]
255fn py_extractors() -> &'static DashMap<String, PyExtractor> {
256 static PY_EXTRACTORS: std::sync::OnceLock<DashMap<String, PyExtractor>> =
257 std::sync::OnceLock::new();
258 PY_EXTRACTORS.get_or_init(DashMap::new)
259}
260
261#[cfg(feature = "python")]
267pub fn register_py_extractor(type_name: &str, extractor: PyExtractor) -> Result<(), anyhow::Error> {
268 let reg = py_extractors();
269 match reg.entry(type_name.to_string()) {
270 Entry::Occupied(_) => {
271 anyhow::bail!(
272 "Custom data type \"{type_name}\" is already registered for Python extraction"
273 );
274 }
275 Entry::Vacant(v) => {
276 v.insert(extractor);
277 Ok(())
278 }
279 }
280}
281
282#[cfg(feature = "python")]
289pub fn ensure_py_extractor_registered(
290 type_name: &str,
291 extractor: PyExtractor,
292) -> Result<(), anyhow::Error> {
293 let reg = py_extractors();
294 reg.entry(type_name.to_string())
295 .or_insert_with(|| extractor);
296 Ok(())
297}
298
299#[cfg(feature = "python")]
302#[must_use]
303pub fn try_extract_from_py(
304 type_name: &str,
305 obj: &pyo3::Bound<'_, pyo3::PyAny>,
306) -> Option<Arc<dyn CustomDataTrait>> {
307 let reg = py_extractors();
308 let entry = reg.get(type_name)?;
309 let extractor = entry.value();
310 extractor(obj)
311}
312
313#[cfg(feature = "python")]
314type RustExtractorFactory = Box<dyn Fn() -> PyExtractor + Send + Sync>;
315
316#[cfg(feature = "python")]
317fn rust_extractor_factories() -> &'static DashMap<String, RustExtractorFactory> {
318 static RUST_EXTRACTOR_FACTORIES: std::sync::OnceLock<DashMap<String, RustExtractorFactory>> =
319 std::sync::OnceLock::new();
320 RUST_EXTRACTOR_FACTORIES.get_or_init(DashMap::new)
321}
322
323#[cfg(feature = "python")]
331pub fn register_rust_extractor_factory(
332 type_name: &str,
333 factory: RustExtractorFactory,
334) -> Result<(), anyhow::Error> {
335 let reg = rust_extractor_factories();
336 match reg.entry(type_name.to_string()) {
337 Entry::Occupied(_) => {
338 anyhow::bail!("Rust extractor factory for \"{type_name}\" is already registered");
339 }
340 Entry::Vacant(v) => {
341 v.insert(factory);
342 Ok(())
343 }
344 }
345}
346
347#[cfg(feature = "python")]
354pub fn ensure_rust_extractor_factory_registered(
355 type_name: &str,
356 factory: RustExtractorFactory,
357) -> Result<(), anyhow::Error> {
358 let reg = rust_extractor_factories();
359 reg.entry(type_name.to_string()).or_insert_with(|| factory);
360 Ok(())
361}
362
363#[cfg(feature = "python")]
370pub fn register_rust_extractor<T>() -> Result<(), anyhow::Error>
371where
372 T: CustomDataTrait + for<'a, 'py> pyo3::FromPyObject<'a, 'py> + Send + Sync + 'static,
373{
374 let type_name = T::type_name_static();
375 let factory: RustExtractorFactory = Box::new(|| {
376 Box::new(|obj: &pyo3::Bound<'_, pyo3::PyAny>| {
377 obj.extract::<T>()
378 .ok()
379 .map(|x| Arc::new(x) as Arc<dyn CustomDataTrait>)
380 })
381 });
382 register_rust_extractor_factory(type_name, factory)
383}
384
385#[cfg(feature = "python")]
392pub fn ensure_rust_extractor_registered<T>() -> Result<(), anyhow::Error>
393where
394 T: CustomDataTrait + for<'a, 'py> pyo3::FromPyObject<'a, 'py> + Send + Sync + 'static,
395{
396 let type_name = T::type_name_static();
397 let factory: RustExtractorFactory = Box::new(|| {
398 Box::new(|obj: &pyo3::Bound<'_, pyo3::PyAny>| {
399 obj.extract::<T>()
400 .ok()
401 .map(|x| Arc::new(x) as Arc<dyn CustomDataTrait>)
402 })
403 });
404 ensure_rust_extractor_factory_registered(type_name, factory)
405}
406
407#[cfg(feature = "python")]
409#[must_use]
410pub fn get_rust_extractor(type_name: &str) -> Option<PyExtractor> {
411 let reg = rust_extractor_factories();
412 let factory_ref = reg.get(type_name)?;
413 Some(factory_ref.value()())
414}
415
416#[cfg(test)]
417mod tests {
418 use std::hash::{DefaultHasher, Hash, Hasher};
419
420 use nautilus_core::UnixNanos;
421 use rstest::rstest;
422 use serde::{Deserialize, Serialize};
423
424 use super::*;
425 use crate::data::{CustomData, custom::register_custom_data_json};
426
427 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
428 struct TestRegCustomData {
429 ts_init: UnixNanos,
430 }
431
432 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
433 #[serde(deny_unknown_fields)]
434 struct StrictRegCustomData {
435 ts_init: UnixNanos,
436 }
437
438 impl crate::data::HasTsInit for TestRegCustomData {
439 fn ts_init(&self) -> UnixNanos {
440 self.ts_init
441 }
442 }
443
444 impl crate::data::custom::CustomDataTrait for TestRegCustomData {
445 fn type_name(&self) -> &'static str {
446 "TestRegCustomData"
447 }
448 fn type_name_static() -> &'static str {
449 "TestRegCustomData"
450 }
451 fn as_any(&self) -> &dyn std::any::Any {
452 self
453 }
454 fn ts_event(&self) -> nautilus_core::UnixNanos {
455 self.ts_init
456 }
457 fn to_json(&self) -> anyhow::Result<String> {
458 Ok(serde_json::to_string(self)?)
459 }
460 fn clone_arc(&self) -> Arc<dyn crate::data::CustomDataTrait> {
461 Arc::new(self.clone())
462 }
463 fn eq_arc(&self, other: &dyn crate::data::CustomDataTrait) -> bool {
464 other.as_any().downcast_ref::<Self>() == Some(self)
465 }
466 fn from_json(
467 value: serde_json::Value,
468 ) -> anyhow::Result<Arc<dyn crate::data::CustomDataTrait>> {
469 let t: Self = serde_json::from_value(value)?;
470 Ok(Arc::new(t))
471 }
472 }
473
474 impl crate::data::HasTsInit for StrictRegCustomData {
475 fn ts_init(&self) -> UnixNanos {
476 self.ts_init
477 }
478 }
479
480 impl crate::data::custom::CustomDataTrait for StrictRegCustomData {
481 fn type_name(&self) -> &'static str {
482 "StrictRegCustomData"
483 }
484 fn type_name_static() -> &'static str {
485 "StrictRegCustomData"
486 }
487 fn as_any(&self) -> &dyn std::any::Any {
488 self
489 }
490 fn ts_event(&self) -> nautilus_core::UnixNanos {
491 self.ts_init
492 }
493 fn to_json(&self) -> anyhow::Result<String> {
494 Ok(serde_json::to_string(self)?)
495 }
496 fn clone_arc(&self) -> Arc<dyn crate::data::CustomDataTrait> {
497 Arc::new(self.clone())
498 }
499 fn eq_arc(&self, other: &dyn crate::data::CustomDataTrait) -> bool {
500 other.as_any().downcast_ref::<Self>() == Some(self)
501 }
502 fn from_json(
503 value: serde_json::Value,
504 ) -> anyhow::Result<Arc<dyn crate::data::CustomDataTrait>> {
505 let t: Self = serde_json::from_value(value)?;
506 Ok(Arc::new(t))
507 }
508 }
509
510 #[rstest]
511 fn json_registry_roundtrip() {
512 let _ = register_custom_data_json::<TestRegCustomData>();
513
514 let data = Data::Custom(CustomData::from_arc(Arc::new(TestRegCustomData {
515 ts_init: UnixNanos::from(100),
516 })));
517
518 let json = serde_json::to_string(&data).unwrap();
519 let back: Data = serde_json::from_str(&json).unwrap();
520
521 match (&data, &back) {
522 (Data::Custom(a), Data::Custom(b)) => {
523 assert_eq!(a.data.type_name(), b.data.type_name());
524 assert_eq!(a.data.ts_init(), b.data.ts_init());
525 }
526 _ => panic!("expected Custom variant"),
527 }
528 }
529
530 #[rstest]
531 fn json_registry_roundtrip_with_deny_unknown_fields() {
532 let _ = register_custom_data_json::<StrictRegCustomData>();
533
534 let data = Data::Custom(CustomData::from_arc(Arc::new(StrictRegCustomData {
535 ts_init: UnixNanos::from(200),
536 })));
537
538 let json = serde_json::to_string(&data).unwrap();
539 let back: Data = serde_json::from_str(&json).unwrap();
540
541 match (&data, &back) {
542 (Data::Custom(a), Data::Custom(b)) => {
543 assert_eq!(a.data.type_name(), b.data.type_name());
544 assert_eq!(a.data.ts_init(), b.data.ts_init());
545 }
546 _ => panic!("expected Custom variant"),
547 }
548 }
549
550 #[rstest]
551 fn data_type_registry_result_hashes_like_equal_values_from_all_routes() {
552 fn hash_data_type(data_type: &DataType) -> u64 {
553 let mut hasher = DefaultHasher::new();
554 data_type.hash(&mut hasher);
555 hasher.finish()
556 }
557
558 let metadata = serde_json::json!({"key": "value"});
559 let constructed = DataType::new(
560 "ExampleType",
561 Some(serde_json::from_value(metadata.clone()).unwrap()),
562 Some("catalog/path".to_string()),
563 );
564 let persistence_json = serde_json::json!({
565 "type_name": constructed.type_name(),
566 "metadata": metadata,
567 "identifier": constructed.identifier(),
568 });
569 let persisted = DataType::from_persistence_json(&persistence_json.to_string()).unwrap();
570 let registry_envelope = serde_json::json!({"data_type": persistence_json});
571 let registered = parse_data_type_from_value(®istry_envelope).unwrap();
572 let deserialization_payload = serde_json::json!({
573 "type_name": constructed.type_name(),
574 "metadata": constructed.metadata(),
575 "topic": constructed.topic(),
576 "hash": constructed.precomputed_hash() ^ u64::MAX,
577 "identifier": constructed.identifier(),
578 });
579 let deserialized: DataType = serde_json::from_value(deserialization_payload).unwrap();
580
581 for data_type in [&deserialized, &persisted, ®istered] {
585 assert_eq!(data_type, &constructed);
586 assert_eq!(hash_data_type(data_type), hash_data_type(&constructed));
587 assert_eq!(data_type.type_name(), constructed.type_name());
588 assert_eq!(data_type.metadata(), constructed.metadata());
589 assert_eq!(data_type.identifier(), constructed.identifier());
590 assert_eq!(data_type.topic(), constructed.topic());
591 }
592 }
593
594 #[rstest]
595 fn ensure_json_deserializer_registered_is_idempotent() {
596 let deserializer: JsonDeserializer = Box::new(|value| {
597 let t: TestRegCustomData = serde_json::from_value(value)?;
598 Ok(Arc::new(t) as Arc<dyn crate::data::CustomDataTrait>)
599 });
600 let r1 = ensure_json_deserializer_registered("IdempotentTestJson", deserializer);
601 assert!(r1.is_ok(), "first registration should succeed");
602 let deserializer2: JsonDeserializer = Box::new(|value| {
603 let t: TestRegCustomData = serde_json::from_value(value)?;
604 Ok(Arc::new(t) as Arc<dyn crate::data::CustomDataTrait>)
605 });
606 let r2 = ensure_json_deserializer_registered("IdempotentTestJson", deserializer2);
607 assert!(
608 r2.is_ok(),
609 "second registration with same type_name should succeed (idempotent)"
610 );
611 }
612
613 #[rstest]
614 fn register_json_deserializer_fails_on_duplicate() {
615 let deserializer: JsonDeserializer = Box::new(|value| {
616 let t: TestRegCustomData = serde_json::from_value(value)?;
617 Ok(Arc::new(t) as Arc<dyn crate::data::CustomDataTrait>)
618 });
619 let r1 = register_json_deserializer("StrictDuplicateTestJson", deserializer);
620 assert!(r1.is_ok());
621 let deserializer2: JsonDeserializer = Box::new(|value| {
622 let t: TestRegCustomData = serde_json::from_value(value)?;
623 Ok(Arc::new(t) as Arc<dyn crate::data::CustomDataTrait>)
624 });
625 let r2 = register_json_deserializer("StrictDuplicateTestJson", deserializer2);
626 assert!(r2.is_err());
627 let err_msg = r2.unwrap_err().to_string();
628 assert!(
629 err_msg.contains("already registered"),
630 "expected 'already registered' in error, found: {err_msg}"
631 );
632 }
633
634 #[rstest]
635 #[cfg(feature = "arrow")]
636 fn ensure_arrow_registered_is_idempotent() {
637 let schema = Arc::new(arrow::datatypes::Schema::empty());
638 let encoder: ArrowEncoder = Box::new(|_| {
639 Ok(arrow::record_batch::RecordBatch::new_empty(Arc::new(
640 arrow::datatypes::Schema::empty(),
641 )))
642 });
643 let decoder: ArrowDecoder = Box::new(|_, _| Ok(Vec::new()));
644
645 let r1 = ensure_arrow_registered("IdempotentTestArrow", schema, encoder, decoder);
646 assert!(r1.is_ok(), "first Arrow registration should succeed");
647
648 let schema2 = Arc::new(arrow::datatypes::Schema::empty());
649 let encoder2: ArrowEncoder = Box::new(|_| {
650 Ok(arrow::record_batch::RecordBatch::new_empty(Arc::new(
651 arrow::datatypes::Schema::empty(),
652 )))
653 });
654 let decoder2: ArrowDecoder = Box::new(|_, _| Ok(Vec::new()));
655
656 let r2 = ensure_arrow_registered("IdempotentTestArrow", schema2, encoder2, decoder2);
657 assert!(
658 r2.is_ok(),
659 "second Arrow registration with same type_name should be idempotent"
660 );
661 }
662}