fakecloud_sagemaker/
state.rs

1//! Account-partitioned, serializable state for the Amazon SageMaker control plane.
2//!
3//! The registry is deliberately schema-light: every named resource family
4//! (models, endpoints, endpoint configs, training jobs, notebook instances,
5//! pipelines, ...) is stored in one uniform two-level map,
6//! `resources[family][id]`, where the stored value is the resource's JSON
7//! record (its persisted attributes plus any minted ARN / id / timestamps).
8//! Keeping every key a plain `String` means the snapshot never depends on the
9//! tuple-key serde adapter and new resource families need no new struct fields.
10//!
11//! Alongside the resource map are:
12//! * `tags` — resource tags keyed by ARN.
13//! * `singletons` — account-scoped singleton values keyed by a stable string
14//!   (e.g. the Service Catalog portfolio status).
15
16use std::collections::BTreeMap;
17use std::sync::Arc;
18
19use parking_lot::RwLock;
20use serde::{Deserialize, Serialize};
21use serde_json::Value;
22
23use fakecloud_core::multi_account::{AccountState, MultiAccountState};
24
25pub const SAGEMAKER_SNAPSHOT_SCHEMA_VERSION: u32 = 1;
26
27/// Per-account SageMaker control-plane state.
28#[derive(Debug, Clone, Default, Serialize, Deserialize)]
29pub struct SageMakerData {
30    /// `family -> id -> record`. The record is the resource's persisted JSON
31    /// attributes plus any minted ARN / id / timestamps.
32    #[serde(default)]
33    pub resources: BTreeMap<String, BTreeMap<String, Value>>,
34    /// Resource tags keyed by ARN.
35    #[serde(default)]
36    pub tags: BTreeMap<String, BTreeMap<String, String>>,
37    /// Account-scoped singleton values keyed by a stable string.
38    #[serde(default)]
39    pub singletons: BTreeMap<String, Value>,
40    /// Monotonic counter used to mint unique ids within an account.
41    #[serde(default)]
42    pub seq: u64,
43}
44
45impl SageMakerData {
46    /// Fetch a resource record by family + id.
47    pub fn get_resource(&self, family: &str, id: &str) -> Option<&Value> {
48        self.resources.get(family).and_then(|m| m.get(id))
49    }
50
51    /// Fetch a mutable resource record by family + id.
52    pub fn get_resource_mut(&mut self, family: &str, id: &str) -> Option<&mut Value> {
53        self.resources.get_mut(family).and_then(|m| m.get_mut(id))
54    }
55
56    /// Insert / replace a resource record.
57    pub fn put_resource(&mut self, family: &str, id: &str, record: Value) {
58        self.resources
59            .entry(family.to_string())
60            .or_default()
61            .insert(id.to_string(), record);
62    }
63
64    /// Remove a resource record, returning it if present.
65    pub fn remove_resource(&mut self, family: &str, id: &str) -> Option<Value> {
66        let removed = self.resources.get_mut(family).and_then(|m| m.remove(id));
67        if let Some(m) = self.resources.get(family) {
68            if m.is_empty() {
69                self.resources.remove(family);
70            }
71        }
72        removed
73    }
74
75    /// Resolve a caller-supplied identifier value to a stored key within a
76    /// family. Matches the direct storage key first, then falls back to a record
77    /// whose *canonical* identifier member — `{Family}Name`, `{Family}Id` or
78    /// `{Family}Arn` — equals the value (so a resource keyed by its Name can
79    /// still be described by the minted Id or ARN the create returned).
80    ///
81    /// The fallback is restricted to the family's own canonical members rather
82    /// than any suffix-matching `*Name` / `*Id` / `*Arn` member, so an unrelated
83    /// member that happens to carry an equal value (e.g. a shared `RoleArn`, or a
84    /// cross-referenced `SourceArn`) on a sibling record cannot mis-resolve to
85    /// the wrong record.
86    pub fn resolve_key(&self, family: &str, value: &str) -> Option<String> {
87        let m = self.resources.get(family)?;
88        if m.contains_key(value) {
89            return Some(value.to_string());
90        }
91        let canonical = [
92            format!("{family}Name"),
93            format!("{family}Id"),
94            format!("{family}Arn"),
95        ];
96        for (k, rec) in m {
97            if let Some(obj) = rec.as_object() {
98                if canonical
99                    .iter()
100                    .any(|cand| obj.get(cand).and_then(Value::as_str) == Some(value))
101                {
102                    return Some(k.clone());
103                }
104            }
105        }
106        None
107    }
108
109    /// All records of a family as `(id, record)` pairs, ordered by id.
110    pub fn list_resource_entries(&self, family: &str) -> Vec<(String, Value)> {
111        self.resources
112            .get(family)
113            .map(|m| m.iter().map(|(k, v)| (k.clone(), v.clone())).collect())
114            .unwrap_or_default()
115    }
116
117    /// Next unique sequence value for id minting.
118    pub fn next_seq(&mut self) -> u64 {
119        self.seq += 1;
120        self.seq
121    }
122}
123
124impl AccountState for SageMakerData {
125    fn new_for_account(_account_id: &str, _region: &str, _endpoint: &str) -> Self {
126        Self::default()
127    }
128}
129
130pub type SharedSageMakerState = Arc<RwLock<MultiAccountState<SageMakerData>>>;
131
132#[derive(Debug, Serialize, Deserialize)]
133pub struct SageMakerSnapshot {
134    pub schema_version: u32,
135    pub accounts: MultiAccountState<SageMakerData>,
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141    use serde_json::json;
142
143    #[test]
144    fn new_account_is_empty() {
145        let data = SageMakerData::new_for_account("000000000000", "us-east-1", "");
146        assert!(data.resources.is_empty());
147        assert!(data.tags.is_empty());
148    }
149
150    #[test]
151    fn resource_round_trips() {
152        let mut d = SageMakerData::default();
153        d.put_resource(
154            "Model",
155            "m1",
156            json!({"ModelName": "m1", "ModelArn": "arn:aws:sagemaker:us-east-1:0:model/m1"}),
157        );
158        assert_eq!(d.get_resource("Model", "m1").unwrap()["ModelName"], "m1");
159        assert_eq!(d.resolve_key("Model", "m1").as_deref(), Some("m1"));
160        // resolve by the minted ARN
161        assert_eq!(
162            d.resolve_key("Model", "arn:aws:sagemaker:us-east-1:0:model/m1")
163                .as_deref(),
164            Some("m1")
165        );
166        assert!(d.remove_resource("Model", "m1").is_some());
167        assert!(d.get_resource("Model", "m1").is_none());
168    }
169
170    #[test]
171    fn resolve_key_ignores_noncanonical_sibling_member() {
172        let mut d = SageMakerData::default();
173        let shared = "arn:aws:iam::0:role/shared";
174        // A record keyed by its name, carrying an unrelated RoleArn member.
175        d.put_resource("Model", "m1", json!({"ModelName": "m1", "RoleArn": shared}));
176        // Resolving by the shared RoleArn must NOT mis-match m1: RoleArn is not a
177        // canonical Model identifier member.
178        assert_eq!(d.resolve_key("Model", shared), None);
179        // The canonical ModelName still resolves.
180        assert_eq!(d.resolve_key("Model", "m1").as_deref(), Some("m1"));
181    }
182}
fakecloud_sagemaker/state.rs

fakecloud_sagemaker/
state.rs
