MySQL on duplicate key update

persistent-mysql/Database/Persist/MySQL.hs

@@ -1,4 +1,5 @@
 {-# LANGUAGE ConstraintKinds #-}
+{-# LANGUAGE GADTs #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE FlexibleContexts #-}
@@ -14,6 +15,9 @@ module Database.Persist.MySQL
     , MySQLBase.defaultSSLInfo
     , MySQLConf(..)
     , mockMigration
+    , insertOnDuplicateKeyUpdate
+    , insertManyOnDuplicateKeyUpdate
+    , SomeField(..)
     ) where
 
 import Control.Arrow
@@ -1014,3 +1018,107 @@ mockMigration mig = do
       result = runReaderT . runWriterT . runWriterT $ mig 
   resp <- result sqlbackend
   mapM_ T.putStrLn $ map snd $ snd resp
+
+-- | MySQL specific 'upsert'. This will prevent multiple queries, when one will
+-- do.
+insertOnDuplicateKeyUpdate
+  :: ( PersistEntityBackend record ~ BaseBackend backend
+     , PersistEntity record
+     , MonadIO m
+     , PersistStore backend
+     , backend ~ SqlBackend
+     )
+  => record
+  -> [Update record]
+  -> SqlPersistT m ()
+insertOnDuplicateKeyUpdate record =
+  insertManyOnDuplicateKeyUpdate [record] []
+
+-- | This wraps values of an Entity's 'EntityField', making them have the same
+-- type. This allows them to be put in lists.
+data SomeField record where
+  SomeField :: EntityField record typ -> SomeField record
+
+-- | Do a bulk insert on the given records in the first parameter. In the event
+-- that a key conflicts with a record currently in the database, the second and
+-- third parameters determine what will happen.
+--
+-- The second parameter is a list of fields to copy from the original value.
+-- This allows you to specify that, when a collision occurs, you'll just update
+-- the value in the database with the field values that you inserted.
+--
+-- The third parameter is a list of updates to perform that are independent of
+-- the value that is provided. You can use this to increment a counter value.
+-- These updates only occur if the original record is present in the database.
+insertManyOnDuplicateKeyUpdate
+  :: ( PersistEntityBackend record ~ SqlBackend
+     , PersistEntity record
+     , MonadIO m
+     )
+  => [record] -- ^ A list of the records you want to insert, or update
+  -> [SomeField record] -- ^ A list of the fields you want to copy over.
+  -> [Update record] -- ^ A list of the updates to apply that aren't dependent on the record being inserted.
+  -> SqlPersistT m ()
+insertManyOnDuplicateKeyUpdate [] _ _ = return ()
+insertManyOnDuplicateKeyUpdate records [] [] = insertMany_ records
+insertManyOnDuplicateKeyUpdate records fieldValues updates =
+  uncurry rawExecute $ mkBulkInsertQuery records fieldValues updates
+
+-- | This creates the query for 'bulkInsertOnDuplicateKeyUpdate'. It will give
+-- garbage results if you don't provide a list of either fields to copy or
+-- fields to update.
+mkBulkInsertQuery
+    :: (PersistEntityBackend record ~ SqlBackend, PersistEntity record)
+    => [record] -- ^ A list of the records you want to insert, or update
+    -> [SomeField record] -- ^ A list of the fields you want to copy over.
+    -> [Update record] -- ^ A list of the updates to apply that aren't dependent on the record being inserted.
+    -> (Text, [PersistValue])
+mkBulkInsertQuery records fieldValues updates =
+    (q, recordValues <> updsValues)
+  where
+    fieldDefs = map (\x -> case x of SomeField rec -> persistFieldDef rec) fieldValues
+    updateFieldNames = map (T.pack . escapeDBName . fieldDB) fieldDefs
+    entityDef' = entityDef records
+    entityFieldNames = map (T.pack . escapeDBName . fieldDB) (entityFields entityDef')
+    tableName = T.pack . escapeDBName . entityDB $ entityDef'
+    recordValues = concatMap (map toPersistValue . toPersistFields) records
+    recordPlaceholders = commaSeparated $ map (parenWrapped . commaSeparated . map (const "?") . toPersistFields) records
+    fieldSets = map (\n -> T.concat [n, "=VALUES(", n, ")"]) updateFieldNames
+    upds = map mkUpdateText updates
+    updsValues = map (\(Update _ val _) -> toPersistValue val) updates
+    q = T.concat
+        [ "INSERT INTO "
+        , tableName
+        , " ("
+        , commaSeparated entityFieldNames
+        , ") "
+        , " VALUES "
+        , recordPlaceholders
+        , " ON DUPLICATE KEY UPDATE "
+        , commaSeparated (fieldSets <> upds)
+        ]
+
+-- | Vendored from @persistent@.
+mkUpdateText :: PersistEntity record => Update record -> Text
+mkUpdateText x =
+  case updateUpdate x of
+    Assign -> n <> "=?"
+    Add -> T.concat [n, "=", n, "+?"]
+    Subtract -> T.concat [n, "=", n, "-?"]
+    Multiply -> T.concat [n, "=", n, "*?"]
+    Divide -> T.concat [n, "=", n, "/?"]
+    BackendSpecificUpdate up ->
+      error . T.unpack $ "BackendSpecificUpdate " <> up <> " not supported"
+  where
+    n = T.pack . escapeDBName . fieldDB . updateFieldDef $ x
+
+commaSeparated :: [Text] -> Text
+commaSeparated = T.intercalate ", "
+
+parenWrapped :: Text -> Text
+parenWrapped t = T.concat ["(", t, ")"]
+
+-- | Gets the 'FieldDef' for an 'Update'. Vendored from @persistent@.
+updateFieldDef :: PersistEntity v => Update v -> FieldDef
+updateFieldDef (Update f _ _) = persistFieldDef f
+updateFieldDef BackendUpdate {} = error "updateFieldDef did not expect BackendUpdate"