harmos_storage/manager/
arrow.rs

1/* ***
2 *
3 * Harmos
4 *
5 * SPDX-License-Identifier: Apache-2.0
6 *
7 * ***
8 */
9
10/* === Modules === */
11
12/* === Re-Exports === */
13
14/* === Imports === */
15
16// Standard Library Imports
17
18// Third Party Imports
19
20// Local Imports
21use super::*;
22
23/* === Constants === */
24
25/* === Macros === */
26
27/* === Type Aliases === */
28
29/* === Structs === */
30
31/* === Enums === */
32
33/* === Traits === */
34
35/* === Implementations === */
36
37impl<C> StorageManager<C> {
38    #[tracing::instrument(name = "storage.manager.arrow.lookup_schema", level = "debug", skip_all)]
39    pub async fn lookup_schema(&self, key: &SchemaKey) -> Result<StorageSchema> {
40        let snapshot = self.snapshot(&self.cursor_saved()).await?;
41        crate::query::schema_at(&self.inner, &snapshot, key).await
42    }
43
44    #[tracing::instrument(name = "storage.manager.arrow.admit_arrow_records", level = "debug", skip_all)]
45    pub async fn admit_arrow_records(&self, scope: Option<&Identifier>, batch: &crate::writes::DecodedBatch) -> Result<StorageCursor> {
46        let schema = self.lookup_schema(&batch.schema).await?;
47        let mut records = crate::writes::convert::records(scope, batch, &schema.schema)?;
48        for record in &mut records {
49            if record.source.as_ref().is_some_and(|source| source != &batch.destination) {
50                return Err(Error::InvalidInput {
51                    field: "record destination",
52                    reason: "record source differs from authorized destination",
53                });
54            }
55            record.source = Some(batch.destination.clone());
56        }
57        self.admit_records(records, None).await
58    }
59
60    #[tracing::instrument(name = "storage.manager.arrow.admit_arrow_streams", level = "debug", skip_all)]
61    pub async fn admit_arrow_streams(&self, scope: Option<&Identifier>, batch: &crate::writes::DecodedBatch) -> Result<StorageCursor> {
62        let schema = self.lookup_schema(&batch.schema).await?;
63        let rows = crate::writes::convert::streams(scope, batch, &schema.schema)?;
64        self.admit_stream(rows, None).await
65    }
66
67    #[tracing::instrument(name = "storage.manager.arrow.admit_arrow_objects", level = "debug", skip_all)]
68    pub async fn admit_arrow_objects(&self, _scope: Option<&Identifier>, batch: &crate::writes::DecodedBatch) -> Result<StorageCursor> {
69        if batch.destination != batch.schema.id {
70            return Err(Error::InvalidInput {
71                field: "object destination",
72                reason: "destination differs from object schema",
73            });
74        }
75        let _writer = tracing::Instrument::instrument(self.inner.object_writes.lock(), tracing::trace_span!("storage.lock.wait")).await;
76        let definition = (self.inner.object_definition)(&batch.schema)?;
77        let graph = definition.graph(&harmos_domain::config::SchemaSettings::default())?;
78        let values = crate::writes::convert::objects(batch, &graph, self.inner.settings.max_batch_bytes())?;
79        let count = std::num::NonZeroUsize::new(values.len()).ok_or(Error::InvalidInput {
80            field: "object batch",
81            reason: "empty atomic batch",
82        })?;
83        let allowance = self.inner.settings.max_batch_bytes();
84        self.check_object_preconditions(batch, &values).await?;
85        let receipts = self
86            .admit_object_group_locked(allowance, count, None, |group| {
87                for value in values {
88                    group.write_value(value)?;
89                }
90                Ok(())
91            })
92            .await?;
93        receipts.last().map(|receipt| receipt.position.clone()).ok_or(Error::InvalidInput {
94            field: "object batch",
95            reason: "accepted batch has no receipt",
96        })
97    }
98
99    #[tracing::instrument(name = "storage.manager.arrow.check_object_preconditions", level = "debug", skip_all)]
100    pub(super) async fn check_object_preconditions(
101        &self, batch: &crate::writes::DecodedBatch, values: &[crate::ObjectValue],
102    ) -> Result<()> {
103        use arrow_array::Array;
104        let Some(metadata) = batch
105            .data
106            .as_record_batch()
107            .column_by_name("_storage")
108            .and_then(|array| array.as_any().downcast_ref::<arrow_array::StructArray>())
109        else {
110            return Err(Error::InvalidInput {
111                field: "object batch",
112                reason: "missing object metadata",
113            });
114        };
115        let Some(expected) = metadata.column_by_name("expected_revision") else {
116            return Ok(());
117        };
118        let expected = expected.as_any().downcast_ref::<arrow_array::UInt64Array>().ok_or(Error::InvalidInput {
119            field: "object revision",
120            reason: "expected_revision must be nullable UInt64",
121        })?;
122        let snapshot = self.snapshot(&self.cursor_saved()).await?;
123        let page_bytes = self.inner.settings.scan_bytes();
124        let store = metadata.column_by_name("store").and_then(|array| array.as_any().downcast_ref::<arrow_array::StringArray>());
125        let epoch = metadata.column_by_name("epoch").and_then(|array| array.as_any().downcast_ref::<arrow_array::StringArray>());
126        for (row, value) in values.iter().enumerate() {
127            if expected.is_null(row) {
128                continue;
129            }
130            if expected.value(row) != 0
131                && (store.is_none_or(|store| store.is_null(row) || store.value(row) != self.inner.primary.history().store.as_str())
132                    || epoch.is_none_or(|epoch| epoch.is_null(row) || epoch.value(row) != self.inner.primary.history().epoch.as_str()))
133            {
134                return Err(Error::InvalidInput {
135                    field: "object revision",
136                    reason: "expected revision belongs to a different storage history",
137                });
138            }
139            let selection = crate::query::ObjectSelection {
140                schema: Some(value.schema.clone()),
141                key: Some(value.item.clone()),
142                versions: crate::query::ObjectVersions::Latest,
143            };
144            let page = crate::ScanPage {
145                after: 0,
146                rows: std::num::NonZeroUsize::MIN,
147                bytes: page_bytes,
148            };
149            let (latest, _) = snapshot.scan::<crate::ObjectEntry>(&selection, page).await?;
150            let revision = latest.first().map_or(0, |entry| entry.version.revision().get());
151            if revision != expected.value(row) {
152                return Err(Error::InvalidInput {
153                    field: "object revision",
154                    reason: "atomic batch revision precondition failed",
155                });
156            }
157        }
158        Ok(())
159    }
160}
161
162/* === Public Functions === */
163
164/* === Private Functions === */
165
166/* === Tests === */
harmos_storage/manager/arrow.rs

harmos_storage/manager/
arrow.rs
