chore: reorganize rust crates

crates/libpk/src/db/mod.rs

@@ -0,0 +1,96 @@
+use fred::clients::RedisPool;
+use sqlx::postgres::{PgConnectOptions, PgPool, PgPoolOptions};
+use std::str::FromStr;
+use tracing::info;
+
+pub mod repository;
+pub mod types;
+
+pub async fn init_redis() -> anyhow::Result<RedisPool> {
+    info!("connecting to redis");
+    let redis = RedisPool::new(
+        fred::types::RedisConfig::from_url_centralized(crate::config.db.data_redis_addr.as_ref())
+            .expect("redis url is invalid"),
+        None,
+        None,
+        Some(Default::default()),
+        10,
+    )?;
+
+    let redis_handle = redis.connect_pool();
+    tokio::spawn(async move { redis_handle });
+
+    Ok(redis)
+}
+
+pub async fn init_data_db() -> anyhow::Result<PgPool> {
+    info!("connecting to database");
+
+    let mut options = PgConnectOptions::from_str(&crate::config.db.data_db_uri)?;
+
+    if let Some(password) = crate::config.db.db_password.clone() {
+        options = options.password(&password);
+    }
+
+    let mut pool = PgPoolOptions::new();
+
+    if let Some(max_conns) = crate::config.db.data_db_max_connections {
+        pool = pool.max_connections(max_conns);
+    }
+
+    if let Some(min_conns) = crate::config.db.data_db_min_connections {
+        pool = pool.min_connections(min_conns);
+    }
+
+    Ok(pool.connect_with(options).await?)
+}
+
+pub async fn init_messages_db() -> anyhow::Result<PgPool> {
+    info!("connecting to messages database");
+
+    let mut options = PgConnectOptions::from_str(
+        &crate::config
+            .db
+            .messages_db_uri
+            .as_ref()
+            .expect("missing messages db uri"),
+    )?;
+
+    if let Some(password) = crate::config.db.db_password.clone() {
+        options = options.password(&password);
+    }
+
+    let mut pool = PgPoolOptions::new();
+
+    if let Some(max_conns) = crate::config.db.data_db_max_connections {
+        pool = pool.max_connections(max_conns);
+    }
+
+    if let Some(min_conns) = crate::config.db.data_db_min_connections {
+        pool = pool.min_connections(min_conns);
+    }
+
+    Ok(pool.connect_with(options).await?)
+}
+
+pub async fn init_stats_db() -> anyhow::Result<PgPool> {
+    info!("connecting to stats database");
+
+    let mut options = PgConnectOptions::from_str(
+        &crate::config
+            .db
+            .stats_db_uri
+            .as_ref()
+            .expect("missing messages db uri"),
+    )?;
+
+    if let Some(password) = crate::config.db.db_password.clone() {
+        options = options.password(&password);
+    }
+
+    Ok(PgPoolOptions::new()
+        .max_connections(1)
+        .min_connections(1)
+        .connect_with(options)
+        .await?)
+}

crates/libpk/src/db/repository/auth.rs

@@ -0,0 +1,20 @@
+pub async fn legacy_token_auth(
+    pool: &sqlx::postgres::PgPool,
+    token: &str,
+) -> anyhow::Result<Option<i32>> {
+    let mut system: Vec<LegacyTokenDbResponse> =
+        sqlx::query_as("select id from systems where token = $1")
+            .bind(token)
+            .fetch_all(pool)
+            .await?;
+    Ok(if let Some(system) = system.pop() {
+        Some(system.id)
+    } else {
+        None
+    })
+}
+
+#[derive(sqlx::FromRow)]
+struct LegacyTokenDbResponse {
+    id: i32,
+}

crates/libpk/src/db/repository/avatars.rs

@@ -0,0 +1,111 @@
+use sqlx::{PgPool, Postgres, Transaction};
+
+use crate::db::types::avatars::*;
+
+pub async fn get_by_id(pool: &PgPool, id: String) -> anyhow::Result<Option<ImageMeta>> {
+    Ok(sqlx::query_as("select * from images where id = $1")
+        .bind(id)
+        .fetch_optional(pool)
+        .await?)
+}
+
+pub async fn get_by_original_url(
+    pool: &PgPool,
+    original_url: &str,
+) -> anyhow::Result<Option<ImageMeta>> {
+    Ok(
+        sqlx::query_as("select * from images where original_url = $1")
+            .bind(original_url)
+            .fetch_optional(pool)
+            .await?,
+    )
+}
+
+pub async fn get_by_attachment_id(
+    pool: &PgPool,
+    attachment_id: u64,
+) -> anyhow::Result<Option<ImageMeta>> {
+    Ok(
+        sqlx::query_as("select * from images where original_attachment_id = $1")
+            .bind(attachment_id as i64)
+            .fetch_optional(pool)
+            .await?,
+    )
+}
+
+pub async fn remove_deletion_queue(pool: &PgPool, attachment_id: u64) -> anyhow::Result<()> {
+    sqlx::query(
+        r#"
+            delete from image_cleanup_jobs
+            where id in (
+                select id from images
+                where original_attachment_id = $1
+            )
+        "#,
+    )
+    .bind(attachment_id as i64)
+    .execute(pool)
+    .await?;
+
+    Ok(())
+}
+
+pub async fn pop_queue(
+    pool: &PgPool,
+) -> anyhow::Result<Option<(Transaction<Postgres>, ImageQueueEntry)>> {
+    let mut tx = pool.begin().await?;
+    let res: Option<ImageQueueEntry> = sqlx::query_as("delete from image_queue where itemid = (select itemid from image_queue order by itemid for update skip locked limit 1) returning *")
+        .fetch_optional(&mut *tx).await?;
+    Ok(res.map(|x| (tx, x)))
+}
+
+pub async fn get_queue_length(pool: &PgPool) -> anyhow::Result<i64> {
+    Ok(sqlx::query_scalar("select count(*) from image_queue")
+        .fetch_one(pool)
+        .await?)
+}
+
+pub async fn get_stats(pool: &PgPool) -> anyhow::Result<Stats> {
+    Ok(sqlx::query_as(
+        "select count(*) as total_images, sum(file_size) as total_file_size from images",
+    )
+    .fetch_one(pool)
+    .await?)
+}
+
+pub async fn add_image(pool: &PgPool, meta: ImageMeta) -> anyhow::Result<bool> {
+    let kind_str = match meta.kind {
+        ImageKind::Avatar => "avatar",
+        ImageKind::Banner => "banner",
+    };
+
+    let res = sqlx::query("insert into images (id, url, content_type, original_url, file_size, width, height, original_file_size, original_type, original_attachment_id, kind, uploaded_by_account, uploaded_by_system, uploaded_at) values ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, (now() at time zone 'utc')) on conflict (id) do nothing")
+        .bind(meta.id)
+        .bind(meta.url)
+        .bind(meta.content_type)
+        .bind(meta.original_url)
+        .bind(meta.file_size)
+        .bind(meta.width)
+        .bind(meta.height)
+        .bind(meta.original_file_size)
+        .bind(meta.original_type)
+        .bind(meta.original_attachment_id)
+        .bind(kind_str)
+        .bind(meta.uploaded_by_account)
+        .bind(meta.uploaded_by_system)
+        .execute(pool).await?;
+    Ok(res.rows_affected() > 0)
+}
+
+pub async fn push_queue(
+    conn: &mut sqlx::PgConnection,
+    url: &str,
+    kind: ImageKind,
+) -> anyhow::Result<()> {
+    sqlx::query("insert into image_queue (url, kind) values ($1, $2)")
+        .bind(url)
+        .bind(kind)
+        .execute(conn)
+        .await?;
+    Ok(())
+}

crates/libpk/src/db/repository/mod.rs

@@ -0,0 +1,7 @@
+mod stats;
+pub use stats::*;
+
+pub mod avatars;
+
+mod auth;
+pub use auth::*;

crates/libpk/src/db/repository/stats.rs

@@ -0,0 +1,26 @@
+pub async fn get_stats(pool: &sqlx::postgres::PgPool) -> anyhow::Result<Counts> {
+    let counts: Counts = sqlx::query_as("select * from info").fetch_one(pool).await?;
+    Ok(counts)
+}
+
+pub async fn insert_stats(
+    pool: &sqlx::postgres::PgPool,
+    table: &str,
+    value: i64,
+) -> anyhow::Result<()> {
+    // danger sql injection
+    sqlx::query(format!("insert into {table} values (now(), $1)").as_str())
+        .bind(value)
+        .execute(pool)
+        .await?;
+    Ok(())
+}
+
+#[derive(serde::Serialize, sqlx::FromRow)]
+pub struct Counts {
+    pub system_count: i64,
+    pub member_count: i64,
+    pub group_count: i64,
+    pub switch_count: i64,
+    pub message_count: i64,
+}

crates/libpk/src/db/types/avatars.rs

@@ -0,0 +1,53 @@
+use serde::{Deserialize, Serialize};
+use sqlx::FromRow;
+use time::OffsetDateTime;
+use uuid::Uuid;
+
+#[derive(FromRow)]
+pub struct ImageMeta {
+    pub id: String,
+    pub kind: ImageKind,
+    pub content_type: String,
+    pub url: String,
+    pub file_size: i32,
+    pub width: i32,
+    pub height: i32,
+    pub uploaded_at: Option<OffsetDateTime>,
+
+    pub original_url: Option<String>,
+    pub original_attachment_id: Option<i64>,
+    pub original_file_size: Option<i32>,
+    pub original_type: Option<String>,
+    pub uploaded_by_account: Option<i64>,
+    pub uploaded_by_system: Option<Uuid>,
+}
+
+#[derive(FromRow, Serialize)]
+pub struct Stats {
+    pub total_images: i64,
+    pub total_file_size: i64,
+}
+
+#[derive(Serialize, Deserialize, Clone, Copy, Debug, sqlx::Type, PartialEq)]
+#[serde(rename_all = "snake_case")]
+#[sqlx(rename_all = "snake_case", type_name = "text")]
+pub enum ImageKind {
+    Avatar,
+    Banner,
+}
+
+impl ImageKind {
+    pub fn size(&self) -> (u32, u32) {
+        match self {
+            Self::Avatar => (512, 512),
+            Self::Banner => (1024, 1024),
+        }
+    }
+}
+
+#[derive(FromRow)]
+pub struct ImageQueueEntry {
+    pub itemid: i32,
+    pub url: String,
+    pub kind: ImageKind,
+}

crates/libpk/src/db/types/mod.rs

@@ -0,0 +1 @@
+pub mod avatars;