如何为alembic revision指定非默认script.py.mako模板生成修订

一、指定非默认script.py.mako模板生成修订 #

Alembic支持直接通过命令行参数指定自定义模板，无需额外配置多套默认模板：

• 执行命令时用--template参数传入模板文件路径，比如项目根目录下有templates/custom_script.py.mako，命令如下：

  alembic revision --autogenerate --message "自定义修订内容" --template templates/custom_script.py.mako
  

• 若要统一管理模板，可在alembic.ini中配置模板目录：

  template_path = alembic/templates
  

  之后直接指定模板文件名即可：

  alembic revision --autogenerate --message "自定义修订内容" --template custom_script.py.mako
  

二、基于命令行标识生成带自定义逻辑的修订 #

如果需要根据命令行传入的标识（如--with-custom-actions）生成包含数据库写入、API调用的修订，可通过修改env.py和模板实现：

1. 传递自定义命令行参数：
   执行命令时通过--分隔自定义参数：

   alembic revision --autogenerate --message "带自定义操作的修订" -- --with-custom-actions
   

2. 在env.py中捕获参数并注入模板上下文：
   修改env.py中的run_migrations_online函数，获取参数后传递给模板：

   import sys
   
   def run_migrations_online():
       # 保留原有配置代码
       # 捕获自定义命令行参数
       enable_custom_actions = "--with-custom-actions" in sys.argv
       # 配置模板上下文变量
       context.configure(
           # 原有配置（如url、target_metadata等）
           template_args={
               "enable_custom_actions": enable_custom_actions
           }
       )
       # 保留原有执行迁移的代码
   

3. 修改script.py.mako模板生成自定义内容：
   在模板中根据上下文变量添加自定义逻辑：

   """${message}"""
   from alembic import op
   import sqlalchemy as sa
   % if enable_custom_actions:
   def run_custom_actions():
       # 示例：写入操作日志到数据库
       conn = op.get_bind()
       conn.execute(sa.text("INSERT INTO migration_logs (revision_id, note) VALUES ('${rev_id}', '执行自定义操作')"))
       # 示例：调用外部API（需提前导入requests）
       # import requests
       # requests.post("https://your-api.com/migration-notify", json={"revision": "${rev_id}"})
   % endif
   
   def upgrade():
       ${upgrades if upgrades else "pass"}
       % if enable_custom_actions:
       run_custom_actions()
       % endif
   
   def downgrade():
       ${downgrades if downgrades else "pass"}
   

   这样，当命令行传入--with-custom-actions时，生成的修订文件会自动包含自定义操作代码。

三、其他自定义修订的方法 #

• 手动编写修订：无需自动生成数据库变更时，可直接用alembic revision --message "手动修订"生成空模板，再手动添加自定义内容。
• 利用修订生成钩子：在env.py中注册revision_create事件钩子，动态修改生成的脚本内容：

  from alembic.runtime.environment import EnvironmentContext
  
  def modify_revision_script(context, revision, directives):
      script = directives[0]
      # 在upgrade函数开头添加自定义代码
      script.body = script.body.replace(
          "def upgrade():",
          "def upgrade():\n    # 自动添加的自定义逻辑"
      )
  
  EnvironmentContext.configure.event.listen("revision_create", modify_revision_script)
  

内容的提问来源于stack exchange，提问作者anonymous anonymous