tandem-memory 0.6.9

Memory storage and embedding utilities for Tandem
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

impl MemoryDatabase {
    fn ensure_chunk_scope_columns(
        &self,
        conn: &Connection,
        table: &str,
        existing_cols: &HashSet<String>,
    ) -> MemoryResult<()> {
        if !existing_cols.contains("owner_org_unit_id") {
            conn.execute(
                &format!("ALTER TABLE {table} ADD COLUMN owner_org_unit_id TEXT"),
                [],
            )?;
        }
        if !existing_cols.contains("tenant_shared") {
            conn.execute(
                &format!("ALTER TABLE {table} ADD COLUMN tenant_shared INTEGER NOT NULL DEFAULT 0"),
                [],
            )?;
        }
        // Per-scope envelope for hosted-KMS encryption (TAN-668). NULL for
        // local/plaintext rows; carries the wrapped DEK + key scope for hosted
        // rows so `content`/`metadata` can be decrypted before use.
        if !existing_cols.contains("crypto_envelope") {
            conn.execute(
                &format!("ALTER TABLE {table} ADD COLUMN crypto_envelope TEXT"),
                [],
            )?;
        }
        self.backfill_chunk_scope_columns(conn, table)
    }

    fn backfill_chunk_scope_columns(&self, conn: &Connection, table: &str) -> MemoryResult<()> {
        let rows = {
            let mut stmt = conn.prepare(&format!(
                "SELECT id, metadata FROM {table}
                 WHERE (owner_org_unit_id IS NULL OR tenant_shared = 0)
                   AND metadata IS NOT NULL
                   AND TRIM(metadata) != ''"
            ))?;
            let rows = stmt
                .query_map([], |row| {
                    Ok((row.get::<_, String>(0)?, row.get::<_, Option<String>>(1)?))
                })?
                .collect::<Result<Vec<_>, _>>()?;
            rows
        };

        for (id, metadata_stored) in rows {
            let Some(metadata_stored) = metadata_stored else {
                continue;
            };
            let metadata_plain = match self.crypto.decrypt_field(&metadata_stored) {
                Ok(metadata_plain) => metadata_plain,
                Err(err) => {
                    tracing::warn!(
                        table = table,
                        chunk_id = id.as_str(),
                        "skipping owner_org_unit_id backfill for unreadable chunk metadata: {}",
                        err
                    );
                    continue;
                }
            };
            let Ok(metadata) = serde_json::from_str::<serde_json::Value>(&metadata_plain) else {
                continue;
            };
            let owner_org_unit_id = owner_org_unit_id_from_metadata(Some(&metadata));
            let tenant_shared = tenant_shared_from_metadata(Some(&metadata));
            if owner_org_unit_id.is_none() && !tenant_shared {
                continue;
            }
            conn.execute(
                &format!(
                    "UPDATE {table}
                     SET owner_org_unit_id = COALESCE(owner_org_unit_id, ?1),
                         tenant_shared = CASE WHEN ?2 = 1 THEN 1 ELSE tenant_shared END
                     WHERE id = ?3"
                ),
                params![owner_org_unit_id.as_deref(), i64::from(tenant_shared), id],
            )?;
        }

        Ok(())
    }
}