refactor(state): share consistent device reads and row assembly

wakey-control-plane/src/state/store/agent_devices/read.rs

@@ -0,0 +1,270 @@
+//! Fleet and single-agent reads share one sequence and a single SQLite read view.
+
+use crate::state::types::*;
+use anyhow::{Context, Result};
+use sqlx::{SqliteConnection, SqlitePool};
+
+mod assembly;
+use assembly::DeviceRows;
+
+/// A read transaction prevents child rows from belonging to a newer snapshot
+/// than their parents. Keep scoped SQL explicit so agent reads use indexed filters.
+pub(super) async fn load(
+    pool: &SqlitePool,
+    agent_id: Option<&str>,
+) -> Result<Vec<AgentDeviceWithChildren>> {
+    let mut tx = pool
+        .begin()
+        .await
+        .context("failed starting device read transaction")?;
+    let rows = DeviceRows {
+        devices: load_devices(&mut tx, agent_id).await?,
+        macs: load_macs(&mut tx, agent_id).await?,
+        ips: load_ips(&mut tx, agent_id).await?,
+        hostnames: load_hostnames(&mut tx, agent_id).await?,
+        endpoints: load_endpoints(&mut tx, agent_id).await?,
+        facts: load_facts(&mut tx, agent_id).await?,
+    };
+    tx.commit()
+        .await
+        .context("failed completing device read transaction")?;
+    Ok(rows.into_devices())
+}
+
+async fn load_devices(
+    connection: &mut SqliteConnection,
+    agent_id: Option<&str>,
+) -> Result<Vec<AgentDeviceRow>> {
+    match agent_id {
+        Some(agent_id) => sqlx::query_as!(
+            AgentDeviceRow,
+            r#"SELECT agent_id, device_key, presence, display_name,
+                    first_seen_unix, last_seen_unix
+             FROM agent_devices
+             WHERE agent_id = ?1
+             ORDER BY device_key"#,
+            agent_id,
+        )
+        .fetch_all(&mut *connection)
+        .await
+        .context("failed listing agent devices for agent"),
+        None => sqlx::query_as!(
+            AgentDeviceRow,
+            r#"SELECT agent_id, device_key, presence, display_name,
+                    first_seen_unix, last_seen_unix
+             FROM agent_devices
+             ORDER BY agent_id, device_key"#,
+        )
+        .fetch_all(&mut *connection)
+        .await
+        .context("failed listing agent devices"),
+    }
+}
+
+async fn load_macs(
+    connection: &mut SqliteConnection,
+    agent_id: Option<&str>,
+) -> Result<Vec<AgentDeviceMacRow>> {
+    match agent_id {
+        Some(agent_id) => sqlx::query_as!(
+            AgentDeviceMacRow,
+            r#"SELECT agent_id as "agent_id!", device_key as "device_key!", mac as "mac!"
+             FROM agent_device_macs
+             WHERE agent_id = ?1"#,
+            agent_id,
+        )
+        .fetch_all(&mut *connection)
+        .await
+        .context("failed listing agent device macs for agent"),
+        None => sqlx::query_as!(
+            AgentDeviceMacRow,
+            r#"SELECT agent_id as "agent_id!", device_key as "device_key!", mac as "mac!"
+             FROM agent_device_macs"#,
+        )
+        .fetch_all(&mut *connection)
+        .await
+        .context("failed listing agent device macs"),
+    }
+}
+
+async fn load_ips(
+    connection: &mut SqliteConnection,
+    agent_id: Option<&str>,
+) -> Result<Vec<AgentDeviceIpRow>> {
+    match agent_id {
+        Some(agent_id) => sqlx::query_as!(
+            AgentDeviceIpRow,
+            r#"SELECT agent_id as "agent_id!", device_key as "device_key!", ip as "ip!"
+             FROM agent_device_ips
+             WHERE agent_id = ?1"#,
+            agent_id,
+        )
+        .fetch_all(&mut *connection)
+        .await
+        .context("failed listing agent device ips for agent"),
+        None => sqlx::query_as!(
+            AgentDeviceIpRow,
+            r#"SELECT agent_id as "agent_id!", device_key as "device_key!", ip as "ip!"
+             FROM agent_device_ips"#,
+        )
+        .fetch_all(&mut *connection)
+        .await
+        .context("failed listing agent device ips"),
+    }
+}
+
+async fn load_hostnames(
+    connection: &mut SqliteConnection,
+    agent_id: Option<&str>,
+) -> Result<Vec<AgentDeviceHostnameRow>> {
+    match agent_id {
+        Some(agent_id) => sqlx::query_as!(
+            AgentDeviceHostnameRow,
+            r#"SELECT agent_id as "agent_id!", device_key as "device_key!", hostname as "hostname!"
+             FROM agent_device_hostnames
+             WHERE agent_id = ?1"#,
+            agent_id,
+        )
+        .fetch_all(&mut *connection)
+        .await
+        .context("failed listing agent device hostnames for agent"),
+        None => sqlx::query_as!(
+            AgentDeviceHostnameRow,
+            r#"SELECT agent_id as "agent_id!", device_key as "device_key!", hostname as "hostname!"
+             FROM agent_device_hostnames"#,
+        )
+        .fetch_all(&mut *connection)
+        .await
+        .context("failed listing agent device hostnames"),
+    }
+}
+
+async fn load_endpoints(
+    connection: &mut SqliteConnection,
+    agent_id: Option<&str>,
+) -> Result<Vec<AgentDeviceEndpointRow>> {
+    match agent_id {
+        Some(agent_id) => sqlx::query_as::<_, AgentDeviceEndpointRow>(
+            r#"SELECT agent_id, device_key, endpoint_key, source, mac, ip,
+                    hostname, interface, presence, first_seen_unix, last_seen_unix
+             FROM agent_device_endpoints
+             WHERE agent_id = ?1"#,
+        )
+        .bind(agent_id)
+        .fetch_all(&mut *connection)
+        .await
+        .context("failed listing agent device endpoints for agent"),
+        None => sqlx::query_as::<_, AgentDeviceEndpointRow>(
+            r#"SELECT agent_id, device_key, endpoint_key, source, mac, ip,
+                    hostname, interface, presence, first_seen_unix, last_seen_unix
+             FROM agent_device_endpoints"#,
+        )
+        .fetch_all(&mut *connection)
+        .await
+        .context("failed listing agent device endpoints"),
+    }
+}
+
+async fn load_facts(
+    connection: &mut SqliteConnection,
+    agent_id: Option<&str>,
+) -> Result<Vec<AgentDeviceFactRow>> {
+    match agent_id {
+        Some(agent_id) => {
+sqlx::query_as!(
+            AgentDeviceFactRow,
+            r#"SELECT agent_id as "agent_id!", device_key as "device_key!", fact_json as "fact_json!"
+             FROM agent_device_facts
+             WHERE agent_id = ?1"#,
+            agent_id,
+        )
+        .fetch_all(&mut *connection)
+        .await
+        .context("failed listing agent device facts for agent")
+        }
+        None => {
+sqlx::query_as!(
+            AgentDeviceFactRow,
+            r#"SELECT agent_id as "agent_id!", device_key as "device_key!", fact_json as "fact_json!"
+             FROM agent_device_facts"#,
+        )
+        .fetch_all(&mut *connection)
+        .await
+        .context("failed listing agent device facts")
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::state::store::helpers::test_helpers::TestStore;
+    use wakey_core::{Device, NeighborEntry, NeighborState};
+
+    fn device(ip: &str) -> Device {
+        Device::from_parts(
+            vec![NeighborEntry {
+                ip: ip.parse().unwrap(),
+                mac: Some("aa:bb:cc:dd:ee:01".parse().unwrap()),
+                dev: Some("br-lan".into()),
+                state: NeighborState::Reachable,
+            }],
+            vec![],
+        )
+    }
+
+    #[tokio::test]
+    async fn scoped_reads_match_fleet_rows_without_leaking_other_agents() {
+        let ts = TestStore::new().await;
+        let store = ts.store();
+        for (agent, ip) in [("a", "192.168.1.10"), ("b", "192.168.1.20")] {
+            store
+                .replace_agent_device_snapshot(agent, &[device(ip)])
+                .await
+                .unwrap();
+        }
+        let fleet = store.list_agent_device_rows().await.unwrap();
+        assert_eq!(fleet.len(), 2);
+        for agent in ["a", "b", "missing"] {
+            let scoped = store.list_agent_device_rows_for_agent(agent).await.unwrap();
+            let expected = fleet
+                .iter()
+                .filter(|row| row.device.agent_id == agent)
+                .cloned()
+                .collect::<Vec<_>>();
+            // Compare every parent/child field, including timestamps and endpoint metadata.
+            assert_eq!(format!("{scoped:?}"), format!("{expected:?}"));
+        }
+    }
+
+    #[tokio::test]
+    async fn child_reads_share_the_parent_snapshot_during_a_concurrent_write() {
+        let ts = TestStore::new().await;
+        let store = ts.store();
+        store
+            .replace_agent_device_snapshot("a", &[device("192.168.1.10")])
+            .await
+            .unwrap();
+        let mut tx = store.pool.begin().await.unwrap();
+        let parents = load_devices(&mut tx, Some("a")).await.unwrap();
+        assert_eq!(parents.len(), 1);
+
+        // In WAL mode the writer can commit while the reader still sees the
+        // original snapshot. No sleeps or timing assumptions are needed.
+        store
+            .replace_agent_device_snapshot("a", &[device("192.168.1.20")])
+            .await
+            .unwrap();
+        let ips = load_ips(&mut tx, Some("a")).await.unwrap();
+        let endpoints = load_endpoints(&mut tx, Some("a")).await.unwrap();
+        assert_eq!(ips[0].ip, "192.168.1.10");
+        assert_eq!(endpoints[0].ip.as_deref(), Some("192.168.1.10"));
+        tx.commit().await.unwrap();
+
+        let current = store.list_agent_device_rows_for_agent("a").await.unwrap();
+        assert_eq!(
+            current[0].ips,
+            ["192.168.1.20".parse::<std::net::IpAddr>().unwrap()]
+        );
+    }
+}

wakey-control-plane/src/state/store/agent_devices/read/assembly.rs

@@ -0,0 +1,198 @@
+//! Attach decoded children to their parent using the full agent/device identity.
+
+use std::collections::BTreeMap;
+
+use crate::state::types::*;
+
+#[derive(Default)]
+pub(super) struct DeviceRows {
+    pub(super) devices: Vec<AgentDeviceRow>,
+    pub(super) macs: Vec<AgentDeviceMacRow>,
+    pub(super) ips: Vec<AgentDeviceIpRow>,
+    pub(super) hostnames: Vec<AgentDeviceHostnameRow>,
+    pub(super) endpoints: Vec<AgentDeviceEndpointRow>,
+    pub(super) facts: Vec<AgentDeviceFactRow>,
+}
+
+type DeviceIndex = BTreeMap<(String, String), AgentDeviceWithChildren>;
+
+impl DeviceRows {
+    pub(super) fn into_devices(self) -> Vec<AgentDeviceWithChildren> {
+        let mut index = self
+            .devices
+            .into_iter()
+            .map(|device| {
+                let key = (device.agent_id.clone(), device.device_key.clone());
+                (
+                    key,
+                    AgentDeviceWithChildren {
+                        device,
+                        macs: Vec::new(),
+                        ips: Vec::new(),
+                        hostnames: Vec::new(),
+                        endpoints: Vec::new(),
+                        facts: Vec::new(),
+                    },
+                )
+            })
+            .collect::<DeviceIndex>();
+        add_macs(&mut index, self.macs);
+        add_ips(&mut index, self.ips);
+        add_hostnames(&mut index, self.hostnames);
+        add_endpoints(&mut index, self.endpoints);
+        add_facts(&mut index, self.facts);
+        // Matches the SQL order: agent_id/device_key for fleet reads, and
+        // device_key within one agent. Child row order is left unchanged.
+        index.into_values().collect()
+    }
+}
+
+fn add_macs(index: &mut DeviceIndex, rows: Vec<AgentDeviceMacRow>) {
+    for row in rows {
+        let mac = match macaddr::MacAddr::try_from(&row) {
+            Ok(mac) => mac,
+            Err(error) => {
+                tracing::warn!(%error, agent_id = %row.agent_id, device_key = %row.device_key,
+                    raw_mac = %row.mac, "failed to parse mac from agent_device_macs row");
+                continue;
+            }
+        };
+        if let Some(device) = index.get_mut(&(row.agent_id, row.device_key)) {
+            device.macs.push(mac);
+        }
+    }
+}
+
+fn add_ips(index: &mut DeviceIndex, rows: Vec<AgentDeviceIpRow>) {
+    for row in rows {
+        let ip = match std::net::IpAddr::try_from(&row) {
+            Ok(ip) => ip,
+            Err(error) => {
+                tracing::warn!(%error, agent_id = %row.agent_id, device_key = %row.device_key,
+                    raw_ip = %row.ip, "failed to parse ip from agent_device_ips row");
+                continue;
+            }
+        };
+        if let Some(device) = index.get_mut(&(row.agent_id, row.device_key)) {
+            device.ips.push(ip);
+        }
+    }
+}
+
+fn add_hostnames(index: &mut DeviceIndex, rows: Vec<AgentDeviceHostnameRow>) {
+    for row in rows {
+        if let Some(device) = index.get_mut(&(row.agent_id, row.device_key)) {
+            device.hostnames.push(row.hostname);
+        }
+    }
+}
+
+fn add_endpoints(index: &mut DeviceIndex, rows: Vec<AgentDeviceEndpointRow>) {
+    for row in rows {
+        if let Some(endpoint) = row.to_endpoint()
+            && let Some(device) = index.get_mut(&(row.agent_id, row.device_key))
+        {
+            device.endpoints.push(endpoint);
+        }
+    }
+}
+
+fn add_facts(index: &mut DeviceIndex, rows: Vec<AgentDeviceFactRow>) {
+    for row in rows {
+        if let Some(device) = index.get_mut(&(row.agent_id, row.device_key)) {
+            device.facts.push(row.fact_json);
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn parent(agent: &str, key: &str) -> AgentDeviceRow {
+        AgentDeviceRow {
+            agent_id: agent.into(),
+            device_key: key.into(),
+            presence: "unknown".into(),
+            display_name: None,
+            first_seen_unix: 1,
+            last_seen_unix: 2,
+        }
+    }
+
+    #[test]
+    fn assembly_keeps_agent_boundaries_order_and_empty_parents() {
+        let rows = DeviceRows {
+            devices: vec![
+                parent("b", "shared"),
+                parent("a", "shared"),
+                parent("a", "empty"),
+            ],
+            hostnames: vec![
+                AgentDeviceHostnameRow {
+                    agent_id: "b".into(),
+                    device_key: "shared".into(),
+                    hostname: "second".into(),
+                },
+                AgentDeviceHostnameRow {
+                    agent_id: "a".into(),
+                    device_key: "shared".into(),
+                    hostname: "first".into(),
+                },
+                AgentDeviceHostnameRow {
+                    agent_id: "missing".into(),
+                    device_key: "shared".into(),
+                    hostname: "orphan".into(),
+                },
+            ],
+            facts: vec![AgentDeviceFactRow {
+                agent_id: "a".into(),
+                device_key: "shared".into(),
+                fact_json: "{\"action\":\"remove\"}".into(),
+            }],
+            ..Default::default()
+        }
+        .into_devices();
+        assert_eq!(rows.len(), 3);
+        assert_eq!(rows[0].device.device_key, "empty");
+        assert!(rows[0].hostnames.is_empty());
+        assert_eq!(rows[1].hostnames, ["first"]);
+        assert_eq!(rows[1].facts, ["{\"action\":\"remove\"}"]);
+        assert_eq!(rows[2].hostnames, ["second"]);
+        assert!(rows[2].facts.is_empty());
+    }
+
+    #[test]
+    fn malformed_children_do_not_hide_valid_children_or_the_parent() {
+        let rows = DeviceRows {
+            devices: vec![parent("a", "pc")],
+            macs: ["not-a-mac", "aa:bb:cc:dd:ee:01"]
+                .into_iter()
+                .map(|mac| AgentDeviceMacRow {
+                    agent_id: "a".into(),
+                    device_key: "pc".into(),
+                    mac: mac.into(),
+                })
+                .collect(),
+            ips: ["not-an-ip", "192.168.1.10"]
+                .into_iter()
+                .map(|ip| AgentDeviceIpRow {
+                    agent_id: "a".into(),
+                    device_key: "pc".into(),
+                    ip: ip.into(),
+                })
+                .collect(),
+            ..Default::default()
+        }
+        .into_devices();
+        assert_eq!(rows.len(), 1);
+        assert_eq!(
+            rows[0].macs,
+            ["aa:bb:cc:dd:ee:01".parse::<macaddr::MacAddr>().unwrap()]
+        );
+        assert_eq!(
+            rows[0].ips,
+            ["192.168.1.10".parse::<std::net::IpAddr>().unwrap()]
+        );
+    }
+}