From 7a4330d1e03a02aef779cdb16b7b8da1088c547f Mon Sep 17 00:00:00 2001
From: Manu Konomi <konomi360@gmail.com>
Date: Wed, 24 Jan 2024 19:23:30 +0100
Subject: [PATCH] Modified "processTable" to add composite primary keys when
 specified in dbml indexes

---
 dbml_sqlite/core.py       | 7 +++++++
 tests/test_dbml_sqlite.py | 2 +-
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/dbml_sqlite/core.py b/dbml_sqlite/core.py
index 07f7e00..dd5d905 100644
--- a/dbml_sqlite/core.py
+++ b/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:
diff --git a/tests/test_dbml_sqlite.py b/tests/test_dbml_sqlite.py
index 5d57123..0a53e95 100644
--- a/tests/test_dbml_sqlite.py
+++ b/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', [], [])