1 {-# LANGUAGE OverloadedStrings #-}
    2 module Databrary.Model.SQL
    3   ( selectQuery
    4   , isUniqueViolation
    5   , isExclusionViolation
    6   , isForeignKeyViolation
    7   , tryUpdateOrInsert
    8   , updateOrInsert
    9   ) where
   10 
   11 import Control.Monad (guard)
   12 import Database.PostgreSQL.Typed.Protocol (PGError(..), pgErrorCode)
   13 import Database.PostgreSQL.Typed.Query (PGQuery)
   14 import qualified Language.Haskell.TH as TH
   15 
   16 import Databrary.Service.DB
   17 import Databrary.Model.SQL.Select
   18 
   19 selectQuery :: Selector -> String -> TH.ExpQ
   20 selectQuery = selectDistinctQuery Nothing
   21 
   22 isUniqueViolation, isExclusionViolation, isForeignKeyViolation :: PGError -> Bool
   23 isUniqueViolation = ("23505" ==) . pgErrorCode
   24 isExclusionViolation e = pgErrorCode e `elem` ["23505","23P01"]
   25 isForeignKeyViolation = ("23503" ==) . pgErrorCode
   26 
   27 tryUpdateOrInsert :: (MonadDB c m, PGQuery q a) => (PGError -> Maybe e) -> q -> q -> m (Either e (Int, [a]))
   28 tryUpdateOrInsert err upd ins = dbTransaction uoi where
   29   err' e
   30     | isUniqueViolation e = Just Nothing
   31     | otherwise = Just <$> err e
   32   uoi = do
   33     u <- dbTryJust err $ dbRunQuery upd
   34     case u of
   35       Right (0, _) -> do
   36         _ <- dbExecuteSimple "SAVEPOINT pre_insert"
   37         i <- dbTryJust err' $ dbRunQuery ins
   38         case i of
   39           Left Nothing -> do
   40             _ <- dbExecuteSimple "ROLLBACK TO SAVEPOINT pre_insert"
   41             uoi
   42           Left (Just e) -> return $ Left e
   43           Right r -> return $ Right r
   44       _ -> return u
   45 
   46 updateOrInsert :: (MonadDB c m, PGQuery q a) => q -> q -> m (Int, [a])
   47 -- updateOrInsert upd ins = either fail return <$> tryUpdateOrInsert (const Nothing) upd ins
   48 updateOrInsert upd ins = dbTransaction uoi where
   49   uoi = do
   50     u@(n, _) <- dbRunQuery upd
   51     if n /= 0
   52       then return u
   53       else do
   54         _ <- dbExecuteSimple "SAVEPOINT pre_insert"
   55         i <- dbTryJust (guard . isUniqueViolation) $ dbRunQuery ins
   56         either (\() -> do
   57           _ <- dbExecuteSimple "ROLLBACK TO SAVEPOINT pre_insert"
   58           uoi)
   59           return i