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