import logging

from alembic.config import Config
from alembic.script import ScriptDirectory
from alembic.environment import EnvironmentContext
from alembic.migration import __name__ as migration_name

def run_alembic_migration(log_handler=None):
  if log_handler:
    logging.getLogger(migration_name).addHandler(log_handler)

  config = Config()
  config.set_main_option("script_location", "data:migrations")
  script = ScriptDirectory.from_config(config)

  def fn(rev, context):
    return script._upgrade_revs('head', rev)

  with EnvironmentContext(config, script, fn=fn, destination_rev='head'):
    script.run_env()