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