1c5a47a0 by Jabis Sevón

initial commit

0 parents
MATTERMOST_DATA_DIRECTORY="/opt/mattermost/data/"
DATABASE_NAME="mattermost"
DATABASE_USER="mattermost"
PGPASSWORD=""
DATABASE_HOST="127.0.0.1"
DATABASE_PORT="5432"
RETENTION_DAYS="730"
FILE_BATCH_SIZE="500"
/data
/target
.env
mmjab
*.sql
*.csv
*.log
\ No newline at end of file
This diff is collapsed. Click to expand it.
[package]
name = "mmjab"
version = "0.1.0"
edition = "2021"
[dependencies]
chrono = "0.4.26"
tokio = { version = "*", features = ["macros", "signal", "rt-multi-thread"] }
tokio-util = { version = "0.7.10", features = ["io"] }
tokio-postgres = { version = "0.7.12", features= ["runtime","array-impls","js","with-bit-vec-0_6", "with-chrono-0_4"] }
log = {version = "0.4.25", features = ["std","serde"] }
anyhow = "1.0.95"
clap = { version = "4.5.27", features = ["env"] }
env_logger = "0.11.6"
dotenv = { version = "0.15.0", features = ["clap", "cli"] }
[profile.release]
lto = true
opt-level = 3
codegen-units = 1
### Mattermost Team version Jabis'd cleanup utilities
\ No newline at end of file
use anyhow::{Result, Context};
use tokio_postgres::{Client, NoTls};
use log::{error, info, warn, trace};
use std::{fs,env};
use std::path::Path;
use chrono::Utc;
use clap::{Arg, Command};
use dotenv::dotenv;
#[tokio::main]
async fn main() {
if env::var("RUST_LOG").is_err() {
env::set_var("RUST_LOG", "info")
}
env_logger::init();
dotenv().ok();
let matches = Command::new("Clean Utility")
.version("1.0")
.about("Cleans old files and database rows based on retention policies")
.arg(
Arg::new("data_dir")
.long("data-dir")
.env("MATTERMOST_DATA_DIRECTORY")
.help("Path to the Mattermost data directory")
.required(true)
)
.arg(
Arg::new("db_name")
.short('n')
.long("db-name")
.env("DATABASE_NAME")
.help("Database name")
.required(true)
)
.arg(
Arg::new("db_user")
.short('u')
.long("db-user")
.env("DATABASE_USER")
.help("Database user")
.required(true)
)
.arg(
Arg::new("db_password")
.short('p')
.long("db-password")
.env("PGPASSWORD")
.help("Database password")
.required(true),
)
.arg(
Arg::new("db_host")
.short('h')
.long("db-host")
.env("DATABASE_HOST")
.help("Database host")
.required(true)
)
.arg(
Arg::new("db_port")
.short('P')
.long("db-port")
.env("DATABASE_PORT")
.help("Database port")
.required(true)
)
.arg(
Arg::new("retention_days")
.short('D')
.long("retention-days")
.env("RETENTION_DAYS")
.help("Number of days to retain data")
.required(true)
)
.arg(
Arg::new("file_batch_size")
.short('b')
.long("file-batch-size")
.env("FILE_BATCH_SIZE")
.help("Batch size for file deletion")
.required(true),
)
.arg(
Arg::new("remove_posts")
.long("remove-posts")
.help("Wipe posts older than timestamp")
.required(false)
)
.arg(
Arg::new("dry_run")
.long("dry-run")
.help("Perform a dry run without making any changes")
.required(false)
)
.get_matches();
let mattermost_data_directory = matches.get_one::<String>("data_dir").unwrap();
let database_name = matches.get_one::<String>("db_name").unwrap();
let database_user = matches.get_one::<String>("db_user").unwrap();
let database_password = matches.get_one::<String>("db_password").unwrap();
let database_host = matches.get_one::<String>("db_host").unwrap();
let database_port = matches.get_one::<String>("db_port").unwrap();
let retention_days = matches.get_one::<String>("retention_days").unwrap();
let file_batch_size = matches.get_one::<String>("file_batch_size").unwrap();
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");
if let Err(err) = clean(
mattermost_data_directory,
database_name,
database_user,
database_password,
database_host,
database_port,
retention_days,
file_batch_size,
remove_posts,
dry_run,
).await {
error!("Cleaning operation failed: {}", err);
} else {
info!("Cleaning operation completed successfully.");
}
}
pub async fn clean(
mattermost_data_directory: &str,
database_name: &str,
database_user: &str,
database_password: &str,
database_host: &str,
database_port: &str,
retention_days: i64,
file_batch_size: usize,
remove_posts: bool,
dry_run: bool,
) -> Result<()> {
validate(
mattermost_data_directory,
database_name,
database_user,
database_host,
retention_days,
file_batch_size,
)?;
let connection_string = format!(
"postgres://{}:{}@{}:{}/{}?sslmode=disable",
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")?;
tokio::spawn(async move {
if let Err(e) = connection.await {
warn!("error happened at spawn {e}");
eprintln!("connection error: {}", e);
}
});
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?;
delete_file_info_rows(&client, millisecond_epoch, dry_run).await?;
if remove_posts {
delete_post_rows(&client, millisecond_epoch, dry_run).await?;
} else {
info!("Skipping posts removal")
}
Ok(())
}
async fn clean_files(
client: &Client,
millisecond_epoch: i64,
mattermost_data_directory: &str,
file_batch_size: usize,
dry_run: bool,
) -> Result<()> {
let mut batch = 0;
let mut more_results = true;
while more_results {
more_results = clean_files_batch(
client,
millisecond_epoch,
mattermost_data_directory,
file_batch_size,
batch,
dry_run,
).await?;
batch += 1;
}
Ok(())
}
async fn clean_files_batch(
client: &Client,
millisecond_epoch: i64,
mattermost_data_directory: &str,
file_batch_size: usize,
batch: usize,
dry_run: bool,
) -> Result<bool> {
let query = "
SELECT path, thumbnailpath, previewpath
FROM fileinfo
WHERE createat < $1
OFFSET $2
LIMIT $3;
";
trace!("Querying: {}",&query);
let offset = (batch * file_batch_size) as i64;
let limit = file_batch_size as i64;
trace!("params: {} {} {}",&millisecond_epoch, &offset, &limit);
let rows = client
.query(query, &[&millisecond_epoch, &offset, &limit])
.await.context("Failed to fetch file info rows")?;
let mut more_results = false;
for row in rows {
more_results = true;
let path: String = row.get("path");
let thumbnail_path: String = row.get("thumbnailpath");
let preview_path: String = row.get("previewpath");
if dry_run {
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")?;
}
}
Ok(more_results)
}
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))?;
trace!("Removed: {:#?} ", &full_path);
num_deleted += 1;
} else {
trace!("Path does not exist: {:#?} ", &full_path);
}
}
}
if num_deleted > 0 {
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<()> {
let query = "
DELETE FROM fileinfo
WHERE createat < $1;
";
trace!("Querying: {}",&query);
trace!("Params: {:#?}",&millisecond_epoch);
if dry_run {
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")?;
info!("Removed {} file information rows", result);
Ok(())
}
async fn delete_post_rows(client: &Client, millisecond_epoch: i64, dry_run: bool) -> Result<()> {
let query = "
DELETE FROM posts
WHERE createat < $1;
";
trace!("Querying: {}",&query);
trace!("Params: {:#?}",&millisecond_epoch);
if dry_run {
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")?;
info!("Removed {} post rows", result);
Ok(())
}
fn validate(
mattermost_data_directory: &str,
database_name: &str,
database_user: &str,
database_host: &str,
retention_days: i64,
file_batch_size: usize,
) -> Result<()> {
if mattermost_data_directory.is_empty()
|| database_name.is_empty()
|| database_user.is_empty()
|| database_host.is_empty()
|| retention_days <= 0
|| file_batch_size == 0
{
anyhow::bail!("Invalid input parameters");
}
Ok(())
}
#!/bin/bash
set -a; source .env; set +a
./mmjab
Styling with Markdown is supported
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!