Modified "processTable" to add composite primary keys when specified in dbml indexes

Author: KoStyle

dbml_sqlite/core.py

@@ -1,3 +1,4 @@
+from ast import IsNot
 import re
 import os
 import uuid
@@ -82,6 +83,7 @@ def processFile(target, emulationMode, tableExists=True, indexExists=True, idxNa
     for table in parsed.tables:
         for index in table.indexes:
             statements.append(processIndex(table, index, idxNameFunc, indexExists=indexExists, join=False))
+            break # We just need a single multi-field primary key
     statements = list(chain.from_iterable(statements))
     if join:
         statements = "".join(statements)
@@ -167,6 +169,11 @@ def processTable(table, emulationMode, tableExists=True, join=True):
         segments.append(processRef(ref, False))
         if j < len(table.refs) - 1:
             segments.append(',\n')
+    if hasattr(table, 'indexes'):
+        for k, index in enumerate(table.indexes):
+            if(index.pk):
+                segments.append(',\n')
+                segments.append(index.sql)
     segments.append('\n);\n')
     segments = list(chain.from_iterable(segments))
     if join:

tests/test_dbml_sqlite.py

@@ -134,7 +134,7 @@ def test_process_enum():
 def test_process_file():
     p = Path('./tests/abc.dbml')
     o = processFile(p, 'full', True, True, MockNameFunc)
-    assert o == 'CREATE TABLE IF NOT EXISTS mytab (\n  name TEXT,\n  phone INTEGER\n);\nCREATE INDEX IF NOT EXISTS _mockname ON mytab (name, phone);\n'
+    assert o == 'CREATE TABLE IF NOT EXISTS mytab (\n  name TEXT,\n  phone INTEGER,\nPRIMARY KEY ("name", "phone")\n);\nCREATE INDEX IF NOT EXISTS _mockname ON mytab (name, phone);\n'
 
 def test_process_index():
     mytab = MockTable('mytab', [], [])