MineSeeker/docker/storage-migrate/migrate.py

#!/usr/bin/env python3
"""Copy the legacy MinIO bucket into RustFS exactly once, without deleting source data."""

import json
import itertools
import os
import sys
import tempfile
import time

import boto3
from botocore.config import Config
from botocore.exceptions import ClientError, EndpointConnectionError


MARKER = "/state/minio-to-rustfs-v2.json"
BUCKET = os.environ.get("STORAGE_BUCKET", "mineseeker")
ACCESS_KEY = os.environ["STORAGE_ACCESS_KEY"]
SECRET_KEY = os.environ["STORAGE_SECRET_KEY"]


def client(endpoint):
    return boto3.client(
        "s3",
        endpoint_url=endpoint,
        aws_access_key_id=ACCESS_KEY,
        aws_secret_access_key=SECRET_KEY,
        region_name="us-east-1",
        config=Config(signature_version="s3v4", s3={"addressing_style": "path"}),
    )


def wait_for_service(s3, label):
    for attempt in range(60):
        try:
            s3.list_buckets()
            print(f"{label} is ready", flush=True)
            return
        except (EndpointConnectionError, ClientError) as error:
            if attempt == 59:
                raise RuntimeError(f"{label} did not become ready: {error}") from error
            time.sleep(2)


def bucket_exists(s3, bucket):
    try:
        s3.head_bucket(Bucket=bucket)
        return True
    except ClientError as error:
        code = error.response.get("Error", {}).get("Code", "")
        status = error.response.get("ResponseMetadata", {}).get("HTTPStatusCode")
        if code in {"NoSuchBucket", "404", "NotFound"} or status == 404:
            return False
        raise


def ensure_bucket(s3, bucket):
    if not bucket_exists(s3, bucket):
        s3.create_bucket(Bucket=bucket)


def reject_versioned_bucket(source, bucket):
    versioning = source.get_bucket_versioning(Bucket=bucket).get("Status")
    if versioning in {"Enabled", "Suspended"}:
        raise RuntimeError(
            f"Bucket {bucket!r} has versioning status {versioning}; "
            "refusing to omit historical object versions or delete markers"
        )


def copy_bucket(source, target, bucket):
    copied = 0
    paginator = source.get_paginator("list_objects_v2")
    for page in paginator.paginate(Bucket=bucket):
        for item in page.get("Contents", []):
            key = item["Key"]
            response = source.get_object(Bucket=bucket, Key=key)
            try:
                # Botocore's request checksum calculation needs a seekable body.
                # Spool to disk beyond 16 MiB so large objects don't consume RAM.
                with tempfile.SpooledTemporaryFile(max_size=16 * 1024 * 1024, mode="w+b") as body:
                    for chunk in response["Body"].iter_chunks(chunk_size=1024 * 1024):
                        if chunk:
                            body.write(chunk)
                    body.seek(0)
                    args = {
                        "Bucket": bucket,
                        "Key": key,
                        "Body": body,
                        "ContentLength": response["ContentLength"],
                        "Metadata": response.get("Metadata", {}),
                    }
                    for field in (
                        "ContentType",
                        "ContentEncoding",
                        "CacheControl",
                        "ContentDisposition",
                        "ContentLanguage",
                        "Expires",
                    ):
                        if response.get(field) is not None:
                            args[field] = response[field]
                    if response.get("WebsiteRedirectLocation") is not None:
                        args["WebsiteRedirectLocation"] = response["WebsiteRedirectLocation"]
                    target.put_object(**args)
            finally:
                response["Body"].close()
            copied += 1
            if copied % 100 == 0:
                print(f"Copied {copied} objects", flush=True)
    return copied


def inventory(s3, bucket):
    paginator = s3.get_paginator("list_objects_v2")
    for page in paginator.paginate(Bucket=bucket):
        for item in page.get("Contents", []):
            yield item["Key"], item["Size"]


def verify_copy(source, target, bucket):
    for source_item, target_item in itertools.zip_longest(
        inventory(source, bucket), inventory(target, bucket)
    ):
        if source_item != target_item:
            raise RuntimeError(
                f"Object inventory mismatch in bucket {bucket!r}: "
                f"source={source_item!r}, target={target_item!r}"
            )


def apply_bucket_policy(source, target, bucket):
    if bucket != BUCKET:
        try:
            response = source.get_bucket_policy(Bucket=bucket)
        except ClientError as error:
            code = error.response.get("Error", {}).get("Code", "")
            status = error.response.get("ResponseMetadata", {}).get("HTTPStatusCode")
            if code in {"NoSuchBucketPolicy", "NoSuchKey", "404", "NotFound"} or status == 404:
                return
            raise
        target.put_bucket_policy(Bucket=bucket, Policy=response["Policy"])
        return

    policy = {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Effect": "Allow",
                "Principal": "*",
                "Action": ["s3:GetObject"],
                "Resource": [
                    f"arn:aws:s3:::{bucket}/media/*",
                    f"arn:aws:s3:::{bucket}/cache/*",
                ],
            }
        ],
    }
    target.put_bucket_policy(Bucket=bucket, Policy=json.dumps(policy))


def migrate_all_buckets(source, target):
    buckets = sorted(bucket["Name"] for bucket in source.list_buckets().get("Buckets", []))
    if BUCKET not in buckets:
        raise RuntimeError(f"Required source bucket {BUCKET!r} does not exist; refusing to mark migration complete")

    totals = {}
    for bucket in buckets:
        reject_versioned_bucket(source, bucket)
        ensure_bucket(target, bucket)
        totals[bucket] = copy_bucket(source, target, bucket)
        verify_copy(source, target, bucket)
        apply_bucket_policy(source, target, bucket)
        print(f"Migrated bucket {bucket!r}: {totals[bucket]} objects", flush=True)
    return totals


def main():
    if os.path.exists(MARKER):
        print("Storage migration already completed; skipping", flush=True)
        return

    source = client(os.environ["SOURCE_ENDPOINT"])
    target = client(os.environ["TARGET_ENDPOINT"])
    wait_for_service(source, "MinIO")
    wait_for_service(target, "RustFS")
    totals = migrate_all_buckets(source, target)

    # Publish the completion marker only after the data and policy are confirmed.
    temporary_marker = f"{MARKER}.tmp"
    with open(temporary_marker, "w", encoding="utf-8") as marker:
        json.dump({"buckets": totals, "objects_copied": sum(totals.values())}, marker)
        marker.write("\n")
        marker.flush()
        os.fsync(marker.fileno())
    os.replace(temporary_marker, MARKER)
    print(
        f"Migration complete: {sum(totals.values())} objects copied across "
        f"{len(totals)} buckets; source data was left untouched",
        flush=True,
    )


if __name__ == "__main__":
    try:
        main()
    except Exception as error:  # Keep the failed container's logs useful to operators.
        print(f"Storage migration failed: {error}", file=sys.stderr, flush=True)
        raise