From aa5cdd47ca3a16e32e94239f4f603e18b016399b Mon Sep 17 00:00:00 2001 From: Eryx-Z Date: Sat, 17 Oct 2026 05:12:55 +0800 Subject: [PATCH] Escape revision messages embedded in triple-quoted docstrings Unescaped messages containing triple quotes could close the module docstring early and leave the rest of the message as executable code when the revision is imported. Backslashes were also interpreted as escape sequences. Escape both before rendering script.py.mako. --- alembic/script/base.py | 8 +++++++- tests/test_script_production.py | 14 ++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/alembic/script/base.py b/alembic/script/base.py index cc136d8b..4b51db96 100644 --- a/alembic/script/base.py +++ b/alembic/script/base.py @@ -45,6 +45,10 @@ _slug_re = re.compile(r"\w+") _default_file_template = "%(rev)s_%(slug)s" +def _escape_docstring_message(message: str) -> str: + """Escape a revision message for embedding in a triple-quoted docstring.""" + return message.replace("\\", "\\\\").replace('"""', '\\"\\"\\"') + class ScriptDirectory: """Provides operations upon an Alembic script directory. @@ -740,7 +744,9 @@ def generate_revision( depends_on=revision.tuple_rev_as_scalar(resolved_depends_on), create_date=create_date, comma=util.format_as_comma, - message=message if message is not None else ("empty message"), + message=_escape_docstring_message( + message if message is not None else "empty message" + ), **kw, ) diff --git a/tests/test_script_production.py b/tests/test_script_production.py index 20360f47..8a91b0be 100644 --- a/tests/test_script_production.py +++ b/tests/test_script_production.py @@ -406,6 +406,20 @@ def test_create_script_basic(self): eq_(rev.down_revision, self.c) assert "some message" in rev.doc + def test_create_script_message_with_triple_quotes(self): + msg = 'add users"""' + "\nprint(\"injected\")\n__doc__ = " + '"""' + rev = command.revision(self.cfg, message=msg) + script = ScriptDirectory.from_config(self.cfg) + rev = script.get_revision(rev.revision) + eq_(rev.doc, msg) + + def test_create_script_message_with_backslashes(self): + msg = r"C:\path\new" + rev = command.revision(self.cfg, message=msg) + script = ScriptDirectory.from_config(self.cfg) + rev = script.get_revision(rev.revision) + eq_(rev.doc, msg) + def test_create_script_splice(self): rev = command.revision( self.cfg, message="some message", head=self.b, splice=True