laser_sdk/fork/
client.rs

1use crate::error::LaserError;
2use crate::fork::{ForkError, ForkInfo, ForkKind};
3use crate::laser::Laser;
4use crate::query::{
5    AGDX_FORK_CREATE_CODE, AGDX_FORK_DELETE_CODE, AGDX_FORK_LIST_CODE, AGDX_FORK_PROMOTE_CODE,
6    AGDX_FORK_PUT_CODE, FORK_OP_VERSION,
7};
8use laser_wire::fork::{
9    ForkCreate, ForkDelete, ForkList, ForkOutcome, ForkPromote, ForkPut, ForkReply,
10};
11use laser_wire::framing::encode_named;
12#[cfg(test)]
13use laser_wire::limits::MAX_FORK_ID_BYTES;
14use serde::Serialize;
15use std::collections::BTreeMap;
16
17impl Laser {
18    /// A handle to one fork by id. Cheap to create, since it borrows the connection.
19    /// Open the fork with [`ForkHandle::create`], then write rows, query it
20    /// (`query(...).fork(id)`), and finally [`promote`](ForkHandle::promote) or
21    /// [`squash`](ForkHandle::squash) it. The id is sent as
22    /// [`resource_name`](Self::resource_name) names it.
23    pub fn fork<'a>(&'a self, fork_id: impl Into<String>) -> ForkHandle<'a> {
24        ForkHandle {
25            laser: self,
26            fork_id: self.resource_name(&fork_id.into()),
27        }
28    }
29
30    /// Every open fork for the authenticated user. A handle that scopes its
31    /// resources to a stream lists only that stream's forks, under the ids the
32    /// caller gave them.
33    pub async fn forks(&self) -> Result<Vec<ForkInfo>, LaserError> {
34        match self
35            .execute_fork(AGDX_FORK_LIST_CODE, &ForkList { v: FORK_OP_VERSION })
36            .await?
37        {
38            ForkOutcome::List(forks) => Ok(self.local_forks(forks)),
39            other => Err(unexpected("list", &other)),
40        }
41    }
42
43    fn local_forks(&self, forks: Vec<ForkInfo>) -> Vec<ForkInfo> {
44        forks
45            .into_iter()
46            .filter_map(|mut fork| {
47                fork.fork_id = self.local_resource_name(&fork.fork_id)?.to_owned();
48                if let Some(parent) = fork.parent.take() {
49                    let local = self.local_resource_name(&parent).map(str::to_owned);
50                    fork.parent = Some(local.unwrap_or(parent));
51                }
52                Some(fork)
53            })
54            .collect()
55    }
56
57    // Send one fork command over the binary connection and decode the reply.
58    // Gated on `forks` (set by the connect-time probe). Without it, `Unsupported`.
59    pub(crate) async fn execute_fork(
60        &self,
61        code: u32,
62        request: &impl Serialize,
63    ) -> Result<ForkOutcome, LaserError> {
64        let capabilities = self.capabilities().await;
65        if !capabilities.forks {
66            return Err(LaserError::unsupported(
67                "fork",
68                "forks are not served by this deployment",
69            ));
70        }
71        // Fail fast on advertised version skew (see `execute_query`): a server
72        // that advertised its accepted fork op version spares the round-trip.
73        if let Some(versions) = capabilities.versions
74            && versions.fork != FORK_OP_VERSION
75        {
76            return Err(ForkError::Version {
77                expected: versions.fork,
78                got: FORK_OP_VERSION,
79            }
80            .into());
81        }
82        let payload = encode_named(request)
83            .map_err(|error| LaserError::Codec(format!("encode fork request: {error}")))?;
84        let payload = self.send_raw_with_response(code, payload).await?;
85        match crate::error::decode_managed_reply::<ForkReply>(&payload)? {
86            ForkReply::Ok(outcome) => Ok(outcome),
87            ForkReply::Err(error) => Err(error.into()),
88            _ => Err(LaserError::Protocol(
89                "fork: unknown reply variant".to_owned(),
90            )),
91        }
92    }
93}
94
95/// A handle to one fork. Build it with [`Laser::fork`].
96pub struct ForkHandle<'a> {
97    laser: &'a Laser,
98    fork_id: String,
99}
100
101impl<'a> ForkHandle<'a> {
102    /// The fork id this handle is bound to, as the caller named it.
103    pub fn id(&self) -> &str {
104        self.laser
105            .local_resource_name(&self.fork_id)
106            .unwrap_or(&self.fork_id)
107    }
108
109    /// The fork id this handle sends, scoped by the connection's
110    /// [`ResourceNaming`](crate::laser::ResourceNaming). Pass it to
111    /// `query(..).fork(..)` or any other request naming this fork.
112    pub fn resource_id(&self) -> &str {
113        &self.fork_id
114    }
115
116    /// Start opening this fork. Choose `.severed()` or `.continuous()` (default
117    /// continuous), optionally narrow the snapshot with `.tables([...])`, then
118    /// `.send().await` for the fork's metadata.
119    pub fn create(&self) -> ForkCreateRequest<'a> {
120        ForkCreateRequest {
121            laser: self.laser,
122            fork_id: self.fork_id.clone(),
123            parent: None,
124            kind: ForkKind::Continuous,
125            tables: Vec::new(),
126        }
127    }
128
129    /// Promote this fork: splice its speculative rows onto the trunk (and apply
130    /// its tombstones), then squash it. Returns the number of rows applied.
131    pub async fn promote(&self) -> Result<usize, LaserError> {
132        match self
133            .laser
134            .execute_fork(
135                AGDX_FORK_PROMOTE_CODE,
136                &ForkPromote {
137                    v: FORK_OP_VERSION,
138                    fork_id: self.fork_id.clone(),
139                },
140            )
141            .await?
142        {
143            ForkOutcome::Promoted { rows } => Ok(rows),
144            other => Err(unexpected("promote", &other)),
145        }
146    }
147
148    /// Squash this fork: discard its speculative rows. Returns `true` when an open
149    /// fork was removed, `false` when none existed.
150    pub async fn squash(&self) -> Result<bool, LaserError> {
151        match self
152            .laser
153            .execute_fork(
154                AGDX_FORK_DELETE_CODE,
155                &ForkDelete {
156                    v: FORK_OP_VERSION,
157                    fork_id: self.fork_id.clone(),
158                },
159            )
160            .await?
161        {
162            ForkOutcome::Deleted(existed) => Ok(existed),
163            other => Err(unexpected("squash", &other)),
164        }
165    }
166
167    /// Write one speculative row into this fork at `(table, partition_id, offset)`.
168    /// Add indexed fields with `.field`, an opaque body with `.payload`, an
169    /// embedding with `.embedding`, or mark `.tombstone()` to hide the trunk row
170    /// at that coordinate from the fork's view. Finish with `.send().await`.
171    pub fn put_row(
172        &self,
173        table: impl Into<String>,
174        partition_id: u32,
175        offset: u64,
176    ) -> ForkPutRequest<'a> {
177        ForkPutRequest {
178            laser: self.laser,
179            fork_id: self.fork_id.clone(),
180            table: self.laser.resource_name(&table.into()),
181            partition_id,
182            offset,
183            projection_id: String::new(),
184            projection_version: 0,
185            fields: BTreeMap::new(),
186            metadata: BTreeMap::new(),
187            payload: None,
188            embedding: None,
189            tombstone: false,
190        }
191    }
192}
193
194/// Fluent builder for [`ForkHandle::create`].
195#[must_use = "call .send().await to create the fork"]
196pub struct ForkCreateRequest<'a> {
197    laser: &'a Laser,
198    fork_id: String,
199    parent: Option<String>,
200    kind: ForkKind,
201    tables: Vec<String>,
202}
203
204impl<'a> ForkCreateRequest<'a> {
205    /// Make a frozen snapshot at the trunk's current offsets (later appends hidden).
206    pub fn severed(mut self) -> Self {
207        self.kind = ForkKind::Severed;
208        self
209    }
210
211    /// Make a live branch that keeps seeing new trunk appends (the default).
212    pub fn continuous(mut self) -> Self {
213        self.kind = ForkKind::Continuous;
214        self
215    }
216
217    /// Record this fork's parent (audit only, the fork still branches off the trunk).
218    pub fn parent(mut self, parent: impl Into<String>) -> Self {
219        self.parent = Some(self.laser.resource_name(&parent.into()));
220        self
221    }
222
223    /// Narrow a severed snapshot to these tables. Empty captures every table.
224    pub fn tables(mut self, tables: impl IntoIterator<Item = impl Into<String>>) -> Self {
225        self.tables = tables
226            .into_iter()
227            .map(|table| self.laser.resource_name(&table.into()))
228            .collect();
229        self
230    }
231
232    /// Open the fork. Returns its metadata.
233    pub async fn send(self) -> Result<ForkInfo, LaserError> {
234        validated_fork_id(&self.fork_id)?;
235        let request = ForkCreate {
236            v: FORK_OP_VERSION,
237            fork_id: self.fork_id,
238            parent: self.parent,
239            kind: self.kind,
240            tables: self.tables,
241        };
242        match self
243            .laser
244            .execute_fork(AGDX_FORK_CREATE_CODE, &request)
245            .await?
246        {
247            ForkOutcome::Created(info) => Ok(info),
248            other => Err(unexpected("create", &other)),
249        }
250    }
251}
252
253/// Fluent builder for [`ForkHandle::put_row`].
254#[must_use = "call .send().await to write the speculative row"]
255pub struct ForkPutRequest<'a> {
256    laser: &'a Laser,
257    fork_id: String,
258    table: String,
259    partition_id: u32,
260    offset: u64,
261    projection_id: String,
262    projection_version: u32,
263    fields: BTreeMap<String, String>,
264    metadata: BTreeMap<String, String>,
265    payload: Option<Vec<u8>>,
266    embedding: Option<Vec<f32>>,
267    tombstone: bool,
268}
269
270impl<'a> ForkPutRequest<'a> {
271    /// Set the projection id/version this speculative row belongs to.
272    pub fn projection(mut self, id: impl Into<String>, version: u32) -> Self {
273        self.projection_id = self.laser.resource_name(&id.into());
274        self.projection_version = version;
275        self
276    }
277
278    /// Add one indexed field (the columns queries filter and order on).
279    pub fn field(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
280        self.fields.insert(name.into(), value.into());
281        self
282    }
283
284    /// Add one metadata header (non-indexed).
285    pub fn metadata(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
286        self.metadata.insert(name.into(), value.into());
287        self
288    }
289
290    /// Attach an opaque payload body.
291    pub fn payload(mut self, payload: impl AsRef<[u8]>) -> Self {
292        self.payload = Some(payload.as_ref().to_vec());
293        self
294    }
295
296    /// Attach an embedding vector. It travels as a JSON array literal, and
297    /// `send` rejects a non-finite component with `LaserError::Invalid`.
298    pub fn embedding(mut self, embedding: impl IntoIterator<Item = f32>) -> Self {
299        self.embedding = Some(embedding.into_iter().collect());
300        self
301    }
302
303    /// Mark this as a tombstone: hide the trunk row at this coordinate from the
304    /// fork's view instead of writing a value.
305    pub fn tombstone(mut self) -> Self {
306        self.tombstone = true;
307        self
308    }
309
310    /// Write the speculative row.
311    pub async fn send(self) -> Result<(), LaserError> {
312        let embedding = self
313            .embedding
314            .as_deref()
315            .map(embedding_literal)
316            .transpose()?;
317        let request = ForkPut {
318            v: FORK_OP_VERSION,
319            fork_id: self.fork_id,
320            table: self.table,
321            partition_id: self.partition_id,
322            offset: self.offset,
323            projection_id: self.projection_id,
324            projection_version: self.projection_version,
325            fields: self.fields,
326            metadata: self.metadata,
327            payload: self.payload,
328            embedding,
329            tombstone: self.tombstone,
330        };
331        match self
332            .laser
333            .execute_fork(AGDX_FORK_PUT_CODE, &request)
334            .await?
335        {
336            ForkOutcome::Written => Ok(()),
337            other => Err(unexpected("put_row", &other)),
338        }
339    }
340}
341
342fn unexpected(op: &str, outcome: &ForkOutcome) -> LaserError {
343    LaserError::Protocol(format!("fork {op}: unexpected outcome {outcome:?}"))
344}
345
346// Reject a malformed fork id before the round-trip, using the wire contract's
347// canonical `validate_fork_id` (non-empty, within `MAX_FORK_ID_BYTES`, and the
348// shared charset safelist) so a client fails fast on exactly what the managed
349// side would reject, and the anti-injection charset rule lives in one place.
350fn validated_fork_id(fork_id: &str) -> Result<(), LaserError> {
351    laser_wire::fork::validate_fork_id(fork_id).map_err(|error| LaserError::Invalid(error.0))
352}
353
354fn embedding_literal(embedding: &[f32]) -> Result<String, LaserError> {
355    if embedding.iter().any(|component| !component.is_finite()) {
356        return Err(LaserError::Invalid(
357            "an embedding component must be a finite number".to_owned(),
358        ));
359    }
360    serde_json::to_string(embedding).map_err(|error| LaserError::Codec(error.to_string()))
361}
362
363#[cfg(test)]
364mod tests {
365    use super::*;
366
367    fn info(fork_id: &str, parent: Option<&str>) -> ForkInfo {
368        ForkInfo {
369            fork_id: fork_id.to_owned(),
370            parent: parent.map(str::to_owned),
371            kind: ForkKind::Continuous,
372            user_id: 1,
373            status: laser_wire::fork::ForkStatus::Open,
374            created_at_micros: 0,
375            row_count: 0,
376        }
377    }
378
379    #[test]
380    fn given_a_default_stream_when_opening_a_fork_then_should_scope_its_names() {
381        let laser = Laser::from_client(crate::iggy::prelude::IggyClient::default())
382            .with_default_stream("acme");
383        let fork = laser.fork("experiment");
384        assert_eq!(fork.resource_id(), "stream:acme/experiment");
385        assert_eq!(fork.id(), "experiment");
386        validated_fork_id(fork.resource_id()).expect("a scoped fork id is valid");
387        let create = fork.create().parent("baseline").tables(["readings"]);
388        assert_eq!(create.parent.as_deref(), Some("stream:acme/baseline"));
389        assert_eq!(create.tables, ["stream:acme/readings"]);
390        let put = fork.put_row("readings", 0, 1).projection("reading.v1", 1);
391        assert_eq!(put.table, "stream:acme/readings");
392        assert_eq!(put.projection_id, "stream:acme/reading.v1");
393    }
394
395    #[test]
396    fn given_listed_forks_when_localized_then_should_keep_only_the_own_stream() {
397        let laser = Laser::from_client(crate::iggy::prelude::IggyClient::default())
398            .with_default_stream("acme");
399        let local = laser.local_forks(vec![
400            info("stream:acme/experiment", Some("stream:acme/baseline")),
401            info("stream:other/experiment", None),
402            info("legacy", None),
403        ]);
404        assert_eq!(local.len(), 1);
405        assert_eq!(local[0].fork_id, "experiment");
406        assert_eq!(local[0].parent.as_deref(), Some("baseline"));
407    }
408
409    #[test]
410    fn given_an_over_long_or_empty_fork_id_when_validated_then_should_reject() {
411        assert!(validated_fork_id("").is_err());
412        assert!(validated_fork_id("experiment-2026-q2").is_ok());
413        let too_long = "f".repeat(MAX_FORK_ID_BYTES + 1);
414        assert!(validated_fork_id(&too_long).is_err());
415        let at_cap = "f".repeat(MAX_FORK_ID_BYTES);
416        assert!(validated_fork_id(&at_cap).is_ok());
417    }
418
419    #[test]
420    fn given_fork_error_when_mapped_then_should_nest_the_typed_error() {
421        let error: LaserError = ForkError::NotFound("f9".to_owned()).into();
422        assert!(matches!(error, LaserError::Fork(ForkError::NotFound(_))));
423        assert!(error.is_not_found());
424        let unsupported: LaserError = ForkError::Unsupported("x".to_owned()).into();
425        assert!(matches!(
426            unsupported,
427            LaserError::Fork(ForkError::Unsupported(_))
428        ));
429        assert!(unsupported.is_unsupported());
430    }
431
432    #[test]
433    fn given_an_embedding_when_encoded_then_should_be_a_json_array_and_reject_non_finite() {
434        assert_eq!(
435            embedding_literal(&[0.5, 1.0]).expect("finite components encode"),
436            "[0.5,1.0]"
437        );
438        assert!(matches!(
439            embedding_literal(&[f32::NAN]),
440            Err(LaserError::Invalid(_))
441        ));
442    }
443}
laser_sdk/fork/client.rs

laser_sdk/fork/
client.rs
