1use std::any::TypeId;
10use std::collections::{BTreeMap, BTreeSet};
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, FlowAlreadyStartedOptions,
17 FlowConfig as ProtoFlowConfig, FlowResetStepMethod, FlowResetType, FlowRetryPolicy,
18 FlowStartOptions, FlowStatus as ProtoFlowStatus, FlowTimeoutPolicy as ProtoFlowTimeoutPolicy,
19 GetFlowSummaryRequest, IdReusePolicy as ProtoIdReusePolicy, InvokeRpcRequest,
20 ListStreamMessagesRequest, ReadStreamRequest, ResetFlowRequest, SearchFlowsRequest,
21 SkipTimerRequest, StartFlowRequest, StepDurability as ProtoStepDurability, StopFlowRequest,
22 StopType as ProtoStopType, TriggerContinueAsNewRequest, UpdateFlowConfigRequest,
23 WaitForAttributeRequest, WaitForFlowRequest, WaitForStepCompletionRequest,
24 WorkerTarget as ProtoWorkerTarget, WriteStreamRequest,
25};
26use tokio::runtime::Runtime;
27use tonic::{Code, Request, transport::Endpoint};
28use uuid::Uuid;
29
30use crate::credential::{CredentialChannel, CredentialInterceptor};
31use crate::sdk_error::{FlowTargetRequirement, ServiceError};
32use crate::server_protocol::{
33 NegotiatedServerProtocolVersion, negotiate_server_protocol, server_info_request_error,
34};
35use crate::stop_flow_options::StopType;
36use crate::time_travel_options::{TimeTravelPoint, TimeTravelStepMethod};
37use crate::value_hydrator::ValueHydrator;
38use crate::value_mapper;
39use crate::wait_options::ClientRequestBudget;
40use crate::worker_dispatcher::{map_flow_timeout_handler_options, map_step_options};
41use crate::{
42 ActiveStepSearchMode, Attribute, AttributeMap, AttributeMatch, BlobCache, ClientOptions, Flow,
43 FlowConfig, FlowErrorType, FlowInfo, FlowResult, FlowStatus, FlowTimeoutPolicy, IdReusePolicy,
44 Registry, RetryPolicy, Rpc, RpcInvokeOptions, SdkError, SdkResult, SearchFlowEntry,
45 SearchFlowsOptions, SearchFlowsPage, StartFlowOptions, StepCompletion, StepDurability,
46 StepExecutionId, StopFlowOptions, Stream, StreamMessage, StreamMessagesPage, TimeTravelOptions,
47 TimerId, Value, WaitForAttributeOptions, WaitForStepCompletionOptions,
48};
49
50pub struct Client {
70 runtime: Runtime,
71 registry: Registry,
72 service: FlowServiceClient<CredentialChannel>,
73 hydrator: ValueHydrator,
74 negotiated_protocol_version: NegotiatedServerProtocolVersion,
75}
76
77impl Client {
78 pub fn new(registry: Registry, blob_cache: Arc<BlobCache>, options: ClientOptions) -> Self {
86 Self::try_new(registry, blob_cache, options)
87 .unwrap_or_else(|error| panic!("cannot create Dex Client: {error}"))
88 }
89
90 pub fn try_new(
97 registry: Registry,
98 blob_cache: Arc<BlobCache>,
99 options: ClientOptions,
100 ) -> SdkResult<Self> {
101 let runtime = Runtime::new().map_err(service_error)?;
102 let endpoint = Endpoint::from_shared(endpoint_address(options.server_address_value()))
103 .map_err(service_error)?;
104 let service = {
105 let _runtime_guard = runtime.enter();
106 FlowServiceClient::new(CredentialChannel::new(
107 endpoint.connect_lazy(),
108 CredentialInterceptor::new(options.credential_value()),
109 ))
110 };
111 let hydrator = ValueHydrator::new(service.clone(), blob_cache);
112 Ok(Self {
113 runtime,
114 registry,
115 service,
116 hydrator,
117 negotiated_protocol_version: NegotiatedServerProtocolVersion::default(),
118 })
119 }
120
121 pub fn start_flow<SomeFlow: Flow>(
130 &self,
131 flow: &SomeFlow,
132 flow_id: &str,
133 input: SomeFlow::StartInput,
134 ) -> SdkResult<String> {
135 self.start_flow_with_options(flow, flow_id, input, StartFlowOptions::new())
136 }
137
138 pub fn start_flow_with_options<SomeFlow: Flow>(
153 &self,
154 flow: &SomeFlow,
155 flow_id: &str,
156 input: SomeFlow::StartInput,
157 options: StartFlowOptions,
158 ) -> SdkResult<String> {
159 require_name(flow_id, "Flow ID")?;
160 let registered = self.registry.flow(flow.flow_type())?;
161 if registered.type_id != TypeId::of::<SomeFlow>() {
162 return Err(SdkError::FlowDefinition {
163 message: format!(
164 "Flow {} does not match the registered Rust type",
165 flow.flow_type()
166 ),
167 });
168 }
169 let (start_step_type, step_input, step_options) = match ®istered.start_step {
170 Some(step) => (
171 step.name.to_string(),
172 Some(value_mapper::encode(&input)?),
173 Some(
174 map_step_options(
175 registered,
176 registered
177 .handler
178 .step_options(step.name)
179 .map_err(sdk_handler_error)?,
180 )
181 .map_err(sdk_handler_error)?,
182 ),
183 ),
184 None => (String::new(), None, None),
185 };
186 let flow_timeout_seconds = optional_seconds(options.timeout)?;
187 let flow_timeout_policy =
188 map_flow_timeout_policy(registered, flow_timeout_seconds, options.timeout_policy)?;
189 let request = StartFlowRequest {
190 flow_id: flow_id.to_string(),
191 flow_type: registered.name.to_string(),
192 flow_timeout_seconds,
193 flow_timeout_policy,
194 start_step_type,
195 step_input,
196 step_options,
197 flow_start_options: Some(self.map_start_options(
198 registered,
199 &options,
200 flow_timeout_seconds,
201 flow_timeout_policy,
202 )?),
203 request_id: options
204 .request_id
205 .clone()
206 .unwrap_or_else(|| Uuid::new_v4().to_string()),
207 };
208 let mut service = self.service.clone();
209 self.runtime.block_on(async {
210 service
211 .start_flow(request)
212 .await
213 .map(|response| response.into_inner().run_id)
214 .map_err(|status| {
215 SdkError::from_status(
216 status,
217 "start_flow",
218 Some(flow_id),
219 FlowTargetRequirement::None,
220 )
221 })
222 })
223 }
224
225 pub fn invoke_rpc<Input: Value, Output: Value>(
239 &self,
240 flow_id: &str,
241 rpc: Rpc<Input, Output>,
242 input: Input,
243 ) -> SdkResult<Output> {
244 self.do_invoke_rpc(flow_id, rpc.name(), &input, &RpcInvokeOptions::new())
245 }
246
247 pub fn invoke_rpc_with_options<Input: Value, Output: Value>(
254 &self,
255 flow_id: &str,
256 rpc: Rpc<Input, Output>,
257 input: Input,
258 options: RpcInvokeOptions,
259 ) -> SdkResult<Output> {
260 self.do_invoke_rpc(flow_id, rpc.name(), &input, &options)
261 }
262
263 pub fn invoke_rpc_without_input<Output: Value>(
267 &self,
268 flow_id: &str,
269 rpc: Rpc<(), Output>,
270 ) -> SdkResult<Output> {
271 self.do_invoke_rpc(flow_id, rpc.name(), &(), &RpcInvokeOptions::new())
272 }
273
274 pub fn invoke_rpc_without_input_with_options<Output: Value>(
280 &self,
281 flow_id: &str,
282 rpc: Rpc<(), Output>,
283 options: RpcInvokeOptions,
284 ) -> SdkResult<Output> {
285 self.do_invoke_rpc(flow_id, rpc.name(), &(), &options)
286 }
287
288 pub fn write_stream<T: Value>(
298 &self,
299 flow_id: &str,
300 stream: &Stream<T>,
301 source: &str,
302 value: T,
303 ) -> SdkResult<()> {
304 require_name(flow_id, "Flow ID")?;
305 require_name(source, "Stream source")?;
306 let flow_type = self.registry.flow_for_stream(stream)?.name.to_string();
307 let request = WriteStreamRequest {
308 flow_id: flow_id.to_string(),
309 flow_type,
310 stream_name: stream.name().to_string(),
311 stream_capacity_bytes: stream.stream_capacity_bytes(),
312 value: Some(value_mapper::encode(&value)?),
313 source: source.to_string(),
314 };
315 self.call_empty(
316 "write_stream",
317 Some(flow_id),
318 FlowTargetRequirement::None,
319 |mut service| async move { service.write_stream(request).await },
320 )
321 }
322
323 pub fn read_stream<T: Value>(
327 &self,
328 flow_id: &str,
329 stream: &Stream<T>,
330 resume_token: &str,
331 ) -> SdkResult<StreamMessage<T>> {
332 self.read_stream_result(flow_id, stream, resume_token, None)
333 }
334
335 pub fn read_stream_with_timeout<T: Value>(
342 &self,
343 flow_id: &str,
344 stream: &Stream<T>,
345 resume_token: &str,
346 timeout: Duration,
347 ) -> SdkResult<StreamMessage<T>> {
348 self.read_stream_result(flow_id, stream, resume_token, Some(timeout))
349 }
350
351 pub fn list_stream_messages<T: Value>(
363 &self,
364 flow_id: &str,
365 stream: &Stream<T>,
366 page_size: i32,
367 before_page_token: &str,
368 ) -> SdkResult<StreamMessagesPage<T>> {
369 require_name(flow_id, "Flow ID")?;
370 if page_size < 1 {
371 return Err(invalid("Stream page size must be positive"));
372 }
373 let flow_type = self.registry.flow_for_stream(stream)?.name.to_string();
374 let request = ListStreamMessagesRequest {
375 flow_id: flow_id.to_string(),
376 flow_type,
377 stream_name: stream.name().to_string(),
378 page_size,
379 before_page_token: before_page_token.to_string(),
380 };
381 let mut service = self.service.clone();
382 let response = self.runtime.block_on(async {
383 service
384 .list_stream_messages(request)
385 .await
386 .map(|response| response.into_inner())
387 .map_err(|status| {
388 SdkError::from_status(
389 status,
390 "list_stream_messages",
391 Some(flow_id),
392 FlowTargetRequirement::None,
393 )
394 })
395 })?;
396 let messages = response
397 .messages
398 .into_iter()
399 .map(|message| {
400 if message.resume_token.is_empty() {
401 return Err(invalid("ListStreamMessages returned an empty resume token"));
402 }
403 let value = message
404 .value
405 .ok_or_else(|| invalid("ListStreamMessages omitted a Stream message Value"))?;
406 let created_time = message.created_time.map(timestamp).ok_or_else(|| {
407 invalid("ListStreamMessages omitted a Stream message creation time")
408 })?;
409 Ok(StreamMessage {
410 value: value_mapper::decode(&value)?,
411 resume_token: message.resume_token,
412 created_time,
413 source: message.source,
414 })
415 })
416 .collect::<SdkResult<Vec<_>>>()?;
417 Ok(StreamMessagesPage {
418 messages,
419 next_page_token: response.next_page_token,
420 })
421 }
422
423 pub fn wait_for_flow(&self, flow_id: &str) -> SdkResult<FlowResult> {
430 self.wait_for_flow_result(flow_id, None)
431 }
432
433 pub fn wait_for_flow_with_timeout(
440 &self,
441 flow_id: &str,
442 timeout: Duration,
443 ) -> SdkResult<FlowResult> {
444 self.wait_for_flow_result(flow_id, Some(timeout))
445 }
446
447 pub fn describe_flow(&self, flow_id: &str) -> SdkResult<FlowInfo> {
453 let mut service = self.service.clone();
454 let response = self.runtime.block_on(async {
455 service
456 .get_flow_summary(GetFlowSummaryRequest {
457 flow_id: flow_id.to_string(),
458 run_id: String::new(),
459 })
460 .await
461 .map(|response| response.into_inner())
462 .map_err(|status| {
463 SdkError::from_status(
464 status,
465 "describe_flow",
466 Some(flow_id),
467 FlowTargetRequirement::Existing,
468 )
469 })
470 })?;
471 let execution = response
472 .flow_execution_id
473 .ok_or_else(|| service_error("GetFlowSummary omitted FlowExecutionID"))?;
474 Ok(FlowInfo {
475 flow_id: execution.flow_id,
476 run_id: execution.run_id,
477 flow_type: response.flow_type,
478 status: map_flow_status(response.flow_status)?,
479 started_at: response
480 .start_time
481 .map_or(SystemTime::UNIX_EPOCH, timestamp),
482 })
483 }
484
485 pub fn search_flows(&self, query: &str, page_size: i32) -> SdkResult<SearchFlowsPage> {
490 self.search_flows_page(query, page_size, "")
491 }
492
493 pub fn search_flows_page(
498 &self,
499 query: &str,
500 page_size: i32,
501 next_page_token: &str,
502 ) -> SdkResult<SearchFlowsPage> {
503 self.search_flows_with_options(query, page_size, next_page_token, SearchFlowsOptions::new())
504 }
505
506 pub fn search_flows_with_options(
508 &self,
509 query: &str,
510 page_size: i32,
511 next_page_token: &str,
512 options: SearchFlowsOptions,
513 ) -> SdkResult<SearchFlowsPage> {
514 if page_size < 0 {
515 return Err(invalid("search page size must not be negative"));
516 }
517 let mut service = self.service.clone();
518 self.runtime.block_on(async {
519 let response = service
520 .search_flows(SearchFlowsRequest {
521 query: query.to_string(),
522 page_size,
523 next_page_token: next_page_token.to_string(),
524 include_continued_as_new: options.include_continued_as_new,
525 })
526 .await
527 .map(|response| response.into_inner())
528 .map_err(|status| {
529 SdkError::from_status(status, "search_flows", None, FlowTargetRequirement::None)
530 })?;
531 let values = response
532 .flow_runs
533 .iter()
534 .flat_map(|entry| {
535 entry.indexed_attributes.iter().map(|attribute| {
536 attribute
537 .value
538 .clone()
539 .map(|value| (entry.flow_id.clone(), value))
540 .ok_or_else(|| invalid("Indexed Attribute has no Value"))
541 })
542 })
543 .collect::<SdkResult<Vec<_>>>()?;
544 let mut values = self.hydrator.hydrate_for_flows(values).await?.into_iter();
545 let mut flows = Vec::with_capacity(response.flow_runs.len());
546 for entry in response.flow_runs {
547 let mut indexed_attributes = BTreeMap::new();
548 for attribute in entry.indexed_attributes {
549 let value = values
550 .next()
551 .ok_or_else(|| invalid("Indexed Attribute hydration count mismatch"))?;
552 indexed_attributes.insert(attribute.key, value_mapper::decode_untyped(&value)?);
553 }
554 flows.push(SearchFlowEntry {
555 flow_id: entry.flow_id,
556 run_id: entry.run_id,
557 flow_type: entry.flow_type,
558 status: map_flow_status(entry.flow_status)?,
559 started_at: entry.start_time.map(timestamp),
560 closed_at: entry.close_time.map(timestamp),
561 indexed_attributes,
562 });
563 }
564 Ok(SearchFlowsPage {
565 flows,
566 next_page_token: response.next_page_token,
567 })
568 })
569 }
570
571 pub fn stop_flow(&self, flow_id: &str, options: StopFlowOptions) -> SdkResult<()> {
577 let stop_type = match options.stop_type {
578 StopType::Cancel => ProtoStopType::Cancel,
579 StopType::Terminate => ProtoStopType::Terminate,
580 StopType::Fail => ProtoStopType::Fail,
581 };
582 self.call_empty(
583 "stop_flow",
584 Some(flow_id),
585 FlowTargetRequirement::Active,
586 |mut service| async move {
587 service
588 .stop_flow(StopFlowRequest {
589 flow_id: flow_id.to_string(),
590 run_id: String::new(),
591 reason: options.reason.unwrap_or_default(),
592 stop_type: stop_type as i32,
593 })
594 .await
595 },
596 )
597 }
598
599 pub fn time_travel(&self, flow_id: &str, options: TimeTravelOptions) -> SdkResult<String> {
608 let mut request = ResetFlowRequest {
609 flow_id: flow_id.to_string(),
610 reset_type: FlowResetType::Unspecified as i32,
611 reason: options.reason.unwrap_or_default(),
612 skip_writes_reapply: options.skip_writes_reapply,
613 ..Default::default()
614 };
615 match options.point {
616 TimeTravelPoint::Beginning => request.reset_type = FlowResetType::Beginning as i32,
617 TimeTravelPoint::HistoryEventTime(time) => {
618 request.reset_type = FlowResetType::HistoryEventTime as i32;
619 request.history_event_time = rfc3339(time)?;
620 }
621 TimeTravelPoint::StepType(step_type) => {
622 request.reset_type = FlowResetType::StepType as i32;
623 request.step_type = step_type.to_string();
624 }
625 TimeTravelPoint::StepExecution(execution, method) => {
626 request.reset_type = FlowResetType::StepExecutionId as i32;
627 request.step_execution_id =
628 format!("{}-{}", execution.step_type, execution.execution_number);
629 request.step_method = match method {
630 TimeTravelStepMethod::WaitFor => FlowResetStepMethod::WaitFor as i32,
631 TimeTravelStepMethod::Execute => FlowResetStepMethod::Execute as i32,
632 };
633 }
634 }
635 let mut service = self.service.clone();
636 self.runtime.block_on(async {
637 service
638 .reset_flow(request)
639 .await
640 .map(|response| response.into_inner().run_id)
641 .map_err(|status| {
642 SdkError::from_status(
643 status,
644 "time_travel",
645 Some(flow_id),
646 FlowTargetRequirement::Existing,
647 )
648 })
649 })
650 }
651
652 pub fn skip_timer(
659 &self,
660 flow_id: &str,
661 step_execution: StepExecutionId,
662 timer: TimerId,
663 ) -> SdkResult<()> {
664 let (timer_condition_id, timer_condition_index) = match timer {
665 TimerId::ConditionId(id) => (id, None),
666 TimerId::ConditionIndex(index) => (
667 String::new(),
668 Some(i32::try_from(index).map_err(|_| invalid("timer index exceeds int32"))?),
669 ),
670 };
671 self.call_empty(
672 "skip_timer",
673 Some(flow_id),
674 FlowTargetRequirement::Active,
675 |mut service| async move {
676 service
677 .skip_timer(SkipTimerRequest {
678 flow_id: flow_id.to_string(),
679 run_id: String::new(),
680 step_execution_id: format!(
681 "{}-{}",
682 step_execution.step_type, step_execution.execution_number
683 ),
684 timer_condition_id,
685 timer_condition_index,
686 })
687 .await
688 },
689 )
690 }
691
692 pub fn wait_for_step_completion(
702 &self,
703 flow_id: &str,
704 step_execution: StepExecutionId,
705 options: WaitForStepCompletionOptions,
706 ) -> SdkResult<()> {
707 let request_budget = ClientRequestBudget::new(options.request_timeout)?;
708 let internal_handler_timeout_seconds = seconds32(options.internal_handler_timeout)?;
709 loop {
710 let request_attempt =
711 request_budget.next_attempt("wait_for_step_completion", flow_id)?;
712 let request_flow_id = flow_id.to_string();
713 let request_step_type = step_execution.step_type.to_string();
714 let request_step_execution_number = step_execution.execution_number.to_string();
715 let request_id = options.request_id.clone();
716 let result = self.call_empty(
717 "wait_for_step_completion",
718 Some(flow_id),
719 FlowTargetRequirement::Active,
720 |mut service| async move {
721 let mut request = Request::new(WaitForStepCompletionRequest {
722 flow_id: request_flow_id,
723 step_type: request_step_type,
724 step_execution_number: request_step_execution_number,
725 request_timeout_seconds: request_attempt.request_timeout_seconds,
726 internal_handler_timeout_seconds,
727 request_id,
728 });
729 if let Some(transport_timeout) = request_attempt.transport_timeout {
730 request.set_timeout(transport_timeout);
731 }
732 service.wait_for_step_completion(request).await
733 },
734 );
735 match result {
736 Err(SdkError::LongPollTimeout { .. }) => continue,
737 Err(error)
738 if error.service_error().is_some_and(|service| {
739 matches!(service.code(), Code::DeadlineExceeded | Code::Cancelled)
740 }) && request_budget.has_deadline() =>
741 {
742 return Err(SdkError::request_timeout(
743 "wait_for_step_completion",
744 flow_id,
745 ));
746 }
747 result => return result,
748 }
749 }
750 }
751
752 pub fn wait_for_attribute_match<T: Value>(
761 &self,
762 flow_id: &str,
763 attribute: &Attribute<T>,
764 attribute_match: AttributeMatch<T>,
765 options: WaitForAttributeOptions,
766 ) -> SdkResult<T> {
767 require_waitable_attribute(attribute.name(), attribute.is_hard_delete_on_null())?;
768 self.wait_for_attribute_value(flow_id, attribute.name(), &attribute_match, options)
769 }
770
771 pub fn wait_for_attribute_map_instance_match<T: Value>(
778 &self,
779 flow_id: &str,
780 attribute: &AttributeMap<T>,
781 instance: &str,
782 attribute_match: AttributeMatch<T>,
783 options: WaitForAttributeOptions,
784 ) -> SdkResult<T> {
785 require_waitable_attribute(attribute.name(), attribute.is_hard_delete_on_null())?;
786 self.wait_for_attribute_value(
787 flow_id,
788 &map_physical_name(attribute.name(), instance)?,
789 &attribute_match,
790 options,
791 )
792 }
793
794 fn wait_for_attribute_value<T: Value>(
795 &self,
796 flow_id: &str,
797 key: &str,
798 attribute_match: &AttributeMatch<T>,
799 options: WaitForAttributeOptions,
800 ) -> SdkResult<T> {
801 let mut encoded_match = attribute_match.encode()?;
802 encoded_match.key = key.to_string();
803 let request_budget = ClientRequestBudget::new(options.request_timeout)?;
804 let internal_handler_timeout_seconds = seconds32(options.internal_handler_timeout)?;
805 let response = loop {
806 let request_attempt =
807 request_budget.next_attempt("wait_for_attribute_match", flow_id)?;
808 let mut service = self.service.clone();
809 let mut request = Request::new(WaitForAttributeRequest {
810 flow_id: flow_id.to_string(),
811 r#match: Some(encoded_match.clone()),
812 request_timeout_seconds: request_attempt.request_timeout_seconds,
813 internal_handler_timeout_seconds,
814 request_id: options.request_id.clone(),
815 });
816 if let Some(transport_timeout) = request_attempt.transport_timeout {
817 request.set_timeout(transport_timeout);
818 }
819 let result = self
820 .runtime
821 .block_on(service.wait_for_attribute(request))
822 .map_err(|status| {
823 SdkError::from_status(
824 status,
825 "wait_for_attribute_match",
826 Some(flow_id),
827 FlowTargetRequirement::Active,
828 )
829 });
830 match result {
831 Ok(response) => break response.into_inner(),
832 Err(SdkError::LongPollTimeout { .. }) => continue,
833 Err(error)
834 if error.service_error().is_some_and(|service| {
835 matches!(service.code(), Code::DeadlineExceeded | Code::Cancelled)
836 }) && request_budget.has_deadline() =>
837 {
838 return Err(SdkError::request_timeout(
839 "wait_for_attribute_match",
840 flow_id,
841 ));
842 }
843 Err(error) => return Err(error),
844 }
845 };
846 let matched_value = response.matched_value.ok_or_else(|| SdkError::Service {
847 service: ServiceError::local(
848 "wait_for_attribute_match",
849 "WaitForAttribute response is incomplete",
850 ),
851 })?;
852 value_mapper::decode(&matched_value)
853 }
854
855 pub fn update_flow_config(&self, flow_id: &str, config: FlowConfig) -> SdkResult<()> {
860 let flow_config = self.map_flow_config(Some(&config))?;
861 self.call_empty(
862 "update_flow_config",
863 Some(flow_id),
864 FlowTargetRequirement::Active,
865 |mut service| async move {
866 service
867 .update_flow_config(UpdateFlowConfigRequest {
868 flow_id: flow_id.to_string(),
869 run_id: String::new(),
870 flow_config: Some(flow_config),
871 })
872 .await
873 },
874 )
875 }
876
877 pub fn trigger_continue_as_new(&self, flow_id: &str) -> SdkResult<()> {
879 self.call_empty(
880 "trigger_continue_as_new",
881 Some(flow_id),
882 FlowTargetRequirement::Active,
883 |mut service| async move {
884 service
885 .trigger_continue_as_new(TriggerContinueAsNewRequest {
886 flow_id: flow_id.to_string(),
887 run_id: String::new(),
888 })
889 .await
890 },
891 )
892 }
893
894 pub fn health_check(&self) -> SdkResult<()> {
898 self.call_empty(
899 "health_check",
900 None,
901 FlowTargetRequirement::None,
902 |mut service| async move { service.health_check(()).await },
903 )
904 }
905
906 fn do_invoke_rpc<Input: Value, Output: Value>(
907 &self,
908 flow_id: &str,
909 rpc_name: &str,
910 input: &Input,
911 invoke_options: &RpcInvokeOptions,
912 ) -> SdkResult<Output> {
913 let (flow, rpc) = self.registry.rpc_with_flow(rpc_name)?;
914 let mut lock_attribute_keys = rpc
915 .locks
916 .iter()
917 .map(|lock| lock.physical_name())
918 .collect::<BTreeSet<_>>();
919 for lock in &invoke_options.lock_attribute_map_instances {
920 let Some(instance) = lock.instance() else {
921 return Err(SdkError::FlowDefinition {
922 message: "RpcInvokeOptions lock must target an AttributeMap instance".into(),
923 });
924 };
925 require_rpc_invoke_definition(
926 flow,
927 lock.attribute_name(),
928 crate::persistence::PersistenceKind::AttributeMap,
929 )?;
930 crate::registry::validate_map_instance(instance).map_err(invalid)?;
931 lock_attribute_keys.insert(lock.physical_name());
932 }
933 let mut load_attribute_map_instances = rpc
934 .load_attribute_maps
935 .iter()
936 .map(|load| map_load_name(&load.name, load.instance.as_deref()))
937 .collect::<BTreeSet<_>>();
938 for load in &invoke_options.load_attribute_map_instances {
939 let instance = require_rpc_invoke_map_load(
940 flow,
941 &load.name,
942 load.instance.as_deref(),
943 crate::persistence::PersistenceKind::AttributeMap,
944 )?;
945 load_attribute_map_instances
946 .insert(crate::registry::physical_name(&load.name, instance));
947 }
948 let mut load_channel_names = rpc
949 .load_channels
950 .iter()
951 .map(|load| load.name.clone())
952 .collect::<Vec<_>>();
953 load_channel_names.sort();
954 let mut load_channel_map_instances = rpc
955 .load_channel_maps
956 .iter()
957 .map(|load| map_load_name(&load.name, load.instance.as_deref()))
958 .collect::<BTreeSet<_>>();
959 for load in &invoke_options.load_channel_map_instances {
960 let instance = require_rpc_invoke_map_load(
961 flow,
962 &load.name,
963 load.instance.as_deref(),
964 crate::persistence::PersistenceKind::ChannelMap,
965 )?;
966 load_channel_map_instances.insert(crate::registry::physical_name(&load.name, instance));
967 }
968 let request = InvokeRpcRequest {
969 flow_id: flow_id.to_string(),
970 run_id: String::new(),
971 rpc_name: rpc_name.to_string(),
972 input: Some(value_mapper::encode(input)?),
973 timeout_seconds: optional_seconds(rpc.timeout)?,
974 lock_attribute_keys: lock_attribute_keys.into_iter().collect(),
975 request_id: Uuid::new_v4().to_string(),
976 is_transactional: rpc.is_transactional,
977 load_attribute_map_instances: load_attribute_map_instances.into_iter().collect(),
978 load_channel_names,
979 load_channel_map_instances: load_channel_map_instances.into_iter().collect(),
980 };
981 let mut service = self.service.clone();
982 let output = self.runtime.block_on(async {
983 service
984 .invoke_rpc(request)
985 .await
986 .map(|response| response.into_inner().output)
987 .map_err(|status| {
988 SdkError::from_status(
989 status,
990 "invoke_rpc",
991 Some(flow_id),
992 FlowTargetRequirement::Active,
993 )
994 })
995 })?;
996 let output = output.ok_or_else(|| invalid("InvokeRPC omitted output"))?;
997 let output = self
998 .runtime
999 .block_on(self.hydrator.hydrate(flow_id, output))?;
1000 value_mapper::decode(&output)
1001 }
1002
1003 fn read_stream_result<T: Value>(
1004 &self,
1005 flow_id: &str,
1006 stream: &Stream<T>,
1007 resume_token: &str,
1008 timeout: Option<Duration>,
1009 ) -> SdkResult<StreamMessage<T>> {
1010 require_name(flow_id, "Flow ID")?;
1011 let flow_type = self.registry.flow_for_stream(stream)?.name.to_string();
1012 let wait_time_seconds = optional_seconds(timeout)?;
1013 let request = ReadStreamRequest {
1014 flow_id: flow_id.to_string(),
1015 flow_type,
1016 stream_name: stream.name().to_string(),
1017 resume_token: resume_token.to_string(),
1018 wait_time_seconds,
1019 };
1020 let mut service = self.service.clone();
1021 let response = self.runtime.block_on(async {
1022 service
1023 .read_stream(request)
1024 .await
1025 .map(|response| response.into_inner())
1026 .map_err(|status| {
1027 SdkError::from_status(
1028 status,
1029 "read_stream",
1030 Some(flow_id),
1031 FlowTargetRequirement::None,
1032 )
1033 })
1034 })?;
1035 let message = response
1036 .message
1037 .ok_or_else(|| invalid("ReadStream omitted its message"))?;
1038 if message.resume_token.is_empty() {
1039 return Err(invalid("ReadStream returned an empty resume token"));
1040 }
1041 let value = message
1042 .value
1043 .ok_or_else(|| invalid("ReadStream omitted its Value"))?;
1044 let created_time = message
1045 .created_time
1046 .map(timestamp)
1047 .ok_or_else(|| invalid("ReadStream omitted its creation time"))?;
1048 Ok(StreamMessage {
1049 value: value_mapper::decode(&value)?,
1050 resume_token: message.resume_token,
1051 created_time,
1052 source: message.source,
1053 })
1054 }
1055
1056 fn wait_for_flow_result(
1057 &self,
1058 flow_id: &str,
1059 timeout: Option<Duration>,
1060 ) -> SdkResult<FlowResult> {
1061 let mut service = self.service.clone();
1062 let response = self.runtime.block_on(async {
1063 service
1064 .wait_for_flow(WaitForFlowRequest {
1065 flow_id: flow_id.to_string(),
1066 run_id: String::new(),
1067 needs_results: true,
1068 wait_time_seconds: optional_seconds(timeout)?,
1069 })
1070 .await
1071 .map(|response| response.into_inner())
1072 .map_err(|status| {
1073 SdkError::from_status(
1074 status,
1075 "wait_for_flow",
1076 Some(flow_id),
1077 FlowTargetRequirement::Existing,
1078 )
1079 })
1080 });
1081 let response = response?;
1082 let mut outputs = Vec::with_capacity(response.results.len());
1083 for result in &response.results {
1084 outputs.push(result.completed_step_output.clone().ok_or_else(|| {
1085 SdkError::ValueMapping {
1086 message: "Step completion output is required".to_string(),
1087 }
1088 })?);
1089 }
1090 let outputs = self
1091 .runtime
1092 .block_on(self.hydrator.hydrate_all(flow_id, outputs))?;
1093 let completions = response
1094 .results
1095 .into_iter()
1096 .zip(outputs)
1097 .map(|(result, output)| {
1098 StepCompletion::new(
1099 result.completed_step_type,
1100 result.completed_step_execution_id,
1101 output,
1102 )
1103 })
1104 .collect();
1105 Ok(FlowResult::new(
1106 map_flow_status(response.flow_status)?,
1107 map_flow_error_type(response.error_type),
1108 (!response.error_message.is_empty()).then_some(response.error_message),
1109 completions,
1110 ))
1111 }
1112
1113 fn map_start_options(
1114 &self,
1115 flow: &crate::registry::RegisteredFlow,
1116 options: &StartFlowOptions,
1117 flow_timeout_seconds: i32,
1118 flow_timeout_policy: i32,
1119 ) -> SdkResult<FlowStartOptions> {
1120 let attributes = options
1121 .attributes
1122 .iter()
1123 .map(|attribute| {
1124 let definition =
1125 flow.persistence
1126 .get(&attribute.attribute_name)
1127 .ok_or_else(|| {
1128 invalid(format!(
1129 "initial Attribute is not registered: {}",
1130 attribute.key
1131 ))
1132 })?;
1133 attribute.encode_attribute_write(definition)
1134 })
1135 .collect::<SdkResult<Vec<_>>>()?;
1136 if let Some(write) = attributes.iter().find(|write| write.is_hard_delete_on_null) {
1137 self.require_hard_delete_attribute_support("start_flow", &write.key)?;
1138 }
1139 let config = options.config_override.as_ref();
1140 let flow_config_override = Some(map_pull_flow_config(config)?);
1141 #[expect(
1142 clippy::needless_update,
1143 reason = "The test preserves forward-compatible construction of SDK options."
1144 )]
1145 let flow_start_options = FlowStartOptions {
1146 id_reuse_policy: match options.id_reuse_policy {
1147 IdReusePolicy::Default => ProtoIdReusePolicy::Unspecified,
1148 IdReusePolicy::AllowIfPreviousFailed => {
1149 ProtoIdReusePolicy::AllowIfPreviousExistsAbnormally
1150 }
1151 IdReusePolicy::AllowIfNotRunning => ProtoIdReusePolicy::AllowIfNoRunning,
1152 IdReusePolicy::Disallow => ProtoIdReusePolicy::DisallowReuse,
1153 IdReusePolicy::TerminateIfRunning => ProtoIdReusePolicy::AllowTerminateIfRunning,
1154 } as i32,
1155 flow_start_delay_seconds: optional_seconds(options.start_delay)?,
1156 retry_policy: options
1157 .retry_policy
1158 .clone()
1159 .map(map_flow_retry)
1160 .transpose()?,
1161 attributes,
1162 flow_config_override,
1163 flow_already_started_options: Some(FlowAlreadyStartedOptions {
1164 ignore_already_started_error: options.ignore_already_started,
1165 }),
1166 timeout_handler_options: map_flow_timeout_handler_options(
1167 flow,
1168 flow_timeout_seconds,
1169 flow_timeout_policy,
1170 options.timeout_handler_options.as_ref(),
1171 )
1172 .map_err(sdk_handler_error)?,
1173 ..Default::default()
1174 };
1175 Ok(flow_start_options)
1176 }
1177
1178 fn require_hard_delete_attribute_support(
1179 &self,
1180 operation: &'static str,
1181 attribute_key: &str,
1182 ) -> SdkResult<()> {
1183 if !self.negotiated_protocol_version.has_negotiated() {
1184 let mut service = self.service.clone();
1185 let server_info = self
1186 .runtime
1187 .block_on(service.get_server_info(()))
1188 .map_err(|status| SdkError::Service {
1189 service: ServiceError::local(operation, server_info_request_error(status)),
1190 })?
1191 .into_inner();
1192 let negotiated_protocol_version =
1193 negotiate_server_protocol(&server_info).map_err(|error| SdkError::Service {
1194 service: ServiceError::local(operation, error),
1195 })?;
1196 self.negotiated_protocol_version
1197 .store(negotiated_protocol_version);
1198 }
1199 self.negotiated_protocol_version
1200 .require_hard_delete_attribute_support(attribute_key)
1201 .map_err(|error| SdkError::Service {
1202 service: ServiceError::local(operation, error),
1203 })
1204 }
1205
1206 fn map_flow_config(&self, config: Option<&FlowConfig>) -> SdkResult<ProtoFlowConfig> {
1207 map_flow_config(config)
1208 }
1209
1210 fn call_empty<Response, Future, Call>(
1211 &self,
1212 operation: &'static str,
1213 flow_id: Option<&str>,
1214 requirement: FlowTargetRequirement,
1215 call: Call,
1216 ) -> SdkResult<()>
1217 where
1218 Call: FnOnce(FlowServiceClient<CredentialChannel>) -> Future,
1219 Future: std::future::Future<Output = Result<tonic::Response<Response>, tonic::Status>>,
1220 {
1221 let service = self.service.clone();
1222 self.runtime
1223 .block_on(call(service))
1224 .map(|_| ())
1225 .map_err(|status| SdkError::from_status(status, operation, flow_id, requirement))
1226 }
1227}
1228
1229pub(crate) fn map_pull_flow_config(config: Option<&FlowConfig>) -> SdkResult<ProtoFlowConfig> {
1230 Ok(ProtoFlowConfig {
1231 worker_target: Some(ProtoWorkerTarget { is_pull_mode: true }),
1232 ..map_flow_config(config)?
1233 })
1234}
1235
1236pub(crate) fn map_flow_config(config: Option<&FlowConfig>) -> SdkResult<ProtoFlowConfig> {
1237 Ok(ProtoFlowConfig {
1238 active_step_search_mode: config
1239 .and_then(|config| config.active_step_search_mode)
1240 .map(|mode| match mode {
1241 ActiveStepSearchMode::All => ProtoSearchMode::EnabledForAll as i32,
1242 ActiveStepSearchMode::WithWaitFor => {
1243 ProtoSearchMode::EnabledForStepsWithWaitFor as i32
1244 }
1245 ActiveStepSearchMode::Disabled => ProtoSearchMode::Disabled as i32,
1246 }),
1247 continue_as_new_threshold: config
1248 .and_then(|config| config.continue_as_new_threshold)
1249 .map(i32::try_from)
1250 .transpose()
1251 .map_err(|_| invalid("continue-as-new threshold exceeds int32"))?,
1252 continue_as_new_page_size_in_bytes: config
1253 .and_then(|config| config.continue_as_new_page_size_bytes)
1254 .map(i32::try_from)
1255 .transpose()
1256 .map_err(|_| invalid("continue-as-new page size exceeds int32"))?,
1257 step_durability: config
1258 .and_then(|config| config.step_durability)
1259 .map(|durability| {
1260 (match durability {
1261 StepDurability::Default => ProtoStepDurability::Unspecified,
1262 StepDurability::Sync => ProtoStepDurability::Sync,
1263 StepDurability::Async => ProtoStepDurability::Async,
1264 }) as i32
1265 }),
1266 worker_target: None,
1267 attribute_store_names: config.and_then(|config| {
1268 config
1269 .attribute_store_names
1270 .as_ref()
1271 .map(|names| AttributeStoreNames {
1272 names: names.clone(),
1273 })
1274 }),
1275 })
1276}
1277
1278pub(crate) fn map_flow_retry(retry: RetryPolicy) -> SdkResult<FlowRetryPolicy> {
1279 Ok(FlowRetryPolicy {
1280 initial_interval_seconds: optional_seconds(retry.initial_interval)?,
1281 backoff_coefficient: retry.backoff_coefficient.unwrap_or_default() as f32,
1282 maximum_interval_seconds: optional_seconds(retry.maximum_interval)?,
1283 maximum_attempts: i32::try_from(retry.maximum_attempts.unwrap_or_default())
1284 .map_err(|_| invalid("maximum attempts exceed int32"))?,
1285 })
1286}
1287
1288pub(crate) fn map_flow_status(status: i32) -> SdkResult<FlowStatus> {
1289 match ProtoFlowStatus::try_from(status).ok() {
1290 Some(ProtoFlowStatus::Running) => Ok(FlowStatus::Running),
1291 Some(ProtoFlowStatus::Completed) => Ok(FlowStatus::Completed),
1292 Some(ProtoFlowStatus::Failed) => Ok(FlowStatus::Failed),
1293 Some(ProtoFlowStatus::ServerSideTimeoutInternalOnly) => {
1294 Ok(FlowStatus::ServerSideTimeoutInternalOnly)
1295 }
1296 Some(ProtoFlowStatus::Terminated) => Ok(FlowStatus::Terminated),
1297 Some(ProtoFlowStatus::Canceled) => Ok(FlowStatus::Canceled),
1298 Some(ProtoFlowStatus::ContinuedAsNew) => Ok(FlowStatus::ContinuedAsNew),
1299 _ => Err(invalid(format!("unknown Flow status {status}"))),
1300 }
1301}
1302
1303pub(crate) fn map_flow_error_type(error_type: i32) -> Option<FlowErrorType> {
1304 use dex_protocol::dex::FlowErrorType as ProtoFlowErrorType;
1305
1306 match ProtoFlowErrorType::try_from(error_type).ok() {
1307 Some(ProtoFlowErrorType::StepDecisionFailingFlow) => {
1308 Some(FlowErrorType::StepDecisionFailed)
1309 }
1310 Some(ProtoFlowErrorType::ClientApiFailingFlow) => Some(FlowErrorType::ClientApiFailed),
1311 Some(ProtoFlowErrorType::WorkerApiFail) => Some(FlowErrorType::WorkerApiFailed),
1312 Some(ProtoFlowErrorType::InvalidUserFlowCode) => Some(FlowErrorType::InvalidUserFlowCode),
1313 Some(ProtoFlowErrorType::Internal) => Some(FlowErrorType::Internal),
1314 Some(ProtoFlowErrorType::FlowTimeout) => Some(FlowErrorType::FlowTimeout),
1315 _ => None,
1316 }
1317}
1318
1319pub(crate) fn map_flow_timeout_policy(
1320 flow: &crate::registry::RegisteredFlow,
1321 timeout_seconds: i32,
1322 policy: FlowTimeoutPolicy,
1323) -> SdkResult<i32> {
1324 if timeout_seconds == 0 {
1325 if policy != FlowTimeoutPolicy::Default {
1326 return Err(invalid("Flow timeout policy requires a positive timeout"));
1327 }
1328 return Ok(ProtoFlowTimeoutPolicy::Unspecified as i32);
1329 }
1330 let resolved = match policy {
1331 FlowTimeoutPolicy::Default if flow.handler.has_timeout_handler() => {
1332 FlowTimeoutPolicy::Handler
1333 }
1334 FlowTimeoutPolicy::Default => FlowTimeoutPolicy::Fail,
1335 configured => configured,
1336 };
1337 if resolved == FlowTimeoutPolicy::Handler && !flow.handler.has_timeout_handler() {
1338 return Err(invalid(format!(
1339 "Flow {} has no timeout handler",
1340 flow.name
1341 )));
1342 }
1343 Ok(match resolved {
1344 FlowTimeoutPolicy::Default => ProtoFlowTimeoutPolicy::Unspecified,
1345 FlowTimeoutPolicy::Fail => ProtoFlowTimeoutPolicy::Fail,
1346 FlowTimeoutPolicy::Cancel => ProtoFlowTimeoutPolicy::Cancel,
1347 FlowTimeoutPolicy::Handler => ProtoFlowTimeoutPolicy::Handler,
1348 } as i32)
1349}
1350
1351fn timestamp(timestamp: prost_types::Timestamp) -> SystemTime {
1352 if timestamp.seconds >= 0 {
1353 SystemTime::UNIX_EPOCH
1354 + Duration::from_secs(timestamp.seconds as u64)
1355 + Duration::from_nanos(timestamp.nanos.max(0) as u64)
1356 } else {
1357 SystemTime::UNIX_EPOCH - Duration::from_secs(timestamp.seconds.unsigned_abs())
1358 }
1359}
1360
1361fn rfc3339(time: SystemTime) -> SdkResult<String> {
1362 let duration = time
1363 .duration_since(SystemTime::UNIX_EPOCH)
1364 .map_err(service_error)?;
1365 let days = (duration.as_secs() / 86_400) as i64;
1366 let seconds = duration.as_secs() % 86_400;
1367 let (year, month, day) = civil_from_days(days);
1368 let hour = seconds / 3_600;
1369 let minute = seconds % 3_600 / 60;
1370 let second = seconds % 60;
1371 Ok(format!(
1372 "{year:04}-{month:02}-{day:02}T{hour:02}:{minute:02}:{second:02}Z"
1373 ))
1374}
1375
1376fn civil_from_days(days: i64) -> (i64, i64, i64) {
1377 let days = days + 719_468;
1378 let era = if days >= 0 { days } else { days - 146_096 } / 146_097;
1379 let day_of_era = days - era * 146_097;
1380 let year_of_era =
1381 (day_of_era - day_of_era / 1_460 + day_of_era / 36_524 - day_of_era / 146_096) / 365;
1382 let mut year = year_of_era + era * 400;
1383 let day_of_year = day_of_era - (365 * year_of_era + year_of_era / 4 - year_of_era / 100);
1384 let month_prime = (5 * day_of_year + 2) / 153;
1385 let day = day_of_year - (153 * month_prime + 2) / 5 + 1;
1386 let month = month_prime + if month_prime < 10 { 3 } else { -9 };
1387 year += i64::from(month <= 2);
1388 (year, month, day)
1389}
1390
1391fn map_load_name(name: &str, instance: Option<&str>) -> String {
1392 match instance {
1393 Some(instance) => crate::registry::physical_name(name, instance),
1394 None => format!("{name}/"),
1395 }
1396}
1397
1398fn require_rpc_invoke_map_load<'a>(
1399 flow: &crate::registry::RegisteredFlow,
1400 name: &str,
1401 instance: Option<&'a str>,
1402 expected_kind: crate::persistence::PersistenceKind,
1403) -> SdkResult<&'a str> {
1404 require_rpc_invoke_definition(flow, name, expected_kind)?;
1405 let instance = instance.ok_or_else(|| SdkError::FlowDefinition {
1406 message: "RpcInvokeOptions requires an exact map instance".to_string(),
1407 })?;
1408 crate::registry::validate_map_instance(instance).map_err(invalid)?;
1409 Ok(instance)
1410}
1411
1412fn require_rpc_invoke_definition(
1413 flow: &crate::registry::RegisteredFlow,
1414 name: &str,
1415 expected_kind: crate::persistence::PersistenceKind,
1416) -> SdkResult<()> {
1417 flow.persistence
1418 .get(name)
1419 .filter(|definition| definition.kind == expected_kind)
1420 .ok_or_else(|| SdkError::FlowDefinition {
1421 message: format!(
1422 "Flow {} does not register the selected map {name}",
1423 flow.name
1424 ),
1425 })?;
1426 Ok(())
1427}
1428
1429fn optional_seconds(duration: Option<Duration>) -> SdkResult<i32> {
1430 duration
1431 .map(seconds32)
1432 .transpose()
1433 .map(Option::unwrap_or_default)
1434}
1435
1436pub(crate) fn seconds32(duration: Duration) -> SdkResult<i32> {
1437 if duration.subsec_nanos() != 0 {
1438 return Err(invalid("Duration must use whole seconds"));
1439 }
1440 i32::try_from(duration.as_secs()).map_err(|_| invalid("Duration exceeds int32"))
1441}
1442
1443fn endpoint_address(address: &str) -> String {
1444 if address.contains("://") {
1445 address.to_string()
1446 } else {
1447 format!("http://{address}")
1448 }
1449}
1450
1451fn require_name(value: &str, kind: &str) -> SdkResult<()> {
1452 if value.is_empty() {
1453 Err(invalid(format!("{kind} is required")))
1454 } else {
1455 Ok(())
1456 }
1457}
1458
1459fn map_physical_name(name: &str, instance: &str) -> SdkResult<String> {
1460 crate::registry::validate_map_instance(instance).map_err(invalid)?;
1461 Ok(crate::registry::physical_name(name, instance))
1462}
1463
1464fn require_waitable_attribute(attribute_name: &str, is_hard_delete_on_null: bool) -> SdkResult<()> {
1465 if is_hard_delete_on_null {
1466 return Err(invalid(format!(
1467 "wait_for_attribute_match does not support hard-delete Attribute {attribute_name}"
1468 )));
1469 }
1470 Ok(())
1471}
1472
1473fn sdk_handler_error(error: impl std::fmt::Display) -> SdkError {
1474 invalid(error.to_string())
1475}
1476
1477pub(crate) fn invalid(message: impl Into<String>) -> SdkError {
1478 SdkError::InvalidArgument {
1479 message: message.into(),
1480 }
1481}
1482
1483fn service_error(error: impl std::fmt::Display) -> SdkError {
1484 SdkError::Service {
1485 service: ServiceError::local("client", error.to_string()),
1486 }
1487}
1488
1489#[cfg(test)]
1490mod tests {
1491 use super::map_flow_config;
1492 use crate::FlowConfig;
1493
1494 #[test]
1495 fn attribute_store_names_preserve_protocol_presence() {
1496 let absent = map_flow_config(Some(&FlowConfig::new())).expect("map absent config");
1497 let named = map_flow_config(Some(
1498 &FlowConfig::new()
1499 .attribute_store_names(vec!["profiles".to_owned(), "audit".to_owned()]),
1500 ))
1501 .expect("map named config");
1502 let disabled = map_flow_config(Some(&FlowConfig::new().attribute_store_names(vec![])))
1503 .expect("map disabled config");
1504
1505 assert_eq!(absent.attribute_store_names, None);
1506 assert_eq!(
1507 named.attribute_store_names.expect("named stores").names,
1508 vec!["profiles", "audit"]
1509 );
1510 assert_eq!(
1511 disabled
1512 .attribute_store_names
1513 .expect("disabled stores")
1514 .names,
1515 Vec::<String>::new()
1516 );
1517 }
1518
1519 mod hard_delete {
1520 use std::convert::Infallible;
1521 use std::sync::atomic::{AtomicU32, Ordering};
1522 use std::sync::{Arc, LazyLock, Mutex};
1523 use std::task::{Context, Poll};
1524
1525 use dex_protocol::dex::ServerInfo;
1526 use tokio::net::TcpListener;
1527 use tokio::sync::oneshot;
1528 use tokio_stream::wrappers::TcpListenerStream;
1529 use tonic::codegen::{BoxFuture, Service, http};
1530 use tonic::server::{NamedService, UnaryService};
1531 use tonic::transport::Server;
1532 use tonic::{Request, Response, Status};
1533
1534 use crate::{
1535 Attribute, AttributeMap, AttributeMatch, BlobCache, BlobCacheConfig, Client,
1536 ClientOptions, Flow, PersistenceSchema, Registry, SdkError, StartFlowOptions,
1537 WaitForAttributeOptions,
1538 };
1539
1540 static API_KEY: LazyLock<Attribute<String>> =
1541 LazyLock::new(|| Attribute::new("api-key").hard_delete_on_null());
1542 static TENANT_KEYS: LazyLock<AttributeMap<String>> =
1543 LazyLock::new(|| AttributeMap::new("tenant-keys").hard_delete_on_null());
1544 static PLAIN_NOTE: LazyLock<Attribute<String>> =
1545 LazyLock::new(|| Attribute::new("plain-note"));
1546
1547 struct HardDeleteClientFlow;
1548
1549 impl Flow for HardDeleteClientFlow {
1550 type StartInput = ();
1551
1552 fn persistence(&self) -> PersistenceSchema {
1553 PersistenceSchema::new()
1554 .attribute(&API_KEY)
1555 .attribute_map(&TENANT_KEYS)
1556 .attribute(&PLAIN_NOTE)
1557 }
1558 }
1559
1560 #[derive(Clone)]
1561 struct OldFlowService {
1562 paths: Arc<Mutex<Vec<String>>>,
1563 current_protocol_version: Arc<AtomicU32>,
1564 }
1565
1566 impl NamedService for OldFlowService {
1567 const NAME: &'static str = "dex.FlowService";
1568 }
1569
1570 impl<B> Service<http::Request<B>> for OldFlowService
1571 where
1572 B: tonic::codegen::Body + Send + 'static,
1573 B::Error: Into<tonic::codegen::StdError> + Send + 'static,
1574 {
1575 type Response = http::Response<tonic::body::Body>;
1576 type Error = Infallible;
1577 type Future = BoxFuture<Self::Response, Self::Error>;
1578
1579 fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
1580 Poll::Ready(Ok(()))
1581 }
1582
1583 fn call(&mut self, request: http::Request<B>) -> Self::Future {
1584 let path = request.uri().path().to_string();
1585 self.paths.lock().expect("paths lock").push(path.clone());
1586 let server_info = GetServerInfo {
1587 current_protocol_version: self.current_protocol_version.load(Ordering::SeqCst),
1588 };
1589 Box::pin(async move {
1590 if path == "/dex.FlowService/GetServerInfo" {
1591 let mut grpc = tonic::server::Grpc::new(tonic_prost::ProstCodec::default());
1592 Ok(grpc.unary(server_info, request).await)
1593 } else {
1594 Ok(Status::unimplemented("fake FlowService").into_http())
1595 }
1596 })
1597 }
1598 }
1599
1600 struct GetServerInfo {
1601 current_protocol_version: u32,
1602 }
1603
1604 impl UnaryService<()> for GetServerInfo {
1605 type Response = ServerInfo;
1606 type Future = BoxFuture<Response<Self::Response>, Status>;
1607
1608 fn call(&mut self, _: Request<()>) -> Self::Future {
1609 let current_protocol_version = self.current_protocol_version;
1610 Box::pin(async move {
1611 Ok(Response::new(ServerInfo {
1612 server_version: "fake".to_string(),
1613 minimum_supported_protocol_version: 1,
1614 current_protocol_version,
1615 server_engine: 0,
1616 }))
1617 })
1618 }
1619 }
1620
1621 fn client_for(address: &str, cache_directory: &tempfile::TempDir) -> Client {
1622 let cache = Arc::new(
1623 BlobCache::open(
1624 BlobCacheConfig::new(cache_directory.path(), 1024 * 1024, 0)
1625 .expect("blob cache config"),
1626 )
1627 .expect("blob cache"),
1628 );
1629 Client::try_new(
1630 Registry::new()
1631 .register(HardDeleteClientFlow)
1632 .expect("register hard-delete Flow"),
1633 cache,
1634 ClientOptions::new().server_address(address),
1635 )
1636 .expect("Client")
1637 }
1638
1639 #[test]
1640 fn wait_for_attribute_rejects_hard_delete_attributes_before_contacting_dex() {
1641 let cache_directory = tempfile::tempdir().expect("blob cache directory");
1642 let client = client_for("127.0.0.1:1", &cache_directory);
1643 let single = client.wait_for_attribute_match(
1644 "flow-1",
1645 &API_KEY,
1646 AttributeMatch::equal_to("key-1".to_string()),
1647 WaitForAttributeOptions::new(),
1648 );
1649 let map = client.wait_for_attribute_map_instance_match(
1650 "flow-1",
1651 &TENANT_KEYS,
1652 "tenant-a",
1653 AttributeMatch::equal_to("key-a".to_string()),
1654 WaitForAttributeOptions::new(),
1655 );
1656 for result in [single, map] {
1657 let Err(SdkError::InvalidArgument { message }) = result else {
1658 panic!("expected a local InvalidArgument");
1659 };
1660 assert!(message.contains("does not support hard-delete Attribute"));
1661 }
1662 }
1663
1664 #[test]
1665 fn start_flow_negotiates_once_and_rejects_hard_delete_writes_to_old_servers() {
1666 let runtime = tokio::runtime::Runtime::new().expect("server runtime");
1667 let listener = runtime
1668 .block_on(TcpListener::bind("127.0.0.1:0"))
1669 .expect("bind fake FlowService");
1670 let address = listener.local_addr().expect("fake server address");
1671 let paths = Arc::new(Mutex::new(Vec::new()));
1672 let current_protocol_version = Arc::new(AtomicU32::new(2));
1673 let (shutdown, stopped) = oneshot::channel::<()>();
1674 runtime.spawn(
1675 Server::builder()
1676 .add_service(OldFlowService {
1677 paths: Arc::clone(&paths),
1678 current_protocol_version: Arc::clone(¤t_protocol_version),
1679 })
1680 .serve_with_incoming_shutdown(TcpListenerStream::new(listener), async {
1681 stopped.await.expect("FlowService shutdown signal");
1682 }),
1683 );
1684 let cache_directory = tempfile::tempdir().expect("blob cache directory");
1685 let client = client_for(&address.to_string(), &cache_directory);
1686
1687 for options in [
1688 StartFlowOptions::new().initial_attribute(&API_KEY, "key-1".to_string()),
1689 StartFlowOptions::new().initial_attribute_map(
1690 &TENANT_KEYS,
1691 "tenant-a",
1692 "key-a".to_string(),
1693 ),
1694 ] {
1695 let error = client
1696 .start_flow_with_options(&HardDeleteClientFlow, "flow-1", (), options)
1697 .expect_err("old Server rejects hard-delete writes");
1698 let SdkError::Service { service } = error else {
1699 panic!("expected a Service error, got {error}");
1700 };
1701 assert!(service.detail().contains("protocol version 3 or later"));
1702 assert!(
1703 service
1704 .detail()
1705 .contains("negotiated protocol version is 2")
1706 );
1707 }
1708 assert_eq!(
1709 vec!["/dex.FlowService/GetServerInfo".to_string()],
1710 *paths.lock().expect("paths lock")
1711 );
1712
1713 current_protocol_version.store(3, Ordering::SeqCst);
1714 let error = client
1715 .start_flow_with_options(
1716 &HardDeleteClientFlow,
1717 "flow-1",
1718 (),
1719 StartFlowOptions::new().initial_attribute(&API_KEY, "key-1".to_string()),
1720 )
1721 .expect_err("cached negotiation still rejects the write");
1722 assert!(
1723 error
1724 .to_string()
1725 .contains("negotiated protocol version is 2")
1726 );
1727 client
1728 .start_flow_with_options(
1729 &HardDeleteClientFlow,
1730 "flow-1",
1731 (),
1732 StartFlowOptions::new().initial_attribute(&PLAIN_NOTE, "note".to_string()),
1733 )
1734 .expect_err("fake StartFlow is unimplemented");
1735 assert_eq!(
1736 vec![
1737 "/dex.FlowService/GetServerInfo".to_string(),
1738 "/dex.FlowService/StartFlow".to_string(),
1739 ],
1740 *paths.lock().expect("paths lock")
1741 );
1742
1743 let fresh_client = client_for(&address.to_string(), &cache_directory);
1744 fresh_client
1745 .start_flow_with_options(
1746 &HardDeleteClientFlow,
1747 "flow-1",
1748 (),
1749 StartFlowOptions::new().initial_attribute(&API_KEY, "key-1".to_string()),
1750 )
1751 .expect_err("fake StartFlow is unimplemented");
1752 assert_eq!(
1753 Some(&"/dex.FlowService/StartFlow".to_string()),
1754 paths.lock().expect("paths lock").last()
1755 );
1756 shutdown.send(()).expect("stop fake FlowService");
1757 }
1758 }
1759}
1760
1761#[cfg(test)]
1762#[path = "client_pull_routing_tests.rs"]
1763mod pull_routing_tests;