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