CamFinder/src/db.rs

use crate::models::ChunkRecord;
use anyhow::{Context, Result};
use std::convert::TryFrom;
use tokio_postgres::{Client, NoTls};

pub struct Database {
    client: Client,
}

impl Database {
    pub async fn open(database_url: impl AsRef<str>) -> Result<Self> {
        let database_url = database_url.as_ref().trim();
        let (client, connection) = tokio_postgres::connect(database_url, NoTls)
            .await
            .with_context(|| format!("failed to connect to postgres at {database_url}"))?;

        tokio::spawn(async move {
            if let Err(err) = connection.await {
                eprintln!("postgres connection error: {err}");
            }
        });

        let db = Self { client };
        db.init().await?;
        Ok(db)
    }

    async fn init(&self) -> Result<()> {
        self.client
            .batch_execute(
                r#"
                CREATE TABLE IF NOT EXISTS open_rtsp_endpoint (
                    id BIGSERIAL PRIMARY KEY,
                    ip INET NOT NULL UNIQUE,
                    port INTEGER NOT NULL,
                    scanned_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
                    last_seen_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
                    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
                    updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
                );

                CREATE INDEX IF NOT EXISTS open_rtsp_endpoint_port_idx
                    ON open_rtsp_endpoint (port);

                CREATE TABLE IF NOT EXISTS scan_chunks (
                    chunk_start BIGINT PRIMARY KEY NOT NULL,
                    chunk_end BIGINT NOT NULL,
                    tested_count BIGINT NOT NULL DEFAULT 0,
                    open_count BIGINT NOT NULL DEFAULT 0,
                    is_complete BOOLEAN NOT NULL DEFAULT FALSE,
                    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
                    updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
                );
                "#,
            )
            .await
            .context("failed to initialize postgres schema")?;
        Ok(())
    }

    pub async fn load_chunk_records(&self) -> Result<Vec<ChunkRecord>> {
        let rows = self
            .client
            .query(
                r#"
                SELECT chunk_start, chunk_end, tested_count, open_count, is_complete
                FROM scan_chunks
                ORDER BY chunk_start ASC
                "#,
                &[],
            )
            .await
            .context("failed to load scan chunks")?;

        rows.into_iter().map(row_to_chunk_record).collect()
    }

    pub async fn load_chunk_record(&self, chunk_start: usize) -> Result<Option<ChunkRecord>> {
        let chunk_start = usize_to_i64(chunk_start, "chunk_start")?;
        let row = self
            .client
            .query_opt(
                r#"
                SELECT chunk_start, chunk_end, tested_count, open_count, is_complete
                FROM scan_chunks
                WHERE chunk_start = $1
                "#,
                &[&chunk_start],
            )
            .await
            .context("failed to load scan chunk")?;

        row.map(row_to_chunk_record).transpose()
    }

    pub async fn save_chunk_record(&self, record: &ChunkRecord) -> Result<()> {
        let chunk_start = usize_to_i64(record.chunk_start, "chunk_start")?;
        let chunk_end = usize_to_i64(record.chunk_end, "chunk_end")?;
        let tested_count = usize_to_i64(record.tested_count, "tested_count")?;
        let open_count = usize_to_i64(record.open_count, "open_count")?;

        self.client
            .execute(
                r#"
                INSERT INTO scan_chunks (
                    chunk_start,
                    chunk_end,
                    tested_count,
                    open_count,
                    is_complete,
                    created_at,
                    updated_at
                )
                VALUES ($1, $2, $3, $4, $5, NOW(), NOW())
                ON CONFLICT (chunk_start) DO UPDATE SET
                    chunk_end = EXCLUDED.chunk_end,
                    tested_count = EXCLUDED.tested_count,
                    open_count = EXCLUDED.open_count,
                    is_complete = EXCLUDED.is_complete,
                    updated_at = NOW()
                "#,
                &[
                    &chunk_start,
                    &chunk_end,
                    &tested_count,
                    &open_count,
                    &record.is_complete,
                ],
            )
            .await
            .context("failed to save scan chunk")?;

        Ok(())
    }

    pub async fn upsert_open_endpoint(&self, ip: std::net::Ipv4Addr, port: u16) -> Result<()> {
        let ip = std::net::IpAddr::V4(ip);
        let port = i32::from(port);

        self.client
            .execute(
                r#"
                INSERT INTO open_rtsp_endpoint (
                    ip,
                    port,
                    scanned_at,
                    last_seen_at,
                    created_at,
                    updated_at
                )
                VALUES ($1, $2, NOW(), NOW(), NOW(), NOW())
                ON CONFLICT (ip) DO UPDATE SET
                    port = EXCLUDED.port,
                    scanned_at = NOW(),
                    last_seen_at = NOW(),
                    updated_at = NOW()
                "#,
                &[&ip, &port],
            )
            .await
            .context("failed to upsert open endpoint")?;

        Ok(())
    }

    pub async fn count_open_endpoints(&self, port: u16) -> Result<usize> {
        let port = i32::from(port);
        let count: i64 = self
            .client
            .query_one(
                "SELECT COUNT(*)::BIGINT FROM open_rtsp_endpoint WHERE port = $1",
                &[&port],
            )
            .await
            .context("failed to count open endpoints")?
            .get(0);

        i64_to_usize(count, "open endpoint count")
    }

    pub async fn summarize_chunks(&self) -> Result<(usize, usize)> {
        let row = self
            .client
            .query_one(
                r#"
                SELECT
                    COALESCE(SUM(tested_count), 0)::BIGINT,
                    COALESCE(SUM(open_count), 0)::BIGINT
                FROM scan_chunks
                "#,
                &[],
            )
            .await
            .context("failed to summarize scan chunks")?;

        let tested = i64_to_usize(row.get::<_, i64>(0), "tested count")?;
        let open = i64_to_usize(row.get::<_, i64>(1), "open count")?;
        Ok((tested, open))
    }
}

fn row_to_chunk_record(row: tokio_postgres::Row) -> Result<ChunkRecord> {
    Ok(ChunkRecord {
        chunk_start: i64_to_usize(row.get::<_, i64>(0), "chunk_start")?,
        chunk_end: i64_to_usize(row.get::<_, i64>(1), "chunk_end")?,
        tested_count: i64_to_usize(row.get::<_, i64>(2), "tested_count")?,
        open_count: i64_to_usize(row.get::<_, i64>(3), "open_count")?,
        is_complete: row.get::<_, bool>(4),
    })
}

fn i64_to_usize(value: i64, field: &str) -> Result<usize> {
    usize::try_from(value).with_context(|| format!("{field} out of range: {value}"))
}

fn usize_to_i64(value: usize, field: &str) -> Result<i64> {
    i64::try_from(value).with_context(|| format!("{field} out of range: {value}"))
}