Skip to main content

floe_core/state/
mod.rs

1use std::collections::{BTreeMap, HashSet};
2use std::fs;
3use std::path::{Path, PathBuf};
4
5use serde::{Deserialize, Serialize};
6use tempfile::NamedTempFile;
7
8use crate::config::{
9    ConfigBase, EntityConfig, IncrementalMode, ResolvedPath, RootConfig, StorageResolver,
10};
11use crate::io::storage::{
12    extensions, local::LocalClient, CloudClient, ConditionalWrite, StorageClient, StoredObject,
13};
14use crate::{ConfigError, FloeResult};
15
16pub const ENTITY_STATE_SCHEMA_V1: &str = "floe.state.file-ingest.v1";
17pub const ENTITY_STATE_SCHEMA_V2: &str = "floe.state.file-ingest.v2";
18pub const ENTITY_STATE_FILENAME: &str = "state.json";
19const STATE_CAS_RETRIES: usize = 5;
20pub const CLAIM_TTL_SECONDS: i64 = 60 * 60;
21
22#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
23pub struct EntityState {
24    pub schema: String,
25    pub entity: String,
26    pub updated_at: Option<String>,
27    #[serde(default)]
28    pub files: BTreeMap<String, EntityFileState>,
29    #[serde(default)]
30    pub claims: BTreeMap<String, EntityFileClaim>,
31}
32
33impl EntityState {
34    pub fn new(entity: impl Into<String>) -> Self {
35        Self {
36            schema: ENTITY_STATE_SCHEMA_V2.to_string(),
37            entity: entity.into(),
38            updated_at: None,
39            files: BTreeMap::new(),
40            claims: BTreeMap::new(),
41        }
42    }
43}
44
45#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
46pub struct EntityFileState {
47    pub processed_at: String,
48    pub size: Option<u64>,
49    pub mtime: Option<String>,
50}
51
52#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
53pub struct EntityFileClaim {
54    pub run_id: String,
55    pub acquired_at: String,
56    pub expires_at: String,
57    pub size: Option<u64>,
58    pub mtime: Option<String>,
59}
60
61#[derive(Debug, Clone)]
62pub struct EntityStateInspection {
63    pub entity_name: String,
64    pub incremental_mode: IncrementalMode,
65    pub path: ResolvedPath,
66    pub state: Option<EntityState>,
67}
68
69pub fn resolve_entity_state_path(
70    resolver: &StorageResolver,
71    entity: &EntityConfig,
72) -> FloeResult<ResolvedPath> {
73    if let Some(state) = &entity.state {
74        if let Some(path) = state.path.as_deref() {
75            let resolved = if is_remote_uri(path) {
76                resolver.resolve_path(
77                    &entity.name,
78                    "entity.state.path",
79                    entity.source.storage.as_deref(),
80                    path,
81                )?
82            } else {
83                resolver.resolve_local_path(path)?
84            };
85            return Ok(resolved);
86        }
87    }
88
89    let resolved_source = resolver.resolve_path(
90        &entity.name,
91        "entity.source.path",
92        entity.source.storage.as_deref(),
93        &entity.source.path,
94    )?;
95    let source_root = derive_source_root(
96        &entity.source.path,
97        &entity.source.format,
98        resolved_source.local_path.as_deref(),
99    );
100    let default_path = join_state_path(&source_root, &entity.name);
101    let resolved = resolver.resolve_path(
102        &entity.name,
103        "entity.state.path",
104        entity.source.storage.as_deref(),
105        &default_path,
106    )?;
107    Ok(resolved)
108}
109
110pub fn read_entity_state(path: &Path) -> FloeResult<Option<EntityState>> {
111    if !path.exists() {
112        return Ok(None);
113    }
114    let payload = fs::read_to_string(path)?;
115    let state = parse_entity_state(payload.as_bytes())?;
116    Ok(Some(state))
117}
118
119fn parse_entity_state(payload: &[u8]) -> FloeResult<EntityState> {
120    let mut state: EntityState = serde_json::from_slice(payload)?;
121    if state.schema == ENTITY_STATE_SCHEMA_V1 {
122        state.schema = ENTITY_STATE_SCHEMA_V2.to_string();
123        state.claims.clear();
124    }
125    Ok(state)
126}
127
128#[derive(Debug, Clone)]
129pub enum EntityStateTarget {
130    Local { path: PathBuf, uri: String },
131    Remote { storage: String, uri: String },
132}
133
134#[derive(Debug, Clone)]
135pub struct LoadedEntityState {
136    pub target: EntityStateTarget,
137    pub state: EntityState,
138    pub version: Option<String>,
139    pub existed: bool,
140}
141
142#[derive(Debug, Clone)]
143pub struct ClaimedEntityState {
144    pub target: EntityStateTarget,
145    pub state: EntityState,
146    pub version: Option<String>,
147}
148
149#[derive(Debug, Clone)]
150pub struct EntityStateClaimOutcome {
151    pub pending_inputs: Vec<crate::io::format::InputFile>,
152    pub claimed_state: Option<ClaimedEntityState>,
153    pub active_claims: Vec<String>,
154    pub already_processed: Vec<(crate::io::format::InputFile, EntityFileState)>,
155}
156
157pub fn claim_entity_inputs(
158    resolver: &StorageResolver,
159    cloud: &mut CloudClient,
160    entity: &EntityConfig,
161    run_id: &str,
162    input_files: Vec<crate::io::format::InputFile>,
163) -> FloeResult<EntityStateClaimOutcome> {
164    if input_files.is_empty() {
165        return Ok(EntityStateClaimOutcome {
166            pending_inputs: Vec::new(),
167            claimed_state: None,
168            active_claims: Vec::new(),
169            already_processed: Vec::new(),
170        });
171    }
172
173    for _ in 0..STATE_CAS_RETRIES {
174        let mut loaded = load_entity_state(resolver, cloud, entity)?;
175        remove_expired_claims(&mut loaded.state);
176        let mut pending_inputs = Vec::new();
177        let mut active_claims = Vec::new();
178        let mut already_processed = Vec::new();
179        let acquired_at = now_rfc3339();
180        let expires_at = rfc3339_after_seconds(CLAIM_TTL_SECONDS);
181
182        for input_file in &input_files {
183            if let Some(recorded) = loaded.state.files.get(&input_file.source_uri) {
184                already_processed.push((input_file.clone(), recorded.clone()));
185                continue;
186            }
187            match loaded.state.claims.get(&input_file.source_uri) {
188                Some(_) => {
189                    active_claims.push(input_file.source_uri.clone());
190                }
191                None => {
192                    loaded.state.claims.insert(
193                        input_file.source_uri.clone(),
194                        EntityFileClaim {
195                            run_id: run_id.to_string(),
196                            acquired_at: acquired_at.clone(),
197                            expires_at: expires_at.clone(),
198                            size: input_file.source_size,
199                            mtime: input_file.source_mtime.clone(),
200                        },
201                    );
202                    pending_inputs.push(input_file.clone());
203                }
204            }
205        }
206
207        if pending_inputs.is_empty() {
208            if active_claims.is_empty() {
209                let _ = persist_loaded_state(cloud, resolver, &loaded)?;
210            }
211            return Ok(EntityStateClaimOutcome {
212                pending_inputs,
213                claimed_state: None,
214                active_claims,
215                already_processed,
216            });
217        }
218
219        loaded.state.schema = ENTITY_STATE_SCHEMA_V2.to_string();
220        loaded.state.updated_at = Some(acquired_at);
221        match persist_loaded_state(cloud, resolver, &loaded)? {
222            Some(version) => {
223                return Ok(EntityStateClaimOutcome {
224                    pending_inputs,
225                    claimed_state: Some(ClaimedEntityState {
226                        target: loaded.target,
227                        state: loaded.state,
228                        version,
229                    }),
230                    active_claims,
231                    already_processed,
232                });
233            }
234            None => continue,
235        }
236    }
237
238    Err(Box::new(ConfigError(format!(
239        "entity.name={} incremental state update conflicted after {STATE_CAS_RETRIES} retries",
240        entity.name
241    ))))
242}
243
244pub fn promote_claimed_entity_state(
245    resolver: &StorageResolver,
246    cloud: &mut CloudClient,
247    entity_name: &str,
248    run_id: &str,
249    claimed: &ClaimedEntityState,
250) -> FloeResult<()> {
251    mutate_claimed_state(resolver, cloud, entity_name, claimed, |state, our_uris| {
252        let processed_at = now_rfc3339();
253        let claimed_files: Vec<String> = state
254            .claims
255            .iter()
256            .filter(|(uri, claim)| claim.run_id == run_id && our_uris.contains(*uri))
257            .map(|(source_uri, _)| source_uri.clone())
258            .collect();
259        for source_uri in claimed_files {
260            if let Some(claim) = state.claims.remove(&source_uri) {
261                state.files.insert(
262                    source_uri,
263                    EntityFileState {
264                        processed_at: processed_at.clone(),
265                        size: claim.size,
266                        mtime: claim.mtime,
267                    },
268                );
269            }
270        }
271        state.updated_at = Some(processed_at);
272    })
273}
274
275pub fn release_claimed_entity_state(
276    resolver: &StorageResolver,
277    cloud: &mut CloudClient,
278    entity_name: &str,
279    run_id: &str,
280    claimed: &ClaimedEntityState,
281) -> FloeResult<()> {
282    mutate_claimed_state(resolver, cloud, entity_name, claimed, |state, our_uris| {
283        state
284            .claims
285            .retain(|uri, claim| !(claim.run_id == run_id && our_uris.contains(uri)));
286        state.updated_at = Some(now_rfc3339());
287    })
288}
289
290pub fn renew_claimed_entity_state(
291    resolver: &StorageResolver,
292    cloud: &mut CloudClient,
293    entity_name: &str,
294    run_id: &str,
295    claimed: &ClaimedEntityState,
296) -> FloeResult<()> {
297    mutate_claimed_state(resolver, cloud, entity_name, claimed, |state, our_uris| {
298        let now = now_rfc3339();
299        let expires_at = rfc3339_after_seconds(CLAIM_TTL_SECONDS);
300        for (uri, claim) in state.claims.iter_mut() {
301            if claim.run_id == run_id && our_uris.contains(uri) {
302                claim.expires_at = expires_at.clone();
303            }
304        }
305        state.updated_at = Some(now);
306    })
307}
308
309fn mutate_claimed_state(
310    resolver: &StorageResolver,
311    cloud: &mut CloudClient,
312    entity_name: &str,
313    claimed: &ClaimedEntityState,
314    mutate: impl Fn(&mut EntityState, &HashSet<String>),
315) -> FloeResult<()> {
316    let our_uris: HashSet<String> = claimed.state.claims.keys().cloned().collect();
317    for attempt in 0..STATE_CAS_RETRIES {
318        let mut loaded = if attempt == 0 {
319            LoadedEntityState {
320                target: claimed.target.clone(),
321                state: claimed.state.clone(),
322                version: claimed.version.clone(),
323                existed: claimed.version.is_some(),
324            }
325        } else {
326            load_target_state_with_entity_name(
327                cloud,
328                resolver,
329                entity_name,
330                claimed.target.clone(),
331            )?
332        };
333        mutate(&mut loaded.state, &our_uris);
334        loaded.state.schema = ENTITY_STATE_SCHEMA_V2.to_string();
335        let persisted = if loaded.state.files.is_empty() && loaded.state.claims.is_empty() {
336            delete_loaded_state(cloud, resolver, &loaded)?
337        } else {
338            persist_loaded_state(cloud, resolver, &loaded)?
339        };
340        if persisted.is_some() {
341            return Ok(());
342        }
343    }
344    Err(Box::new(ConfigError(format!(
345        "entity.name={} incremental state update conflicted after {STATE_CAS_RETRIES} retries",
346        entity_name
347    ))))
348}
349
350pub fn inspect_entity_state_with_base(
351    config_path: &Path,
352    config_base: ConfigBase,
353    entity_name: &str,
354) -> FloeResult<EntityStateInspection> {
355    let config = crate::load_config(config_path)?;
356    inspect_entity_state(&config, config_base, entity_name)
357}
358
359pub fn inspect_entity_state(
360    config: &RootConfig,
361    config_base: ConfigBase,
362    entity_name: &str,
363) -> FloeResult<EntityStateInspection> {
364    let (entity, path) = resolve_entity_state_target(config, config_base.clone(), entity_name)?;
365    let resolver = StorageResolver::new(config, config_base)?;
366    let target = state_target_from_resolved(&path)?;
367    let mut cloud = CloudClient::new();
368    let loaded = load_target_state_with_resolver(&mut cloud, &resolver, entity, target)?;
369    let state = loaded.existed.then_some(loaded.state);
370
371    Ok(EntityStateInspection {
372        entity_name: entity.name.clone(),
373        incremental_mode: entity.incremental_mode,
374        path,
375        state,
376    })
377}
378
379pub fn reset_entity_state_with_base(
380    config_path: &Path,
381    config_base: ConfigBase,
382    entity_name: &str,
383) -> FloeResult<bool> {
384    let config = crate::load_config(config_path)?;
385    let (entity, path) = resolve_entity_state_target(&config, config_base.clone(), entity_name)?;
386    let target = state_target_from_resolved(&path)?;
387    match target {
388        EntityStateTarget::Local { path, .. } => {
389            if path.exists() {
390                fs::remove_file(&path)?;
391                return Ok(true);
392            }
393            Ok(false)
394        }
395        EntityStateTarget::Remote { storage, uri } => {
396            let mut cloud = CloudClient::new();
397            let resolver = StorageResolver::new(&config, config_base)?;
398            let client = cloud.client_for_context(
399                &resolver,
400                &storage,
401                &format!("entity.name={}", entity.name),
402            )?;
403            let Some(object) = client.read_object(&uri)? else {
404                return Ok(false);
405            };
406            match client.delete_object_conditional(&uri, Some(&object.version))? {
407                ConditionalWrite::Written { .. } => Ok(true),
408                ConditionalWrite::Conflict => Err(Box::new(ConfigError(format!(
409                    "entity.name={} remote state changed while resetting: {}",
410                    entity.name, uri
411                )))),
412            }
413        }
414    }
415}
416
417fn resolve_entity_state_target<'a>(
418    config: &'a RootConfig,
419    config_base: ConfigBase,
420    entity_name: &str,
421) -> FloeResult<(&'a EntityConfig, ResolvedPath)> {
422    let entity = config
423        .entities
424        .iter()
425        .find(|entity| entity.name == entity_name)
426        .ok_or_else(|| {
427            Box::new(ConfigError(format!("entity not found: {entity_name}")))
428                as Box<dyn std::error::Error + Send + Sync>
429        })?;
430    let resolver = StorageResolver::new(config, config_base)?;
431    let path = resolve_entity_state_path(&resolver, entity)?;
432    Ok((entity, path))
433}
434
435pub fn write_entity_state_atomic(path: &Path, state: &EntityState) -> FloeResult<()> {
436    let parent = path.parent().ok_or_else(|| {
437        Box::new(ConfigError(format!(
438            "state path has no parent directory: {}",
439            path.display()
440        ))) as Box<dyn std::error::Error + Send + Sync>
441    })?;
442    fs::create_dir_all(parent)?;
443
444    let mut temp = NamedTempFile::new_in(parent)?;
445    serde_json::to_writer_pretty(temp.as_file_mut(), state)?;
446    temp.as_file_mut().sync_all()?;
447    temp.persist(path).map_err(|err| err.error)?;
448    Ok(())
449}
450
451fn load_entity_state(
452    resolver: &StorageResolver,
453    cloud: &mut CloudClient,
454    entity: &EntityConfig,
455) -> FloeResult<LoadedEntityState> {
456    let resolved = resolve_entity_state_path(resolver, entity)?;
457    let target = state_target_from_resolved(&resolved)?;
458    load_target_state_with_resolver(cloud, resolver, entity, target)
459}
460
461fn load_target_state_with_resolver(
462    cloud: &mut CloudClient,
463    resolver: &StorageResolver,
464    entity: &EntityConfig,
465    target: EntityStateTarget,
466) -> FloeResult<LoadedEntityState> {
467    load_target_state_with_entity_name(cloud, resolver, &entity.name, target)
468}
469
470fn load_target_state_with_entity_name(
471    cloud: &mut CloudClient,
472    resolver: &StorageResolver,
473    entity_name: &str,
474    target: EntityStateTarget,
475) -> FloeResult<LoadedEntityState> {
476    match target {
477        EntityStateTarget::Local { path, uri } => {
478            let object = LocalClient::new().read_object(&uri)?;
479            let (state, version, existed) = resolve_loaded_state(entity_name, object)?;
480            Ok(LoadedEntityState {
481                target: EntityStateTarget::Local { path, uri },
482                state,
483                version,
484                existed,
485            })
486        }
487        EntityStateTarget::Remote { storage, uri } => {
488            let client = cloud.client_for_context(
489                resolver,
490                &storage,
491                &format!("entity.name={entity_name}"),
492            )?;
493            let object = client.read_object(&uri)?;
494            let (state, version, existed) = resolve_loaded_state(entity_name, object)?;
495            Ok(LoadedEntityState {
496                target: EntityStateTarget::Remote { storage, uri },
497                state,
498                version,
499                existed,
500            })
501        }
502    }
503}
504
505fn resolve_loaded_state(
506    entity_name: &str,
507    object: Option<StoredObject>,
508) -> FloeResult<(EntityState, Option<String>, bool)> {
509    match object {
510        Some(object) => Ok((
511            validate_entity_state_name(entity_name, parse_entity_state(&object.body)?)?,
512            Some(object.version),
513            true,
514        )),
515        None => Ok((EntityState::new(entity_name), None, false)),
516    }
517}
518
519fn with_state_client<R, F>(
520    cloud: &mut CloudClient,
521    resolver: &StorageResolver,
522    target: &EntityStateTarget,
523    f: F,
524) -> FloeResult<R>
525where
526    F: FnOnce(&str, &dyn StorageClient) -> FloeResult<R>,
527{
528    match target {
529        EntityStateTarget::Local { uri, .. } => f(uri, &LocalClient::new()),
530        EntityStateTarget::Remote { uri, storage } => {
531            let client = cloud.client_for_context(resolver, storage, "entity state")?;
532            f(uri, client)
533        }
534    }
535}
536
537fn conditional_write_to_version(cw: ConditionalWrite) -> Option<Option<String>> {
538    match cw {
539        ConditionalWrite::Written { version } => Some(Some(version)),
540        ConditionalWrite::Conflict => None,
541    }
542}
543
544fn persist_loaded_state(
545    cloud: &mut CloudClient,
546    resolver: &StorageResolver,
547    loaded: &LoadedEntityState,
548) -> FloeResult<Option<Option<String>>> {
549    let mut state = loaded.state.clone();
550    state.schema = ENTITY_STATE_SCHEMA_V2.to_string();
551    let body = serde_json::to_vec_pretty(&state)?;
552    let version = loaded.version.as_deref();
553    let cw = with_state_client(cloud, resolver, &loaded.target, |uri, client| {
554        client.write_object_conditional(uri, version, &body)
555    })?;
556    Ok(conditional_write_to_version(cw))
557}
558
559fn delete_loaded_state(
560    cloud: &mut CloudClient,
561    resolver: &StorageResolver,
562    loaded: &LoadedEntityState,
563) -> FloeResult<Option<Option<String>>> {
564    let version = loaded.version.as_deref();
565    let cw = with_state_client(cloud, resolver, &loaded.target, |uri, client| {
566        client.delete_object_conditional(uri, version)
567    })?;
568    Ok(conditional_write_to_version(cw))
569}
570
571fn state_target_from_resolved(resolved: &ResolvedPath) -> FloeResult<EntityStateTarget> {
572    if let Some(path) = &resolved.local_path {
573        return Ok(EntityStateTarget::Local {
574            path: path.clone(),
575            uri: resolved.uri.clone(),
576        });
577    }
578    if is_remote_uri(&resolved.uri) {
579        return Ok(EntityStateTarget::Remote {
580            storage: resolved.storage.clone(),
581            uri: resolved.uri.clone(),
582        });
583    }
584    Err(Box::new(ConfigError(format!(
585        "state path is neither local nor supported remote: {}",
586        resolved.uri
587    ))))
588}
589
590fn remove_expired_claims(state: &mut EntityState) {
591    let now = time::OffsetDateTime::now_utc();
592    state.claims.retain(|_, claim| {
593        time::OffsetDateTime::parse(
594            &claim.expires_at,
595            &time::format_description::well_known::Rfc3339,
596        )
597        .map(|expires_at| expires_at > now)
598        .unwrap_or(false)
599    });
600}
601
602fn rfc3339_offset(seconds: i64) -> String {
603    (time::OffsetDateTime::now_utc() + time::Duration::seconds(seconds))
604        .format(&time::format_description::well_known::Rfc3339)
605        .unwrap_or_else(|_| crate::report::now_rfc3339())
606}
607
608fn now_rfc3339() -> String {
609    rfc3339_offset(0)
610}
611
612fn rfc3339_after_seconds(seconds: i64) -> String {
613    rfc3339_offset(seconds)
614}
615
616fn join_state_path(source_root: &str, entity_name: &str) -> String {
617    if source_root.is_empty() || source_root == "." {
618        format!(".floe/state/{entity_name}/{ENTITY_STATE_FILENAME}")
619    } else {
620        format!(
621            "{}/.floe/state/{entity_name}/{ENTITY_STATE_FILENAME}",
622            source_root.trim_end_matches(is_path_separator)
623        )
624    }
625}
626
627fn derive_source_root(
628    raw_path: &str,
629    source_format: &str,
630    resolved_local_path: Option<&Path>,
631) -> String {
632    let trimmed = raw_path.trim_end_matches(is_path_separator);
633    if trimmed.is_empty() {
634        return String::new();
635    }
636
637    if let Some(prefix) = prefix_before_first_glob(trimmed) {
638        if prefix.is_empty() {
639            return String::new();
640        }
641        if prefix.ends_with(is_path_separator) {
642            return prefix.trim_end_matches(is_path_separator).to_string();
643        }
644        return parent_like(prefix)
645            .unwrap_or(prefix)
646            .trim_end_matches(is_path_separator)
647            .to_string();
648    }
649
650    if let Some(path) = resolved_local_path.filter(|path| path.exists()) {
651        if path.is_dir() {
652            return trimmed.to_string();
653        }
654        if path.is_file() {
655            return parent_like(trimmed)
656                .unwrap_or(trimmed)
657                .trim_end_matches(is_path_separator)
658                .to_string();
659        }
660    }
661
662    if matches_source_file_suffix(trimmed, source_format) {
663        return parent_like(trimmed)
664            .unwrap_or(trimmed)
665            .trim_end_matches(is_path_separator)
666            .to_string();
667    }
668
669    trimmed.to_string()
670}
671
672fn prefix_before_first_glob(value: &str) -> Option<&str> {
673    let index = value.find(['*', '?', '['])?;
674    Some(&value[..index])
675}
676
677fn matches_source_file_suffix(value: &str, source_format: &str) -> bool {
678    let Some(segment) = value.rsplit(is_path_separator).next() else {
679        return false;
680    };
681    let segment = segment.to_ascii_lowercase();
682
683    extensions::suffixes_for_format(source_format)
684        .map(|suffixes| {
685            suffixes
686                .iter()
687                .any(|suffix| segment.ends_with(&suffix.to_ascii_lowercase()))
688        })
689        .unwrap_or(false)
690}
691
692fn parent_like(value: &str) -> Option<&str> {
693    value.rfind(is_path_separator).map(|index| {
694        if index == 0 {
695            &value[..1]
696        } else {
697            &value[..index]
698        }
699    })
700}
701
702fn is_path_separator(ch: char) -> bool {
703    ch == '/' || ch == '\\'
704}
705
706fn is_remote_uri(value: &str) -> bool {
707    value.starts_with("s3://") || value.starts_with("gs://") || value.starts_with("abfs://")
708}
709
710pub fn validate_entity_state(entity: &EntityConfig, state: EntityState) -> FloeResult<EntityState> {
711    validate_entity_state_name(&entity.name, state)
712}
713
714fn validate_entity_state_name(entity_name: &str, state: EntityState) -> FloeResult<EntityState> {
715    if state.schema != ENTITY_STATE_SCHEMA_V1 && state.schema != ENTITY_STATE_SCHEMA_V2 {
716        return Err(Box::new(ConfigError(format!(
717            "entity.name={} state schema mismatch: expected {} or {}, got {}",
718            entity_name, ENTITY_STATE_SCHEMA_V1, ENTITY_STATE_SCHEMA_V2, state.schema
719        ))));
720    }
721
722    if state.entity != entity_name {
723        return Err(Box::new(ConfigError(format!(
724            "entity.name={} state entity mismatch: expected {}, got {}",
725            entity_name, entity_name, state.entity
726        ))));
727    }
728
729    Ok(state)
730}