nautilus_model/data/
registry.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! Registries for custom data: JSON (de)serialization and Arrow encode/decode.
17//!
18//! Mirrors Python's `register_custom_data_class` surface in `custom.py`.
19//! The registry only stores type name -> callbacks for lookup; each type provides
20//! its own deserialize/encode/decode via the trait or registration.
21
22use std::sync::Arc;
23
24#[cfg(feature = "arrow")]
25use arrow::{
26    datatypes::{DataType as ArrowDataType, Field, Schema},
27    record_batch::RecordBatch,
28};
29use dashmap::{DashMap, mapref::entry::Entry};
30use nautilus_core::Params;
31#[cfg(feature = "python")]
32use pyo3::types::PyAnyMethods;
33
34use crate::data::{CustomData, CustomDataTrait, Data, DataType};
35
36pub type JsonDeserializer =
37    Box<dyn Fn(serde_json::Value) -> Result<Arc<dyn CustomDataTrait>, anyhow::Error> + Send + Sync>;
38#[cfg(feature = "arrow")]
39pub type ArrowEncoder =
40    Box<dyn Fn(&[Arc<dyn CustomDataTrait>]) -> Result<RecordBatch, anyhow::Error> + Send + Sync>;
41#[cfg(feature = "arrow")]
42pub type ArrowDecoder = Box<
43    dyn Fn(
44            &std::collections::HashMap<String, String>,
45            RecordBatch,
46        ) -> Result<Vec<Data>, anyhow::Error>
47        + Send
48        + Sync,
49>;
50
51/// Validates that a custom Arrow write schema contains no unsupported opaque byte fields.
52///
53/// `allow_binary` is reserved for the Rust macro's documented `Vec<u8>` exemption. Python
54/// schemas cannot prove that provenance and must pass `false`.
55///
56/// # Errors
57///
58/// Returns an error if an unsupported opaque byte field is found. The error names the first such
59/// field.
60#[cfg(feature = "arrow")]
61pub fn validate_custom_arrow_schema(
62    type_name: &str,
63    schema: &Schema,
64    allow_binary: bool,
65) -> anyhow::Result<()> {
66    for field in schema.fields() {
67        validate_custom_arrow_field(type_name, field, allow_binary)?;
68    }
69    Ok(())
70}
71
72#[cfg(feature = "arrow")]
73fn validate_custom_arrow_field(
74    type_name: &str,
75    field: &Field,
76    allow_binary: bool,
77) -> anyhow::Result<()> {
78    match field.data_type() {
79        ArrowDataType::Binary if allow_binary => Ok(()),
80        ArrowDataType::Binary
81        | ArrowDataType::LargeBinary
82        | ArrowDataType::BinaryView
83        | ArrowDataType::FixedSizeBinary(_) => anyhow::bail!(
84            "custom write schema `{type_name}` contains opaque byte field `{}`: {}",
85            field.name(),
86            field.data_type(),
87        ),
88        ArrowDataType::List(child)
89        | ArrowDataType::LargeList(child)
90        | ArrowDataType::ListView(child)
91        | ArrowDataType::LargeListView(child)
92        | ArrowDataType::FixedSizeList(child, _)
93        | ArrowDataType::Map(child, _) => {
94            validate_custom_arrow_field(type_name, child, allow_binary)
95        }
96        ArrowDataType::Struct(children) => {
97            for child in children {
98                validate_custom_arrow_field(type_name, child, allow_binary)?;
99            }
100            Ok(())
101        }
102        ArrowDataType::Dictionary(_, value) => {
103            let child = Field::new(field.name(), value.as_ref().clone(), field.is_nullable());
104            validate_custom_arrow_field(type_name, &child, allow_binary)
105        }
106        _ => Ok(()),
107    }
108}
109
110struct Registries {
111    json: DashMap<String, JsonDeserializer>,
112    #[cfg(feature = "arrow")]
113    arrow: DashMap<String, (Arc<Schema>, ArrowEncoder, ArrowDecoder)>,
114}
115
116fn registries() -> &'static Registries {
117    static REGISTRIES: std::sync::OnceLock<Registries> = std::sync::OnceLock::new();
118    REGISTRIES.get_or_init(|| Registries {
119        json: DashMap::new(),
120        #[cfg(feature = "arrow")]
121        arrow: DashMap::new(),
122    })
123}
124
125/// Registers a JSON deserializer for the given custom data type name.
126/// When `Data::deserialize` sees this type name, it will call this function.
127///
128/// # Errors
129/// Returns an error if the type is already registered.
130pub fn register_json_deserializer(
131    type_name: &str,
132    deserializer: JsonDeserializer,
133) -> Result<(), anyhow::Error> {
134    let reg = registries();
135    match reg.json.entry(type_name.to_string()) {
136        Entry::Occupied(_) => {
137            anyhow::bail!("Custom data type \"{type_name}\" is already registered for JSON");
138        }
139        Entry::Vacant(v) => {
140            v.insert(deserializer);
141            Ok(())
142        }
143    }
144}
145
146/// Registers a JSON deserializer for the given custom data type name if not already registered.
147/// If the type is already registered, returns `Ok(())` without overwriting (idempotent).
148/// Use this where repeated registration can occur (e.g. module init).
149///
150/// Does not return an error (idempotent insert into `DashMap`).
151#[allow(
152    clippy::missing_errors_doc,
153    reason = "idempotent registration only inserts into a DashMap"
154)]
155pub fn ensure_json_deserializer_registered(
156    type_name: &str,
157    deserializer: JsonDeserializer,
158) -> Result<(), anyhow::Error> {
159    let reg = registries();
160    reg.json
161        .entry(type_name.to_string())
162        .or_insert_with(|| deserializer);
163    Ok(())
164}
165
166/// Parses a "`data_type`" JSON object into `DataType` (`type_name`, metadata, identifier).
167fn parse_data_type_from_value(value: &serde_json::Value) -> Option<DataType> {
168    let obj = value.get("data_type")?.as_object()?;
169    let type_name = obj.get("type_name")?.as_str()?;
170    let metadata = obj.get("metadata").and_then(|m| {
171        if m.is_null() {
172            None
173        } else {
174            let p: Params = serde_json::from_value(m.clone()).ok()?;
175            if p.is_empty() { None } else { Some(p) }
176        }
177    });
178    let identifier = obj
179        .get("identifier")
180        .and_then(|v| v.as_str())
181        .map(String::from);
182    Some(DataType::new(type_name, metadata, identifier))
183}
184
185/// Parses the canonical `CustomData` JSON envelope `{ type, data_type, payload }` and returns
186/// the payload value to pass to the registered type deserializer. Does not depend on
187/// user payload field names.
188fn parse_envelope_payload(value: &serde_json::Value) -> Result<serde_json::Value, anyhow::Error> {
189    let payload = value
190        .get("payload")
191        .cloned()
192        .ok_or_else(|| anyhow::anyhow!("CustomData JSON missing 'payload' field"))?;
193    Ok(payload)
194}
195
196/// Looks up and runs the JSON deserializer for the given type name.
197/// Returns `None` if the type is not registered.
198///
199/// # Errors
200/// Returns an error if the deserializer fails.
201pub fn deserialize_custom_from_json(
202    type_name: &str,
203    value: &serde_json::Value,
204) -> Result<Option<Data>, anyhow::Error> {
205    let reg = registries();
206    let deserializer_ref = match reg.json.get(type_name) {
207        Some(d) => d,
208        None => return Ok(None),
209    };
210    let data_type = parse_data_type_from_value(value);
211    let payload = parse_envelope_payload(value)?;
212    let arc = deserializer_ref.value()(payload)?;
213    let custom = match data_type {
214        Some(dt) => CustomData::new(arc, dt),
215        None => CustomData::from_arc(arc),
216    };
217    Ok(Some(Data::Custom(custom)))
218}
219
220/// Registers Arrow schema, encoder, and decoder for the given custom data type name.
221///
222/// # Errors
223/// Returns an error if the type is already registered for Arrow.
224#[cfg(feature = "arrow")]
225pub fn register_arrow(
226    type_name: &str,
227    schema: Arc<Schema>,
228    encoder: ArrowEncoder,
229    decoder: ArrowDecoder,
230) -> Result<(), anyhow::Error> {
231    let reg = registries();
232    match reg.arrow.entry(type_name.to_string()) {
233        Entry::Occupied(_) => {
234            anyhow::bail!("Custom data type \"{type_name}\" is already registered for Arrow");
235        }
236        Entry::Vacant(v) => {
237            v.insert((schema, encoder, decoder));
238            Ok(())
239        }
240    }
241}
242
243/// Registers Arrow schema, encoder, and decoder for the given custom data type name if not already
244/// registered. If the type is already registered, returns `Ok(())` without overwriting (idempotent).
245/// Use this where repeated registration can occur (e.g. module init).
246///
247/// Does not return an error (idempotent insert into `DashMap`).
248#[allow(
249    clippy::missing_errors_doc,
250    reason = "idempotent registration only inserts into a DashMap"
251)]
252#[cfg(feature = "arrow")]
253pub fn ensure_arrow_registered(
254    type_name: &str,
255    schema: Arc<Schema>,
256    encoder: ArrowEncoder,
257    decoder: ArrowDecoder,
258) -> Result<(), anyhow::Error> {
259    let reg = registries();
260    reg.arrow
261        .entry(type_name.to_string())
262        .or_insert_with(|| (schema, encoder, decoder));
263    Ok(())
264}
265
266/// Returns the Arrow schema for the given custom type name, if registered.
267#[must_use]
268#[cfg(feature = "arrow")]
269pub fn get_arrow_schema(type_name: &str) -> Option<Arc<Schema>> {
270    let reg = registries();
271    reg.arrow
272        .get(type_name)
273        .map(|entry| Arc::clone(&entry.value().0))
274}
275
276/// Encodes a slice of custom data trait objects to a `RecordBatch` using the registered encoder.
277///
278/// # Errors
279/// Returns an error if the type is not registered or encoding fails.
280#[cfg(feature = "arrow")]
281pub fn encode_custom_to_arrow(
282    type_name: &str,
283    items: &[Arc<dyn CustomDataTrait>],
284) -> Result<Option<RecordBatch>, anyhow::Error> {
285    let reg = registries();
286    let entry = match reg.arrow.get(type_name) {
287        Some(e) => e,
288        None => return Ok(None),
289    };
290    let encoder = &entry.value().1;
291    encoder(items).map(Some)
292}
293
294/// Decodes a `RecordBatch` into `Vec<Data>` using the registered decoder.
295///
296/// # Errors
297/// Returns an error if the type is not registered or decoding fails.
298#[expect(
299    clippy::implicit_hasher,
300    reason = "callers always use the default hasher"
301)]
302#[cfg(feature = "arrow")]
303pub fn decode_custom_from_arrow(
304    type_name: &str,
305    metadata: &std::collections::HashMap<String, String>,
306    record_batch: RecordBatch,
307) -> Result<Option<Vec<Data>>, anyhow::Error> {
308    let reg = registries();
309    let entry = match reg.arrow.get(type_name) {
310        Some(e) => e,
311        None => return Ok(None),
312    };
313    let decoder = &entry.value().2;
314    decoder(metadata, record_batch).map(Some)
315}
316
317#[cfg(feature = "python")]
318pub type PyExtractor = Box<
319    dyn for<'a> Fn(&pyo3::Bound<'a, pyo3::PyAny>) -> Option<Arc<dyn CustomDataTrait>> + Send + Sync,
320>;
321
322#[cfg(feature = "python")]
323fn py_extractors() -> &'static DashMap<String, PyExtractor> {
324    static PY_EXTRACTORS: std::sync::OnceLock<DashMap<String, PyExtractor>> =
325        std::sync::OnceLock::new();
326    PY_EXTRACTORS.get_or_init(DashMap::new)
327}
328
329/// Registers a `PyExtractor` for the given custom data type name.
330/// Used by `CustomData` constructor to convert Python objects to `Arc<dyn CustomDataTrait>`.
331///
332/// # Errors
333/// Returns an error if the type is already registered.
334#[cfg(feature = "python")]
335pub fn register_py_extractor(type_name: &str, extractor: PyExtractor) -> Result<(), anyhow::Error> {
336    let reg = py_extractors();
337    match reg.entry(type_name.to_string()) {
338        Entry::Occupied(_) => {
339            anyhow::bail!(
340                "Custom data type \"{type_name}\" is already registered for Python extraction"
341            );
342        }
343        Entry::Vacant(v) => {
344            v.insert(extractor);
345            Ok(())
346        }
347    }
348}
349
350/// Registers a `PyExtractor` for the given custom data type name if not already registered.
351/// If the type is already registered, returns `Ok(())` without overwriting (idempotent).
352/// Use this where repeated registration can occur (e.g. module init).
353///
354/// Does not return an error (idempotent insert into `DashMap`).
355#[allow(
356    clippy::missing_errors_doc,
357    reason = "idempotent registration only inserts into a DashMap"
358)]
359#[cfg(feature = "python")]
360pub fn ensure_py_extractor_registered(
361    type_name: &str,
362    extractor: PyExtractor,
363) -> Result<(), anyhow::Error> {
364    let reg = py_extractors();
365    reg.entry(type_name.to_string())
366        .or_insert_with(|| extractor);
367    Ok(())
368}
369
370/// Tries to extract `Arc<dyn CustomDataTrait>` from a Python object using the registered extractor.
371/// Returns None if no extractor is registered or extraction fails.
372#[cfg(feature = "python")]
373#[must_use]
374pub fn try_extract_from_py(
375    type_name: &str,
376    obj: &pyo3::Bound<'_, pyo3::PyAny>,
377) -> Option<Arc<dyn CustomDataTrait>> {
378    let reg = py_extractors();
379    let entry = reg.get(type_name)?;
380    let extractor = entry.value();
381    extractor(obj)
382}
383
384#[cfg(feature = "python")]
385type RustExtractorFactory = Box<dyn Fn() -> PyExtractor + Send + Sync>;
386
387#[cfg(feature = "python")]
388fn rust_extractor_factories() -> &'static DashMap<String, RustExtractorFactory> {
389    static RUST_EXTRACTOR_FACTORIES: std::sync::OnceLock<DashMap<String, RustExtractorFactory>> =
390        std::sync::OnceLock::new();
391    RUST_EXTRACTOR_FACTORIES.get_or_init(DashMap::new)
392}
393
394#[cfg(feature = "python")]
395type RustClassCheck = Box<dyn Fn(&pyo3::Bound<'_, pyo3::PyAny>) -> bool + Send + Sync>;
396
397#[cfg(feature = "python")]
398fn rust_class_checks() -> &'static DashMap<String, RustClassCheck> {
399    static RUST_CLASS_CHECKS: std::sync::OnceLock<DashMap<String, RustClassCheck>> =
400        std::sync::OnceLock::new();
401    RUST_CLASS_CHECKS.get_or_init(DashMap::new)
402}
403
404/// Registers a factory that produces a `PyExtractor` for the given type name.
405/// Crates (e.g. persistence) call this at load time for each Rust custom data type.
406/// When `register_custom_data_class(cls)` is called with that type's class, the factory is invoked
407/// and the extractor is registered in the main `PyExtractor` registry.
408///
409/// # Errors
410/// Returns an error if the type name is already registered.
411#[cfg(feature = "python")]
412pub fn register_rust_extractor_factory(
413    type_name: &str,
414    factory: RustExtractorFactory,
415) -> Result<(), anyhow::Error> {
416    let reg = rust_extractor_factories();
417    match reg.entry(type_name.to_string()) {
418        Entry::Occupied(_) => {
419            anyhow::bail!("Rust extractor factory for \"{type_name}\" is already registered");
420        }
421        Entry::Vacant(v) => {
422            v.insert(factory);
423            Ok(())
424        }
425    }
426}
427
428/// Registers a factory that produces a `PyExtractor` for the given type name if not already
429/// registered. If the type is already registered, returns `Ok(())` without overwriting (idempotent).
430/// Use this where repeated registration can occur (e.g. module load).
431///
432/// Does not return an error (idempotent insert into `DashMap`).
433#[allow(
434    clippy::missing_errors_doc,
435    reason = "idempotent registration only inserts into a DashMap"
436)]
437#[cfg(feature = "python")]
438pub fn ensure_rust_extractor_factory_registered(
439    type_name: &str,
440    factory: RustExtractorFactory,
441) -> Result<(), anyhow::Error> {
442    let reg = rust_extractor_factories();
443    reg.entry(type_name.to_string()).or_insert_with(|| factory);
444    Ok(())
445}
446
447/// Registers a Rust custom data type for Python extraction. Call once per type at module load
448/// (e.g. in the persistence PyO3 module). Uses [`register_rust_extractor_factory`] with a
449/// factory that builds the extractor for `T`.
450///
451/// # Errors
452/// Returns an error if the type name is already registered.
453#[cfg(feature = "python")]
454pub fn register_rust_extractor<T>() -> Result<(), anyhow::Error>
455where
456    T: CustomDataTrait
457        + pyo3::PyTypeInfo
458        + for<'a, 'py> pyo3::FromPyObject<'a, 'py>
459        + Send
460        + Sync
461        + 'static,
462{
463    let type_name = T::type_name_static();
464    ensure_rust_class_registered::<T>(type_name);
465
466    let factory: RustExtractorFactory = Box::new(|| {
467        Box::new(|obj: &pyo3::Bound<'_, pyo3::PyAny>| {
468            obj.extract::<T>()
469                .ok()
470                .map(|x| Arc::new(x) as Arc<dyn CustomDataTrait>)
471        })
472    });
473    register_rust_extractor_factory(type_name, factory)
474}
475
476/// Registers a Rust custom data type for Python extraction if not already registered.
477/// If the type is already registered, returns `Ok(())` without overwriting (idempotent).
478/// Use this where repeated registration can occur (e.g. module load).
479///
480/// Does not return an error (idempotent insert into `DashMap`).
481#[allow(
482    clippy::missing_errors_doc,
483    reason = "idempotent registration only inserts into a DashMap"
484)]
485#[cfg(feature = "python")]
486pub fn ensure_rust_extractor_registered<T>() -> Result<(), anyhow::Error>
487where
488    T: CustomDataTrait
489        + pyo3::PyTypeInfo
490        + for<'a, 'py> pyo3::FromPyObject<'a, 'py>
491        + Send
492        + Sync
493        + 'static,
494{
495    let type_name = T::type_name_static();
496    ensure_rust_class_registered::<T>(type_name);
497
498    let factory: RustExtractorFactory = Box::new(|| {
499        Box::new(|obj: &pyo3::Bound<'_, pyo3::PyAny>| {
500            obj.extract::<T>()
501                .ok()
502                .map(|x| Arc::new(x) as Arc<dyn CustomDataTrait>)
503        })
504    });
505    ensure_rust_extractor_factory_registered(type_name, factory)
506}
507
508/// Calls the registered factory for the given type name and returns the extractor, if any.
509#[cfg(feature = "python")]
510#[must_use]
511pub fn get_rust_extractor(type_name: &str) -> Option<PyExtractor> {
512    let reg = rust_extractor_factories();
513    let factory_ref = reg.get(type_name)?;
514    Some(factory_ref.value()())
515}
516
517// Returns `true` when a Rust type registered with a different PyO3 class uses `type_name`
518#[cfg(feature = "python")]
519#[must_use]
520pub(crate) fn conflicts_with_rust_class(
521    type_name: &str,
522    data_class: &pyo3::Bound<'_, pyo3::PyAny>,
523) -> bool {
524    rust_class_checks()
525        .get(type_name)
526        .is_some_and(|check| !check.value()(data_class))
527}
528
529#[cfg(feature = "python")]
530fn ensure_rust_class_registered<T: pyo3::PyTypeInfo>(type_name: &str) {
531    rust_class_checks()
532        .entry(type_name.to_string())
533        .or_insert_with(|| Box::new(|data_class| data_class.is(T::type_object(data_class.py()))));
534}
535
536#[cfg(test)]
537mod tests {
538    use std::hash::{DefaultHasher, Hash, Hasher};
539
540    use nautilus_core::UnixNanos;
541    use rstest::rstest;
542    use serde::{Deserialize, Serialize};
543
544    use super::*;
545    use crate::data::{CustomData, custom::register_custom_data_json};
546
547    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
548    struct TestRegCustomData {
549        ts_init: UnixNanos,
550    }
551
552    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
553    #[serde(deny_unknown_fields)]
554    struct StrictRegCustomData {
555        ts_init: UnixNanos,
556    }
557
558    impl crate::data::HasTsInit for TestRegCustomData {
559        fn ts_init(&self) -> UnixNanos {
560            self.ts_init
561        }
562    }
563
564    impl crate::data::custom::CustomDataTrait for TestRegCustomData {
565        fn type_name(&self) -> &'static str {
566            "TestRegCustomData"
567        }
568        fn type_name_static() -> &'static str {
569            "TestRegCustomData"
570        }
571        fn as_any(&self) -> &dyn std::any::Any {
572            self
573        }
574        fn ts_event(&self) -> nautilus_core::UnixNanos {
575            self.ts_init
576        }
577        fn to_json(&self) -> anyhow::Result<String> {
578            Ok(serde_json::to_string(self)?)
579        }
580        fn clone_arc(&self) -> Arc<dyn crate::data::CustomDataTrait> {
581            Arc::new(self.clone())
582        }
583        fn eq_arc(&self, other: &dyn crate::data::CustomDataTrait) -> bool {
584            other.as_any().downcast_ref::<Self>() == Some(self)
585        }
586        fn from_json(
587            value: serde_json::Value,
588        ) -> anyhow::Result<Arc<dyn crate::data::CustomDataTrait>> {
589            let t: Self = serde_json::from_value(value)?;
590            Ok(Arc::new(t))
591        }
592    }
593
594    impl crate::data::HasTsInit for StrictRegCustomData {
595        fn ts_init(&self) -> UnixNanos {
596            self.ts_init
597        }
598    }
599
600    impl crate::data::custom::CustomDataTrait for StrictRegCustomData {
601        fn type_name(&self) -> &'static str {
602            "StrictRegCustomData"
603        }
604        fn type_name_static() -> &'static str {
605            "StrictRegCustomData"
606        }
607        fn as_any(&self) -> &dyn std::any::Any {
608            self
609        }
610        fn ts_event(&self) -> nautilus_core::UnixNanos {
611            self.ts_init
612        }
613        fn to_json(&self) -> anyhow::Result<String> {
614            Ok(serde_json::to_string(self)?)
615        }
616        fn clone_arc(&self) -> Arc<dyn crate::data::CustomDataTrait> {
617            Arc::new(self.clone())
618        }
619        fn eq_arc(&self, other: &dyn crate::data::CustomDataTrait) -> bool {
620            other.as_any().downcast_ref::<Self>() == Some(self)
621        }
622        fn from_json(
623            value: serde_json::Value,
624        ) -> anyhow::Result<Arc<dyn crate::data::CustomDataTrait>> {
625            let t: Self = serde_json::from_value(value)?;
626            Ok(Arc::new(t))
627        }
628    }
629
630    #[rstest]
631    fn json_registry_roundtrip() {
632        let _ = register_custom_data_json::<TestRegCustomData>();
633
634        let data = Data::Custom(CustomData::from_arc(Arc::new(TestRegCustomData {
635            ts_init: UnixNanos::from(100),
636        })));
637
638        let json = serde_json::to_string(&data).unwrap();
639        let back: Data = serde_json::from_str(&json).unwrap();
640
641        match (&data, &back) {
642            (Data::Custom(a), Data::Custom(b)) => {
643                assert_eq!(a.data.type_name(), b.data.type_name());
644                assert_eq!(a.data.ts_init(), b.data.ts_init());
645            }
646            _ => panic!("expected Custom variant"),
647        }
648    }
649
650    #[rstest]
651    fn json_registry_roundtrip_with_deny_unknown_fields() {
652        let _ = register_custom_data_json::<StrictRegCustomData>();
653
654        let data = Data::Custom(CustomData::from_arc(Arc::new(StrictRegCustomData {
655            ts_init: UnixNanos::from(200),
656        })));
657
658        let json = serde_json::to_string(&data).unwrap();
659        let back: Data = serde_json::from_str(&json).unwrap();
660
661        match (&data, &back) {
662            (Data::Custom(a), Data::Custom(b)) => {
663                assert_eq!(a.data.type_name(), b.data.type_name());
664                assert_eq!(a.data.ts_init(), b.data.ts_init());
665            }
666            _ => panic!("expected Custom variant"),
667        }
668    }
669
670    #[rstest]
671    fn data_type_registry_result_hashes_like_equal_values_from_all_routes() {
672        fn hash_data_type(data_type: &DataType) -> u64 {
673            let mut hasher = DefaultHasher::new();
674            data_type.hash(&mut hasher);
675            hasher.finish()
676        }
677
678        let metadata = serde_json::json!({"key": "value"});
679        let constructed = DataType::new(
680            "ExampleType",
681            Some(serde_json::from_value(metadata.clone()).unwrap()),
682            Some("catalog/path".to_string()),
683        );
684        let persistence_json = serde_json::json!({
685            "type_name": constructed.type_name(),
686            "metadata": metadata,
687            "identifier": constructed.identifier(),
688        });
689        let persisted = DataType::from_persistence_json(&persistence_json.to_string()).unwrap();
690        let registry_envelope = serde_json::json!({"data_type": persistence_json});
691        let registered = parse_data_type_from_value(&registry_envelope).unwrap();
692        let deserialization_payload = serde_json::json!({
693            "type_name": constructed.type_name(),
694            "metadata": constructed.metadata(),
695            "topic": constructed.topic(),
696            "hash": constructed.precomputed_hash() ^ u64::MAX,
697            "identifier": constructed.identifier(),
698        });
699        let deserialized: DataType = serde_json::from_value(deserialization_payload).unwrap();
700
701        // Eq and Hash both observe `topic` alone, so they cannot detect a route that drops or
702        // rewrites the other fields. Assert those explicitly as well, or a registry regression
703        // that lost `identifier` would still satisfy this test.
704        for data_type in [&deserialized, &persisted, &registered] {
705            assert_eq!(data_type, &constructed);
706            assert_eq!(hash_data_type(data_type), hash_data_type(&constructed));
707            assert_eq!(data_type.type_name(), constructed.type_name());
708            assert_eq!(data_type.metadata(), constructed.metadata());
709            assert_eq!(data_type.identifier(), constructed.identifier());
710            assert_eq!(data_type.topic(), constructed.topic());
711        }
712    }
713
714    #[rstest]
715    fn ensure_json_deserializer_registered_is_idempotent() {
716        let deserializer: JsonDeserializer = Box::new(|value| {
717            let t: TestRegCustomData = serde_json::from_value(value)?;
718            Ok(Arc::new(t) as Arc<dyn crate::data::CustomDataTrait>)
719        });
720        let r1 = ensure_json_deserializer_registered("IdempotentTestJson", deserializer);
721        assert!(r1.is_ok(), "first registration should succeed");
722        let deserializer2: JsonDeserializer = Box::new(|value| {
723            let t: TestRegCustomData = serde_json::from_value(value)?;
724            Ok(Arc::new(t) as Arc<dyn crate::data::CustomDataTrait>)
725        });
726        let r2 = ensure_json_deserializer_registered("IdempotentTestJson", deserializer2);
727        assert!(
728            r2.is_ok(),
729            "second registration with same type_name should succeed (idempotent)"
730        );
731    }
732
733    #[rstest]
734    fn register_json_deserializer_fails_on_duplicate() {
735        let deserializer: JsonDeserializer = Box::new(|value| {
736            let t: TestRegCustomData = serde_json::from_value(value)?;
737            Ok(Arc::new(t) as Arc<dyn crate::data::CustomDataTrait>)
738        });
739        let r1 = register_json_deserializer("StrictDuplicateTestJson", deserializer);
740        assert!(r1.is_ok());
741        let deserializer2: JsonDeserializer = Box::new(|value| {
742            let t: TestRegCustomData = serde_json::from_value(value)?;
743            Ok(Arc::new(t) as Arc<dyn crate::data::CustomDataTrait>)
744        });
745        let r2 = register_json_deserializer("StrictDuplicateTestJson", deserializer2);
746        assert!(r2.is_err());
747        let err_msg = r2.unwrap_err().to_string();
748        assert!(
749            err_msg.contains("already registered"),
750            "expected 'already registered' in error, found: {err_msg}"
751        );
752    }
753
754    #[rstest]
755    #[cfg(feature = "arrow")]
756    fn ensure_arrow_registered_is_idempotent() {
757        let schema = Arc::new(arrow::datatypes::Schema::empty());
758        let encoder: ArrowEncoder = Box::new(|_| {
759            Ok(arrow::record_batch::RecordBatch::new_empty(Arc::new(
760                arrow::datatypes::Schema::empty(),
761            )))
762        });
763        let decoder: ArrowDecoder = Box::new(|_, _| Ok(Vec::new()));
764
765        let r1 = ensure_arrow_registered("IdempotentTestArrow", schema, encoder, decoder);
766        assert!(r1.is_ok(), "first Arrow registration should succeed");
767
768        let schema2 = Arc::new(arrow::datatypes::Schema::empty());
769        let encoder2: ArrowEncoder = Box::new(|_| {
770            Ok(arrow::record_batch::RecordBatch::new_empty(Arc::new(
771                arrow::datatypes::Schema::empty(),
772            )))
773        });
774        let decoder2: ArrowDecoder = Box::new(|_, _| Ok(Vec::new()));
775
776        let r2 = ensure_arrow_registered("IdempotentTestArrow", schema2, encoder2, decoder2);
777        assert!(
778            r2.is_ok(),
779            "second Arrow registration with same type_name should be idempotent"
780        );
781    }
782
783    #[rstest]
784    #[cfg(feature = "arrow")]
785    fn python_custom_arrow_schema_rejects_opaque_bytes() {
786        let schema = Schema::new(vec![Field::new("payload", ArrowDataType::Binary, false)]);
787
788        let error = validate_custom_arrow_schema("PythonPayload", &schema, false).unwrap_err();
789
790        assert_eq!(
791            error.to_string(),
792            "custom write schema `PythonPayload` contains opaque byte field `payload`: Binary",
793        );
794    }
795}
nautilus_model/data/registry.rs

nautilus_model/data/
registry.rs
