import logging
import os
from functools import wraps


def data_migration(migration):
    @wraps(migration)
    def wrapped(*args, **kwargs):
        if not os.getenv("BILLING_PROXY_RUN_DATA_MIGRATIONS", False):
            logging.getLogger("alembic.runtime.migration").warning(
                "Skipping data migration"
            )
            return

        migration(*args, **kwargs)

    return wrapped
