Added scaffoded project

src/db.rs

@@ -0,0 +1,218 @@
+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 = ip.to_string();
+        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::inet, $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}"))
+}