From fac4c03b45abe9dd3e5b2713dad7f1649992f6e2 Mon Sep 17 00:00:00 2001
From: Zach Smith <zach@driver.xyz>
Date: Fri, 2 Feb 2018 11:25:31 -0800
Subject: Apply fix to foreign and primary keys

---
 src/main/scala/SchemaParser.scala | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/src/main/scala/SchemaParser.scala b/src/main/scala/SchemaParser.scala
index 5ea61bf..e7ef031 100644
--- a/src/main/scala/SchemaParser.scala
+++ b/src/main/scala/SchemaParser.scala
@@ -19,12 +19,18 @@ object ModelTransformation {
 
   def citextNoLength(dbModel: m.Model): m.Model =
     dbModel.copy(
-      tables = dbModel.tables.map(
-        table =>
-          table.copy(
-            columns = table.columns.map(citextColumnNoLength),
-            indices = table.indices.map(index => index.copy(columns = index.columns.map(citextColumnNoLength)))
-        )))
+      tables = dbModel.tables.map(table =>
+        table.copy(
+          primaryKey = table.primaryKey.map(pk => pk.copy(columns = pk.columns.map(citextColumnNoLength))),
+          columns = table.columns.map(citextColumnNoLength),
+          indices = table.indices.map(index => index.copy(columns = index.columns.map(citextColumnNoLength))),
+          foreignKeys = table.foreignKeys.map { fk =>
+            fk.copy(
+              referencedColumns = fk.referencedColumns.map(citextColumnNoLength),
+              referencingColumns = fk.referencingColumns.map(citextColumnNoLength)
+            )
+          }
+      )))
 
   def references(
       dbModel: m.Model,
-- 
cgit v1.2.3