dataplane_sdk/sdk/
internal.rs

1//  Copyright (c) 2026 Metaform Systems, Inc
2//
3//  This program and the accompanying materials are made available under the
4//  terms of the Apache License, Version 2.0 which is available at
5//  https://www.apache.org/licenses/LICENSE-2.0
6//
7//    SPDX-License-Identifier: Apache-2.0
8//
9//    Contributors:
10//         Metaform Systems, Inc. - initial API and implementation
11//
12
13use crate::{
14    core::{
15        db::{
16            control_plane::ControlPlaneRepo,
17            data_flow::DataFlowRepo,
18            tx::{Transaction, TransactionalContext},
19        },
20        error::{DbError, HandlerError},
21        handler::DataFlowHandler,
22        model::{
23            data_address::DataAddress,
24            data_flow::{DataFlow, DataFlowState, DataFlowType, TransitionError},
25            messages::{
26                DataFlowPrepareMessage, DataFlowResumeMessage, DataFlowStartMessage,
27                DataFlowStartedNotificationMessage, DataFlowStatusMessage,
28                DataFlowStatusResponseMessage,
29            },
30        },
31    },
32    error::{SdkError, SdkResult},
33};
34
35pub struct DataPlaneSdkInternal<C>
36where
37    C: TransactionalContext,
38{
39    pub(crate) ctx: C,
40    pub(crate) repo: Box<dyn DataFlowRepo<Transaction = C::Transaction>>,
41    pub(crate) control_plane_repo: Box<dyn ControlPlaneRepo<Transaction = C::Transaction>>,
42    pub(crate) handler: Box<dyn DataFlowHandler<Transaction = C::Transaction>>,
43    pub(crate) client: reqwest::Client,
44}
45
46impl<C> DataPlaneSdkInternal<C>
47where
48    C: TransactionalContext,
49{
50    pub async fn start(
51        &self,
52        participant_context_id: &str,
53        control_plane_id: &str,
54        req: DataFlowStartMessage,
55    ) -> SdkResult<DataFlowStatusMessage> {
56        let mut flow = DataFlow::builder()
57            .id(req.process_id)
58            .counter_party_id(req.counter_party_id)
59            .maybe_data_address(req.data_address)
60            .participant_context_id(participant_context_id)
61            .state(DataFlowState::Initiating)
62            .metadata(req.metadata)
63            .participant_id(req.participant_id)
64            .dataspace_context(req.dataspace_context)
65            .dataset_id(req.dataset_id)
66            .agreement_id(req.agreement_id)
67            .control_plane_id(control_plane_id)
68            .labels(req.labels)
69            .profile(req.profile)
70            .kind(DataFlowType::Provider)
71            .build();
72
73        if self.handler.can_handle(&flow).await? {
74            let mut tx = self.ctx.begin().await?;
75            let response = self.handler.on_start(&mut tx, &flow).await?;
76
77            match response.state {
78                DataFlowState::Starting => flow.transition_to_starting()?,
79                DataFlowState::Started => flow.transition_to_started()?,
80                _ => {
81                    return Err(SdkError::Handler(HandlerError::NotSupported(
82                        "Handler can only transition to Starting or Started state".to_string(),
83                    )));
84                }
85            }
86
87            self.repo.create(&mut tx, &flow).await?;
88            tx.commit().await?;
89
90            Ok(response)
91        } else {
92            Err(SdkError::Handler(HandlerError::NotSupported(
93                "Data flow handler cannot handle this flow".to_string(),
94            )))
95        }
96    }
97
98    pub async fn prepare(
99        &self,
100        participant_context_id: &str,
101        control_plane_id: &str,
102        req: DataFlowPrepareMessage,
103    ) -> SdkResult<DataFlowStatusMessage> {
104        let mut flow = DataFlow::builder()
105            .id(req.process_id)
106            .counter_party_id(req.counter_party_id)
107            .participant_context_id(participant_context_id)
108            .state(DataFlowState::Initiating)
109            .metadata(req.metadata)
110            .participant_id(req.participant_id)
111            .dataspace_context(req.dataspace_context)
112            .dataset_id(req.dataset_id)
113            .agreement_id(req.agreement_id)
114            .control_plane_id(control_plane_id)
115            .labels(req.labels)
116            .profile(req.profile)
117            .kind(DataFlowType::Consumer)
118            .build();
119
120        if self.handler.can_handle(&flow).await? {
121            let mut tx = self.ctx.begin().await?;
122            let response = self.handler.on_prepare(&mut tx, &flow).await?;
123
124            match response.state {
125                DataFlowState::Preparing => flow.transition_to_preparing()?,
126                DataFlowState::Prepared => flow.transition_to_prepared()?,
127                _ => {
128                    return Err(SdkError::Handler(HandlerError::NotSupported(format!(
129                        "Handler can only transition to Preparing or Prepared state: current state {:?}",
130                        response.state
131                    ))));
132                }
133            }
134            self.repo.create(&mut tx, &flow).await?;
135            tx.commit().await?;
136
137            Ok(response)
138        } else {
139            Err(SdkError::Handler(HandlerError::NotSupported(
140                "Data flow handler cannot handle this flow".to_string(),
141            )))
142        }
143    }
144
145    pub async fn terminate(
146        &self,
147        _ctx: &str,
148        flow_id: &str,
149        reason: Option<String>,
150    ) -> SdkResult<()> {
151        let mut tx = self.ctx.begin().await?;
152        let mut flow = self
153            .repo
154            .fetch_by_id(&mut tx, flow_id)
155            .await?
156            .ok_or_else(|| DbError::NotFound(flow_id.to_string()))?;
157
158        flow.transition_to_terminated(reason)?;
159        self.repo.update(&mut tx, &flow).await?;
160
161        self.handler.on_terminate(&mut tx, &flow).await?;
162
163        tx.commit().await?;
164        Ok(())
165    }
166
167    pub async fn started(
168        &self,
169        _ctx: &str,
170        flow_id: &str,
171        msg: DataFlowStartedNotificationMessage,
172    ) -> SdkResult<()> {
173        let mut tx = self.ctx.begin().await?;
174        let mut flow = self
175            .repo
176            .fetch_by_id(&mut tx, flow_id)
177            .await?
178            .ok_or_else(|| DbError::NotFound(flow_id.to_string()))?;
179
180        flow.data_address = msg.data_address;
181
182        self.handler.on_started(&mut tx, &flow).await?;
183
184        flow.transition_to_started()?;
185        self.repo.update(&mut tx, &flow).await?;
186
187        tx.commit().await?;
188        Ok(())
189    }
190
191    pub async fn completed(&self, _ctx: &str, flow_id: &str) -> SdkResult<()> {
192        let mut tx = self.ctx.begin().await?;
193        let mut flow = self
194            .repo
195            .fetch_by_id(&mut tx, flow_id)
196            .await?
197            .ok_or_else(|| DbError::NotFound(flow_id.to_string()))?;
198
199        flow.transition_to_completed()?;
200
201        self.handler.on_completed(&mut tx, &flow).await?;
202
203        self.repo.update(&mut tx, &flow).await?;
204
205        tx.commit().await?;
206        Ok(())
207    }
208
209    pub async fn suspend(
210        &self,
211        _ctx: &str,
212        flow_id: &str,
213        reason: Option<String>,
214    ) -> SdkResult<()> {
215        let mut tx = self.ctx.begin().await?;
216
217        let mut flow = self
218            .repo
219            .fetch_by_id(&mut tx, flow_id)
220            .await?
221            .ok_or_else(|| DbError::NotFound(flow_id.to_string()))?;
222
223        flow.transition_to_suspended(reason)?;
224        self.repo.update(&mut tx, &flow).await?;
225
226        self.handler.on_suspend(&mut tx, &flow).await?;
227
228        tx.commit().await?;
229
230        Ok(())
231    }
232
233    pub async fn resume(
234        &self,
235        _ctx: &str,
236        flow_id: &str,
237        msg: DataFlowResumeMessage,
238    ) -> SdkResult<DataFlowStatusMessage> {
239        let mut tx = self.ctx.begin().await?;
240
241        let mut flow = self
242            .repo
243            .fetch_by_id(&mut tx, flow_id)
244            .await?
245            .ok_or_else(|| DbError::NotFound(flow_id.to_string()))?;
246
247        flow.data_address = msg.data_address;
248
249        let response = self.handler.on_resume(&mut tx, &flow).await?;
250
251        match response.state {
252            DataFlowState::Starting => flow.transition_to_starting()?,
253            DataFlowState::Started => flow.transition_to_started()?,
254            _ => {
255                return Err(SdkError::Handler(HandlerError::NotSupported(
256                    "Handler can only transition to Starting or Started state".to_string(),
257                )));
258            }
259        }
260
261        self.repo.update(&mut tx, &flow).await?;
262
263        tx.commit().await?;
264
265        Ok(response)
266    }
267    pub async fn status(
268        &self,
269        _ctx: &str,
270        flow_id: &str,
271    ) -> SdkResult<DataFlowStatusResponseMessage> {
272        let mut tx = self.ctx.begin().await?;
273
274        let flow = self
275            .repo
276            .fetch_by_id(&mut tx, flow_id)
277            .await?
278            .ok_or_else(|| DbError::NotFound(flow_id.to_string()))?;
279
280        tx.commit().await?;
281
282        Ok(DataFlowStatusResponseMessage::builder()
283            .data_flow_id(flow.id)
284            .state(flow.state)
285            .build())
286    }
287
288    /// Notifies the control plane that the data flow has been prepared, by
289    /// POSTing to the URL of the flow's referenced control plane (resolved from
290    /// `control_plane_id`). See the Data Plane Signaling "Control Plane
291    /// Endpoint" section.
292    pub async fn notify_prepared(
293        &self,
294        ctx: &str,
295        flow_id: &str,
296        data_address: Option<DataAddress>,
297    ) -> SdkResult<()> {
298        self.send_callback(
299            ctx,
300            flow_id,
301            "prepared",
302            data_address.clone(),
303            None,
304            move |flow| {
305                flow.data_address = data_address;
306                flow.transition_to_prepared()
307            },
308        )
309        .await
310    }
311
312    /// Notifies the control plane that the data flow has started.
313    pub async fn notify_started(
314        &self,
315        ctx: &str,
316        flow_id: &str,
317        data_address: Option<DataAddress>,
318    ) -> SdkResult<()> {
319        self.send_callback(
320            ctx,
321            flow_id,
322            "started",
323            data_address.clone(),
324            None,
325            |flow| {
326                flow.data_address = data_address;
327                flow.transition_to_started()
328            },
329        )
330        .await
331    }
332
333    /// Notifies the control plane that the data flow has completed.
334    pub async fn notify_completed(&self, ctx: &str, flow_id: &str) -> SdkResult<()> {
335        self.send_callback(ctx, flow_id, "completed", None, None, |flow| {
336            flow.transition_to_completed()
337        })
338        .await
339    }
340
341    /// Notifies the control plane that the data flow has errored. The optional
342    /// `error` is forwarded as the `error` field of the status message.
343    pub async fn notify_errored(
344        &self,
345        ctx: &str,
346        flow_id: &str,
347        error: Option<String>,
348    ) -> SdkResult<()> {
349        self.send_callback(ctx, flow_id, "errored", None, error.clone(), move |flow| {
350            flow.transition_to_terminated(error)
351        })
352        .await
353    }
354
355    async fn send_callback<CB>(
356        &self,
357        _ctx: &str,
358        flow_id: &str,
359        operation: &str,
360        data_address: Option<DataAddress>,
361        error: Option<String>,
362        op: CB,
363    ) -> SdkResult<()>
364    where
365        CB: FnOnce(&mut DataFlow) -> Result<(), TransitionError>,
366    {
367        let mut tx = self.ctx.begin().await?;
368
369        let mut flow = self
370            .fetch_by_id(&mut tx, flow_id)
371            .await?
372            .ok_or_else(|| DbError::NotFound(flow_id.to_string()))?;
373
374        let control_plane = self
375            .control_plane_repo
376            .fetch_by_id(&mut tx, &flow.control_plane_id)
377            .await?
378            .ok_or_else(|| DbError::NotFound(flow.control_plane_id.clone()))?;
379
380        op(&mut flow)?;
381
382        let msg = DataFlowStatusMessage::builder()
383            .data_flow_id(flow.id.clone())
384            .maybe_data_address(data_address)
385            .state(flow.state.clone())
386            .maybe_error(error)
387            .build();
388
389        let url = format!(
390            "{}/transfers/{}/dataflow/{}",
391            control_plane.url.trim_end_matches('/'),
392            flow.id,
393            operation
394        );
395
396        let resp = self.client.post(&url).json(&msg).send().await?;
397        if !resp.status().is_success() {
398            let status = resp.status().as_u16();
399            let body = resp.text().await.unwrap_or_default();
400            return Err(SdkError::NotificationStatus { status, body });
401        }
402
403        self.repo.update(&mut tx, &flow).await?;
404        tx.commit().await?;
405
406        Ok(())
407    }
408
409    pub async fn fetch_by_id(
410        &self,
411        tx: &mut C::Transaction,
412        flow_id: &str,
413    ) -> SdkResult<Option<DataFlow>> {
414        self.repo.fetch_by_id(tx, flow_id).await.map(Ok)?
415    }
416
417    pub fn ctx(&self) -> &C {
418        &self.ctx
419    }
420}
dataplane_sdk/sdk/internal.rs

dataplane_sdk/sdk/
internal.rs
