fix: quote tablegroup and project name on dbml render #46

pydbml/renderer/dbml/default/project.py

@@ -4,6 +4,7 @@
 from pydbml.classes import Project
 from pydbml.renderer.dbml.default.renderer import DefaultDBMLRenderer
 from pydbml.renderer.dbml.default.utils import comment_to_dbml
+from pydbml.tools import doublequote_string
 
 
 def render_items(items: Dict[str, str]) -> str:
@@ -19,7 +20,8 @@ def render_items(items: Dict[str, str]) -> str:
 @DefaultDBMLRenderer.renderer_for(Project)
 def render_project(model: Project) -> str:
     result = comment_to_dbml(model.comment) if model.comment else ''
-    result += f'Project "{model.name}" {{\n'
+    quoted_name = doublequote_string(model.name)
+    result += f'Project {quoted_name} {{\n'
     result += render_items(model.items)
     if model.note:
         result += indent(DefaultDBMLRenderer.render(model.note), '    ') + '\n'

pydbml/renderer/dbml/default/table_group.py

@@ -4,12 +4,14 @@
 from pydbml.renderer.dbml.default.renderer import DefaultDBMLRenderer
 from pydbml.renderer.dbml.default.table import get_full_name_for_dbml
 from pydbml.renderer.dbml.default.utils import comment_to_dbml
+from pydbml.tools import doublequote_string
 
 
 @DefaultDBMLRenderer.renderer_for(TableGroup)
 def render_table_group(model: TableGroup) -> str:
     result = comment_to_dbml(model.comment) if model.comment else ''
-    result += f'TableGroup {model.name}'
+    quoted_name = doublequote_string(model.name)
+    result += f'TableGroup {quoted_name}'
     if model.color:
         result += f' [color: {model.color}]'
     result += ' {\n'

pydbml/tools.py

@@ -27,6 +27,14 @@ def strip_empty_lines(source: str) -> str:
     return pattern.sub('\g<content>', source)
 
 
+def doublequote_string(source: str) -> str:
+    """Safely wrap a single-line string in double quotes"""
+    if '\n' in source:
+        raise ValueError(f'Multiline strings are not allowed: {source!r}')
+    result = source.strip('"').replace('"', '\\"')
+    return f'"{result}"'
+
+
 def remove_indentation(source: str) -> str:
     if not source:
         return source

test/test_data/integration1.dbml

@@ -38,7 +38,7 @@ Ref {
     "Employees"."favorite_book_id" > "books"."id"
 }
 
-TableGroup Unanimate {
+TableGroup "Unanimate" {
     "books"
     "countries"
 }

test/test_renderer/test_dbml/test_table_group.py

@@ -12,7 +12,7 @@ def test_simple(table1: Table, table2: Table, table3: Table) -> None:
             items=[table1, table2, table3],
         )
         expected = (
-            'TableGroup mygroup {\n'
+            'TableGroup "mygroup" {\n'
             '    "products"\n'
             '    "products"\n'
             '    "orders"\n'
@@ -31,7 +31,7 @@ def test_full(table1: Table, table2: Table, table3: Table) -> None:
         )
         expected = (
             '// My comment\n'
-            'TableGroup mygroup [color: #FFF] {\n'
+            'TableGroup "mygroup" [color: #FFF] {\n'
             '    "products"\n'
             '    "products"\n'
             '    "orders"\n'

test/test_tools.py

@@ -1,7 +1,9 @@
 from unittest import TestCase
 
+import pytest
+
 from pydbml.classes import Note
-from pydbml.tools import remove_indentation
+from pydbml.tools import remove_indentation, doublequote_string
 from pydbml.renderer.sql.default.utils import comment_to_sql
 from pydbml.tools import indent
 from pydbml.renderer.dbml.default.utils import note_option_to_dbml, comment_to_dbml
@@ -10,103 +12,120 @@
 
 class TestCommentToDBML(TestCase):
     def test_comment(self) -> None:
-        oneline = 'comment'
-        self.assertEqual(f'// {oneline}\n', comment_to_dbml(oneline))
+        oneline = "comment"
+        self.assertEqual(f"// {oneline}\n", comment_to_dbml(oneline))
 
-        expected = \
-'''// 
+        expected = """// 
 // line1
 // line2
 // line3
 // 
-'''
-        source = '\nline1\nline2\nline3\n'
+"""
+        source = "\nline1\nline2\nline3\n"
         self.assertEqual(comment_to_dbml(source), expected)
 
 
 class TestCommentToSQL(TestCase):
     def test_comment(self) -> None:
-        oneline = 'comment'
-        self.assertEqual(f'-- {oneline}\n', comment_to_sql(oneline))
+        oneline = "comment"
+        self.assertEqual(f"-- {oneline}\n", comment_to_sql(oneline))
 
-        expected = \
-'''-- 
+        expected = """-- 
 -- line1
 -- line2
 -- line3
 -- 
-'''
-        source = '\nline1\nline2\nline3\n'
+"""
+        source = "\nline1\nline2\nline3\n"
         self.assertEqual(comment_to_sql(source), expected)
 
 
 class TestNoteOptionToDBML(TestCase):
     def test_oneline(self) -> None:
-        note = Note('one line note')
+        note = Note("one line note")
         self.assertEqual(f"note: 'one line note'", note_option_to_dbml(note))
 
     def test_oneline_with_quote(self) -> None:
-        note = Note('one line\'d note')
+        note = Note("one line'd note")
         self.assertEqual(f"note: 'one line\\'d note'", note_option_to_dbml(note))
 
     def test_multiline(self) -> None:
-        note = Note('line1\nline2\nline3')
+        note = Note("line1\nline2\nline3")
         expected = "note: '''line1\nline2\nline3'''"
         self.assertEqual(expected, note_option_to_dbml(note))
 
     def test_multiline_with_quotes(self) -> None:
-        note = Note('line1\n\'\'\'line2\nline3')
+        note = Note("line1\n'''line2\nline3")
         expected = "note: '''line1\n\\'''line2\nline3'''"
         self.assertEqual(expected, note_option_to_dbml(note))
 
 
 class TestIndent(TestCase):
     def test_empty(self) -> None:
-        self.assertEqual(indent(''), '')
+        self.assertEqual(indent(""), "")
 
     def test_nonempty(self) -> None:
-        oneline = 'one line text'
-        self.assertEqual(indent(oneline), f'    {oneline}')
-        source = 'line1\nline2\nline3'
-        expected = '    line1\n    line2\n    line3'
+        oneline = "one line text"
+        self.assertEqual(indent(oneline), f"    {oneline}")
+        source = "line1\nline2\nline3"
+        expected = "    line1\n    line2\n    line3"
         self.assertEqual(indent(source), expected)
-        expected2 = '  line1\n  line2\n  line3'
+        expected2 = "  line1\n  line2\n  line3"
         self.assertEqual(indent(source, 2), expected2)
 
 
 class TestStripEmptyLines(TestCase):
     def test_empty(self) -> None:
-        source = ''
+        source = ""
         self.assertEqual(strip_empty_lines(source), source)
 
     def test_no_empty_lines(self) -> None:
-        source = 'line1\n\n\nline2'
+        source = "line1\n\n\nline2"
         self.assertEqual(strip_empty_lines(source), source)
 
     def test_empty_lines(self) -> None:
-        stripped = '   line1\n\n line2'
-        source = f'\n \n   \n\t \t \n  \n{stripped}\n\n\n   \n \t \n\t \n   \n'
+        stripped = "   line1\n\n line2"
+        source = f"\n \n   \n\t \t \n  \n{stripped}\n\n\n   \n \t \n\t \n   \n"
         self.assertEqual(strip_empty_lines(source), stripped)
 
     def test_one_empty_line(self) -> None:
-        stripped = '   line1\n\n line2'
-        source = f'\n{stripped}'
+        stripped = "   line1\n\n line2"
+        source = f"\n{stripped}"
         self.assertEqual(strip_empty_lines(source), stripped)
-        source = f'{stripped}\n'
+        source = f"{stripped}\n"
         self.assertEqual(strip_empty_lines(source), stripped)
 
     def test_end(self) -> None:
-        stripped = '   line1\n\n line2'
-        source = f'\n{stripped}\n   '
+        stripped = "   line1\n\n line2"
+        source = f"\n{stripped}\n   "
         self.assertEqual(strip_empty_lines(source), stripped)
 
 
 class TestRemoveIndentation(TestCase):
     def test_empty(self) -> None:
-        source = ''
+        source = ""
         self.assertEqual(remove_indentation(source), source)
 
     def test_not_empty(self) -> None:
-        source = '    line1\n     line2'
-        expected = 'line1\n line2'
+        source = "    line1\n     line2"
+        expected = "line1\n line2"
         self.assertEqual(remove_indentation(source), expected)
+
+
+class TestDoublequoteString:
+    @staticmethod
+    @pytest.mark.parametrize(
+        "source,expected",
+        [
+            ("Test string", '"Test string"'),
+            ('String with "quotes"!', '"String with \\"quotes\\"!"'),
+            ('"Quoted string"', '"Quoted string"'),
+        ],
+    )
+    def test_oneline(source: str, expected: str) -> None:
+        assert doublequote_string(source) == expected
+
+    @staticmethod
+    def test_multiline() -> None:
+        with pytest.raises(ValueError):
+            doublequote_string('line1\nline2')