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