From f0e8d23dd298e10da7659216782b6537b478fa80 Mon Sep 17 00:00:00 2001
From: "Adam C. Baker" <adam.durandal@gmail.com>
Date: Tue, 7 Oct 2014 23:52:51 -0700
Subject: [PATCH] getting foreign key data with a column.

---
 src/PgStructure.hs           | 18 +++++++++---------
 test/Unit/PgStructureSpec.hs |  8 +++++++-
 2 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/src/PgStructure.hs b/src/PgStructure.hs
index 53bc70096..42d5a6c45 100644
--- a/src/PgStructure.hs
+++ b/src/PgStructure.hs
@@ -1,8 +1,6 @@
--- {{{ Imports
-
+{-# OPTIONS_GHC -fno-warn-incomplete-patterns #-}
 {-# LANGUAGE OverloadedStrings #-}
 
-
 module PgStructure where
 
 import Data.Functor ( (<$>) )
@@ -19,8 +17,6 @@ import Database.HDBC.PostgreSQL
 
 import Data.Aeson ((.=))
 
--- }}}
-
 data Table = Table {
   tableSchema :: String
 , tableName :: String
@@ -69,6 +65,7 @@ data Column = Column {
 , colMaxLen :: Maybe Int
 , colPrecision :: Maybe Int
 , colDefault :: Maybe String
+, colFK :: Maybe ForeignKey
 } deriving (Show)
 
 instance JSON.ToJSON Column where
@@ -120,11 +117,15 @@ columns s t conn = do
         \  from information_schema.columns\
         \ where table_schema = ?\
         \   and table_name = ?" [toSql s, toSql t]
-  return $ mapMaybe mkColumn r
+  fks <- foreignKeys s t conn
+  let lookupFK (_:_:name:_) = Map.lookup (fromSql name) fks
+      lookupFK _ = Nothing
+  let cols = zipWith ($) (map mkColumn r) (map lookupFK r)
+  return cols
 
   where
-    mkColumn [schema, table, name, pos, nullable, colT, updatable, maxlen, precision, defVal] =
-      Just $ Column (fromSql schema)
+    --TODO: handle failed pattern match with an appropriate exception
+    mkColumn [schema, table, name, pos, nullable, colT, updatable, maxlen, precision, defVal] = Column (fromSql schema)
         (fromSql table)
         (fromSql name)
         (fromSql pos)
@@ -134,7 +135,6 @@ columns s t conn = do
         (fromSql maxlen)
         (fromSql precision)
         (fromSql defVal)
-    mkColumn _ = Nothing
 
 printTables :: String -> Connection -> IO BL.ByteString
 printTables schema conn = JSON.encode <$> tables schema conn
diff --git a/test/Unit/PgStructureSpec.hs b/test/Unit/PgStructureSpec.hs
index 9045dc044..e0a177b14 100644
--- a/test/Unit/PgStructureSpec.hs
+++ b/test/Unit/PgStructureSpec.hs
@@ -16,12 +16,18 @@ spec = around dbWithSchema $ beforeWith setRole $ do
       map tableName ts `shouldBe` ["auto_incrementing_pk","compound_pk",
         "has_fk","items","menagerie","no_pk", "simple_pk"]
 
-  describe "columns" $
+  describe "columns" $ do
     it "responds with each column for the table" $ \conn -> do
       cs <- columns "1" "auto_incrementing_pk" conn
       map colName cs `shouldBe` ["id","nullable_string","non_nullable_string",
         "inserted_at"]
 
+    it "includes foreign key data" $ \conn -> do
+      cs <- columns "1" "has_fk" conn
+      map colFK cs `shouldBe` [Nothing,
+        Just $ ForeignKey "auto_incrementing_pk" "id",
+        Just $ ForeignKey "simple_pk" "k"]
+
   describe "foreignKeys" $
     it "has a description of the foreign key columns" $ \conn ->
       foreignKeys "1" "has_fk" conn `shouldReturn` M.fromList [