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

dataplane_sdk/sdk/
internal.rs
