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_serializable_type` and `register_arrow` 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::{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
63/// Registers a JSON deserializer for the given custom data type name.
64/// When `Data::deserialize` sees this type name, it will call this function.
65///
66/// # Errors
67/// Returns an error if the type is already registered.
68pub 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
84/// Registers a JSON deserializer for the given custom data type name if not already registered.
85/// If the type is already registered, returns `Ok(())` without overwriting (idempotent).
86/// Use this where repeated registration can occur (e.g. module init).
87///
88/// # Errors
89/// Does not return an error (idempotent insert into `DashMap`).
90pub 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
101/// Parses a "`data_type`" JSON object into `DataType` (`type_name`, metadata, identifier).
102fn 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
120/// Parses the canonical `CustomData` JSON envelope `{ type, data_type, payload }` and returns
121/// the payload value to pass to the registered type deserializer. Does not depend on
122/// user payload field names.
123fn 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
131/// Looks up and runs the JSON deserializer for the given type name.
132/// Returns `None` if the type is not registered.
133///
134/// # Errors
135/// Returns an error if the deserializer fails.
136pub 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/// Registers Arrow schema, encoder, and decoder for the given custom data type name.
156///
157/// # Errors
158/// Returns an error if the type is already registered for Arrow.
159#[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/// Registers Arrow schema, encoder, and decoder for the given custom data type name if not already
179/// registered. If the type is already registered, returns `Ok(())` without overwriting (idempotent).
180/// Use this where repeated registration can occur (e.g. module init).
181///
182/// # Errors
183/// Does not return an error (idempotent insert into `DashMap`).
184#[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/// Returns the Arrow schema for the given custom type name, if registered.
199#[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/// Encodes a slice of custom data trait objects to a `RecordBatch` using the registered encoder.
209///
210/// # Errors
211/// Returns an error if the type is not registered or encoding fails.
212#[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/// Decodes a `RecordBatch` into `Vec<Data>` using the registered decoder.
227///
228/// # Errors
229/// Returns an error if the type is not registered or decoding fails.
230#[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/// Registers a `PyExtractor` for the given custom data type name.
262/// Used by `CustomData` constructor to convert Python objects to `Arc<dyn CustomDataTrait>`.
263///
264/// # Errors
265/// Returns an error if the type is already registered.
266#[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/// Registers a `PyExtractor` for the given custom data type name if not already registered.
283/// If the type is already registered, returns `Ok(())` without overwriting (idempotent).
284/// Use this where repeated registration can occur (e.g. module init).
285///
286/// # Errors
287/// Does not return an error (idempotent insert into `DashMap`).
288#[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/// Tries to extract `Arc<dyn CustomDataTrait>` from a Python object using the registered extractor.
300/// Returns None if no extractor is registered or extraction fails.
301#[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/// Registers a factory that produces a `PyExtractor` for the given type name.
324/// Crates (e.g. persistence) call this at load time for each Rust custom data type.
325/// When `register_custom_data_class(cls)` is called with that type's class, the factory is invoked
326/// and the extractor is registered in the main `PyExtractor` registry.
327///
328/// # Errors
329/// Returns an error if the type name is already registered.
330#[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/// Registers a factory that produces a `PyExtractor` for the given type name if not already
348/// registered. If the type is already registered, returns `Ok(())` without overwriting (idempotent).
349/// Use this where repeated registration can occur (e.g. module load).
350///
351/// # Errors
352/// Does not return an error (idempotent insert into `DashMap`).
353#[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/// Registers a Rust custom data type for Python extraction. Call once per type at module load
364/// (e.g. in the persistence PyO3 module). Uses [`register_rust_extractor_factory`] with a
365/// factory that builds the extractor for `T`.
366///
367/// # Errors
368/// Returns an error if the type name is already registered.
369#[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/// Registers a Rust custom data type for Python extraction if not already registered.
386/// If the type is already registered, returns `Ok(())` without overwriting (idempotent).
387/// Use this where repeated registration can occur (e.g. module load).
388///
389/// # Errors
390/// Does not return an error (idempotent insert into `DashMap`).
391#[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/// Calls the registered factory for the given type name and returns the extractor, if any.
408#[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(&registry_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        // Eq and Hash both observe `topic` alone, so they cannot detect a route that drops or
582        // rewrites the other fields. Assert those explicitly as well, or a registry regression
583        // that lost `identifier` would still satisfy this test.
584        for data_type in [&deserialized, &persisted, &registered] {
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}