1use std::any::TypeId;
10use std::collections::BTreeMap;
11use std::sync::Arc;
12use std::time::{Duration, SystemTime};
13
14use dex_protocol::dex::flow_service_client::FlowServiceClient;
15use dex_protocol::dex::{
16 ActiveStepSearchMode as ProtoSearchMode, AttributeStoreNames, AttributeWrite,
17 FlowAlreadyStartedOptions, FlowConfig as ProtoFlowConfig, FlowResetStepMethod, FlowResetType,
18 FlowRetryPolicy, FlowStartOptions, FlowStatus as ProtoFlowStatus,
19 FlowTimeoutPolicy as ProtoFlowTimeoutPolicy, GetAttributesRequest, GetFlowSummaryRequest,
20 IdReusePolicy as ProtoIdReusePolicy, InvokeRpcRequest, PublishToChannelRequest,
21 ReadStreamRequest, ResetFlowRequest, SearchFlowsRequest, SetAttributesRequest,
22 SkipTimerRequest, StartFlowRequest, StepDurability as ProtoStepDurability, StopFlowRequest,
23 StopType as ProtoStopType, TriggerContinueAsNewRequest, UpdateFlowConfigRequest,
24 WaitForAttributeCondition, WaitForAttributeEqual, WaitForAttributeRequest, WaitForFlowRequest,
25 WaitForStepCompletionRequest, WorkerTarget as ProtoWorkerTarget, WriteStreamRequest,
26 wait_for_attribute_condition,
27};
28use tokio::runtime::Runtime;
29use tonic::transport::Endpoint;
30use uuid::Uuid;
31
32use crate::sdk_error::{FlowTargetRequirement, ServiceError};
33use crate::stop_flow_options::StopType;
34use crate::time_travel_options::{TimeTravelPoint, TimeTravelStepMethod};
35use crate::value_hydrator::ValueHydrator;
36use crate::value_mapper;
37use crate::worker_dispatcher::map_step_options;
38use crate::{
39 ActiveStepSearchMode, Attribute, AttributeMap, BlobCache, Channel, ChannelMap, ClientOptions,
40 Flow, FlowConfig, FlowErrorType, FlowInfo, FlowResult, FlowStatus, FlowTimeoutPolicy,
41 IdReusePolicy, Registry, RetryPolicy, Rpc, SdkError, SdkResult, SearchFlowEntry,
42 SearchFlowsPage, StartFlowOptions, StepCompletion, StepDurability, StepExecutionId,
43 StopFlowOptions, Stream, StreamMessage, TimeTravelOptions, TimerId, Value, WorkerTarget,
44};
45
46pub struct Client {
66 runtime: Runtime,
67 registry: Registry,
68 service: FlowServiceClient<tonic::transport::Channel>,
69 hydrator: ValueHydrator,
70 options: ClientOptions,
71}
72
73impl Client {
74 pub fn new(registry: Registry, blob_cache: Arc<BlobCache>, options: ClientOptions) -> Self {
82 Self::try_new(registry, blob_cache, options)
83 .unwrap_or_else(|error| panic!("cannot create Dex Client: {error}"))
84 }
85
86 pub fn try_new(
93 registry: Registry,
94 blob_cache: Arc<BlobCache>,
95 options: ClientOptions,
96 ) -> SdkResult<Self> {
97 let runtime = Runtime::new().map_err(service_error)?;
98 let endpoint = Endpoint::from_shared(endpoint_address(options.server_address_value()))
99 .map_err(service_error)?;
100 let service = {
101 let _runtime_guard = runtime.enter();
102 FlowServiceClient::new(endpoint.connect_lazy())
103 };
104 let hydrator = ValueHydrator::new(service.clone(), blob_cache);
105 Ok(Self {
106 runtime,
107 registry,
108 service,
109 hydrator,
110 options,
111 })
112 }
113
114 pub fn start_flow<SomeFlow: Flow>(
123 &self,
124 flow: &SomeFlow,
125 flow_id: &str,
126 input: SomeFlow::StartInput,
127 ) -> SdkResult<String> {
128 self.start_flow_with_options(flow, flow_id, input, StartFlowOptions::new())
129 }
130
131 pub fn start_flow_with_options<SomeFlow: Flow>(
140 &self,
141 flow: &SomeFlow,
142 flow_id: &str,
143 input: SomeFlow::StartInput,
144 options: StartFlowOptions,
145 ) -> SdkResult<String> {
146 require_name(flow_id, "Flow ID")?;
147 let registered = self.registry.flow(flow.flow_type())?;
148 if registered.type_id != TypeId::of::<SomeFlow>() {
149 return Err(SdkError::FlowDefinition {
150 message: format!(
151 "Flow {} does not match the registered Rust type",
152 flow.flow_type()
153 ),
154 });
155 }
156 let (start_step_type, step_input, step_options) = match ®istered.start_step {
157 Some(step) => (
158 step.name.to_string(),
159 Some(value_mapper::encode(&input)?),
160 Some(
161 map_step_options(
162 registered,
163 registered
164 .handler
165 .step_options(step.name)
166 .map_err(sdk_handler_error)?,
167 )
168 .map_err(sdk_handler_error)?,
169 ),
170 ),
171 None => (String::new(), None, None),
172 };
173 let flow_timeout_seconds = optional_seconds(options.timeout)?;
174 let flow_timeout_policy =
175 map_flow_timeout_policy(registered, flow_timeout_seconds, options.timeout_policy)?;
176 let request = StartFlowRequest {
177 flow_id: flow_id.to_string(),
178 flow_type: registered.name.to_string(),
179 flow_timeout_seconds,
180 flow_timeout_policy,
181 start_step_type,
182 step_input,
183 step_options,
184 flow_start_options: Some(self.map_start_options(registered, &options)?),
185 request_id: options
186 .request_id
187 .clone()
188 .unwrap_or_else(|| Uuid::new_v4().to_string()),
189 };
190 let mut service = self.service.clone();
191 self.runtime.block_on(async {
192 service
193 .start_flow(request)
194 .await
195 .map(|response| response.into_inner().run_id)
196 .map_err(|status| {
197 SdkError::from_status(
198 status,
199 "start_flow",
200 Some(flow_id),
201 FlowTargetRequirement::None,
202 )
203 })
204 })
205 }
206
207 pub fn invoke_rpc<Input: Value, Output: Value>(
214 &self,
215 flow_id: &str,
216 rpc: Rpc<Input, Output>,
217 input: Input,
218 ) -> SdkResult<Output> {
219 self.do_invoke_rpc(flow_id, rpc.name(), &input)
220 }
221
222 pub fn invoke_rpc_without_input<Output: Value>(
226 &self,
227 flow_id: &str,
228 rpc: Rpc<(), Output>,
229 ) -> SdkResult<Output> {
230 self.do_invoke_rpc(flow_id, rpc.name(), &())
231 }
232
233 pub fn get_attribute<T: Value>(
237 &self,
238 flow_id: &str,
239 attribute: &Attribute<T>,
240 ) -> SdkResult<Option<T>> {
241 self.get_attribute_value(flow_id, attribute.name())
242 }
243
244 pub fn get_attribute_map_instance<T: Value>(
248 &self,
249 flow_id: &str,
250 attribute: &AttributeMap<T>,
251 instance: &str,
252 ) -> SdkResult<Option<T>> {
253 self.get_attribute_value(
254 flow_id,
255 &crate::registry::physical_name(attribute.name(), instance),
256 )
257 }
258
259 pub fn set_attribute<T: Value>(
265 &self,
266 flow_id: &str,
267 attribute: &Attribute<T>,
268 value: T,
269 ) -> SdkResult<()> {
270 self.set_attribute_value(
271 flow_id,
272 attribute.name(),
273 &value,
274 attribute.index().map(|index| index.proto_config(false)),
275 attribute.sync_config(),
276 )
277 }
278
279 pub fn set_attribute_map_instance<T: Value>(
285 &self,
286 flow_id: &str,
287 attribute: &AttributeMap<T>,
288 instance: &str,
289 value: T,
290 ) -> SdkResult<()> {
291 self.set_attribute_value(
292 flow_id,
293 &crate::registry::physical_name(attribute.name(), instance),
294 &value,
295 attribute.index().map(|index| index.proto_config(true)),
296 attribute.sync_config(),
297 )
298 }
299
300 pub fn publish<T: Value>(
302 &self,
303 flow_id: &str,
304 channel: &Channel<T>,
305 value: T,
306 ) -> SdkResult<()> {
307 self.publish_values(flow_id, channel.name(), [value])
308 }
309
310 pub fn publish_many<T: Value>(
312 &self,
313 flow_id: &str,
314 channel: &Channel<T>,
315 values: impl IntoIterator<Item = T>,
316 ) -> SdkResult<()> {
317 self.publish_values(flow_id, channel.name(), values)
318 }
319
320 pub fn publish_map<T: Value>(
322 &self,
323 flow_id: &str,
324 channel: &ChannelMap<T>,
325 instance: &str,
326 values: impl IntoIterator<Item = T>,
327 ) -> SdkResult<()> {
328 self.publish_values(
329 flow_id,
330 &crate::registry::physical_name(channel.name(), instance),
331 values,
332 )
333 }
334
335 pub fn write_stream<T: Value>(
345 &self,
346 flow_id: &str,
347 stream: &Stream<T>,
348 source: &str,
349 value: T,
350 ) -> SdkResult<()> {
351 require_name(flow_id, "Flow ID")?;
352 require_name(source, "Stream source")?;
353 let flow_type = self.registry.flow_for_stream(stream)?.name.to_string();
354 let request = WriteStreamRequest {
355 flow_id: flow_id.to_string(),
356 flow_type,
357 stream_name: stream.name().to_string(),
358 stream_capacity_bytes: stream.stream_capacity_bytes(),
359 value: Some(value_mapper::encode(&value)?),
360 source: source.to_string(),
361 };
362 self.call_empty(
363 "write_stream",
364 Some(flow_id),
365 FlowTargetRequirement::None,
366 |mut service| async move { service.write_stream(request).await },
367 )
368 }
369
370 pub fn read_stream<T: Value>(
374 &self,
375 flow_id: &str,
376 stream: &Stream<T>,
377 resume_token: &str,
378 ) -> SdkResult<StreamMessage<T>> {
379 self.read_stream_result(flow_id, stream, resume_token, None)
380 }
381
382 pub fn read_stream_with_timeout<T: Value>(
389 &self,
390 flow_id: &str,
391 stream: &Stream<T>,
392 resume_token: &str,
393 timeout: Duration,
394 ) -> SdkResult<StreamMessage<T>> {
395 self.read_stream_result(flow_id, stream, resume_token, Some(timeout))
396 }
397
398 pub fn wait_for_flow(&self, flow_id: &str) -> SdkResult<FlowResult> {
405 self.wait_for_flow_result(flow_id, None)
406 }
407
408 pub fn wait_for_flow_with_timeout(
415 &self,
416 flow_id: &str,
417 timeout: Duration,
418 ) -> SdkResult<FlowResult> {
419 self.wait_for_flow_result(flow_id, Some(timeout))
420 }
421
422 pub fn describe_flow(&self, flow_id: &str) -> SdkResult<FlowInfo> {
428 let mut service = self.service.clone();
429 let response = self.runtime.block_on(async {
430 service
431 .get_flow_summary(GetFlowSummaryRequest {
432 flow_id: flow_id.to_string(),
433 run_id: String::new(),
434 })
435 .await
436 .map(|response| response.into_inner())
437 .map_err(|status| {
438 SdkError::from_status(
439 status,
440 "describe_flow",
441 Some(flow_id),
442 FlowTargetRequirement::Existing,
443 )
444 })
445 })?;
446 let execution = response
447 .flow_execution_id
448 .ok_or_else(|| service_error("GetFlowSummary omitted FlowExecutionID"))?;
449 Ok(FlowInfo {
450 flow_id: execution.flow_id,
451 run_id: execution.run_id,
452 flow_type: response.flow_type,
453 status: map_flow_status(response.flow_status)?,
454 started_at: response
455 .start_time
456 .map_or(SystemTime::UNIX_EPOCH, timestamp),
457 })
458 }
459
460 pub fn search_flows(&self, query: &str, page_size: i32) -> SdkResult<SearchFlowsPage> {
465 self.search_flows_page(query, page_size, "")
466 }
467
468 pub fn search_flows_page(
473 &self,
474 query: &str,
475 page_size: i32,
476 next_page_token: &str,
477 ) -> SdkResult<SearchFlowsPage> {
478 if page_size < 0 {
479 return Err(invalid("search page size must not be negative"));
480 }
481 let mut service = self.service.clone();
482 self.runtime.block_on(async {
483 let response = service
484 .search_flows(SearchFlowsRequest {
485 query: query.to_string(),
486 page_size,
487 next_page_token: next_page_token.to_string(),
488 })
489 .await
490 .map(|response| response.into_inner())
491 .map_err(|status| {
492 SdkError::from_status(status, "search_flows", None, FlowTargetRequirement::None)
493 })?;
494 let values = response
495 .flow_runs
496 .iter()
497 .flat_map(|entry| entry.indexed_attributes.iter())
498 .map(|attribute| {
499 attribute
500 .value
501 .clone()
502 .ok_or_else(|| invalid("Indexed Attribute has no Value"))
503 })
504 .collect::<SdkResult<Vec<_>>>()?;
505 let mut values = self.hydrator.hydrate_all(values).await?.into_iter();
506 let mut flows = Vec::with_capacity(response.flow_runs.len());
507 for entry in response.flow_runs {
508 let mut indexed_attributes = BTreeMap::new();
509 for attribute in entry.indexed_attributes {
510 let value = values
511 .next()
512 .ok_or_else(|| invalid("Indexed Attribute hydration count mismatch"))?;
513 indexed_attributes.insert(attribute.key, value_mapper::decode_untyped(&value)?);
514 }
515 flows.push(SearchFlowEntry {
516 flow_id: entry.flow_id,
517 run_id: entry.run_id,
518 flow_type: entry.flow_type,
519 status: map_flow_status(entry.flow_status)?,
520 started_at: entry.start_time.map(timestamp),
521 closed_at: entry.close_time.map(timestamp),
522 indexed_attributes,
523 });
524 }
525 Ok(SearchFlowsPage {
526 flows,
527 next_page_token: response.next_page_token,
528 })
529 })
530 }
531
532 pub fn stop_flow(&self, flow_id: &str, options: StopFlowOptions) -> SdkResult<()> {
538 let stop_type = match options.stop_type {
539 StopType::Cancel => ProtoStopType::Cancel,
540 StopType::Terminate => ProtoStopType::Terminate,
541 StopType::Fail => ProtoStopType::Fail,
542 };
543 self.call_empty(
544 "stop_flow",
545 Some(flow_id),
546 FlowTargetRequirement::Active,
547 |mut service| async move {
548 service
549 .stop_flow(StopFlowRequest {
550 flow_id: flow_id.to_string(),
551 run_id: String::new(),
552 reason: options.reason.unwrap_or_default(),
553 stop_type: stop_type as i32,
554 })
555 .await
556 },
557 )
558 }
559
560 pub fn time_travel(&self, flow_id: &str, options: TimeTravelOptions) -> SdkResult<String> {
569 let mut request = ResetFlowRequest {
570 flow_id: flow_id.to_string(),
571 run_id: String::new(),
572 reset_type: FlowResetType::Unspecified as i32,
573 reason: options.reason.unwrap_or_default(),
574 history_event_time: String::new(),
575 step_type: String::new(),
576 step_execution_id: String::new(),
577 skip_writes_reapply: options.skip_writes_reapply,
578 step_method: FlowResetStepMethod::Unspecified as i32,
579 };
580 match options.point {
581 TimeTravelPoint::Beginning => request.reset_type = FlowResetType::Beginning as i32,
582 TimeTravelPoint::HistoryEventTime(time) => {
583 request.reset_type = FlowResetType::HistoryEventTime as i32;
584 request.history_event_time = rfc3339(time)?;
585 }
586 TimeTravelPoint::StepType(step_type) => {
587 request.reset_type = FlowResetType::StepType as i32;
588 request.step_type = step_type.to_string();
589 }
590 TimeTravelPoint::StepExecution(execution, method) => {
591 request.reset_type = FlowResetType::StepExecutionId as i32;
592 request.step_execution_id =
593 format!("{}-{}", execution.step_type, execution.execution_number);
594 request.step_method = match method {
595 TimeTravelStepMethod::WaitFor => FlowResetStepMethod::WaitFor as i32,
596 TimeTravelStepMethod::Execute => FlowResetStepMethod::Execute as i32,
597 };
598 }
599 }
600 let mut service = self.service.clone();
601 self.runtime.block_on(async {
602 service
603 .reset_flow(request)
604 .await
605 .map(|response| response.into_inner().run_id)
606 .map_err(|status| {
607 SdkError::from_status(
608 status,
609 "time_travel",
610 Some(flow_id),
611 FlowTargetRequirement::Existing,
612 )
613 })
614 })
615 }
616
617 pub fn skip_timer(
624 &self,
625 flow_id: &str,
626 step_execution: StepExecutionId,
627 timer: TimerId,
628 ) -> SdkResult<()> {
629 let (timer_condition_id, timer_condition_index) = match timer {
630 TimerId::ConditionId(id) => (id, None),
631 TimerId::ConditionIndex(index) => (
632 String::new(),
633 Some(i32::try_from(index).map_err(|_| invalid("timer index exceeds int32"))?),
634 ),
635 };
636 self.call_empty(
637 "skip_timer",
638 Some(flow_id),
639 FlowTargetRequirement::Active,
640 |mut service| async move {
641 service
642 .skip_timer(SkipTimerRequest {
643 flow_id: flow_id.to_string(),
644 run_id: String::new(),
645 step_execution_id: format!(
646 "{}-{}",
647 step_execution.step_type, step_execution.execution_number
648 ),
649 timer_condition_id,
650 timer_condition_index,
651 })
652 .await
653 },
654 )
655 }
656
657 pub fn wait_for_step_completion(
664 &self,
665 flow_id: &str,
666 step_execution: StepExecutionId,
667 timeout: Duration,
668 ) -> SdkResult<()> {
669 let wait_time_seconds = seconds32(timeout)?;
670 self.call_empty(
671 "wait_for_step_completion",
672 Some(flow_id),
673 FlowTargetRequirement::Active,
674 |mut service| async move {
675 service
676 .wait_for_step_completion(WaitForStepCompletionRequest {
677 flow_id: flow_id.to_string(),
678 step_type: step_execution.step_type.to_string(),
679 step_execution_number: step_execution.execution_number.to_string(),
680 wait_time_seconds,
681 request_id: Uuid::new_v4().to_string(),
682 })
683 .await
684 },
685 )
686 }
687
688 pub fn wait_for_attribute_equal<T: Value>(
694 &self,
695 flow_id: &str,
696 attribute: &Attribute<T>,
697 expected: T,
698 timeout: Duration,
699 ) -> SdkResult<()> {
700 self.wait_for_attribute_value(flow_id, attribute.name(), &expected, timeout)
701 }
702
703 pub fn wait_for_attribute_map_instance_equal<T: Value>(
708 &self,
709 flow_id: &str,
710 attribute: &AttributeMap<T>,
711 instance: &str,
712 expected: T,
713 timeout: Duration,
714 ) -> SdkResult<()> {
715 self.wait_for_attribute_value(
716 flow_id,
717 &crate::registry::physical_name(attribute.name(), instance),
718 &expected,
719 timeout,
720 )
721 }
722
723 fn wait_for_attribute_value<T: Value>(
724 &self,
725 flow_id: &str,
726 key: &str,
727 expected: &T,
728 timeout: Duration,
729 ) -> SdkResult<()> {
730 let value = value_mapper::encode(expected)?;
731 if !matches!(
732 value.kind,
733 Some(dex_protocol::dex::value::Kind::StringValue(_))
734 | Some(dex_protocol::dex::value::Kind::BoolValue(_))
735 | Some(dex_protocol::dex::value::Kind::IntValue(_))
736 | Some(dex_protocol::dex::value::Kind::DoubleValue(_))
737 ) {
738 return Err(invalid(
739 "wait_for_attribute_equal supports only string, boolean, or number values",
740 ));
741 }
742 let wait_time_seconds = seconds32(timeout)?;
743 self.call_empty(
744 "wait_for_attribute_equal",
745 Some(flow_id),
746 FlowTargetRequirement::Active,
747 |mut service| async move {
748 service
749 .wait_for_attribute(WaitForAttributeRequest {
750 flow_id: flow_id.to_string(),
751 run_id: String::new(),
752 condition: Some(WaitForAttributeCondition {
753 kind: Some(wait_for_attribute_condition::Kind::Equal(
754 WaitForAttributeEqual {
755 key: key.to_string(),
756 value: Some(value),
757 },
758 )),
759 }),
760 wait_time_seconds,
761 request_id: Uuid::new_v4().to_string(),
762 })
763 .await
764 },
765 )
766 }
767
768 pub fn update_flow_config(&self, flow_id: &str, config: FlowConfig) -> SdkResult<()> {
772 let flow_config = self.map_flow_config(Some(&config))?;
773 self.call_empty(
774 "update_flow_config",
775 Some(flow_id),
776 FlowTargetRequirement::Active,
777 |mut service| async move {
778 service
779 .update_flow_config(UpdateFlowConfigRequest {
780 flow_id: flow_id.to_string(),
781 run_id: String::new(),
782 flow_config: Some(flow_config),
783 })
784 .await
785 },
786 )
787 }
788
789 pub fn trigger_continue_as_new(&self, flow_id: &str) -> SdkResult<()> {
791 self.call_empty(
792 "trigger_continue_as_new",
793 Some(flow_id),
794 FlowTargetRequirement::Active,
795 |mut service| async move {
796 service
797 .trigger_continue_as_new(TriggerContinueAsNewRequest {
798 flow_id: flow_id.to_string(),
799 run_id: String::new(),
800 })
801 .await
802 },
803 )
804 }
805
806 pub fn health_check(&self) -> SdkResult<()> {
810 self.call_empty(
811 "health_check",
812 None,
813 FlowTargetRequirement::None,
814 |mut service| async move { service.health_check(()).await },
815 )
816 }
817
818 fn do_invoke_rpc<Input: Value, Output: Value>(
819 &self,
820 flow_id: &str,
821 rpc_name: &str,
822 input: &Input,
823 ) -> SdkResult<Output> {
824 let rpc = self.registry.rpc(rpc_name)?;
825 let request = InvokeRpcRequest {
826 flow_id: flow_id.to_string(),
827 run_id: String::new(),
828 rpc_name: rpc_name.to_string(),
829 input: Some(value_mapper::encode(input)?),
830 timeout_seconds: optional_seconds(rpc.timeout)?,
831 lock_attribute_keys: rpc.locks.iter().map(|lock| lock.physical_name()).collect(),
832 request_id: Uuid::new_v4().to_string(),
833 };
834 let mut service = self.service.clone();
835 let output = self.runtime.block_on(async {
836 service
837 .invoke_rpc(request)
838 .await
839 .map(|response| response.into_inner().output)
840 .map_err(|status| {
841 SdkError::from_status(
842 status,
843 "invoke_rpc",
844 Some(flow_id),
845 FlowTargetRequirement::Active,
846 )
847 })
848 })?;
849 let output = output.ok_or_else(|| invalid("InvokeRPC omitted output"))?;
850 let output = self.runtime.block_on(self.hydrator.hydrate(output))?;
851 value_mapper::decode(&output)
852 }
853
854 fn get_attribute_value<T: Value>(&self, flow_id: &str, key: &str) -> SdkResult<Option<T>> {
855 let mut service = self.service.clone();
856 let response = self.runtime.block_on(async {
857 service
858 .get_attributes(GetAttributesRequest {
859 flow_id: flow_id.to_string(),
860 run_id: String::new(),
861 keys: vec![key.to_string()],
862 all_keys: false,
863 })
864 .await
865 .map(|response| response.into_inner())
866 .map_err(|status| {
867 SdkError::from_status(
868 status,
869 "get_attribute",
870 Some(flow_id),
871 FlowTargetRequirement::Existing,
872 )
873 })
874 })?;
875 let Some(entry) = response.attributes.into_iter().next() else {
876 return Ok(None);
877 };
878 let value = entry
879 .value
880 .ok_or_else(|| invalid("GetAttributes returned an empty Value"))?;
881 let value = self.runtime.block_on(self.hydrator.hydrate(value))?;
882 value_mapper::decode(&value).map(Some)
883 }
884
885 fn set_attribute_value<T: Value>(
886 &self,
887 flow_id: &str,
888 key: &str,
889 value: &T,
890 index_config: Option<dex_protocol::dex::IndexConfig>,
891 sync_config: Option<dex_protocol::dex::AttributeSyncConfig>,
892 ) -> SdkResult<()> {
893 let write = AttributeWrite {
894 key: key.to_string(),
895 value: Some(value_mapper::encode(value)?),
896 index_config,
897 sync_config,
898 };
899 self.call_empty(
900 "set_attribute",
901 Some(flow_id),
902 FlowTargetRequirement::Active,
903 |mut service| async move {
904 service
905 .set_attributes(SetAttributesRequest {
906 flow_id: flow_id.to_string(),
907 run_id: String::new(),
908 attributes: vec![write],
909 request_id: Uuid::new_v4().to_string(),
910 })
911 .await
912 },
913 )
914 }
915
916 fn publish_values<T: Value>(
917 &self,
918 flow_id: &str,
919 channel_name: &str,
920 values: impl IntoIterator<Item = T>,
921 ) -> SdkResult<()> {
922 let messages = values
923 .into_iter()
924 .map(|value| {
925 Ok(dex_protocol::dex::ChannelMessage {
926 channel_name: channel_name.to_string(),
927 value: Some(value_mapper::encode(&value)?),
928 })
929 })
930 .collect::<SdkResult<Vec<_>>>()?;
931 self.call_empty(
932 "publish",
933 Some(flow_id),
934 FlowTargetRequirement::Active,
935 |mut service| async move {
936 service
937 .publish_to_channel(PublishToChannelRequest {
938 flow_id: flow_id.to_string(),
939 run_id: String::new(),
940 messages,
941 })
942 .await
943 },
944 )
945 }
946
947 fn read_stream_result<T: Value>(
948 &self,
949 flow_id: &str,
950 stream: &Stream<T>,
951 resume_token: &str,
952 timeout: Option<Duration>,
953 ) -> SdkResult<StreamMessage<T>> {
954 require_name(flow_id, "Flow ID")?;
955 let flow_type = self.registry.flow_for_stream(stream)?.name.to_string();
956 let wait_time_seconds = optional_seconds(timeout)?;
957 let request = ReadStreamRequest {
958 flow_id: flow_id.to_string(),
959 flow_type,
960 stream_name: stream.name().to_string(),
961 resume_token: resume_token.to_string(),
962 wait_time_seconds,
963 };
964 let mut service = self.service.clone();
965 let response = self.runtime.block_on(async {
966 service
967 .read_stream(request)
968 .await
969 .map(|response| response.into_inner())
970 .map_err(|status| {
971 SdkError::from_status(
972 status,
973 "read_stream",
974 Some(flow_id),
975 FlowTargetRequirement::None,
976 )
977 })
978 })?;
979 let message = response
980 .message
981 .ok_or_else(|| invalid("ReadStream omitted its message"))?;
982 if message.resume_token.is_empty() {
983 return Err(invalid("ReadStream returned an empty resume token"));
984 }
985 let value = message
986 .value
987 .ok_or_else(|| invalid("ReadStream omitted its Value"))?;
988 let created_time = message
989 .created_time
990 .map(timestamp)
991 .ok_or_else(|| invalid("ReadStream omitted its creation time"))?;
992 Ok(StreamMessage {
993 value: value_mapper::decode(&value)?,
994 resume_token: message.resume_token,
995 created_time,
996 source: message.source,
997 })
998 }
999
1000 fn wait_for_flow_result(
1001 &self,
1002 flow_id: &str,
1003 timeout: Option<Duration>,
1004 ) -> SdkResult<FlowResult> {
1005 let mut service = self.service.clone();
1006 let response = self.runtime.block_on(async {
1007 service
1008 .wait_for_flow(WaitForFlowRequest {
1009 flow_id: flow_id.to_string(),
1010 run_id: String::new(),
1011 needs_results: true,
1012 wait_time_seconds: optional_seconds(timeout)?,
1013 })
1014 .await
1015 .map(|response| response.into_inner())
1016 .map_err(|status| {
1017 SdkError::from_status(
1018 status,
1019 "wait_for_flow",
1020 Some(flow_id),
1021 FlowTargetRequirement::Existing,
1022 )
1023 })
1024 });
1025 let response = response?;
1026 let mut outputs = Vec::with_capacity(response.results.len());
1027 for result in &response.results {
1028 outputs.push(result.completed_step_output.clone().ok_or_else(|| {
1029 SdkError::ValueMapping {
1030 message: "Step completion output is required".to_string(),
1031 }
1032 })?);
1033 }
1034 let outputs = self.runtime.block_on(self.hydrator.hydrate_all(outputs))?;
1035 let completions = response
1036 .results
1037 .into_iter()
1038 .zip(outputs)
1039 .map(|(result, output)| {
1040 StepCompletion::new(
1041 result.completed_step_type,
1042 result.completed_step_execution_id,
1043 output,
1044 )
1045 })
1046 .collect();
1047 Ok(FlowResult::new(
1048 map_flow_status(response.flow_status)?,
1049 map_flow_error_type(response.error_type),
1050 (!response.error_message.is_empty()).then_some(response.error_message),
1051 completions,
1052 ))
1053 }
1054
1055 fn map_start_options(
1056 &self,
1057 flow: &crate::registry::RegisteredFlow,
1058 options: &StartFlowOptions,
1059 ) -> SdkResult<FlowStartOptions> {
1060 let attributes = options
1061 .attributes
1062 .iter()
1063 .map(|attribute| {
1064 let logical_name = attribute.key.split('/').next().unwrap_or(&attribute.key);
1065 if !flow.persistence.contains_key(logical_name) {
1066 return Err(invalid(format!(
1067 "initial Attribute is not registered: {}",
1068 attribute.key
1069 )));
1070 }
1071 Ok(AttributeWrite {
1072 key: attribute.key.clone(),
1073 value: Some(attribute.value.encode()?),
1074 index_config: attribute.index_config.clone(),
1075 sync_config: attribute.sync_config,
1076 })
1077 })
1078 .collect::<SdkResult<Vec<_>>>()?;
1079 let config = options.config_override.as_ref();
1080 let flow_config_override =
1081 if config.is_some() || self.options.worker_target_value().is_some() {
1082 Some(self.map_flow_config(config)?)
1083 } else {
1084 None
1085 };
1086 #[allow(clippy::needless_update)]
1087 let flow_start_options = FlowStartOptions {
1088 id_reuse_policy: match options.id_reuse_policy {
1089 IdReusePolicy::Default => ProtoIdReusePolicy::Unspecified,
1090 IdReusePolicy::AllowIfPreviousFailed => {
1091 ProtoIdReusePolicy::AllowIfPreviousExistsAbnormally
1092 }
1093 IdReusePolicy::AllowIfNotRunning => ProtoIdReusePolicy::AllowIfNoRunning,
1094 IdReusePolicy::Disallow => ProtoIdReusePolicy::DisallowReuse,
1095 IdReusePolicy::TerminateIfRunning => ProtoIdReusePolicy::AllowTerminateIfRunning,
1096 } as i32,
1097 flow_start_delay_seconds: optional_seconds(options.start_delay)?,
1098 retry_policy: options
1099 .retry_policy
1100 .clone()
1101 .map(map_flow_retry)
1102 .transpose()?,
1103 attributes,
1104 flow_config_override,
1105 flow_already_started_options: Some(FlowAlreadyStartedOptions {
1106 ignore_already_started_error: options.ignore_already_started,
1107 }),
1108 ..Default::default()
1109 };
1110 Ok(flow_start_options)
1111 }
1112
1113 fn map_flow_config(&self, config: Option<&FlowConfig>) -> SdkResult<ProtoFlowConfig> {
1114 map_flow_config(config, self.options.worker_target_value())
1115 }
1116
1117 fn call_empty<Response, Future, Call>(
1118 &self,
1119 operation: &'static str,
1120 flow_id: Option<&str>,
1121 requirement: FlowTargetRequirement,
1122 call: Call,
1123 ) -> SdkResult<()>
1124 where
1125 Call: FnOnce(FlowServiceClient<tonic::transport::Channel>) -> Future,
1126 Future: std::future::Future<Output = Result<tonic::Response<Response>, tonic::Status>>,
1127 {
1128 let service = self.service.clone();
1129 self.runtime
1130 .block_on(call(service))
1131 .map(|_| ())
1132 .map_err(|status| SdkError::from_status(status, operation, flow_id, requirement))
1133 }
1134}
1135
1136pub(crate) fn map_flow_config(
1137 config: Option<&FlowConfig>,
1138 default_target: Option<&WorkerTarget>,
1139) -> SdkResult<ProtoFlowConfig> {
1140 let target = config
1141 .and_then(|config| config.worker_target.as_ref())
1142 .or(default_target);
1143 Ok(ProtoFlowConfig {
1144 active_step_search_mode: config
1145 .and_then(|config| config.active_step_search_mode)
1146 .map(|mode| match mode {
1147 ActiveStepSearchMode::All => ProtoSearchMode::EnabledForAll as i32,
1148 ActiveStepSearchMode::WithWaitFor => {
1149 ProtoSearchMode::EnabledForStepsWithWaitFor as i32
1150 }
1151 ActiveStepSearchMode::Disabled => ProtoSearchMode::Disabled as i32,
1152 }),
1153 continue_as_new_threshold: config
1154 .and_then(|config| config.continue_as_new_threshold)
1155 .map(i32::try_from)
1156 .transpose()
1157 .map_err(|_| invalid("continue-as-new threshold exceeds int32"))?,
1158 continue_as_new_page_size_in_bytes: config
1159 .and_then(|config| config.continue_as_new_page_size_bytes)
1160 .map(i32::try_from)
1161 .transpose()
1162 .map_err(|_| invalid("continue-as-new page size exceeds int32"))?,
1163 step_durability: config
1164 .and_then(|config| config.step_durability)
1165 .map(|durability| {
1166 (match durability {
1167 StepDurability::Default => ProtoStepDurability::Unspecified,
1168 StepDurability::Sync => ProtoStepDurability::Sync,
1169 StepDurability::Async => ProtoStepDurability::Async,
1170 }) as i32
1171 }),
1172 worker_target: target.map(map_worker_target),
1173 attribute_store_names: config.and_then(|config| {
1174 config
1175 .attribute_store_names
1176 .as_ref()
1177 .map(|names| AttributeStoreNames {
1178 names: names.clone(),
1179 })
1180 }),
1181 })
1182}
1183
1184pub(crate) fn map_flow_retry(retry: RetryPolicy) -> SdkResult<FlowRetryPolicy> {
1185 Ok(FlowRetryPolicy {
1186 initial_interval_seconds: optional_seconds(retry.initial_interval)?,
1187 backoff_coefficient: retry.backoff_coefficient.unwrap_or_default() as f32,
1188 maximum_interval_seconds: optional_seconds(retry.maximum_interval)?,
1189 maximum_attempts: i32::try_from(retry.maximum_attempts.unwrap_or_default())
1190 .map_err(|_| invalid("maximum attempts exceed int32"))?,
1191 })
1192}
1193
1194fn map_worker_target(target: &WorkerTarget) -> ProtoWorkerTarget {
1195 ProtoWorkerTarget {
1196 address: target.address().to_string(),
1197 is_headless_address: target.is_headless(),
1198 }
1199}
1200
1201pub(crate) fn map_flow_status(status: i32) -> SdkResult<FlowStatus> {
1202 match ProtoFlowStatus::try_from(status).ok() {
1203 Some(ProtoFlowStatus::Running) => Ok(FlowStatus::Running),
1204 Some(ProtoFlowStatus::Completed) => Ok(FlowStatus::Completed),
1205 Some(ProtoFlowStatus::Failed) => Ok(FlowStatus::Failed),
1206 Some(ProtoFlowStatus::ServerSideTimeoutInternalOnly) => {
1207 Ok(FlowStatus::ServerSideTimeoutInternalOnly)
1208 }
1209 Some(ProtoFlowStatus::Terminated) => Ok(FlowStatus::Terminated),
1210 Some(ProtoFlowStatus::Canceled) => Ok(FlowStatus::Canceled),
1211 Some(ProtoFlowStatus::ContinuedAsNew) => Ok(FlowStatus::ContinuedAsNew),
1212 _ => Err(invalid(format!("unknown Flow status {status}"))),
1213 }
1214}
1215
1216pub(crate) fn map_flow_error_type(error_type: i32) -> Option<FlowErrorType> {
1217 use dex_protocol::dex::FlowErrorType as ProtoFlowErrorType;
1218
1219 match ProtoFlowErrorType::try_from(error_type).ok() {
1220 Some(ProtoFlowErrorType::StepDecisionFailingFlow) => {
1221 Some(FlowErrorType::StepDecisionFailed)
1222 }
1223 Some(ProtoFlowErrorType::ClientApiFailingFlow) => Some(FlowErrorType::ClientApiFailed),
1224 Some(ProtoFlowErrorType::WorkerApiFail) => Some(FlowErrorType::WorkerApiFailed),
1225 Some(ProtoFlowErrorType::InvalidUserFlowCode) => Some(FlowErrorType::InvalidUserFlowCode),
1226 Some(ProtoFlowErrorType::Internal) => Some(FlowErrorType::Internal),
1227 Some(ProtoFlowErrorType::FlowTimeout) => Some(FlowErrorType::FlowTimeout),
1228 _ => None,
1229 }
1230}
1231
1232pub(crate) fn map_flow_timeout_policy(
1233 flow: &crate::registry::RegisteredFlow,
1234 timeout_seconds: i32,
1235 policy: FlowTimeoutPolicy,
1236) -> SdkResult<i32> {
1237 if timeout_seconds == 0 {
1238 if policy != FlowTimeoutPolicy::Default {
1239 return Err(invalid("Flow timeout policy requires a positive timeout"));
1240 }
1241 return Ok(ProtoFlowTimeoutPolicy::Unspecified as i32);
1242 }
1243 let resolved = match policy {
1244 FlowTimeoutPolicy::Default if flow.handler.has_timeout_handler() => {
1245 FlowTimeoutPolicy::Handler
1246 }
1247 FlowTimeoutPolicy::Default => FlowTimeoutPolicy::Fail,
1248 configured => configured,
1249 };
1250 if resolved == FlowTimeoutPolicy::Handler && !flow.handler.has_timeout_handler() {
1251 return Err(invalid(format!(
1252 "Flow {} has no timeout handler",
1253 flow.name
1254 )));
1255 }
1256 Ok(match resolved {
1257 FlowTimeoutPolicy::Default => ProtoFlowTimeoutPolicy::Unspecified,
1258 FlowTimeoutPolicy::Fail => ProtoFlowTimeoutPolicy::Fail,
1259 FlowTimeoutPolicy::Cancel => ProtoFlowTimeoutPolicy::Cancel,
1260 FlowTimeoutPolicy::Handler => ProtoFlowTimeoutPolicy::Handler,
1261 } as i32)
1262}
1263
1264fn timestamp(timestamp: prost_types::Timestamp) -> SystemTime {
1265 if timestamp.seconds >= 0 {
1266 SystemTime::UNIX_EPOCH
1267 + Duration::from_secs(timestamp.seconds as u64)
1268 + Duration::from_nanos(timestamp.nanos.max(0) as u64)
1269 } else {
1270 SystemTime::UNIX_EPOCH - Duration::from_secs(timestamp.seconds.unsigned_abs())
1271 }
1272}
1273
1274fn rfc3339(time: SystemTime) -> SdkResult<String> {
1275 let duration = time
1276 .duration_since(SystemTime::UNIX_EPOCH)
1277 .map_err(service_error)?;
1278 let days = (duration.as_secs() / 86_400) as i64;
1279 let seconds = duration.as_secs() % 86_400;
1280 let (year, month, day) = civil_from_days(days);
1281 let hour = seconds / 3_600;
1282 let minute = seconds % 3_600 / 60;
1283 let second = seconds % 60;
1284 Ok(format!(
1285 "{year:04}-{month:02}-{day:02}T{hour:02}:{minute:02}:{second:02}Z"
1286 ))
1287}
1288
1289fn civil_from_days(days: i64) -> (i64, i64, i64) {
1290 let days = days + 719_468;
1291 let era = if days >= 0 { days } else { days - 146_096 } / 146_097;
1292 let day_of_era = days - era * 146_097;
1293 let year_of_era =
1294 (day_of_era - day_of_era / 1_460 + day_of_era / 36_524 - day_of_era / 146_096) / 365;
1295 let mut year = year_of_era + era * 400;
1296 let day_of_year = day_of_era - (365 * year_of_era + year_of_era / 4 - year_of_era / 100);
1297 let month_prime = (5 * day_of_year + 2) / 153;
1298 let day = day_of_year - (153 * month_prime + 2) / 5 + 1;
1299 let month = month_prime + if month_prime < 10 { 3 } else { -9 };
1300 year += i64::from(month <= 2);
1301 (year, month, day)
1302}
1303
1304fn optional_seconds(duration: Option<Duration>) -> SdkResult<i32> {
1305 duration
1306 .map(seconds32)
1307 .transpose()
1308 .map(Option::unwrap_or_default)
1309}
1310
1311fn seconds32(duration: Duration) -> SdkResult<i32> {
1312 if duration.subsec_nanos() != 0 {
1313 return Err(invalid("Duration must use whole seconds"));
1314 }
1315 i32::try_from(duration.as_secs()).map_err(|_| invalid("Duration exceeds int32"))
1316}
1317
1318fn endpoint_address(address: &str) -> String {
1319 if address.contains("://") {
1320 address.to_string()
1321 } else {
1322 format!("http://{address}")
1323 }
1324}
1325
1326fn require_name(value: &str, kind: &str) -> SdkResult<()> {
1327 if value.is_empty() {
1328 Err(invalid(format!("{kind} is required")))
1329 } else {
1330 Ok(())
1331 }
1332}
1333
1334fn sdk_handler_error(error: impl std::fmt::Display) -> SdkError {
1335 invalid(error.to_string())
1336}
1337
1338fn invalid(message: impl Into<String>) -> SdkError {
1339 SdkError::InvalidArgument {
1340 message: message.into(),
1341 }
1342}
1343
1344fn service_error(error: impl std::fmt::Display) -> SdkError {
1345 SdkError::Service {
1346 service: ServiceError::local("client", error.to_string()),
1347 }
1348}
1349
1350#[cfg(test)]
1351mod tests {
1352 use super::map_flow_config;
1353 use crate::FlowConfig;
1354
1355 #[test]
1356 fn attribute_store_names_preserve_protocol_presence() {
1357 let absent = map_flow_config(Some(&FlowConfig::new()), None).expect("map absent config");
1358 let named = map_flow_config(
1359 Some(
1360 &FlowConfig::new()
1361 .attribute_store_names(vec!["profiles".to_owned(), "audit".to_owned()]),
1362 ),
1363 None,
1364 )
1365 .expect("map named config");
1366 let disabled =
1367 map_flow_config(Some(&FlowConfig::new().attribute_store_names(vec![])), None)
1368 .expect("map disabled config");
1369
1370 assert_eq!(absent.attribute_store_names, None);
1371 assert_eq!(
1372 named.attribute_store_names.expect("named stores").names,
1373 vec!["profiles", "audit"]
1374 );
1375 assert_eq!(
1376 disabled
1377 .attribute_store_names
1378 .expect("disabled stores")
1379 .names,
1380 Vec::<String>::new()
1381 );
1382 }
1383}