@@ -0,0 +1,172 @@
1+module Control.Carrier.Squeal
2+ ( DBConnection,
3+ runSquealWithConn,
4+ runSquealWithConnRethrow,
5+ runSquealWithConn',
6+ SquealC (..),
7+ getSquealPool,
8+ runSqueal',
9+ runSqueal,
10+ runSquealPool,
11+ runSquealRethrow,
12+ SquealPool (..),
13+ module Control.Algebra,
14+ )
15+where
16+
17+import Control.Algebra
18+import Control.Carrier.Orphans ()
19+import Control.Effect.Squeal
20+import Control.Monad.IO.Unlift
21+import Data.Functor
22+import qualified Squeal.PostgreSQL as Sq
23+import UnliftIO
24+import UnliftIO.Pool
25+
26+type DBConnection (schemas :: SchemasType) = K Connection schemas
27+
28+newtype SquealC schemas m k = SquealC {unSquealC :: DBConnection schemas -> m k}
29+
30+instance Functor m => Functor (SquealC schemas m) where
31+ fmap f (SquealC mk) = SquealC $ (fmap . fmap) f mk
32+
33+instance Applicative m => Applicative (SquealC schemas m) where
34+ pure x = SquealC $ \_ -> pure x
35+ (SquealC mklhs) <*> (SquealC mkrhs) = SquealC $ \r -> mklhs r <*> mkrhs r
36+
37+instance Monad m => Monad (SquealC schemas m) where
38+ (SquealC mk) >>= f = SquealC $ \r -> mk r >>= (runSquealWithConn' r . f)
39+
40+instance MonadIO m => MonadIO (SquealC schemas m) where
41+ liftIO = SquealC . const . liftIO
42+
43+instance MonadUnliftIO m => MonadUnliftIO (SquealC schemas m) where
44+ withRunInIO inner = SquealC $ \r -> withRunInIO (\f -> inner (f . runSquealWithConn' r))
45+
46+instance (MonadUnliftIO m, Algebra sig m) => Algebra (Squeal schemas :+: sig) (SquealC schemas m) where
47+ alg (L (ManipulateParams man x mk)) = SquealC $ \r -> do
48+ res <- flip evalPQ r $ Sq.manipulateParams man x
49+ runSquealWithConn' r $ mk res
50+ alg (L (TraversePrepared man x mk)) = SquealC $ \r -> do
51+ res <- flip evalPQ r $ Sq.traversePrepared man x
52+ runSquealWithConn' r $ mk res
53+ alg (L (TraversePrepared_ man x mk)) = SquealC $ \r -> do
54+ flip evalPQ r $ Sq.traversePrepared_ man x
55+ runSquealWithConn' r mk
56+ alg (R other) = SquealC $ \r -> alg . hmap (runSquealWithConn' r) $ other
57+
58+-- | Run a squeal session using the given database connection without a transaction and without any error handling. You probably shouldn't use this.
59+runSquealWithConn' :: DBConnection schemas -> SquealC schemas m k -> m k
60+runSquealWithConn' r (SquealC mk) = mk r
61+
62+-- | Run a squeal session using the given database connection, transaction mode and error handler.
63+runSquealWithConn ::
64+ MonadUnliftIO m =>
65+ DBConnection schemas ->
66+ Maybe TransactionMode ->
67+ (SquealException -> m k) ->
68+ SquealC schemas m k ->
69+ m k
70+runSquealWithConn db tr er mk =
71+ handleSqueal er $ maybe id (transactionallyRetry' db) tr (runSquealWithConn' db mk)
72+ where
73+ transactionallyRetry' ::
74+ (MonadUnliftIO m) =>
75+ DBConnection schemas ->
76+ TransactionMode ->
77+ m x ->
78+ m x
79+ transactionallyRetry' conn mode tx = mask $ \restore ->
80+ loop . try $ do
81+ x <- restore tx
82+ flip evalPQ conn $ Sq.manipulate_ commit
83+ return x
84+ where
85+ loop attempt = do
86+ flip evalPQ conn $ Sq.manipulate_ $ begin mode
87+ attempt >>= \case
88+ Left (PQException (PQState _ (Just "40001") _)) -> do
89+ flip evalPQ conn $ Sq.manipulate_ rollback
90+ loop attempt
91+ Left err -> do
92+ flip evalPQ conn $ Sq.manipulate_ rollback
93+ throwIO err
94+ Right x -> return x
95+
96+-- | Run a squeal session using the given database connection and transaction mode. Errors will not be handled.
97+runSquealWithConnRethrow ::
98+ MonadUnliftIO m =>
99+ DBConnection schemas ->
100+ Maybe TransactionMode ->
101+ SquealC schemas m k ->
102+ m k
103+runSquealWithConnRethrow db tr = runSquealWithConn db tr throwIO
104+
105+newtype SquealPoolC schemas m k = SquealPoolC {unSquealPoolC :: Pool (DBConnection schemas) -> m k}
106+
107+instance Functor m => Functor (SquealPoolC schemas m) where
108+ fmap f (SquealPoolC mk) = SquealPoolC $ (fmap . fmap) f mk
109+
110+instance Applicative m => Applicative (SquealPoolC schemas m) where
111+ pure x = SquealPoolC $ const $ pure x
112+ (SquealPoolC mklhs) <*> (SquealPoolC mkrhs) = SquealPoolC $ \r -> mklhs r <*> mkrhs r
113+
114+instance Monad m => Monad (SquealPoolC schemas m) where
115+ (SquealPoolC mk) >>= f = SquealPoolC $ \r -> mk r >>= (($ r) . unSquealPoolC . f)
116+
117+instance MonadIO m => MonadIO (SquealPoolC schemas m) where
118+ liftIO = SquealPoolC . const . liftIO
119+
120+instance MonadUnliftIO m => MonadUnliftIO (SquealPoolC schemas m) where
121+ withRunInIO inner = SquealPoolC $ \r -> withRunInIO (\f -> inner (f . ($ r) . unSquealPoolC))
122+
123+runSquealPool :: Pool (DBConnection schemas) -> SquealPoolC schemas m k -> m k
124+runSquealPool conn (SquealPoolC f) = f conn
125+
126+instance Algebra sig m => Algebra (SquealPool schemas :+: sig) (SquealPoolC schemas m) where
127+ alg (L (GetSquealPool mk)) = SquealPoolC $ \r -> runSquealPool r $ mk r
128+ alg (R other) = SquealPoolC $ \r -> alg . hmap (runSquealPool r) $ other
129+
130+newtype SquealPool schemas m k = GetSquealPool (Pool (DBConnection schemas) -> m k)
131+
132+instance Functor m => Functor (SquealPool schemas m) where
133+ fmap f (GetSquealPool mk) = GetSquealPool ((fmap . fmap) f mk)
134+
135+instance HFunctor (SquealPool schemas) where
136+ hmap f (GetSquealPool mk) = GetSquealPool (fmap f mk)
137+
138+instance Effect (SquealPool schemas) where
139+ thread ctx f (GetSquealPool mk) = GetSquealPool $ \y -> f (ctx $> mk y)
140+
141+getSquealPool :: Has (SquealPool schemas) sig m => m (Pool (DBConnection schemas))
142+getSquealPool = send $ GetSquealPool pure
143+
144+-- | Run a squeal session picking a database connection from the connection pool without a transaction and without any error handling. You probably shouldn't use this.
145+runSqueal' ::
146+ (MonadUnliftIO m, Has (SquealPool schemas) sig m) =>
147+ SquealC schemas m k ->
148+ m k
149+runSqueal' = runSqueal Nothing throwIO
150+
151+-- | Run a squeal session picking a database connection from the connection pool with the given transaction mode and error handler.
152+runSqueal ::
153+ (MonadUnliftIO m, Has (SquealPool schemas) sig m) =>
154+ Maybe TransactionMode ->
155+ (SquealException -> m k) ->
156+ SquealC schemas m k ->
157+ m k
158+runSqueal tr er mk = do
159+ pool <- getSquealPool
160+ withResource pool $ \db ->
161+ runSquealWithConn db tr er mk
162+
163+-- | Run a squeal session picking a database connection from the connection pool with the given transaction mode. Errors will not be handled.
164+runSquealRethrow ::
165+ (MonadUnliftIO m, Has (SquealPool schemas) sig m) =>
166+ Maybe TransactionMode ->
167+ SquealC schemas m k ->
168+ m k
169+runSquealRethrow tr mk = do
170+ pool <- getSquealPool
171+ withResource pool $ \db ->
172+ runSquealWithConn db tr throwIO mk