Merge branch 'main' into rust-command-parser

lib/libpk/Cargo.toml

@@ -5,17 +5,21 @@ edition = "2021"
 
 [dependencies]
 anyhow = { workspace = true } 
-config = "0.13.3"
+config = "0.14.0"
 fred = { workspace = true }
 gethostname = "0.4.1"
 lazy_static = { workspace = true }
 metrics = { workspace = true }
-metrics-exporter-prometheus = { version = "0.11.0", default-features = false, features = ["tokio", "http-listener", "tracing"] }
+metrics-exporter-prometheus = { version = "0.15.3", default-features = false, features = ["tokio", "http-listener", "tracing"] }
 serde = { workspace = true }
 sqlx = { workspace = true }
+time = { workspace = true }
 tokio = { workspace = true }
 tracing = { workspace = true }
+tracing-gelf = "0.7.1"
 tracing-subscriber = { workspace = true}
+twilight-model = { workspace = true }
+uuid = { workspace = true }
 
 prost = { workspace = true }
 prost-types = { workspace = true }

lib/libpk/src/_config.rs

@@ -3,11 +3,27 @@ use lazy_static::lazy_static;
 use serde::Deserialize;
 use std::sync::Arc;
 
+use twilight_model::id::{marker::UserMarker, Id};
+
+#[derive(Clone, Deserialize, Debug)]
+pub struct ClusterSettings {
+    pub node_id: u32,
+    pub total_shards: u32,
+    pub total_nodes: u32,
+}
+
 #[derive(Deserialize, Debug)]
 pub struct DiscordConfig {
-    pub client_id: u32,
+    pub client_id: Id<UserMarker>,
     pub bot_token: String,
     pub client_secret: String,
+    pub max_concurrency: u32,
+    #[serde(default)]
+    pub cluster: Option<ClusterSettings>,
+    pub api_base_url: Option<String>,
+
+    #[serde(default = "_default_api_addr")]
+    pub cache_api_addr: String,
 }
 
 #[derive(Deserialize, Debug)]
@@ -24,7 +40,7 @@ fn _default_api_addr() -> String {
     "0.0.0.0:5000".to_string()
 }
 
-#[derive(Deserialize, Debug)]
+#[derive(Deserialize, Clone, Debug)]
 pub struct ApiConfig {
     #[serde(default = "_default_api_addr")]
     pub addr: String,
@@ -38,27 +54,69 @@ pub struct ApiConfig {
     pub temp_token2: Option<String>,
 }
 
+#[derive(Deserialize, Clone, Debug)]
+pub struct AvatarsConfig {
+    pub s3: S3Config,
+    pub cdn_url: String,
+
+    #[serde(default)]
+    pub migrate_worker_count: u32,
+}
+
+#[derive(Deserialize, Clone, Debug)]
+pub struct S3Config {
+    pub bucket: String,
+    pub application_id: String,
+    pub application_key: String,
+    pub endpoint: String,
+}
+
 fn _metrics_default() -> bool {
     false
 }
+fn _json_log_default() -> bool {
+    false
+}
 
 #[derive(Deserialize, Debug)]
 pub struct PKConfig {
     pub db: DatabaseConfig,
 
-    pub discord: DiscordConfig,
-    pub api: ApiConfig,
+    #[serde(default)]
+    pub discord: Option<DiscordConfig>,
+    #[serde(default)]
+    pub api: Option<ApiConfig>,
+    #[serde(default)]
+    pub avatars: Option<AvatarsConfig>,
 
     #[serde(default = "_metrics_default")]
     pub run_metrics_server: bool,
 
-    pub(crate) gelf_log_url: Option<String>,
+    #[serde(default = "_json_log_default")]
+    pub(crate) json_log: bool,
+}
+
+impl PKConfig {
+    pub fn api(self) -> ApiConfig {
+        self.api.expect("missing api config")
+    }
+
+    pub fn discord_config(self) -> DiscordConfig {
+        self.discord.expect("missing discord config")
+    }
 }
 
 lazy_static! {
     #[derive(Debug)]
-    pub static ref CONFIG: Arc<PKConfig> = Arc::new(Config::builder()
-    .add_source(config::Environment::with_prefix("pluralkit").separator("__"))
-    .build().unwrap()
-    .try_deserialize::<PKConfig>().unwrap());
+    pub static ref CONFIG: Arc<PKConfig> = {
+        if let Ok(var) = std::env::var("NOMAD_ALLOC_INDEX")
+            && std::env::var("pluralkit__discord__cluster__total_nodes").is_ok() {
+            std::env::set_var("pluralkit__discord__cluster__node_id", var);
+        }
+
+        Arc::new(Config::builder()
+        .add_source(config::Environment::with_prefix("pluralkit").separator("__"))
+        .build().unwrap()
+        .try_deserialize::<PKConfig>().unwrap())
+    };
 }

lib/libpk/src/db/mod.rs

@@ -4,6 +4,7 @@ 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");

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

@@ -0,0 +1,87 @@
+use sqlx::{PgPool, Postgres, Transaction};
+
+use crate::db::types::avatars::*;
+
+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 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(())
+}

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

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

lib/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,
+}

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

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

lib/libpk/src/lib.rs

@@ -1,30 +1,25 @@
-use gethostname::gethostname;
+#![feature(let_chains)]
 use metrics_exporter_prometheus::PrometheusBuilder;
-use tracing_subscriber::{prelude::__tracing_subscriber_SubscriberExt, EnvFilter, Registry};
+use tracing_subscriber::EnvFilter;
 
 pub mod db;
 pub mod proto;
+pub mod util;
 
 pub mod _config;
 pub use crate::_config::CONFIG as config;
 
 pub fn init_logging(component: &str) -> anyhow::Result<()> {
-    let subscriber = Registry::default()
-        .with(EnvFilter::from_default_env())
-        .with(tracing_subscriber::fmt::layer());
-
-    if let Some(gelf_url) = &config.gelf_log_url {
-        let gelf_logger = tracing_gelf::Logger::builder()
-            .additional_field("component", component)
-            .additional_field("hostname", gethostname().to_str());
-        let mut conn_handle = gelf_logger
-            .init_udp_with_subscriber(gelf_url, subscriber)
-            .unwrap();
-        tokio::spawn(async move { conn_handle.connect().await });
+    // todo: fix component
+    if config.json_log {
+        tracing_subscriber::fmt()
+            .json()
+            .with_env_filter(EnvFilter::from_default_env())
+            .init();
     } else {
-        // gelf_logger internally sets the global subscriber
-        tracing::subscriber::set_global_default(subscriber)
-            .expect("unable to set global subscriber");
+        tracing_subscriber::fmt()
+            .with_env_filter(EnvFilter::from_default_env())
+            .init();
     }
 
     Ok(())

lib/libpk/src/util/mod.rs

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

lib/libpk/src/util/redis.rs

@@ -0,0 +1,15 @@
+use fred::error::RedisError;
+
+pub trait RedisErrorExt<T> {
+    fn to_option_or_error(self) -> Result<Option<T>, RedisError>;
+}
+
+impl<T> RedisErrorExt<T> for Result<T, RedisError> {
+    fn to_option_or_error(self) -> Result<Option<T>, RedisError> {
+        match self {
+            Ok(v) => Ok(Some(v)),
+            Err(error) if error.is_not_found() => Ok(None),
+            Err(error) => Err(error),
+        }
+    }
+}