return a map of foreign keys

src/PgStructure.hs

@@ -11,8 +11,8 @@ import Data.Maybe (mapMaybe)
 import Control.Applicative ( (<*>) )
 
 import qualified Data.ByteString.Lazy as BL
-
 import qualified Data.Aeson as JSON
+import qualified Data.Map as Map
 
 import Database.HDBC hiding (colType, colNullable)
 import Database.HDBC.PostgreSQL
@@ -37,10 +37,10 @@ toBool :: String -> Bool
 toBool = (== "YES")
 
 data ForeignKey = ForeignKey {
-  fkCol::String, fkTableReferred::String, fkColReferred::String
+  fkTable::String, fkCol::String
 } deriving (Eq, Show)
 
-foreignKeys :: String -> String -> Connection -> IO [ForeignKey]
+foreignKeys :: String -> String -> Connection -> IO (Map.Map String ForeignKey)
 foreignKeys schema table conn = do
   r <- quickQuery conn
     "select kcu.column_name, ccu.table_name AS foreign_table_name,\
@@ -53,8 +53,10 @@ foreignKeys schema table conn = do
     \where constraint_type = 'FOREIGN KEY' \
     \  and tc.table_name=? and tc.table_schema = ? \
     \order by kcu.column_name" (map toSql [table, schema])
-  return [ForeignKey col reftab refcol | [col, reftab, refcol] <-
+  return $ foldl addKey Map.empty $ map (map fromSql) r
-    map (map fromSql) r]
+  where
+    addKey m [col, ftab, fcol] = Map.insert col (ForeignKey ftab fcol) m
+    addKey m _ = m --should never happen
 
 data Column = Column {
   colSchema :: String

test/Unit/PgStructureSpec.hs

@@ -6,6 +6,7 @@ import PgStructure (Table(..), tables, Column(..), columns, ForeignKey(..),
 
 import Database.HDBC (quickQuery)
 import SpecHelper(dbWithSchema)
+import qualified Data.Map as M;
 
 spec :: Spec
 spec = around dbWithSchema $ beforeWith setRole $ do
@@ -23,10 +24,7 @@ spec = around dbWithSchema $ beforeWith setRole $ do
 
   describe "foreignKeys" $
     it "has a description of the foreign key columns" $ \conn ->
-      foreignKeys "1" "has_fk" conn `shouldReturn` [
+      foreignKeys "1" "has_fk" conn `shouldReturn` M.fromList [
-        ForeignKey { fkCol="auto_inc_fk",
+        ("auto_inc_fk", ForeignKey {fkTable="auto_incrementing_pk", fkCol="id"}),
-                   fkTableReferred="auto_incrementing_pk", fkColReferred="id"},
+        ("simple_fk", ForeignKey { fkTable="simple_pk", fkCol="k"})]
-        ForeignKey { fkCol="simple_fk", fkTableReferred="simple_pk",
-                   fkColReferred="k"}
-        ]
   where setRole conn = quickQuery conn "set role dbapi_test" [] >> return conn