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