Refactored API code into its own module

src/api/mod.rs

@@ -1 +1,241 @@
+use crate::api::models::{
+    AddImage, AlbumQuery, CreateAlbum, ImageQuery, ImageSort, PicContext, PicOxError, Response,
+};
+use crate::config::PicOxConfig;
+use crate::model::album::Album;
+use crate::model::api_key::ApiKey;
+use crate::model::image::{Image, ImageData};
+use crate::state::Context;
+use crate::storage_manager::StorageManager;
+use axum::body::Bytes;
+use axum::extract::{DefaultBodyLimit, Multipart, Path, Query, Request, State};
+use axum::http::HeaderMap;
+use axum::middleware::Next;
+use axum::response::IntoResponse;
+use axum::routing::{get, post};
+use axum::{middleware, Json, Router};
+use j_db::database::Database;
+use j_db::model::JdbModel;
+use log::info;
+use rand::prelude::SliceRandom;
+use rand::thread_rng;
+use std::sync::Arc;
+use tokio::sync::RwLock;
 
+pub mod models;
+
+async fn create_album(
+    State(context): State<PicContext>,
+    headers: HeaderMap,
+    Json(album): Json<CreateAlbum>,
+) -> Result<Response<Album>, PicOxError> {
+    let user_id = get_user_id_from_headers(&headers)?;
+    let new_album = Album::new(&album.album_name, Vec::new(), user_id);
+
+    info!(
+        "Creating new album '{}pub pub pub ' for user {}",
+        album.album_name, user_id
+    );
+    let new_album = context.db.insert(new_album)?;
+
+    Ok(Response(new_album))
+}
+
+async fn get_album(
+    album_id: Path<u64>,
+    State(context): State<PicContext>,
+) -> Result<Response<Album>, PicOxError> {
+    let album = context.db.get::<Album>(*album_id)?;
+
+    Ok(Response(album))
+}
+
+async fn query_images(
+    image_query: Query<ImageQuery>,
+    State(context): State<PicContext>,
+) -> Result<Response<Vec<Image>>, PicOxError> {
+    let album_id = if let Some(album) = &image_query.album {
+        Some(
+            Album::find_album_by_query(
+                &context.db,
+                AlbumQuery {
+                    album_name: Some(album.to_string()),
+                },
+            )
+            .ok_or(PicOxError::AlbumNotFound)?
+            .id()
+            .unwrap(),
+        )
+    } else {
+        None
+    };
+
+    let mut images: Vec<Image> = context
+        .db
+        .filter(|_, img: &Image| {
+            if let Some(album_id) = album_id {
+                if img.album != album_id {
+                    return false;
+                }
+            }
+
+            if !image_query.tags.is_empty() {
+                let mut found = false;
+                for tag in &image_query.tags {
+                    if img.tags.contains(tag) {
+                        found = true;
+                        break;
+                    }
+                }
+
+                if !found {
+                    return false;
+                }
+            }
+
+            true
+        })?
+        .collect();
+
+    match image_query.order {
+        ImageSort::Random => {
+            images.shuffle(&mut thread_rng());
+        }
+        ImageSort::DateAscending => {
+            images.sort_by(|img_a, img_b| img_a.create_date.cmp(&img_b.create_date))
+        }
+        ImageSort::DateDescending => {
+            images.sort_by(|img_a, img_b| img_b.create_date.cmp(&img_a.create_date))
+        }
+        ImageSort::None => {}
+    }
+
+    if images.len() > image_query.limit {
+        images.drain(image_query.limit..);
+    }
+
+    Ok(Response(images))
+}
+
+fn get_user_id_from_headers(headers: &HeaderMap) -> Result<u64, PicOxError> {
+    let user = headers.get("user").ok_or(PicOxError::NoUserInHeader)?;
+    let user_str = user.to_str().unwrap();
+    let user: u64 = user_str.parse().unwrap();
+
+    Ok(user)
+}
+
+async fn add_image(
+    State(context): State<PicContext>,
+    headers: HeaderMap,
+    mut img_data: Multipart,
+) -> Result<Response<Image>, PicOxError> {
+    let mut data: Vec<u8> = Vec::new();
+    let mut metadata: Option<AddImage> = None;
+    let mut file_name = None;
+    while let Some(field) = img_data.next_field().await.unwrap() {
+        let field_name = field.name();
+        if let Some(field_name) = field_name {
+            if field_name == "metadata" {
+                let metadata_json = field.text().await.unwrap();
+                metadata = Some(serde_json::from_str(&metadata_json).unwrap())
+            } else if field_name == "img_data" {
+                file_name = Some(field.file_name().unwrap().to_string());
+                let file_segment = field.bytes().await.unwrap_or(Bytes::new());
+                data.extend_from_slice(file_segment.as_ref());
+            }
+        }
+    }
+
+    let user = headers.get("user").unwrap().to_str().unwrap();
+    let user: u64 = user.parse().unwrap();
+
+    let mut album =
+        Album::find_album_by_query(&context.db, metadata.clone().unwrap().album).unwrap();
+
+    let mut store_manager = context.store_manager.write().await;
+    let img = store_manager
+        .store_img(
+            &context.db,
+            None,
+            ImageData::Bytes(data),
+            &file_name.unwrap(),
+            user,
+            album.id().unwrap(),
+        )
+        .await?;
+
+    info!(
+        "Creating new image id={} at {:?}",
+        img.id().unwrap(),
+        img.storage_location
+    );
+
+    album.images.push(img.id().unwrap());
+
+    context.db.insert::<Album>(album)?;
+
+    Ok(Response(img))
+}
+
+async fn query_album(
+    album_query: Query<AlbumQuery>,
+    State(context): State<PicContext>,
+) -> Result<Response<Option<Album>>, PicOxError> {
+    let resp = Album::find_album_by_query(&context.db, album_query.0);
+    Ok(Response(resp))
+}
+
+async fn check_token_header(
+    State(context): State<PicContext>,
+    mut request: Request,
+    next: Next,
+) -> Result<impl IntoResponse, PicOxError> {
+    let headers = request.headers();
+
+    if let Some(token) = headers.get("token") {
+        if let Some(api_key) = ApiKey::find_api_key_by_token(&context.db, token.to_str().unwrap())?
+        {
+            info!(
+                "Authenticated user {}: '{}'",
+                api_key.id().unwrap(),
+                api_key.description
+            );
+            request
+                .headers_mut()
+                .insert("user", api_key.id().unwrap().into());
+            return Ok(next.run(request).await);
+        }
+    }
+
+    Err(PicOxError::TokenInvalid)
+}
+
+pub async fn run_picox(db: Database, config: PicOxConfig) {
+    let store_manager = StorageManager::new(config.storage_config.clone());
+
+    let context = Context {
+        db,
+        config: config.clone(),
+        store_manager: RwLock::new(store_manager),
+    };
+
+    let context = Arc::new(context);
+
+    let app = Router::new()
+        .route("/api/image/", post(add_image))
+        .layer(DefaultBodyLimit::max(1024 * 1024 * 1024))
+        .route("/api/album/create", post(create_album))
+        .layer(middleware::from_fn_with_state(
+            context.clone(),
+            check_token_header,
+        ))
+        .route("/api/album/:id", get(get_album))
+        .route("/api/album/", get(query_album))
+        .route("/api/image/", get(query_images))
+        .with_state(context);
+
+    let listener = tokio::net::TcpListener::bind(&config.host).await.unwrap();
+    info!("Serving at {}", config.host);
+    axum::serve(listener, app).await.unwrap();
+}

src/api/models.rs

@@ -0,0 +1,118 @@
+use crate::state::Context;
+use crate::storage_manager::StoreError;
+use axum::http::StatusCode;
+use axum::response::IntoResponse;
+use axum::Json;
+use axum_macros::FromRequest;
+use serde::{Deserialize, Serialize};
+use std::sync::Arc;
+
+pub type PicContext = Arc<Context>;
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CreateAlbum {
+    pub album_name: String,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct AlbumQuery {
+    pub album_name: Option<String>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct AddImage {
+    pub album: AlbumQuery,
+    pub tags: Vec<String>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, Default)]
+pub enum ImageSort {
+    Random,
+    DateAscending,
+    DateDescending,
+    #[default]
+    None,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ImageQuery {
+    pub album: Option<String>,
+    #[serde(default)]
+    pub tags: Vec<String>,
+    #[serde(default)]
+    pub order: ImageSort,
+    #[serde(default)]
+    pub limit: usize,
+}
+
+#[derive(FromRequest)]
+#[from_request(via(axum::Json), rejection(PicOxError))]
+pub struct Response<T>(pub T);
+
+impl<T> IntoResponse for Response<T>
+where
+    Json<T>: IntoResponse,
+{
+    fn into_response(self) -> axum::response::Response {
+        Json(self.0).into_response()
+    }
+}
+
+#[allow(dead_code)]
+pub enum PicOxError {
+    StoreError(StoreError),
+    DbError(j_db::error::JDbError),
+    AlbumNotFound,
+    ImageNotFound,
+    TokenInvalid,
+    NoUserInHeader,
+}
+
+impl From<StoreError> for PicOxError {
+    fn from(value: StoreError) -> Self {
+        Self::StoreError(value)
+    }
+}
+
+impl From<j_db::error::JDbError> for PicOxError {
+    fn from(value: j_db::error::JDbError) -> Self {
+        Self::DbError(value)
+    }
+}
+
+#[derive(Serialize)]
+pub struct ErrorResponse {
+    pub message: String,
+}
+
+impl IntoResponse for PicOxError {
+    fn into_response(self) -> axum::response::Response {
+        let (status, message) = match self {
+            PicOxError::StoreError(err) => match err {
+                StoreError::InvalidFile => (StatusCode::BAD_REQUEST, err.to_string()),
+                StoreError::OutOfStorage => (StatusCode::INSUFFICIENT_STORAGE, err.to_string()),
+                StoreError::ImageTooBig => (StatusCode::UNAUTHORIZED, err.to_string()),
+                StoreError::IOError(_) => (
+                    StatusCode::INTERNAL_SERVER_ERROR,
+                    "IO Error Has Occurred!".to_string(),
+                ),
+            },
+            PicOxError::DbError(err) => (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()),
+            PicOxError::AlbumNotFound => (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "No album found!".to_string(),
+            ),
+            PicOxError::ImageNotFound => (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "Image not found".to_string(),
+            ),
+            PicOxError::TokenInvalid => (StatusCode::UNAUTHORIZED, "Token is invalid".to_string()),
+            PicOxError::NoUserInHeader => (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "User not found in header".to_string(),
+            ),
+        };
+
+        (status, Response(ErrorResponse { message })).into_response()
+    }
+}

src/main.rs

@@ -5,32 +5,12 @@ mod state;
 mod storage_manager;
 
 use crate::config::PicOxConfig;
-use crate::model::album::Album;
 use crate::model::api_key::{ApiKey, ApiPermissions};
-use crate::model::image::{Image, ImageData};
-use crate::state::Context;
-use crate::storage_manager::{StorageManager, StoreError};
-use axum::body::Bytes;
-use axum::extract::{DefaultBodyLimit, Multipart, Path, Query, Request, State};
-use axum::http::{HeaderMap, StatusCode};
-use axum::middleware::Next;
-use axum::response::IntoResponse;
-use axum::routing::{get, post};
-use axum::{middleware, Json, Router};
-use axum_macros::FromRequest;
 use base64::Engine;
 use j_db::database::Database;
-use j_db::model::JdbModel;
 use log::info;
-use rand::seq::SliceRandom;
-use rand::thread_rng;
-use serde::{Deserialize, Serialize};
 use std::path::PathBuf;
-use std::sync::Arc;
 use structopt::StructOpt;
-use tokio::sync::RwLock;
-
-type PicContext = Arc<Context>;
 
 #[derive(StructOpt, Debug, Clone)]
 #[structopt(about = "PicOx Commands")]
@@ -65,319 +45,6 @@ struct Args {
     pub command: SubCommands,
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
-struct CreateAlbum {
-    pub album_name: String,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-struct AlbumQuery {
-    pub album_name: Option<String>,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-struct AddImage {
-    pub album: AlbumQuery,
-    pub tags: Vec<String>,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize, Default)]
-pub enum ImageSort {
-    Random,
-    DateAscending,
-    DateDescending,
-    #[default]
-    None,
-}
-
-#[derive(Debug, Clone, Serialize, Deserialize)]
-struct ImageQuery {
-    pub album: Option<String>,
-    #[serde(default)]
-    pub tags: Vec<String>,
-    #[serde(default)]
-    pub order: ImageSort,
-    #[serde(default)]
-    pub limit: usize,
-}
-
-#[derive(FromRequest)]
-#[from_request(via(axum::Json), rejection(PicOxError))]
-struct Response<T>(T);
-
-impl<T> IntoResponse for Response<T>
-where
-    Json<T>: IntoResponse,
-{
-    fn into_response(self) -> axum::response::Response {
-        Json(self.0).into_response()
-    }
-}
-
-#[allow(dead_code)]
-pub enum PicOxError {
-    StoreError(StoreError),
-    DbError(j_db::error::JDbError),
-    AlbumNotFound,
-    ImageNotFound,
-    TokenInvalid,
-    NoUserInHeader,
-}
-
-impl From<StoreError> for PicOxError {
-    fn from(value: StoreError) -> Self {
-        Self::StoreError(value)
-    }
-}
-
-impl From<j_db::error::JDbError> for PicOxError {
-    fn from(value: j_db::error::JDbError) -> Self {
-        Self::DbError(value)
-    }
-}
-
-#[derive(Serialize)]
-struct ErrorResponse {
-    pub message: String,
-}
-
-impl IntoResponse for PicOxError {
-    fn into_response(self) -> axum::response::Response {
-        let (status, message) = match self {
-            PicOxError::StoreError(err) => match err {
-                StoreError::InvalidFile => (StatusCode::BAD_REQUEST, err.to_string()),
-                StoreError::OutOfStorage => (StatusCode::INSUFFICIENT_STORAGE, err.to_string()),
-                StoreError::ImageTooBig => (StatusCode::UNAUTHORIZED, err.to_string()),
-                StoreError::IOError(_) => (
-                    StatusCode::INTERNAL_SERVER_ERROR,
-                    "IO Error Has Occurred!".to_string(),
-                ),
-            },
-            PicOxError::DbError(err) => (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()),
-            PicOxError::AlbumNotFound => (
-                StatusCode::INTERNAL_SERVER_ERROR,
-                "No album found!".to_string(),
-            ),
-            PicOxError::ImageNotFound => (
-                StatusCode::INTERNAL_SERVER_ERROR,
-                "Image not found".to_string(),
-            ),
-            PicOxError::TokenInvalid => (StatusCode::UNAUTHORIZED, "Token is invalid".to_string()),
-            PicOxError::NoUserInHeader => (
-                StatusCode::INTERNAL_SERVER_ERROR,
-                "User not found in header".to_string(),
-            ),
-        };
-
-        (status, Response(ErrorResponse { message })).into_response()
-    }
-}
-
-async fn create_album(
-    State(context): State<PicContext>,
-    headers: HeaderMap,
-    Json(album): Json<CreateAlbum>,
-) -> Result<Response<Album>, PicOxError> {
-    let user_id = get_user_id_from_headers(&headers)?;
-    let new_album = Album::new(&album.album_name, Vec::new(), user_id);
-
-    info!(
-        "Creating new album '{}' for user {}",
-        album.album_name, user_id
-    );
-    let new_album = context.db.insert(new_album)?;
-
-    Ok(Response(new_album))
-}
-
-async fn get_album(
-    album_id: Path<u64>,
-    State(context): State<PicContext>,
-) -> Result<Response<Album>, PicOxError> {
-    let album = context.db.get::<Album>(*album_id)?;
-
-    Ok(Response(album))
-}
-
-async fn query_images(
-    image_query: Query<ImageQuery>,
-    State(context): State<PicContext>,
-) -> Result<Response<Vec<Image>>, PicOxError> {
-    let album_id = if let Some(album) = &image_query.album {
-        Some(
-            Album::find_album_by_query(
-                &context.db,
-                AlbumQuery {
-                    album_name: Some(album.to_string()),
-                },
-            )
-            .ok_or(PicOxError::AlbumNotFound)?
-            .id()
-            .unwrap(),
-        )
-    } else {
-        None
-    };
-
-    let mut images: Vec<Image> = context
-        .db
-        .filter(|_, img: &Image| {
-            if let Some(album_id) = album_id {
-                if img.album != album_id {
-                    return false;
-                }
-            }
-
-            if !image_query.tags.is_empty() {
-                let mut found = false;
-                for tag in &image_query.tags {
-                    if img.tags.contains(tag) {
-                        found = true;
-                        break;
-                    }
-                }
-
-                if !found {
-                    return false;
-                }
-            }
-
-            true
-        })?
-        .collect();
-
-    match image_query.order {
-        ImageSort::Random => {
-            images.shuffle(&mut thread_rng());
-        }
-        ImageSort::DateAscending => {
-            images.sort_by(|img_a, img_b| img_a.create_date.cmp(&img_b.create_date))
-        }
-        ImageSort::DateDescending => {
-            images.sort_by(|img_a, img_b| img_b.create_date.cmp(&img_a.create_date))
-        }
-        ImageSort::None => {}
-    }
-
-    if images.len() > image_query.limit {
-        images.drain(image_query.limit..);
-    }
-
-    Ok(Response(images))
-}
-
-fn get_user_id_from_headers(headers: &HeaderMap) -> Result<u64, PicOxError> {
-    let user = headers.get("user").ok_or(PicOxError::NoUserInHeader)?;
-    let user_str = user.to_str().unwrap();
-    let user: u64 = user_str.parse().unwrap();
-
-    Ok(user)
-}
-
-async fn add_image(
-    State(context): State<PicContext>,
-    headers: HeaderMap,
-    mut img_data: Multipart,
-) -> Result<Response<Image>, PicOxError> {
-    let mut data: Vec<u8> = Vec::new();
-    let mut metadata: Option<AddImage> = None;
-    let mut file_name = None;
-    while let Some(field) = img_data.next_field().await.unwrap() {
-        let field_name = field.name();
-        if let Some(field_name) = field_name {
-            if field_name == "metadata" {
-                let metadata_json = field.text().await.unwrap();
-                metadata = Some(serde_json::from_str(&metadata_json).unwrap())
-            } else if field_name == "img_data" {
-                file_name = Some(field.file_name().unwrap().to_string());
-                let file_segment = field.bytes().await.unwrap_or(Bytes::new());
-                data.extend_from_slice(file_segment.as_ref());
-            }
-        }
-    }
-
-    let user = headers.get("user").unwrap().to_str().unwrap();
-    let user: u64 = user.parse().unwrap();
-
-    let mut album =
-        Album::find_album_by_query(&context.db, metadata.clone().unwrap().album).unwrap();
-
-    let mut store_manager = context.store_manager.write().await;
-    let img = store_manager
-        .store_img(
-            &context.db,
-            None,
-            ImageData::Bytes(data),
-            &file_name.unwrap(),
-            user,
-            album.id().unwrap(),
-        )
-        .await?;
-
-    album.images.push(img.id().unwrap());
-
-    context.db.insert::<Album>(album)?;
-
-    Ok(Response(img))
-}
-
-async fn query_album(
-    album_query: Query<AlbumQuery>,
-    State(context): State<PicContext>,
-) -> Result<Response<Option<Album>>, PicOxError> {
-    let resp = Album::find_album_by_query(&context.db, album_query.0);
-    Ok(Response(resp))
-}
-
-async fn check_token_header(
-    State(context): State<PicContext>,
-    mut request: Request,
-    next: Next,
-) -> Result<impl IntoResponse, PicOxError> {
-    let headers = request.headers();
-
-    if let Some(token) = headers.get("token") {
-        if let Some(api_key) = ApiKey::find_api_key_by_token(&context.db, token.to_str().unwrap())?
-        {
-            request
-                .headers_mut()
-                .insert("user", api_key.id().unwrap().into());
-            return Ok(next.run(request).await);
-        }
-    }
-
-    Err(PicOxError::TokenInvalid)
-}
-
-async fn run_picox(db: Database, config: PicOxConfig) {
-    let store_manager = StorageManager::new(config.storage_config.clone());
-
-    let context = Context {
-        db,
-        config: config.clone(),
-        store_manager: RwLock::new(store_manager),
-    };
-
-    let context = Arc::new(context);
-
-    let app = Router::new()
-        .route("/api/image/", post(add_image))
-        .layer(DefaultBodyLimit::max(1024 * 1024 * 1024))
-        .route("/api/album/create", post(create_album))
-        .layer(middleware::from_fn_with_state(
-            context.clone(),
-            check_token_header,
-        ))
-        .route("/api/album/:id", get(get_album))
-        .route("/api/album/", get(query_album))
-        .route("/api/image/", get(query_images))
-        .with_state(context);
-
-    let listener = tokio::net::TcpListener::bind(&config.host).await.unwrap();
-    info!("Serving at {}", config.host);
-    axum::serve(listener, app).await.unwrap();
-}
-
 #[tokio::main]
 async fn main() {
     let args = Args::from_args();
@@ -391,7 +58,7 @@ async fn main() {
 
     match args.command {
         SubCommands::Start => {
-            run_picox(db, config).await;
+            api::run_picox(db, config).await;
         }
         SubCommands::CreateKey {
             description,

src/model/album.rs

@@ -1,4 +1,4 @@
-use crate::AlbumQuery;
+use crate::api::models::AlbumQuery;
 use chrono::{DateTime, Utc};
 use j_db::database::Database;
 use j_db::model::JdbModel;