surreal_sync_core/
id_columns.rs

1//! Primary-key / record-ID column helpers shared across sources.
2//!
3//! - [`apply_id_column_overrides`] — optional per-table PK overrides after schema discovery
4//! - [`flatten_composite_id`] — Array IDs → Text joined with a separator (for transforms / relations)
5
6use crate::schema::{ColumnDefinition, DatabaseSchema, TableDefinition};
7use crate::types::Type;
8use crate::values::Value;
9use std::collections::HashMap;
10use thiserror::Error;
11
12/// Per-table ordered ID / primary-key column names.
13///
14/// Keys are table names; values are ordered column lists. A single-column list
15/// clears `composite_primary_key`; two or more set it.
16pub type IdColumnOverrides = HashMap<String, Vec<String>>;
17
18/// Errors from ID-column override parsing / application.
19#[derive(Debug, Error)]
20pub enum IdColumnsError {
21    #[error("{0}")]
22    Message(String),
23}
24
25impl IdColumnsError {
26    fn msg(s: impl Into<String>) -> Self {
27        Self::Message(s.into())
28    }
29}
30
31/// Parse CLI-style overrides: `table=col1,col2` (repeatable) or a bare
32/// `col1,col2` list applied to `default_table` when provided.
33///
34/// Empty entries are ignored. Column names are trimmed; empty column names error.
35pub fn parse_id_column_overrides(
36    entries: &[String],
37    default_table: Option<&str>,
38) -> Result<IdColumnOverrides, IdColumnsError> {
39    let mut out = IdColumnOverrides::new();
40    for entry in entries {
41        let entry = entry.trim();
42        if entry.is_empty() {
43            continue;
44        }
45        let (table, cols_str) = if let Some((table, cols)) = entry.split_once('=') {
46            (table.trim().to_string(), cols)
47        } else if let Some(table) = default_table {
48            (table.to_string(), entry)
49        } else {
50            return Err(IdColumnsError::msg(format!(
51                "id-columns entry '{entry}' must be table=col1,col2 \
52                 (or pass a default table for bare col1,col2 lists)"
53            )));
54        };
55        if table.is_empty() {
56            return Err(IdColumnsError::msg(format!(
57                "id-columns entry '{entry}' has an empty table name"
58            )));
59        }
60        let cols: Vec<String> = cols_str
61            .split(',')
62            .map(|s| s.trim().to_string())
63            .filter(|s| !s.is_empty())
64            .collect();
65        if cols.is_empty() {
66            return Err(IdColumnsError::msg(format!(
67                "id-columns entry '{entry}' has no column names"
68            )));
69        }
70        out.insert(table, cols);
71    }
72    Ok(out)
73}
74
75/// Apply per-table ID column overrides onto a discovered [`DatabaseSchema`].
76///
77/// For each overridden table:
78/// - Sets `primary_key` to the first column (type looked up from existing columns,
79///   or [`Type::Text`] if missing).
80/// - Moves former PK / other columns so all non-first PK columns remain in `columns`.
81/// - Sets `composite_primary_key` when `cols.len() > 1`, otherwise clears it.
82pub fn apply_id_column_overrides(
83    schema: &mut DatabaseSchema,
84    overrides: &IdColumnOverrides,
85) -> Result<(), IdColumnsError> {
86    for (table_name, cols) in overrides {
87        let table = schema.get_table_mut(table_name).ok_or_else(|| {
88            IdColumnsError::msg(format!(
89                "id-columns override for unknown table '{table_name}'"
90            ))
91        })?;
92        apply_override_to_table(table, cols)?;
93    }
94    Ok(())
95}
96
97fn apply_override_to_table(
98    table: &mut TableDefinition,
99    cols: &[String],
100) -> Result<(), IdColumnsError> {
101    if cols.is_empty() {
102        return Err(IdColumnsError::msg(format!(
103            "id-columns override for table '{}' must list at least one column",
104            table.name
105        )));
106    }
107
108    // Gather all existing columns including current PK.
109    let mut by_name: HashMap<String, ColumnDefinition> = HashMap::new();
110    by_name.insert(table.primary_key.name.clone(), table.primary_key.clone());
111    for c in &table.columns {
112        by_name.insert(c.name.clone(), c.clone());
113    }
114
115    let mut pk_defs = Vec::with_capacity(cols.len());
116    for name in cols {
117        let def = by_name
118            .remove(name)
119            .unwrap_or_else(|| ColumnDefinition::new(name.clone(), Type::Text));
120        pk_defs.push(def);
121    }
122
123    table.primary_key = pk_defs[0].clone();
124    // Remaining columns: leftover non-PK columns plus extra PK parts (so fields
125    // stay available on the row when sources exclude only the synthetic id).
126    let mut other: Vec<ColumnDefinition> = by_name.into_values().collect();
127    other.extend(pk_defs.into_iter().skip(1));
128    other.sort_by(|a, b| a.name.cmp(&b.name));
129    table.columns = other;
130
131    if cols.len() > 1 {
132        table.composite_primary_key = Some(cols.to_vec());
133    } else {
134        table.composite_primary_key = None;
135    }
136    Ok(())
137}
138
139/// Flatten a composite (Array) ID into a Text ID joined with `separator`.
140///
141/// Scalar IDs are returned unchanged. Used by the `flatten_id` transform and
142/// by relation-edge ID flattening.
143pub fn flatten_composite_id(id: Value, separator: &str) -> Value {
144    match id {
145        Value::Array { elements, .. } => {
146            let parts: Vec<String> = elements.into_iter().map(stringify_id_part).collect();
147            Value::Text(parts.join(separator))
148        }
149        other => other,
150    }
151}
152
153/// Stringify one composite-key part for colon/separator joining.
154pub fn stringify_id_part(v: Value) -> String {
155    match v {
156        Value::Int8 { value, .. } => value.to_string(),
157        Value::Int16(n) => n.to_string(),
158        Value::Int32(n) => n.to_string(),
159        Value::Int64(n) => n.to_string(),
160        Value::Text(s) => s,
161        Value::VarChar { value, .. } | Value::Char { value, .. } => value,
162        Value::Uuid(u) => u.to_string(),
163        Value::Ulid(u) => u.to_string(),
164        Value::Bool(b) => b.to_string(),
165        Value::Null => String::new(),
166        other => format!("{other:?}"),
167    }
168}
169
170/// Build a record ID from ordered field values (single → scalar, multi → Array).
171pub fn build_composite_record_id(parts: Vec<Value>) -> Value {
172    match parts.len() {
173        0 => Value::Null,
174        1 => parts.into_iter().next().unwrap(),
175        _ => Value::Array {
176            elements: parts,
177            element_type: Box::new(Type::Text),
178        },
179    }
180}
181
182#[cfg(test)]
183mod tests {
184    use super::*;
185
186    fn sample_schema() -> DatabaseSchema {
187        DatabaseSchema::new(vec![TableDefinition::new(
188            "ledger",
189            ColumnDefinition::new("a", Type::Int32),
190            vec![
191                ColumnDefinition::new("b", Type::Int32),
192                ColumnDefinition::new("amount", Type::Int64),
193            ],
194        )])
195    }
196
197    #[test]
198    fn parse_table_eq_form() {
199        let o =
200            parse_id_column_overrides(&["ledger=a,b".into(), "tags=tag_id".into()], None).unwrap();
201        assert_eq!(o["ledger"], vec!["a", "b"]);
202        assert_eq!(o["tags"], vec!["tag_id"]);
203    }
204
205    #[test]
206    fn parse_bare_with_default_table() {
207        let o = parse_id_column_overrides(&["a,b".into()], Some("ledger")).unwrap();
208        assert_eq!(o["ledger"], vec!["a", "b"]);
209    }
210
211    #[test]
212    fn apply_sets_composite() {
213        let mut schema = sample_schema();
214        let overrides = parse_id_column_overrides(&["ledger=a,b".into()], None).unwrap();
215        apply_id_column_overrides(&mut schema, &overrides).unwrap();
216        let t = schema.get_table("ledger").unwrap();
217        assert_eq!(t.primary_key.name, "a");
218        assert_eq!(
219            t.composite_primary_key.as_deref(),
220            Some(["a".to_string(), "b".to_string()].as_slice())
221        );
222        assert_eq!(t.primary_key_column_names(), vec!["a", "b"]);
223    }
224
225    #[test]
226    fn flatten_joins_with_separator() {
227        let id = Value::Array {
228            elements: vec![Value::Int32(4), Value::Text("x".into())],
229            element_type: Box::new(Type::Text),
230        };
231        assert_eq!(flatten_composite_id(id, ":"), Value::Text("4:x".into()));
232    }
233
234    #[test]
235    fn build_composite_array() {
236        let id = build_composite_record_id(vec![Value::Int32(1), Value::Int32(2)]);
237        assert!(matches!(id, Value::Array { elements, .. } if elements.len() == 2));
238    }
239}
surreal_sync_core/id_columns.rs

surreal_sync_core/
id_columns.rs
