diff --git a/src/PgStructure.hs b/src/PgStructure.hs index a6c25f157..53bc70096 100644 --- a/src/PgStructure.hs +++ b/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] <- - map (map fromSql) r] + return $ foldl addKey Map.empty $ 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 diff --git a/test/Unit/PgStructureSpec.hs b/test/Unit/PgStructureSpec.hs index bf9480b4d..9045dc044 100644 --- a/test/Unit/PgStructureSpec.hs +++ b/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` [ - ForeignKey { fkCol="auto_inc_fk", - fkTableReferred="auto_incrementing_pk", fkColReferred="id"}, - ForeignKey { fkCol="simple_fk", fkTableReferred="simple_pk", - fkColReferred="k"} - ] + foreignKeys "1" "has_fk" conn `shouldReturn` M.fromList [ + ("auto_inc_fk", ForeignKey {fkTable="auto_incrementing_pk", fkCol="id"}), + ("simple_fk", ForeignKey { fkTable="simple_pk", fkCol="k"})] where setRole conn = quickQuery conn "set role dbapi_test" [] >> return conn