Use Postgres' RETURNING capability to optimize insertMany.

persistent-mysql/Database/Persist/MySQL.hs

@@ -103,6 +103,7 @@ open' ci logFunc = do
         { connPrepare    = prepare' conn
         , connStmtMap    = smap
         , connInsertSql  = insertSql'
+        , connInsertManySql = Nothing
         , connClose      = MySQL.close conn
         , connMigrateSql = migrate' ci
         , connBegin      = const $ MySQL.execute_ conn "start transaction" >> return ()

persistent-mysql/persistent-mysql.cabal

@@ -1,5 +1,5 @@
 name:            persistent-mysql
-version:         2.1.3.1
+version:         2.1.3.2
 license:         MIT
 license-file:    LICENSE
 author:          Felipe Lessa <felipe.lessa@gmail.com>, Michael Snoyman
@@ -31,7 +31,7 @@ library
                    , mysql-simple          >= 0.2.2.3  && < 0.3
                    , mysql                 >= 0.1.1.3  && < 0.2
                    , blaze-builder
-                   , persistent            >= 2.1      && < 3
+                   , persistent            >= 2.1.7    && < 3
                    , containers            >= 0.2
                    , bytestring            >= 0.9
                    , text                  >= 0.11.0.6

persistent-postgresql/ChangeLog.md

@@ -1,3 +1,7 @@
+## 2.1.6.1
+
+* Optimize the `insertMany` function to insert all rows and retrieve their keys in one SQL query. [#407](https://github.com/yesodweb/persistent/pull/407)
+
 ## 2.1.6
 
 * Postgresql exceptions [#353](https://github.com/yesodweb/persistent/issues/353)

persistent-postgresql/Database/Persist/Postgresql.hs

@@ -147,6 +147,7 @@ openSimpleConn logFunc conn = do
         { connPrepare    = prepare' conn
         , connStmtMap    = smap
         , connInsertSql  = insertSql'
+        , connInsertManySql = Just insertManySql'
         , connClose      = PG.close conn
         , connMigrateSql = migrate'
         , connBegin      = const $ PG.begin    conn
@@ -188,6 +189,21 @@ insertSql' ent vals =
        Just _pdef -> ISRManyKeys sql vals
        Nothing -> ISRSingle (sql <> " RETURNING " <> escape (fieldDB (entityId ent)))
 
+-- | SQL for inserting multiple rows at once and returning their primary keys.
+insertManySql' :: EntityDef -> [[PersistValue]] -> InsertSqlResult
+insertManySql' ent valss =
+  let sql = T.concat
+                [ "INSERT INTO "
+                , escape (entityDB ent)
+                , "("
+                , T.intercalate "," $ map (escape . fieldDB) $ entityFields ent
+                , ") VALUES ("
+                , T.intercalate "),(" $ replicate (length valss) $ T.intercalate "," $ map (const "?") (entityFields ent)
+                , ") RETURNING "
+                , T.intercalate ", " $ dbIdColumnsEsc escape ent
+                ]
+  in ISRSingle sql
+
 execute' :: PG.Connection -> PG.Query -> [PersistValue] -> IO Int64
 execute' conn query vals = PG.execute conn query (map P vals)
 

persistent-postgresql/persistent-postgresql.cabal

@@ -1,5 +1,5 @@
 name:            persistent-postgresql
-version:         2.1.6
+version:         2.1.6.1
 license:         MIT
 license-file:    LICENSE
 author:          Felipe Lessa, Michael Snoyman <michael@snoyman.com>
@@ -19,7 +19,7 @@ library
                    , transformers          >= 0.2.1
                    , postgresql-simple     >= 0.4.0    && < 0.5
                    , postgresql-libpq      >= 0.6.1    && < 0.10
-                   , persistent            >= 2.1.5    && < 3
+                   , persistent            >= 2.1.7    && < 3
                    , containers            >= 0.2
                    , bytestring            >= 0.9
                    , text                  >= 0.7

persistent-sqlite/Database/Persist/Sqlite.hs

@@ -92,6 +92,7 @@ wrapConnectionWal enableWal conn logFunc = do
         { connPrepare = prepare' conn
         , connStmtMap = smap
         , connInsertSql = insertSql'
+        , connInsertManySql = Nothing
         , connClose = Sqlite.close conn
         , connMigrateSql = migrate'
         , connBegin = helper "BEGIN"

persistent-sqlite/persistent-sqlite.cabal

@@ -1,5 +1,5 @@
 name:            persistent-sqlite
-version:         2.1.4.2
+version:         2.1.4.3
 license:         MIT
 license-file:    LICENSE
 author:          Michael Snoyman <michael@snoyman.com>
@@ -25,7 +25,7 @@ library
     build-depends:   base                    >= 4.6         && < 5
                    , bytestring              >= 0.9.1
                    , transformers            >= 0.2.1
-                   , persistent              >= 2.1       && < 3
+                   , persistent              >= 2.1.7       && < 3
                    , monad-control           >= 0.2
                    , containers              >= 0.2
                    , text                    >= 0.7

persistent-test/CompositeTest.hs

@@ -262,6 +262,18 @@ specs = describe "composite" $
       let [Entity newkca1 newca2] = xs
       matchCitizenAddressK kca1 @== matchCitizenAddressK newkca1
       ca1 @== newca2
+    it "insertMany" $ db $ do
+      [kp1, kp2] <- insertMany [p1, p2]
+      newp1 <- get kp1
+      newp1 @== Just p1
+
+      newp2 <- get kp2
+      newp2 @== Just p2
+    it "rawSql instance" $ db $ do
+      key <- insert p1
+      keyFromRaw <- rawSql "SELECT name, name2, age FROM test_parent LIMIT 1" []
+      [key] @== keyFromRaw
+
 #endif
 
 matchK :: (PersistField a, PersistEntity record) => Key record -> Either Text a

persistent-test/PrimaryTest.hs

@@ -43,5 +43,9 @@ specs = describe "primary key reference" $ do
     kf <- insert $ Foo "name"
     kb <- insert $ Bar kf
     return ()
+  it "uses RawSql for a Primary key" $ db $ do
+    key <- insert $ Foo "name"
+    keyFromRaw <- rawSql "SELECT name FROM foo LIMIT 1" []
+    [key] @== keyFromRaw
 #  endif
 #endif

persistent/ChangeLog.md

@@ -1,3 +1,7 @@
+## 2.1.7
+
+* Add a `RawSql` instance for `Key`. This allows selecting primary keys using functions like `rawSql`. [#407](https://github.com/yesodweb/persistent/pull/407)
+
 ## 2.1.6
 
 Important! If persistent-template is not upgraded to 2.1.3.3

persistent/Database/Persist/Class/PersistStore.hs

@@ -80,31 +80,38 @@ class
             => val -> ReaderT backend m ()
     insert_ val = insert val >> return ()
 
-    -- | Create multiple records in the database.
+    -- | Create multiple records in the database and return their 'Key's.
     --
-    -- If you don't need the inserted @Key@s, use 'insertMany_'
+    -- If you don't need the inserted 'Key's, use 'insertMany_'.
     --
-    -- Some backends use the slow default implementation of
-    -- @mapM insert@
+    -- The MongoDB and PostgreSQL backends insert all records and
+    -- retrieve their keys in one database query.
+    --
+    -- The SQLite and MySQL backends use the slow, default implementation of
+    -- @mapM insert@.
     insertMany :: (MonadIO m, backend ~ PersistEntityBackend val, PersistEntity val)
                => [val] -> ReaderT backend m [Key val]
     insertMany = mapM insert
 
-    -- | Same as 'insertMany', but doesn't return any @Key@s.
+    -- | Same as 'insertMany', but doesn't return any 'Key's.
+    --
+    -- The MongoDB, PostgreSQL, and MySQL backends insert all records in
+    -- one database query.
     --
-    -- SQL backends currently use an efficient implementation for this
-    -- unlike 'insertMany'.
+    -- The SQLite backend inserts rows one-by-one.
     insertMany_ :: (MonadIO m, backend ~ PersistEntityBackend val, PersistEntity val)
                 => [val] -> ReaderT backend m ()
     insertMany_ x = insertMany x >> return ()
 
-    -- | Same as 'insertMany_', but takes an 'Entity' instead of just a record
+    -- | Same as 'insertMany_', but takes an 'Entity' instead of just a record.
     --
     -- Useful when migrating data from one entity to another
-    -- and want to preserve ids
+    -- and want to preserve ids.
+    --
+    -- The MongoDB backend inserts all the entities in one database query.
     --
-    -- Some backends use the slow default implementation of
-    -- @mapM insertKey@
+    -- The SQL backends use the slow, default implementation of
+    -- @mapM_ insertKey@.
     insertEntityMany :: (MonadIO m, backend ~ PersistEntityBackend val, PersistEntity val)
                      => [Entity val] -> ReaderT backend m ()
     insertEntityMany = mapM_ (\(Entity k record) -> insertKey k record)

persistent/Database/Persist/Sql/Class.hs

@@ -58,6 +58,13 @@ instance PersistField a => RawSql (Single a) where
     rawSqlProcessRow [pv]  = Single <$> fromPersistValue pv
     rawSqlProcessRow _     = Left $ pack "RawSql (Single a): wrong number of columns."
 
+instance (PersistEntity a, PersistEntityBackend a ~ SqlBackend) => RawSql (Key a) where
+  rawSqlCols _ key         = (length $ keyToValues key, [])
+  rawSqlColCountReason key = "The primary key is composed of "
+                             ++ (show $ length $ keyToValues key)
+                             ++ " columns"
+  rawSqlProcessRow         = keyFromValues
+
 instance (PersistEntity a, PersistEntityBackend a ~ SqlBackend) => RawSql (Entity a) where
     rawSqlCols escape = ((+1) . length . entityFields &&& process) . entityDef . Just . entityVal
         where

persistent/Database/Persist/Sql/Orphan/PersistStore.hs

@@ -187,6 +187,21 @@ instance PersistStore SqlBackend where
         t = entityDef $ Just val
         vals = map toPersistValue $ toPersistFields val
 
+    insertMany [] = return []
+    insertMany vals = do
+        conn <- ask
+
+        case connInsertManySql conn of
+            Nothing -> mapM insert vals
+            Just insertManyFn -> do
+                case insertManyFn ent valss of
+                    ISRSingle sql -> rawSql sql (concat valss)
+                    _ -> error "ISRSingle is expected from the connInsertManySql function"
+                where
+                    ent = entityDef vals
+                    valss = map (map toPersistValue . toPersistFields) vals
+
+
     insertMany_ [] = return ()
     insertMany_ vals = do
         conn <- ask
@@ -200,7 +215,7 @@ instance PersistStore SqlBackend where
                 , ")"
                 ]
 
-        -- SQLite support is only in later versions
+        -- SQLite only supports multi-row inserts in 3.7.11+ (see https://www.sqlite.org/releaselog/3_7_11.html).
         if connRDBMS conn == "sqlite"
             then mapM_ insert vals
             else rawExecute sql (concat valss)

persistent/Database/Persist/Sql/Types.hs

@@ -43,6 +43,7 @@ data SqlBackend = SqlBackend
     { connPrepare :: Text -> IO Statement
     -- | table name, column names, id name, either 1 or 2 statements to run
     , connInsertSql :: EntityDef -> [PersistValue] -> InsertSqlResult
+    , connInsertManySql :: Maybe (EntityDef -> [[PersistValue]] -> InsertSqlResult) -- ^ SQL for inserting many rows and returning their primary keys, for backends that support this functioanlity. If 'Nothing', rows will be inserted one-at-a-time using 'connInsertSql'.
     , connStmtMap :: IORef (Map Text Statement)
     , connClose :: IO ()
     , connMigrateSql

persistent/persistent.cabal

@@ -1,5 +1,5 @@
 name:            persistent
-version:         2.1.6
+version:         2.1.7
 license:         MIT
 license-file:    LICENSE
 author:          Michael Snoyman <michael@snoyman.com>