Update driver interface to get composite primary key infos.

relational-query-HDBC/src/Database/HDBC/Query/TH.hs

@@ -21,8 +21,8 @@ module Database.HDBC.Query.TH (
   defineTableFromDB
   ) where
 
-import Data.Maybe (listToMaybe, maybeToList)
-import Data.List (elemIndex)
+import Data.Maybe (listToMaybe, isJust, catMaybes)
+import qualified Data.Map as Map
 
 import Database.HDBC (IConnection, SqlValue)
 
@@ -78,15 +78,24 @@ defineTableFromDB connect drv scm tbl derives = do
         withConnectionIO connect
         (\conn ->  do
             (cols, notNullIdxs) <- getFields drv conn scm tbl
-            mayPrimaryKey       <- getPrimaryKey drv conn scm tbl
+            primCols            <- getPrimaryKey drv conn scm tbl
 
-            mayPrimaryIdx <- case mayPrimaryKey of
-              Just key -> case elemIndex key $ map fst cols of
-                Nothing -> do putLog $ "defineTableFromDB: fail to find index of pkey - " ++ key ++ ". Something wrong!!"
-                              return   Nothing
-                Just ix ->    return $ Just ix
-              Nothing  ->     return   Nothing
-            return (cols, notNullIdxs, mayPrimaryIdx) )
+            let colIxMap = Map.fromList $ zip [c | (c, _) <- cols] [(0 :: Int) .. ]
+                lookup' k = do
+                  case Map.lookup k colIxMap of
+                    Just i  -> return $ Just i
+                    Nothing -> do
+                      putLog $ "defineTableFromDB: fail to find index of pkey - " ++ k ++ ". Something wrong!!"
+                      return Nothing
 
-  (cols, notNullIdxs, mayPrimaryIdx) <- runIO getDBinfo
-  defineTableDefault scm tbl cols derives (maybeToList mayPrimaryIdx) (listToMaybe notNullIdxs)
+            primaryIxs <- case primCols of
+              _:_ -> do
+                founds <- mapM lookup' primCols
+                if all isJust founds
+                  then return $ catMaybes founds
+                  else return []
+              []  ->   return []
+            return (cols, notNullIdxs, primaryIxs) )
+
+  (cols, notNullIdxs, primaryIxs) <- runIO getDBinfo
+  defineTableDefault scm tbl cols derives primaryIxs (listToMaybe notNullIdxs)

relational-query-HDBC/src/Database/HDBC/Schema/Driver.hs

@@ -38,15 +38,15 @@ data Driver conn =
                      -> IO ([(String, TypeQ)], [Int]) --  Action to get column name and Haskell type pairs and not-null columns index.
 
     -- | Get primary key column name.
-  , getPrimaryKey :: conn              --  Connection to query system catalog
-                  -> String            --  Schema name string
-                  -> String            --  Table name string
-                  -> IO (Maybe String) --  Action to get primary key column name
+  , getPrimaryKey :: conn          --  Connection to query system catalog
+                  -> String        --  Schema name string
+                  -> String        --  Table name string
+                  -> IO ([String]) --  Action to get column names of primary key
   }
 
 -- | Empty definition of 'Driver'
 emptyDriver :: IConnection conn => Driver conn
-emptyDriver =  Driver [] (\_ _ _ _ -> return ([],[])) (\_ _ _ -> return Nothing)
+emptyDriver =  Driver [] (\_ _ _ _ -> return ([],[])) (\_ _ _ -> return [])
 
 -- | Helper function to call 'getFieldsWithMap' using 'typeMap' of 'Driver'.
 getFields :: IConnection conn => Driver conn -> conn -> String -> String -> IO ([(String, TypeQ)], [Int])

relational-query-HDBC/src/Database/HDBC/Schema/IBMDB2.hs

@@ -27,7 +27,7 @@ import qualified Language.Haskell.TH.Lib.Extra as TH
 
 import Database.HDBC (IConnection, SqlValue)
 
-import Database.HDBC.Record.Query (runQuery', listToUnique)
+import Database.HDBC.Record.Query (runQuery')
 import Database.HDBC.Record.Persistable ()
 
 import Database.Record.TH (defineRecordWithSqlTypeDefaultFromDefined)
@@ -60,16 +60,15 @@ getPrimaryKey' :: IConnection conn
               => conn
               -> String
               -> String
-              -> IO (Maybe String)
+              -> IO [String]
 getPrimaryKey' conn scm' tbl' = do
   let tbl = map toUpper tbl'
       scm = map toUpper scm'
-  mayPrim <- runQuery' conn (scm, tbl) primaryKeyQuerySQL
-             >>= listToUnique
-  let mayPrimaryKey = normalizeColumn `fmap` mayPrim
-  putLog $ "getPrimaryKey: primary key = " ++ show mayPrimaryKey
+  primCols <- runQuery' conn (scm, tbl) primaryKeyQuerySQL
+  let primaryKeyCols = normalizeColumn `fmap` primCols
+  putLog $ "getPrimaryKey: primary key = " ++ show primaryKeyCols
 
-  return mayPrimaryKey
+  return primaryKeyCols
 
 getFields' :: IConnection conn
           => TypeMap

relational-query-HDBC/src/Database/HDBC/Schema/PostgreSQL.hs

@@ -28,7 +28,7 @@ import Database.HDBC (IConnection, SqlValue)
 import Database.Record.TH (defineRecordWithSqlTypeDefaultFromDefined)
 import qualified Database.Relational.Query.Table as Table
 
-import Database.HDBC.Record.Query (runQuery', listToUnique)
+import Database.HDBC.Record.Query (runQuery')
 import Database.HDBC.Record.Persistable ()
 
 import Database.Relational.Schema.PostgreSQL
@@ -60,13 +60,15 @@ getPrimaryKey' :: IConnection conn
               => conn
               -> String
               -> String
-              -> IO (Maybe String)
+              -> IO [String]
 getPrimaryKey' conn scm' tbl' = do
   let scm = map toLower scm'
       tbl = map toLower tbl'
-  mayPrim <- runQuery' conn (scm, tbl) primaryKeyQuerySQL
-             >>= listToUnique
-  return $ normalizeColumn `fmap` mayPrim
+  primCols <- runQuery' conn (scm, tbl) primaryKeyQuerySQL
+  let primaryKeyCols = normalizeColumn `fmap` primCols
+  putLog $ "getPrimaryKey: primary key = " ++ show primaryKeyCols
+
+  return primaryKeyCols
 
 getFields' :: IConnection conn
           => TypeMap