Script to move datasets from one hierarchical store to another

[alpapan]

Galaxy version
"version_major": "23.2", "version_minor": "2.dev0"

Describe the bug
Not a bug but I thought I share this for galaxy admins not using distributed storage.
Please let me know if this is not the appropriate place or you prefer a pull request.

Once one sets up multiple hierarchical stores (e.g. due to space running out), this is a script that will move files that are above certain size and haven't been updated for a certain number of days (and are not purged).

It relies on the UUID of the datasets.

It has a dry run function.

Feel free to comment/edit it if you think it is appropriate.

One idea I was thinking was to edit this to search for purged files only (of certain size) and cron move them to a slow disk as they await deletion (only issue there is that if someone unpurged them then they would stay at the slow disk... just an idea)

thank you,
alexie

#!/usr/bin/env python3

## NB: script has been written with some AI assistance

## example run
## move_large_datasets_to_store.py --galaxy-config config/galaxy.yml 
## --galaxy-data-dir /old_dir/database/objects/
## --target-data-dir /new_dir/database/objects/
## --file-size-threshold-gb 50 --dry-run

import os
import shutil
import psycopg2
import yaml
from argparse import ArgumentParser
from datetime import datetime

def get_db_connection(galaxy_config_path):
    with open(galaxy_config_path, 'r') as config_file:
        config = yaml.safe_load(config_file)
    return config.get('galaxy', {}).get('database_connection')

def format_uuid(uuid):
    return f"{uuid[:8]}-{uuid[8:12]}-{uuid[12:16]}-{uuid[16:20]}-{uuid[20:]}"

def construct_dataset_path(dataset_uuid):
    subdir = os.path.join(dataset_uuid[0], dataset_uuid[1], dataset_uuid[2])
    filename = f"dataset_{dataset_uuid}.dat"
    return os.path.join(subdir, filename)

def gb_to_bytes(gb_size):
    return gb_size * (1024 ** 3)

def bytes_to_gb(bytes_size):
    return bytes_size / (1024 ** 3)

def move_datasets(galaxy_config_path, galaxy_data_dir, target_data_dir, file_size_threshold_gb, days_threshold, purged_only, dry_run=False):
    galaxy_db_conn = get_db_connection(galaxy_config_path)
    
    if not galaxy_db_conn:
        print("Error: Unable to find 'database_connection' in galaxy.yml")
        return

    file_size_threshold_bytes = gb_to_bytes(file_size_threshold_gb)

    try:
        conn = psycopg2.connect(galaxy_db_conn)
        cursor = conn.cursor()

        purged_condition = f"AND purged = {str(purged_only).lower()}"

        query = f"""
        SELECT uuid, file_size 
        FROM dataset 
        WHERE update_time < now()::DATE-{days_threshold} 
        AND file_size > {file_size_threshold_bytes} 
        AND deleted = false 
        {purged_condition} 
        ORDER BY file_size
        """
        
        cursor.execute(query)
        datasets = cursor.fetchall()

        total_size_gb = 0
        total_files = len(datasets)
        moved_files = 0

        for uuid, file_size in datasets:
            formatted_uuid = format_uuid(uuid)
            dataset_path = construct_dataset_path(formatted_uuid)
            current_path = os.path.join(galaxy_data_dir, dataset_path)
            target_path = os.path.join(target_data_dir, dataset_path)

            # Check if the dataset already exists in the target directory
            if os.path.exists(target_path):
                print(f"Info: File {target_path} already exists in target directory. Skipping.")
                continue

            # Check if the dataset exists in the source directory
            if not os.path.exists(current_path):
                print(f"Warning: Source file {current_path} does not exist in galaxy data directory.")
                continue

            file_size_gb = bytes_to_gb(file_size)
            total_size_gb += file_size_gb

            if dry_run:
                print(f"Would move dataset {formatted_uuid} from {current_path} to {target_path}. Size: {file_size_gb:.2f} GB")
            else:
                if not os.path.exists(os.path.dirname(target_path)):
                    os.makedirs(os.path.dirname(target_path), exist_ok=True)
                
                try:
                    shutil.copy2(current_path, target_path)
                    os.remove(current_path)
                    print(f"Moved dataset {formatted_uuid} to {target_path}. Size: {file_size_gb:.2f} GB")
                    moved_files += 1
                except IOError as e:
                    print(f"Error moving file {formatted_uuid}: {e}")

        print(f"Total size of {total_files} files processed: {total_size_gb:.2f} GB")
        print(f"Number of files moved: {moved_files}")

        cursor.close()
        conn.close()

    except psycopg2.Error as e:
        print(f"Database error: {e}")

if __name__ == "__main__":
    parser = ArgumentParser(description="Move Galaxy datasets older than a specified number of days and larger than a specified size to another object store.")
    parser.add_argument("--galaxy-config", required=True, help="Path to the galaxy.yml configuration file")
    parser.add_argument("--galaxy-data-dir", required=True, help="Path to the current Galaxy data directory")
    parser.add_argument("--target-data-dir", required=True, help="Path to the target data directory")
    parser.add_argument("--file-size-threshold-gb", type=float, default=10.0, help="File size threshold in gigabytes (default: 10.0)")
    parser.add_argument("--days-threshold", type=int, default=7, help="Number of days to search for update_time (default: 7, 0 is allowed)")
    parser.add_argument("--purged-only", action="store_true", help="Include only purged datasets in the query")
    parser.add_argument("--dry-run", action="store_true", help="Perform a dry run without moving files")

    args = parser.parse_args()
    move_datasets(args.galaxy_config, args.galaxy_data_dir, args.target_data_dir, args.file_size_threshold_gb, args.days_threshold, args.purged_only, args.dry_run)