Skip to main content

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) -> anyhow::Result<Option<DataType>> {
168    let Some(obj) = value.get("data_type").and_then(|v| v.as_object()) else {
169        return Ok(None);
170    };
171
172    let Some(type_name) = obj.get("type_name").and_then(|v| v.as_str()) else {
173        return Ok(None);
174    };
175
176    let metadata = obj.get("metadata").and_then(|m| {
177        if m.is_null() {
178            None
179        } else {
180            let p: Params = serde_json::from_value(m.clone()).ok()?;
181            if p.is_empty() { None } else { Some(p) }
182        }
183    });
184    let identifier = obj
185        .get("identifier")
186        .and_then(|v| v.as_str())
187        .map(String::from);
188    DataType::try_new(type_name, metadata, identifier).map(Some)
189}
190
191/// Parses the canonical `CustomData` JSON envelope `{ type, data_type, payload }` and returns
192/// the payload value to pass to the registered type deserializer. Does not depend on
193/// user payload field names.
194fn parse_envelope_payload(value: &serde_json::Value) -> Result<serde_json::Value, anyhow::Error> {
195    let payload = value
196        .get("payload")
197        .cloned()
198        .ok_or_else(|| anyhow::anyhow!("CustomData JSON missing 'payload' field"))?;
199    Ok(payload)
200}
201
202/// Looks up and runs the JSON deserializer for the given type name.
203/// Returns `None` if the type is not registered.
204///
205/// # Errors
206///
207/// Returns an error if metadata contains the reserved `identifier` key, the payload is missing,
208/// or the deserializer fails.
209pub fn deserialize_custom_from_json(
210    type_name: &str,
211    value: &serde_json::Value,
212) -> Result<Option<Data>, anyhow::Error> {
213    let reg = registries();
214    let deserializer_ref = match reg.json.get(type_name) {
215        Some(d) => d,
216        None => return Ok(None),
217    };
218
219    let data_type = parse_data_type_from_value(value)?;
220    let payload = parse_envelope_payload(value)?;
221    let arc = deserializer_ref.value()(payload)?;
222    let custom = match data_type {
223        Some(dt) => CustomData::new(arc, dt),
224        None => CustomData::from_arc(arc),
225    };
226    Ok(Some(Data::Custom(custom)))
227}
228
229/// Registers Arrow schema, encoder, and decoder for the given custom data type name.
230///
231/// # Errors
232/// Returns an error if the type is already registered for Arrow.
233#[cfg(feature = "arrow")]
234pub fn register_arrow(
235    type_name: &str,
236    schema: Arc<Schema>,
237    encoder: ArrowEncoder,
238    decoder: ArrowDecoder,
239) -> Result<(), anyhow::Error> {
240    let reg = registries();
241    match reg.arrow.entry(type_name.to_string()) {
242        Entry::Occupied(_) => {
243            anyhow::bail!("Custom data type \"{type_name}\" is already registered for Arrow");
244        }
245        Entry::Vacant(v) => {
246            v.insert((schema, encoder, decoder));
247            Ok(())
248        }
249    }
250}
251
252/// Registers Arrow schema, encoder, and decoder for the given custom data type name if not already
253/// registered. If the type is already registered, returns `Ok(())` without overwriting (idempotent).
254/// Use this where repeated registration can occur (e.g. module init).
255///
256/// Does not return an error (idempotent insert into `DashMap`).
257#[allow(
258    clippy::missing_errors_doc,
259    reason = "idempotent registration only inserts into a DashMap"
260)]
261#[cfg(feature = "arrow")]
262pub fn ensure_arrow_registered(
263    type_name: &str,
264    schema: Arc<Schema>,
265    encoder: ArrowEncoder,
266    decoder: ArrowDecoder,
267) -> Result<(), anyhow::Error> {
268    let reg = registries();
269    reg.arrow
270        .entry(type_name.to_string())
271        .or_insert_with(|| (schema, encoder, decoder));
272    Ok(())
273}
274
275/// Returns the Arrow schema for the given custom type name, if registered.
276#[must_use]
277#[cfg(feature = "arrow")]
278pub fn get_arrow_schema(type_name: &str) -> Option<Arc<Schema>> {
279    let reg = registries();
280    reg.arrow
281        .get(type_name)
282        .map(|entry| Arc::clone(&entry.value().0))
283}
284
285/// Encodes a slice of custom data trait objects to a `RecordBatch` using the registered encoder.
286///
287/// # Errors
288/// Returns an error if the type is not registered or encoding fails.
289#[cfg(feature = "arrow")]
290pub fn encode_custom_to_arrow(
291    type_name: &str,
292    items: &[Arc<dyn CustomDataTrait>],
293) -> Result<Option<RecordBatch>, anyhow::Error> {
294    let reg = registries();
295    let entry = match reg.arrow.get(type_name) {
296        Some(e) => e,
297        None => return Ok(None),
298    };
299    let encoder = &entry.value().1;
300    encoder(items).map(Some)
301}
302
303/// Decodes a `RecordBatch` into `Vec<Data>` using the registered decoder.
304///
305/// # Errors
306/// Returns an error if the type is not registered or decoding fails.
307#[expect(
308    clippy::implicit_hasher,
309    reason = "callers always use the default hasher"
310)]
311#[cfg(feature = "arrow")]
312pub fn decode_custom_from_arrow(
313    type_name: &str,
314    metadata: &std::collections::HashMap<String, String>,
315    record_batch: RecordBatch,
316) -> Result<Option<Vec<Data>>, anyhow::Error> {
317    let reg = registries();
318    let entry = match reg.arrow.get(type_name) {
319        Some(e) => e,
320        None => return Ok(None),
321    };
322    let decoder = &entry.value().2;
323    decoder(metadata, record_batch).map(Some)
324}
325
326#[cfg(feature = "python")]
327pub type PyExtractor = Box<
328    dyn for<'a> Fn(&pyo3::Bound<'a, pyo3::PyAny>) -> Option<Arc<dyn CustomDataTrait>> + Send + Sync,
329>;
330
331#[cfg(feature = "python")]
332fn py_extractors() -> &'static DashMap<String, PyExtractor> {
333    static PY_EXTRACTORS: std::sync::OnceLock<DashMap<String, PyExtractor>> =
334        std::sync::OnceLock::new();
335    PY_EXTRACTORS.get_or_init(DashMap::new)
336}
337
338/// Registers a `PyExtractor` for the given custom data type name.
339/// Used by `CustomData` constructor to convert Python objects to `Arc<dyn CustomDataTrait>`.
340///
341/// # Errors
342/// Returns an error if the type is already registered.
343#[cfg(feature = "python")]
344pub fn register_py_extractor(type_name: &str, extractor: PyExtractor) -> Result<(), anyhow::Error> {
345    let reg = py_extractors();
346    match reg.entry(type_name.to_string()) {
347        Entry::Occupied(_) => {
348            anyhow::bail!(
349                "Custom data type \"{type_name}\" is already registered for Python extraction"
350            );
351        }
352        Entry::Vacant(v) => {
353            v.insert(extractor);
354            Ok(())
355        }
356    }
357}
358
359/// Registers a `PyExtractor` for the given custom data type name if not already registered.
360/// If the type is already registered, returns `Ok(())` without overwriting (idempotent).
361/// Use this where repeated registration can occur (e.g. module init).
362///
363/// Does not return an error (idempotent insert into `DashMap`).
364#[allow(
365    clippy::missing_errors_doc,
366    reason = "idempotent registration only inserts into a DashMap"
367)]
368#[cfg(feature = "python")]
369pub fn ensure_py_extractor_registered(
370    type_name: &str,
371    extractor: PyExtractor,
372) -> Result<(), anyhow::Error> {
373    let reg = py_extractors();
374    reg.entry(type_name.to_string())
375        .or_insert_with(|| extractor);
376    Ok(())
377}
378
379/// Tries to extract `Arc<dyn CustomDataTrait>` from a Python object using the registered extractor.
380/// Returns None if no extractor is registered or extraction fails.
381#[cfg(feature = "python")]
382#[must_use]
383pub fn try_extract_from_py(
384    type_name: &str,
385    obj: &pyo3::Bound<'_, pyo3::PyAny>,
386) -> Option<Arc<dyn CustomDataTrait>> {
387    let reg = py_extractors();
388    let entry = reg.get(type_name)?;
389    let extractor = entry.value();
390    extractor(obj)
391}
392
393#[cfg(feature = "python")]
394type RustExtractorFactory = Box<dyn Fn() -> PyExtractor + Send + Sync>;
395
396#[cfg(feature = "python")]
397fn rust_extractor_factories() -> &'static DashMap<String, RustExtractorFactory> {
398    static RUST_EXTRACTOR_FACTORIES: std::sync::OnceLock<DashMap<String, RustExtractorFactory>> =
399        std::sync::OnceLock::new();
400    RUST_EXTRACTOR_FACTORIES.get_or_init(DashMap::new)
401}
402
403#[cfg(feature = "python")]
404type RustClassCheck = Box<dyn Fn(&pyo3::Bound<'_, pyo3::PyAny>) -> bool + Send + Sync>;
405
406#[cfg(feature = "python")]
407fn rust_class_checks() -> &'static DashMap<String, RustClassCheck> {
408    static RUST_CLASS_CHECKS: std::sync::OnceLock<DashMap<String, RustClassCheck>> =
409        std::sync::OnceLock::new();
410    RUST_CLASS_CHECKS.get_or_init(DashMap::new)
411}
412
413/// Registers a factory that produces a `PyExtractor` for the given type name.
414/// Crates (e.g. persistence) call this at load time for each Rust custom data type.
415/// When `register_custom_data_class(cls)` is called with that type's class, the factory is invoked
416/// and the extractor is registered in the main `PyExtractor` registry.
417///
418/// # Errors
419/// Returns an error if the type name is already registered.
420#[cfg(feature = "python")]
421pub fn register_rust_extractor_factory(
422    type_name: &str,
423    factory: RustExtractorFactory,
424) -> Result<(), anyhow::Error> {
425    let reg = rust_extractor_factories();
426    match reg.entry(type_name.to_string()) {
427        Entry::Occupied(_) => {
428            anyhow::bail!("Rust extractor factory for \"{type_name}\" is already registered");
429        }
430        Entry::Vacant(v) => {
431            v.insert(factory);
432            Ok(())
433        }
434    }
435}
436
437/// Registers a factory that produces a `PyExtractor` for the given type name if not already
438/// registered. If the type is already registered, returns `Ok(())` without overwriting (idempotent).
439/// Use this where repeated registration can occur (e.g. module load).
440///
441/// Does not return an error (idempotent insert into `DashMap`).
442#[allow(
443    clippy::missing_errors_doc,
444    reason = "idempotent registration only inserts into a DashMap"
445)]
446#[cfg(feature = "python")]
447pub fn ensure_rust_extractor_factory_registered(
448    type_name: &str,
449    factory: RustExtractorFactory,
450) -> Result<(), anyhow::Error> {
451    let reg = rust_extractor_factories();
452    reg.entry(type_name.to_string()).or_insert_with(|| factory);
453    Ok(())
454}
455
456/// Registers a Rust custom data type for Python extraction. Call once per type at module load
457/// (e.g. in the persistence PyO3 module). Uses [`register_rust_extractor_factory`] with a
458/// factory that builds the extractor for `T`.
459///
460/// # Errors
461/// Returns an error if the type name is already registered.
462#[cfg(feature = "python")]
463pub fn register_rust_extractor<T>() -> Result<(), anyhow::Error>
464where
465    T: CustomDataTrait
466        + pyo3::PyTypeInfo
467        + for<'a, 'py> pyo3::FromPyObject<'a, 'py>
468        + Send
469        + Sync
470        + 'static,
471{
472    let type_name = T::type_name_static();
473    ensure_rust_class_registered::<T>(type_name);
474
475    let factory: RustExtractorFactory = Box::new(|| {
476        Box::new(|obj: &pyo3::Bound<'_, pyo3::PyAny>| {
477            obj.extract::<T>()
478                .ok()
479                .map(|x| Arc::new(x) as Arc<dyn CustomDataTrait>)
480        })
481    });
482    register_rust_extractor_factory(type_name, factory)
483}
484
485/// Registers a Rust custom data type for Python extraction if not already registered.
486/// If the type is already registered, returns `Ok(())` without overwriting (idempotent).
487/// Use this where repeated registration can occur (e.g. module load).
488///
489/// Does not return an error (idempotent insert into `DashMap`).
490#[allow(
491    clippy::missing_errors_doc,
492    reason = "idempotent registration only inserts into a DashMap"
493)]
494#[cfg(feature = "python")]
495pub fn ensure_rust_extractor_registered<T>() -> Result<(), anyhow::Error>
496where
497    T: CustomDataTrait
498        + pyo3::PyTypeInfo
499        + for<'a, 'py> pyo3::FromPyObject<'a, 'py>
500        + Send
501        + Sync
502        + 'static,
503{
504    let type_name = T::type_name_static();
505    ensure_rust_class_registered::<T>(type_name);
506
507    let factory: RustExtractorFactory = Box::new(|| {
508        Box::new(|obj: &pyo3::Bound<'_, pyo3::PyAny>| {
509            obj.extract::<T>()
510                .ok()
511                .map(|x| Arc::new(x) as Arc<dyn CustomDataTrait>)
512        })
513    });
514    ensure_rust_extractor_factory_registered(type_name, factory)
515}
516
517/// Calls the registered factory for the given type name and returns the extractor, if any.
518#[cfg(feature = "python")]
519#[must_use]
520pub fn get_rust_extractor(type_name: &str) -> Option<PyExtractor> {
521    let reg = rust_extractor_factories();
522    let factory_ref = reg.get(type_name)?;
523    Some(factory_ref.value()())
524}
525
526// Returns `true` when a Rust type registered with a different PyO3 class uses `type_name`
527#[cfg(feature = "python")]
528#[must_use]
529pub(crate) fn conflicts_with_rust_class(
530    type_name: &str,
531    data_class: &pyo3::Bound<'_, pyo3::PyAny>,
532) -> bool {
533    rust_class_checks()
534        .get(type_name)
535        .is_some_and(|check| !check.value()(data_class))
536}
537
538#[cfg(feature = "python")]
539fn ensure_rust_class_registered<T: pyo3::PyTypeInfo>(type_name: &str) {
540    rust_class_checks()
541        .entry(type_name.to_string())
542        .or_insert_with(|| Box::new(|data_class| data_class.is(T::type_object(data_class.py()))));
543}
544
545#[cfg(test)]
546mod tests {
547    use std::{
548        hash::{DefaultHasher, Hash, Hasher},
549        sync::atomic::{AtomicBool, Ordering},
550    };
551
552    use nautilus_core::UnixNanos;
553    use rstest::rstest;
554    use serde::{Deserialize, Serialize};
555
556    use super::*;
557    use crate::data::{
558        CustomData,
559        custom::{ensure_custom_data_json_registered, register_custom_data_json},
560    };
561
562    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
563    struct TestRegCustomData {
564        ts_init: UnixNanos,
565    }
566
567    #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
568    #[serde(deny_unknown_fields)]
569    struct StrictRegCustomData {
570        ts_init: UnixNanos,
571    }
572
573    impl crate::data::HasTsInit for TestRegCustomData {
574        fn ts_init(&self) -> UnixNanos {
575            self.ts_init
576        }
577    }
578
579    impl crate::data::custom::CustomDataTrait for TestRegCustomData {
580        fn type_name(&self) -> &'static str {
581            "TestRegCustomData"
582        }
583        fn type_name_static() -> &'static str {
584            "TestRegCustomData"
585        }
586        fn as_any(&self) -> &dyn std::any::Any {
587            self
588        }
589        fn ts_event(&self) -> nautilus_core::UnixNanos {
590            self.ts_init
591        }
592        fn to_json(&self) -> anyhow::Result<String> {
593            Ok(serde_json::to_string(self)?)
594        }
595        fn clone_arc(&self) -> Arc<dyn crate::data::CustomDataTrait> {
596            Arc::new(self.clone())
597        }
598        fn eq_arc(&self, other: &dyn crate::data::CustomDataTrait) -> bool {
599            other.as_any().downcast_ref::<Self>() == Some(self)
600        }
601        fn from_json(
602            value: serde_json::Value,
603        ) -> anyhow::Result<Arc<dyn crate::data::CustomDataTrait>> {
604            let t: Self = serde_json::from_value(value)?;
605            Ok(Arc::new(t))
606        }
607    }
608
609    impl crate::data::HasTsInit for StrictRegCustomData {
610        fn ts_init(&self) -> UnixNanos {
611            self.ts_init
612        }
613    }
614
615    impl crate::data::custom::CustomDataTrait for StrictRegCustomData {
616        fn type_name(&self) -> &'static str {
617            "StrictRegCustomData"
618        }
619        fn type_name_static() -> &'static str {
620            "StrictRegCustomData"
621        }
622        fn as_any(&self) -> &dyn std::any::Any {
623            self
624        }
625        fn ts_event(&self) -> nautilus_core::UnixNanos {
626            self.ts_init
627        }
628        fn to_json(&self) -> anyhow::Result<String> {
629            Ok(serde_json::to_string(self)?)
630        }
631        fn clone_arc(&self) -> Arc<dyn crate::data::CustomDataTrait> {
632            Arc::new(self.clone())
633        }
634        fn eq_arc(&self, other: &dyn crate::data::CustomDataTrait) -> bool {
635            other.as_any().downcast_ref::<Self>() == Some(self)
636        }
637        fn from_json(
638            value: serde_json::Value,
639        ) -> anyhow::Result<Arc<dyn crate::data::CustomDataTrait>> {
640            let t: Self = serde_json::from_value(value)?;
641            Ok(Arc::new(t))
642        }
643    }
644
645    #[rstest]
646    fn json_registry_roundtrip() {
647        let _ = register_custom_data_json::<TestRegCustomData>();
648
649        let data = Data::Custom(CustomData::from_arc(Arc::new(TestRegCustomData {
650            ts_init: UnixNanos::from(100),
651        })));
652
653        let json = serde_json::to_string(&data).unwrap();
654        let back: Data = serde_json::from_str(&json).unwrap();
655
656        match (&data, &back) {
657            (Data::Custom(a), Data::Custom(b)) => {
658                assert_eq!(a.data.type_name(), b.data.type_name());
659                assert_eq!(a.data.ts_init(), b.data.ts_init());
660            }
661            _ => panic!("expected Custom variant"),
662        }
663    }
664
665    #[rstest]
666    fn json_registry_roundtrip_with_deny_unknown_fields() {
667        let _ = register_custom_data_json::<StrictRegCustomData>();
668
669        let data = Data::Custom(CustomData::from_arc(Arc::new(StrictRegCustomData {
670            ts_init: UnixNanos::from(200),
671        })));
672
673        let json = serde_json::to_string(&data).unwrap();
674        let back: Data = serde_json::from_str(&json).unwrap();
675
676        match (&data, &back) {
677            (Data::Custom(a), Data::Custom(b)) => {
678                assert_eq!(a.data.type_name(), b.data.type_name());
679                assert_eq!(a.data.ts_init(), b.data.ts_init());
680            }
681            _ => panic!("expected Custom variant"),
682        }
683    }
684
685    #[rstest]
686    #[case("ReservedMetadataString", serde_json::json!("shadow"))]
687    #[case("ReservedMetadataNull", serde_json::Value::Null)]
688    fn json_registry_rejects_reserved_metadata_before_callback(
689        #[case] type_name: &str,
690        #[case] value: serde_json::Value,
691    ) {
692        let called = Arc::new(AtomicBool::new(false));
693        let called_by_callback = called.clone();
694        register_json_deserializer(
695            type_name,
696            Box::new(move |_| {
697                called_by_callback.store(true, Ordering::SeqCst);
698                anyhow::bail!("invalid metadata reached callback")
699            }),
700        )
701        .unwrap();
702
703        let custom = CustomData::from_arc(Arc::new(TestRegCustomData {
704            ts_init: UnixNanos::from(317),
705        }));
706
707        let mut envelope = serde_json::to_value(custom).unwrap();
708        envelope["type"] = serde_json::json!(type_name);
709        envelope["data_type"]["type_name"] = serde_json::json!(type_name);
710        envelope["data_type"]["metadata"]["identifier"] = value;
711
712        let error = deserialize_custom_from_json(type_name, &envelope).unwrap_err();
713
714        assert_eq!(
715            error.to_string(),
716            "`identifier` is reserved for DataType catalog identity"
717        );
718        assert!(!called.load(Ordering::SeqCst));
719    }
720
721    #[rstest]
722    #[case::missing(None)]
723    #[case::null(Some(serde_json::Value::Null))]
724    #[case::empty(Some(serde_json::json!({})))]
725    #[case::populated(Some(serde_json::json!({"source": "feed", "level": 3})))]
726    fn json_registry_preserves_envelope_metadata(#[case] metadata: Option<serde_json::Value>) {
727        ensure_custom_data_json_registered::<StrictRegCustomData>().unwrap();
728        let expected_metadata: Option<Params> = metadata
729            .clone()
730            .filter(|v| !v.is_null() && v.as_object().is_some_and(|m| !m.is_empty()))
731            .map(|v| serde_json::from_value(v).unwrap());
732
733        let data_type = DataType::new(
734            "StrictRegCustomData",
735            expected_metadata.clone(),
736            Some("catalog/317".to_string()),
737        );
738
739        let expected = CustomData::new(
740            Arc::new(StrictRegCustomData {
741                ts_init: UnixNanos::from(317),
742            }),
743            data_type,
744        );
745
746        let mut envelope = serde_json::to_value(&expected).unwrap();
747        if let Some(metadata) = metadata {
748            envelope["data_type"]["metadata"] = metadata;
749        } else {
750            envelope["data_type"]
751                .as_object_mut()
752                .unwrap()
753                .remove("metadata");
754        }
755
756        let actual = deserialize_custom_from_json("StrictRegCustomData", &envelope)
757            .unwrap()
758            .unwrap();
759
760        let Data::Custom(actual) = actual else {
761            panic!("expected Custom variant");
762        };
763
764        assert_eq!(
765            serde_json::to_value(&actual).unwrap(),
766            serde_json::to_value(&expected).unwrap()
767        );
768        assert_eq!(actual.data_type.metadata(), expected_metadata.as_ref());
769        assert_eq!(actual.data_type.identifier(), Some("catalog/317"));
770    }
771
772    #[rstest]
773    #[case::missing(None)]
774    #[case::null(Some(serde_json::Value::Null))]
775    #[case::missing_type_name(Some(serde_json::json!({"metadata": {"source": "feed"}})))]
776    fn json_registry_preserves_missing_descriptor_fallback(
777        #[case] descriptor: Option<serde_json::Value>,
778    ) {
779        ensure_custom_data_json_registered::<StrictRegCustomData>().unwrap();
780
781        let expected = CustomData::from_arc(Arc::new(StrictRegCustomData {
782            ts_init: UnixNanos::from(317),
783        }));
784
785        let mut envelope = serde_json::to_value(&expected).unwrap();
786        if let Some(descriptor) = descriptor {
787            envelope["data_type"] = descriptor;
788        } else {
789            envelope.as_object_mut().unwrap().remove("data_type");
790        }
791
792        let actual = deserialize_custom_from_json("StrictRegCustomData", &envelope)
793            .unwrap()
794            .unwrap();
795
796        assert_eq!(
797            serde_json::to_value(actual).unwrap(),
798            serde_json::to_value(expected).unwrap()
799        );
800    }
801
802    #[rstest]
803    fn data_type_registry_result_hashes_like_equal_values_from_all_routes() {
804        fn hash_data_type(data_type: &DataType) -> u64 {
805            let mut hasher = DefaultHasher::new();
806            data_type.hash(&mut hasher);
807            hasher.finish()
808        }
809
810        let metadata = serde_json::json!({"key": "value"});
811        let constructed = DataType::new(
812            "ExampleType",
813            Some(serde_json::from_value(metadata.clone()).unwrap()),
814            Some("catalog/path".to_string()),
815        );
816        let persistence_json = serde_json::json!({
817            "type_name": constructed.type_name(),
818            "metadata": metadata,
819            "identifier": constructed.identifier(),
820        });
821        let persisted = DataType::from_persistence_json(&persistence_json.to_string()).unwrap();
822        let registry_envelope = serde_json::json!({"data_type": persistence_json});
823        let registered = parse_data_type_from_value(&registry_envelope)
824            .unwrap()
825            .unwrap();
826        let deserialization_payload = serde_json::json!({
827            "type_name": constructed.type_name(),
828            "metadata": constructed.metadata(),
829            "topic": constructed.topic(),
830            "hash": constructed.precomputed_hash() ^ u64::MAX,
831            "identifier": constructed.identifier(),
832        });
833        let deserialized: DataType = serde_json::from_value(deserialization_payload).unwrap();
834
835        // Eq and Hash both observe `topic` alone, so they cannot detect a route that drops or
836        // rewrites the other fields. Assert those explicitly as well, or a registry regression
837        // that lost `identifier` would still satisfy this test.
838        for data_type in [&deserialized, &persisted, &registered] {
839            assert_eq!(data_type, &constructed);
840            assert_eq!(hash_data_type(data_type), hash_data_type(&constructed));
841            assert_eq!(data_type.type_name(), constructed.type_name());
842            assert_eq!(data_type.metadata(), constructed.metadata());
843            assert_eq!(data_type.identifier(), constructed.identifier());
844            assert_eq!(data_type.topic(), constructed.topic());
845        }
846    }
847
848    #[rstest]
849    fn ensure_json_deserializer_registered_is_idempotent() {
850        let deserializer: JsonDeserializer = Box::new(|value| {
851            let t: TestRegCustomData = serde_json::from_value(value)?;
852            Ok(Arc::new(t) as Arc<dyn crate::data::CustomDataTrait>)
853        });
854        let r1 = ensure_json_deserializer_registered("IdempotentTestJson", deserializer);
855        assert!(r1.is_ok(), "first registration should succeed");
856        let deserializer2: JsonDeserializer = Box::new(|value| {
857            let t: TestRegCustomData = serde_json::from_value(value)?;
858            Ok(Arc::new(t) as Arc<dyn crate::data::CustomDataTrait>)
859        });
860        let r2 = ensure_json_deserializer_registered("IdempotentTestJson", deserializer2);
861        assert!(
862            r2.is_ok(),
863            "second registration with same type_name should succeed (idempotent)"
864        );
865    }
866
867    #[rstest]
868    fn register_json_deserializer_fails_on_duplicate() {
869        let deserializer: JsonDeserializer = Box::new(|value| {
870            let t: TestRegCustomData = serde_json::from_value(value)?;
871            Ok(Arc::new(t) as Arc<dyn crate::data::CustomDataTrait>)
872        });
873        let r1 = register_json_deserializer("StrictDuplicateTestJson", deserializer);
874        assert!(r1.is_ok());
875        let deserializer2: JsonDeserializer = Box::new(|value| {
876            let t: TestRegCustomData = serde_json::from_value(value)?;
877            Ok(Arc::new(t) as Arc<dyn crate::data::CustomDataTrait>)
878        });
879        let r2 = register_json_deserializer("StrictDuplicateTestJson", deserializer2);
880        assert!(r2.is_err());
881        let err_msg = r2.unwrap_err().to_string();
882        assert!(
883            err_msg.contains("already registered"),
884            "expected 'already registered' in error, found: {err_msg}"
885        );
886    }
887
888    #[rstest]
889    #[cfg(feature = "arrow")]
890    fn ensure_arrow_registered_is_idempotent() {
891        let schema = Arc::new(arrow::datatypes::Schema::empty());
892        let encoder: ArrowEncoder = Box::new(|_| {
893            Ok(arrow::record_batch::RecordBatch::new_empty(Arc::new(
894                arrow::datatypes::Schema::empty(),
895            )))
896        });
897        let decoder: ArrowDecoder = Box::new(|_, _| Ok(Vec::new()));
898
899        let r1 = ensure_arrow_registered("IdempotentTestArrow", schema, encoder, decoder);
900        assert!(r1.is_ok(), "first Arrow registration should succeed");
901
902        let schema2 = Arc::new(arrow::datatypes::Schema::empty());
903        let encoder2: ArrowEncoder = Box::new(|_| {
904            Ok(arrow::record_batch::RecordBatch::new_empty(Arc::new(
905                arrow::datatypes::Schema::empty(),
906            )))
907        });
908        let decoder2: ArrowDecoder = Box::new(|_, _| Ok(Vec::new()));
909
910        let r2 = ensure_arrow_registered("IdempotentTestArrow", schema2, encoder2, decoder2);
911        assert!(
912            r2.is_ok(),
913            "second Arrow registration with same type_name should be idempotent"
914        );
915    }
916
917    #[rstest]
918    #[cfg(feature = "arrow")]
919    fn python_custom_arrow_schema_rejects_opaque_bytes() {
920        let schema = Schema::new(vec![Field::new("payload", ArrowDataType::Binary, false)]);
921
922        let error = validate_custom_arrow_schema("PythonPayload", &schema, false).unwrap_err();
923
924        assert_eq!(
925            error.to_string(),
926            "custom write schema `PythonPayload` contains opaque byte field `payload`: Binary",
927        );
928    }
929}