Insert Select With Conflict for postgres

src/Database/Esqueleto/PostgreSQL.hs

@@ -1,6 +1,7 @@
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE OverloadedStrings
-           , GADTs, CPP
+           , GADTs, CPP, Rank2Types
+           , ScopedTypeVariables
  #-}
 -- | This module contain PostgreSQL-specific functions.
 --
@@ -20,6 +21,8 @@ module Database.Esqueleto.PostgreSQL
   , random_
   , upsert
   , upsertBy
+  , insertSelectWithConflict
+  , insertSelectWithConflictCount
   -- * Internal
   , unsafeSqlAggregateFunction
   ) where
@@ -33,15 +36,18 @@ import           Database.Esqueleto.Internal.Language         hiding (random_)
 import           Database.Esqueleto.Internal.PersistentImport hiding (upsert, upsertBy)
 import           Database.Esqueleto.Internal.Sql
 import           Database.Esqueleto.Internal.Internal         (EsqueletoError(..), CompositeKeyError(..), 
-                                                              UnexpectedCaseError(..), SetClause)
+                                                              UnexpectedCaseError(..), SetClause, Ident(..),
+                                                              uncommas)
 import           Database.Persist.Class                       (OnlyOneUniqueKey)
 import           Data.List.NonEmpty                           ( NonEmpty( (:|) ) )
+import           Data.Int                                     (Int64)
+import           Data.Proxy                                   (Proxy(..))
 import           Control.Arrow                                ((***), first)
 import           Control.Exception                            (Exception, throw, throwIO)
+import           Control.Monad                                (void)
 import           Control.Monad.IO.Class                       (MonadIO (..))
 import qualified Control.Monad.Trans.Reader                   as R
 
-
 -- | (@random()@) Split out into database specific modules
 -- because MySQL uses `rand()`.
 --
@@ -203,14 +209,81 @@ upsertBy uniqueKey record updates = do
     updatesText conn = first builderToText $ renderUpdates conn updates
     handler conn f = fmap head $ uncurry rawSql $
       (***) (f entDef (uDef :| [])) addVals $ updatesText conn
-    renderUpdates :: SqlBackend
-        -> [SqlExpr (Update val)]
-        -> (TLB.Builder, [PersistValue])
-    renderUpdates conn = uncommas' . concatMap renderUpdate
-        where
-          mk :: SqlExpr (Value ()) -> [(TLB.Builder, [PersistValue])]
-          mk (ERaw _ f)        = [f info]
-          mk (ECompositeKey _) = throw (CompositeKeyErr MakeSetError) -- FIXME
-          renderUpdate :: SqlExpr (Update val) -> [(TLB.Builder, [PersistValue])]
-          renderUpdate (ESet f) = mk (f undefined) -- second parameter of f is always unused
-          info = (projectBackend conn, initialIdentState)
+
+-- | Render postgres updates to be use in a SET clause.
+renderUpdates :: (BackendCompatible SqlBackend backend) => 
+    backend
+    -> [SqlExpr (Update val)]
+    -> (TLB.Builder, [PersistValue])
+renderUpdates conn = uncommas' . concatMap renderUpdate
+    where
+      mk :: SqlExpr (Value ()) -> [(TLB.Builder, [PersistValue])]
+      mk (ERaw _ f)        = [f info]
+      mk (ECompositeKey _) = throw (CompositeKeyErr MakeSetError) -- FIXME
+      renderUpdate :: SqlExpr (Update val) -> [(TLB.Builder, [PersistValue])]
+      renderUpdate (ESet f) = mk (f undefined) -- second parameter of f is always unused
+      info = (projectBackend conn, initialIdentState)
+
+-- | Inserts into a table the results of a query similar to 'insertSelect' but allows
+-- to update values that violate a constraint during insertions.
+--
+-- Example of usage:
+--
+-- @
+-- insertSelectWithConflict 
+--   (SomeFooUnique undefined)
+--   (from $ \b ->
+--     return $ Foo <# (b ^. BarNum)
+--   )
+--   (\current excluded ->
+--     [FooNum =. (current ^. FooNum) +. (excluded ^. FooNum)]
+--   )
+-- @
+--
+-- Inserts to table Foo all Bar.num values and in case of conflict SomeFooUnique, 
+-- the conflicting value is updated to the current plus the excluded.
+insertSelectWithConflict :: (MonadIO m, PersistEntity val) => 
+  Unique val
+  -- ^ Uniqueness constraint, this is used just to get the name of the postgres constraint, the value(s) is(are) never used, so if you have a unique "MyUnique 0", "MyUnique undefined" will work as well.
+  -> SqlQuery (SqlExpr (Insertion val)) 
+  -- ^ Insert query.
+  -> (SqlExpr (Entity val) -> SqlExpr (Entity val) -> [SqlExpr (Update val)])
+  -- ^ A list of updates to be applied in case of the constraint being violated. The expression takes the current and excluded value to produce the updates.
+  -> SqlWriteT m ()
+insertSelectWithConflict unique query = void . insertSelectWithConflictCount unique query
+
+-- | Same as 'insertSelectWithConflict' but returns the number of rows affected.
+insertSelectWithConflictCount :: forall val m. (MonadIO m, PersistEntity val) => 
+  Unique val
+  -> SqlQuery (SqlExpr (Insertion val)) 
+  -> (SqlExpr (Entity val) -> SqlExpr (Entity val) -> [SqlExpr (Update val)])
+  -> SqlWriteT m Int64
+insertSelectWithConflictCount unique query conflictQuery = do
+  conn <- R.ask
+  uncurry rawExecuteCount $
+    combine
+      (toRawSql INSERT_INTO (conn, initialIdentState) (fmap EInsertFinal query))
+      (conflict conn)
+  where
+    proxy :: Proxy val
+    proxy = Proxy
+    updates = conflictQuery entCurrent entExcluded
+    combine (tlb1,vals1) (tlb2,vals2) = (builderToText (tlb1 `mappend` tlb2), vals1 ++ vals2)
+    entExcluded = EEntity $ I "excluded"
+    tableName = unDBName . entityDB . entityDef
+    entCurrent = EEntity $ I (tableName proxy)
+    -- there must be a better way to get the constrain name from a unique, make this not a list search
+    uniqueDef = head . filter ((==) (persistUniqueToFieldNames unique) . uniqueFields) . entityUniques . entityDef $ proxy
+    constraint = TLB.fromText . unDBName . uniqueDBName $ uniqueDef
+    renderedUpdates :: (BackendCompatible SqlBackend backend) => backend -> (TLB.Builder, [PersistValue])
+    renderedUpdates conn = renderUpdates conn updates
+    conflict conn = (foldr1 mappend ([
+        TLB.fromText "ON CONFLICT ON CONSTRAINT \"",
+        constraint,
+        TLB.fromText "\" DO "
+      ] ++ if null updates then [TLB.fromText "NOTHING"] else [      
+        TLB.fromText "UPDATE SET ",
+        updatesTLB
+      ]),values)
+      where
+        (updatesTLB,values) = renderedUpdates conn

test/Common/Test.hs

@@ -52,6 +52,7 @@ module Common.Test
     , Circle (..)
     , Numbers (..)
     , OneUnique(..)
+    , Unique(..)
     ) where
 
 import Control.Monad (forM_, replicateM, replicateM_, void)

test/PostgreSQL/Test.hs

@@ -978,6 +978,50 @@ testUpsert =
       u3e <- EP.upsert u3 [OneUniqueName =. val "fifth"]
       liftIO $ entityVal u3e `shouldBe` u1{oneUniqueName="fifth"}
 
+testInsertSelectWithConflict :: Spec
+testInsertSelectWithConflict =
+  describe "insertSelectWithConflict test" $ do
+    it "Should do Nothing when no updates set" $ run $ do
+      _ <- insert p1
+      _ <- insert p2
+      _ <- insert p3
+      n1 <- EP.insertSelectWithConflictCount (UniqueValue undefined) (
+          from $ \p -> return $ OneUnique <# val "test" <&> (p ^. PersonFavNum)
+        )
+        (\current excluded -> [])
+      uniques1 <- select $ from $ \u -> return u
+      n2 <- EP.insertSelectWithConflictCount (UniqueValue undefined) (
+          from $ \p -> return $ OneUnique <# val "test" <&> (p ^. PersonFavNum)
+        )
+        (\current excluded -> [])
+      uniques2 <- select $ from $ \u -> return u
+      liftIO $ n1 `shouldBe` 3
+      liftIO $ n2 `shouldBe` 0
+      let test = map (OneUnique "test" . personFavNum) [p1,p2,p3]
+      liftIO $ map entityVal uniques1 `shouldBe` test
+      liftIO $ map entityVal uniques2 `shouldBe` test
+    it "Should update a value if given an update on conflict" $ run $ do
+        _ <- insert p1
+        _ <- insert p2
+        _ <- insert p3
+        -- Note, have to sum 4 so that the update does not conflicts again with another row.
+        n1 <- EP.insertSelectWithConflictCount (UniqueValue undefined) (
+            from $ \p -> return $ OneUnique <# val "test" <&> (p ^. PersonFavNum)
+          )
+          (\current excluded -> [OneUniqueValue =. val 4 +. (current ^. OneUniqueValue) +. (excluded ^. OneUniqueValue)])
+        uniques1 <- select $ from $ \u -> return u
+        n2 <- EP.insertSelectWithConflictCount (UniqueValue undefined) (
+            from $ \p -> return $ OneUnique <# val "test" <&> (p ^. PersonFavNum)
+          )
+          (\current excluded -> [OneUniqueValue =. val 4 +. (current ^. OneUniqueValue) +. (excluded ^. OneUniqueValue)])
+        uniques2 <- select $ from $ \u -> return u
+        liftIO $ n1 `shouldBe` 3
+        liftIO $ n2 `shouldBe` 3
+        let test = map (OneUnique "test" . personFavNum) [p1,p2,p3]
+            test2 = map (OneUnique "test" . (+4) . (*2) . personFavNum) [p1,p2,p3]
+        liftIO $ map entityVal uniques1 `shouldBe` test
+        liftIO $ map entityVal uniques2 `shouldBe` test2
+
 type JSONValue = Maybe (JSONB A.Value)
 
 createSaneSQL :: (PersistField a) => SqlExpr (Value a) -> T.Text -> [PersistValue] -> IO ()
@@ -1051,6 +1095,7 @@ main = do
       testPostgresqlTextFunctions
       testInsertUniqueViolation
       testUpsert
+      testInsertSelectWithConflict
       describe "PostgreSQL JSON tests" $ do
         -- NOTE: We only clean the table once, so we
         -- can use its contents across all JSON tests