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