667656cf by Jabis Sevón

fmt

1 parent 1c5a47a0
Showing 1 changed file with 85 additions and 37 deletions
use anyhow::{Result, Context};
use tokio_postgres::{Client, NoTls};
use log::{error, info, warn, trace};
use std::{fs,env};
use std::path::Path;
use anyhow::{Context, Result};
use chrono::Utc;
use clap::{Arg, Command};
use dotenv::dotenv;
use log::{error, info, trace, warn};
use std::path::Path;
use std::{env, fs};
use tokio_postgres::{Client, NoTls};
#[tokio::main]
async fn main() {
......@@ -22,7 +22,7 @@ async fn main() {
.long("data-dir")
.env("MATTERMOST_DATA_DIRECTORY")
.help("Path to the Mattermost data directory")
.required(true)
.required(true),
)
.arg(
Arg::new("db_name")
......@@ -30,7 +30,7 @@ async fn main() {
.long("db-name")
.env("DATABASE_NAME")
.help("Database name")
.required(true)
.required(true),
)
.arg(
Arg::new("db_user")
......@@ -38,7 +38,7 @@ async fn main() {
.long("db-user")
.env("DATABASE_USER")
.help("Database user")
.required(true)
.required(true),
)
.arg(
Arg::new("db_password")
......@@ -54,7 +54,7 @@ async fn main() {
.long("db-host")
.env("DATABASE_HOST")
.help("Database host")
.required(true)
.required(true),
)
.arg(
Arg::new("db_port")
......@@ -62,7 +62,7 @@ async fn main() {
.long("db-port")
.env("DATABASE_PORT")
.help("Database port")
.required(true)
.required(true),
)
.arg(
Arg::new("retention_days")
......@@ -70,7 +70,7 @@ async fn main() {
.long("retention-days")
.env("RETENTION_DAYS")
.help("Number of days to retain data")
.required(true)
.required(true),
)
.arg(
Arg::new("file_batch_size")
......@@ -84,13 +84,13 @@ async fn main() {
Arg::new("remove_posts")
.long("remove-posts")
.help("Wipe posts older than timestamp")
.required(false)
.required(false),
)
.arg(
Arg::new("dry_run")
.long("dry-run")
.help("Perform a dry run without making any changes")
.required(false)
.required(false),
)
.get_matches();
......@@ -105,9 +105,13 @@ async fn main() {
let remove_posts = matches.contains_id("remove_posts");
let dry_run = matches.contains_id("dry_run");
let retention_days = retention_days.parse::<i64>().expect("fucking hell retention");
let file_batch_size = file_batch_size.parse::<usize>().expect("fucking hell retention");
//let file_batch_size = file_batch_size.parse::<usize>().expect("fucking hell retention");
let retention_days = retention_days
.parse::<i64>()
.expect("fucking hell retention");
let file_batch_size = file_batch_size
.parse::<usize>()
.expect("fucking hell batch size");
if let Err(err) = clean(
mattermost_data_directory,
database_name,
......@@ -119,7 +123,9 @@ async fn main() {
file_batch_size,
remove_posts,
dry_run,
).await {
)
.await
{
error!("Cleaning operation failed: {}", err);
} else {
info!("Cleaning operation completed successfully.");
......@@ -152,7 +158,9 @@ pub async fn clean(
database_user, database_password, database_host, database_port, database_name
);
trace!("Connection string: {}", &connection_string);
let (client, connection) = tokio_postgres::connect(&connection_string, NoTls).await.context("Failed to connect to the database")?;
let (client, connection) = tokio_postgres::connect(&connection_string, NoTls)
.await
.context("Failed to connect to the database")?;
tokio::spawn(async move {
if let Err(e) = connection.await {
......@@ -163,7 +171,14 @@ pub async fn clean(
info!("Connection established: OK");
let millisecond_epoch = (Utc::now() - chrono::Duration::days(retention_days)).timestamp_millis();
clean_files(&client, millisecond_epoch, mattermost_data_directory, file_batch_size, dry_run).await?;
clean_files(
&client,
millisecond_epoch,
mattermost_data_directory,
file_batch_size,
dry_run,
)
.await?;
delete_file_info_rows(&client, millisecond_epoch, dry_run).await?;
if remove_posts {
delete_post_rows(&client, millisecond_epoch, dry_run).await?;
......@@ -192,7 +207,8 @@ async fn clean_files(
file_batch_size,
batch,
dry_run,
).await?;
)
.await?;
batch += 1;
}
......@@ -214,13 +230,14 @@ async fn clean_files_batch(
OFFSET $2
LIMIT $3;
";
trace!("Querying: {}",&query);
trace!("Querying: {}", &query);
let offset = (batch * file_batch_size) as i64;
let limit = file_batch_size as i64;
trace!("params: {} {} {}",&millisecond_epoch, &offset, &limit);
trace!("params: {} {} {}", &millisecond_epoch, &offset, &limit);
let rows = client
.query(query, &[&millisecond_epoch, &offset, &limit])
.await.context("Failed to fetch file info rows")?;
.await
.context("Failed to fetch file info rows")?;
let mut more_results = false;
......@@ -231,23 +248,38 @@ async fn clean_files_batch(
let preview_path: String = row.get("previewpath");
if dry_run {
info!("[DRY RUN] Would remove: {:?}, {:?}, {:?}", path, thumbnail_path, preview_path);
info!(
"[DRY RUN] Would remove: {:?}, {:?}, {:?}",
path, thumbnail_path, preview_path
);
} else {
remove_files(mattermost_data_directory, &path, &thumbnail_path, &preview_path).context("Failed to remove files")?;
remove_files(
mattermost_data_directory,
&path,
&thumbnail_path,
&preview_path,
)
.context("Failed to remove files")?;
}
}
Ok(more_results)
}
fn remove_files(base_dir: &str, path: &str, thumbnail_path: &str, preview_path: &str) -> Result<()> {
fn remove_files(
base_dir: &str,
path: &str,
thumbnail_path: &str,
preview_path: &str,
) -> Result<()> {
let files = [path, thumbnail_path, preview_path];
let mut num_deleted = 0;
for file in files {
if !file.is_empty() {
let full_path = Path::new(base_dir).join(file);
if full_path.exists() {
fs::remove_file(full_path.clone()).context(format!("Failed to delete file: {:?}", &full_path))?;
fs::remove_file(full_path.clone())
.context(format!("Failed to delete file: {:?}", &full_path))?;
trace!("Removed: {:#?} ", &full_path);
num_deleted += 1;
} else {
......@@ -256,25 +288,35 @@ fn remove_files(base_dir: &str, path: &str, thumbnail_path: &str, preview_path:
}
}
if num_deleted > 0 {
info!("Deleted: {} files. Main file: {}",num_deleted,path);
info!("Deleted: {} files. Main file: {}", num_deleted, path);
} else {
trace!("No files to be deleted");
}
Ok(())
}
async fn delete_file_info_rows(client: &Client, millisecond_epoch: i64, dry_run: bool) -> Result<()> {
async fn delete_file_info_rows(
client: &Client,
millisecond_epoch: i64,
dry_run: bool,
) -> Result<()> {
let query = "
DELETE FROM fileinfo
WHERE createat < $1;
";
trace!("Querying: {}",&query);
trace!("Params: {:#?}",&millisecond_epoch);
trace!("Querying: {}", &query);
trace!("Params: {:#?}", &millisecond_epoch);
if dry_run {
info!("[DRY RUN] Would delete file info rows older than {}", millisecond_epoch);
info!(
"[DRY RUN] Would delete file info rows older than {}",
millisecond_epoch
);
return Ok(());
}
let result = client.execute(query, &[&millisecond_epoch]).await.context("Failed to delete file info rows")?;
let result = client
.execute(query, &[&millisecond_epoch])
.await
.context("Failed to delete file info rows")?;
info!("Removed {} file information rows", result);
Ok(())
}
......@@ -284,13 +326,19 @@ async fn delete_post_rows(client: &Client, millisecond_epoch: i64, dry_run: bool
DELETE FROM posts
WHERE createat < $1;
";
trace!("Querying: {}",&query);
trace!("Params: {:#?}",&millisecond_epoch);
trace!("Querying: {}", &query);
trace!("Params: {:#?}", &millisecond_epoch);
if dry_run {
info!("[DRY RUN] Would delete post rows older than {}", millisecond_epoch);
info!(
"[DRY RUN] Would delete post rows older than {}",
millisecond_epoch
);
return Ok(());
}
let result = client.execute(query, &[&millisecond_epoch]).await.context("Failed to delete post rows")?;
let result = client
.execute(query, &[&millisecond_epoch])
.await
.context("Failed to delete post rows")?;
info!("Removed {} post rows", result);
Ok(())
}
......
Styling with Markdown is supported
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!