HomeServer/src/db.rs

use std::path::Path;

use anyhow::Result;
use rusqlite::{Connection, OptionalExtension, ToSql};

use crate::devices::{DeviceWithName, Devices, DevicesWithName};

pub struct DataBase {
    sql: Connection,
}

impl DataBase {
    const VERSION_0_1_0: &'static str = "0.1.0";

    pub async fn new(path: impl AsRef<Path>) -> Result<Self> {
        let me = Self {
            sql: Connection::open(path)?,
        };

        me.generate_tables()?;
        me.init()?;

        Ok(me)
    }

    fn generate_tables(&self) -> Result<()> {
        self.sql.execute(
            "CREATE TABLE IF NOT EXISTS meta (
                id      INTEGER PRIMARY KEY,
                version INTEGER NOT NULL
            )",
            [],
        )?;

        self.sql.execute(
            "CREATE TABLE IF NOT EXISTS devices (
                id      INTEGER PRIMARY KEY,
                device  VARCHAR(60) NOT NULL,
                type    VARCHAR(30) NOT NULL,
                control INTEGER NOT NULL,
                name    VARCHAR(80)
            )",
            [],
        )?;

        self.sql.execute(
            "CREATE TABLE IF NOT EXISTS data (
                id          INTEGER PRIMARY KEY,
                time        BIGINT NOT NULL,
                name        VARCHAR(30) NOT NULL,
                value       REAL NOT NULL,
                device_id   INTEGER NOT NULL,
                FOREIGN KEY(device_id) REFERENCES devices(id)
            )",
            [],
        )?;

        self.sql.execute(
            "CREATE TABLE IF NOT EXISTS credentials (
                id          INTEGER PRIMARY KEY,
                key         VARCHAR(60) NOT NULL,
                device_id   BIGINT NOT NULL,
                cred        VARCHAR(256) NOT NULL
            )",
            [],
        )?;

        Ok(())
    }

    pub fn version(&self) -> Result<String> {
        Ok(self
            .sql
            .query_row("SELECT version FROM meta WHERE id=1", [], |row| row.get(0))?)
    }

    fn init(&self) -> Result<()> {
        if self.version().is_err() {
            self.sql.execute(
                "INSERT INTO meta (version)
                VALUES (?1)",
                &[Self::VERSION_0_1_0],
            )?;
        }

        Ok(())
    }

    pub fn register_devices(&self, devices: &Devices) -> Result<()> {
        for (device, control) in devices.plugs.iter() {
            self.sql.execute(
                &format!(
                    "INSERT INTO devices (device, type, control)
                    SELECT \"{device}\", \"plug\", ?1
                    WHERE
                        NOT EXISTS (
                            SELECT device
                            FROM devices
                            WHERE device=\"{device}\"
                        )
                "
                ),
                &[control],
            )?;

            let ctl = if *control { 1 } else { 0 };

            self.sql.execute(
                &format!(
                    "
                        UPDATE devices
                        SET control=\"{ctl}\"
                        WHERE device=\"{device}\"
                    "
                ),
                [],
            )?;
        }

        for device in devices.thermostat.iter() {
            self.sql.execute(
                &format!(
                    "INSERT INTO devices (device, type, control)
                    SELECT \"{device}\", \"thermostat\", false
                    WHERE
                        NOT EXISTS (
                            SELECT device
                            FROM devices
                            WHERE device=\"{device}\"
                        )
                "
                ),
                [],
            )?;
        }

        for device in devices.thermometer.iter() {
            self.sql.execute(
                &format!(
                    "INSERT INTO devices (device, type, control)
                    SELECT \"{device}\", \"thermometer\", false
                    WHERE
                        NOT EXISTS (
                            SELECT device
                            FROM devices
                            WHERE device=\"{device}\"
                        )
                "
                ),
                [],
            )?;
        }

        Ok(())
    }

    pub fn write(&self, device_name: &str, time: u64, name: &str, value: f32) -> Result<()> {
        let params: &[&dyn ToSql] = &[&time, &value];

        self.sql.execute(
            &format!(
                "INSERT INTO data (time, name, value, device_id)
                VALUES (?1, \"{name}\", ?2, (SELECT id FROM devices WHERE device=\"{device_name}\") )"
            ),
            params,
        )?;

        Ok(())
    }

    pub fn devices(&self) -> Result<DevicesWithName> {
        let mut devices = DevicesWithName::default();

        for row in self
            .sql
            .prepare(&format!(
                "
                    SELECT device, type, name, control
                    FROM devices
                "
            ))?
            .query_map([], |row| {
                Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?))
            })?
        {
            let (device, dev_type, name, control): (String, String, Option<String>, i32) = row?;

            match dev_type.as_str() {
                "plug" => devices.plugs.push(DeviceWithName {
                    id: device,
                    desc: name,
                    toggle: control != 0,
                }),
                "thermostat" => devices.thermostat.push(DeviceWithName {
                    id: device,
                    desc: name,
                    toggle: control != 0,
                }),
                "thermometer" => devices.temperature_and_humidity.push(DeviceWithName {
                    id: device,
                    desc: name,
                    toggle: control != 0,
                }),

                _ => panic!(),
            }
        }

        Ok(devices)
    }

    pub fn device_exists(&self, device_id: &str) -> Result<bool> {
        Ok(self
            .sql
            .prepare(&format!(
                "
                    SELECT *
                    FROM devices
                    WHERE device=\"{device_id}\"
                "
            ))?
            .exists([])?)
    }

    pub fn change_device_name(&self, device: &str, description: &str) -> Result<()> {
        self.sql.execute(
            &format!(
                "
                    UPDATE devices
                    SET name=\"{description}\"
                    WHERE device=\"{device}\"
                "
            ),
            [],
        )?;

        Ok(())
    }

    pub fn read(&self, device: &str) -> Result<Vec<(u64, f32)>> {
        self._read(&format!(
            "
            SELECT data.time, data.value
            FROM data
            INNER JOIN devices
                ON data.device_id=devices.id
            WHERE devices.device=\"{device}\"
        "
        ))
    }

    pub fn read_range(&self, device: &str, start: u64, end: u64) -> Result<Vec<(u64, f32)>> {
        self._read(&format!(
            "
            SELECT data.time, data.value
            FROM data
            INNER JOIN devices
                ON data.device_id=devices.id
            WHERE devices.device=\"{device}\"
                AND data.time>={start}
                AND data.time<{end}
        "
        ))
    }

    fn _read(&self, query: &str) -> Result<Vec<(u64, f32)>> {
        self.sql
            .prepare(query)?
            .query_map([], |row| Ok((row.get(0)?, row.get(1)?)))?
            .map(|row| {
                let (time, watts) = row?;

                Ok((time, watts))
            })
            .collect()
    }

    pub fn write_credential(&self, key: &str, device_id: u64, credential: &str) -> Result<()> {
        if self
            .sql
            .prepare(&format!(
                "
                SELECT cred
                FROM credentials
                WHERE key=\"{key}\" AND device_id={device_id}
            "
            ))?
            .exists([])?
        {
            self.sql.execute(
                &format!(
                    "
                        UPDATE credentials
                        SET cred=\"{credential}\"
                        WHERE key=\"{key}\" AND device_id={device_id}
                    "
                ),
                [],
            )?;
        } else {
            self.sql.execute(
                &format!(
                    "INSERT INTO crendetials (key, device_id, cred)
                    VALUES (\"{key}\", {device_id}, \"{credential}\")"
                ),
                [],
            )?;
        }

        Ok(())
    }

    pub fn read_credential(&self, key: &str, device_id: u64) -> Result<Option<String>> {
        Ok(self
            .sql
            .prepare(&format!(
                "
                SELECT cred
                FROM credentials
                WHERE key=\"{key}\" AND device_id={device_id}
            "
            ))?
            .query_row([], |row| Ok(row.get(0)?))
            .optional()?)
    }
}

#[cfg(test)]
mod test {
    use std::fs;

    use anyhow::Result;

    use crate::devices::Devices;

    use super::DataBase;

    #[tokio::test]
    async fn test_connection() -> Result<()> {
        let db = DataBase::new("connection_test.db").await?;
        assert_eq!(DataBase::VERSION_0_1_0, db.version()?);

        fs::remove_file("connection_test.db")?;

        Ok(())
    }

    #[tokio::test]
    async fn test_startup() -> Result<()> {
        let db = DataBase::new("startup_test.db").await?;

        db.register_devices(&Devices {
            plugs: vec![("test".to_string(), true)],
            thermostat: Vec::new(),
            thermometer: Vec::new(),
        })?;

        fs::remove_file("startup_test.db")?;

        Ok(())
    }

    #[tokio::test]
    async fn test_write() -> Result<()> {
        let db = DataBase::new("write_test.db").await?;

        let device_name = "test";

        db.register_devices(&Devices {
            plugs: vec![(device_name.to_string(), true)],
            thermostat: Vec::new(),
            thermometer: Vec::new(),
        })?;

        db.write(device_name, 0, "watts", 5.5)?;

        let device_descriptor = "udo";
        db.change_device_name(device_name, device_descriptor)?;

        let devices = db.devices()?;

        assert_eq!(devices.plugs[0].desc.as_ref().unwrap(), device_descriptor);
        assert_eq!(devices.plugs[0].id, device_name);

        fs::remove_file("write_test.db")?;

        Ok(())
    }
}