{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | An optional persistent session store. Initialize the schema once before
-- starting workers, then share a namespace across workers using the same
-- session data type. The caller owns the connection pool.
module Web.Spock.Session.Postgresql
  ( PostgresqlTx,
    PostgresqlSessionCfg (..),
    defaultPostgresqlSessionCfg,
    SessionDecodeError (..),
    initializePostgresqlSessions,
    newPostgresqlSessionStore,
  ) where

import Control.Concurrent (threadDelay)
import Control.Exception (Exception, catch, throwIO)
import Control.Monad (forM_, unless, void)
import Control.Monad.Reader (ReaderT (..))
import Data.Aeson (FromJSON, ToJSON, eitherDecodeStrict', encode)
import qualified Data.ByteString.Lazy as LBS
import Data.Pool (Pool, withResource)
import Data.Text (Text)
import qualified Data.Text.Encoding as T
import Data.Time (UTCTime)
import Database.PostgreSQL.Simple
import Database.PostgreSQL.Simple.Transaction
import Web.Spock.Config
import Web.Spock.Internal.SessionManager (Session (..))

-- | Deliberately has no MonadIO instance: the whole transaction can be retried,
-- so callbacks must not perform external side effects.
newtype PostgresqlTx a = PostgresqlTx (ReaderT Connection IO a)
  deriving ((forall a b. (a -> b) -> PostgresqlTx a -> PostgresqlTx b)
-> (forall a b. a -> PostgresqlTx b -> PostgresqlTx a)
-> Functor PostgresqlTx
forall a b. a -> PostgresqlTx b -> PostgresqlTx a
forall a b. (a -> b) -> PostgresqlTx a -> PostgresqlTx b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> PostgresqlTx a -> PostgresqlTx b
fmap :: forall a b. (a -> b) -> PostgresqlTx a -> PostgresqlTx b
$c<$ :: forall a b. a -> PostgresqlTx b -> PostgresqlTx a
<$ :: forall a b. a -> PostgresqlTx b -> PostgresqlTx a
Functor, Functor PostgresqlTx
Functor PostgresqlTx =>
(forall a. a -> PostgresqlTx a)
-> (forall a b.
    PostgresqlTx (a -> b) -> PostgresqlTx a -> PostgresqlTx b)
-> (forall a b c.
    (a -> b -> c)
    -> PostgresqlTx a -> PostgresqlTx b -> PostgresqlTx c)
-> (forall a b. PostgresqlTx a -> PostgresqlTx b -> PostgresqlTx b)
-> (forall a b. PostgresqlTx a -> PostgresqlTx b -> PostgresqlTx a)
-> Applicative PostgresqlTx
forall a. a -> PostgresqlTx a
forall a b. PostgresqlTx a -> PostgresqlTx b -> PostgresqlTx a
forall a b. PostgresqlTx a -> PostgresqlTx b -> PostgresqlTx b
forall a b.
PostgresqlTx (a -> b) -> PostgresqlTx a -> PostgresqlTx b
forall a b c.
(a -> b -> c) -> PostgresqlTx a -> PostgresqlTx b -> PostgresqlTx c
forall (f :: * -> *).
Functor f =>
(forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
$cpure :: forall a. a -> PostgresqlTx a
pure :: forall a. a -> PostgresqlTx a
$c<*> :: forall a b.
PostgresqlTx (a -> b) -> PostgresqlTx a -> PostgresqlTx b
<*> :: forall a b.
PostgresqlTx (a -> b) -> PostgresqlTx a -> PostgresqlTx b
$cliftA2 :: forall a b c.
(a -> b -> c) -> PostgresqlTx a -> PostgresqlTx b -> PostgresqlTx c
liftA2 :: forall a b c.
(a -> b -> c) -> PostgresqlTx a -> PostgresqlTx b -> PostgresqlTx c
$c*> :: forall a b. PostgresqlTx a -> PostgresqlTx b -> PostgresqlTx b
*> :: forall a b. PostgresqlTx a -> PostgresqlTx b -> PostgresqlTx b
$c<* :: forall a b. PostgresqlTx a -> PostgresqlTx b -> PostgresqlTx a
<* :: forall a b. PostgresqlTx a -> PostgresqlTx b -> PostgresqlTx a
Applicative, Applicative PostgresqlTx
Applicative PostgresqlTx =>
(forall a b.
 PostgresqlTx a -> (a -> PostgresqlTx b) -> PostgresqlTx b)
-> (forall a b. PostgresqlTx a -> PostgresqlTx b -> PostgresqlTx b)
-> (forall a. a -> PostgresqlTx a)
-> Monad PostgresqlTx
forall a. a -> PostgresqlTx a
forall a b. PostgresqlTx a -> PostgresqlTx b -> PostgresqlTx b
forall a b.
PostgresqlTx a -> (a -> PostgresqlTx b) -> PostgresqlTx b
forall (m :: * -> *).
Applicative m =>
(forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
$c>>= :: forall a b.
PostgresqlTx a -> (a -> PostgresqlTx b) -> PostgresqlTx b
>>= :: forall a b.
PostgresqlTx a -> (a -> PostgresqlTx b) -> PostgresqlTx b
$c>> :: forall a b. PostgresqlTx a -> PostgresqlTx b -> PostgresqlTx b
>> :: forall a b. PostgresqlTx a -> PostgresqlTx b -> PostgresqlTx b
$creturn :: forall a. a -> PostgresqlTx a
return :: forall a. a -> PostgresqlTx a
Monad)

data PostgresqlSessionCfg = PostgresqlSessionCfg
  { -- | Independent applications must use different namespaces.
    PostgresqlSessionCfg -> Text
psc_namespace :: Text,
    -- | Maximum retries for serialization failures or deadlocks. Other errors
    -- propagate immediately. Exhaustion propagates the last SQL exception.
    PostgresqlSessionCfg -> Int
psc_maxRetries :: Int
  }

defaultPostgresqlSessionCfg :: PostgresqlSessionCfg
defaultPostgresqlSessionCfg :: PostgresqlSessionCfg
defaultPostgresqlSessionCfg = Text -> Int -> PostgresqlSessionCfg
PostgresqlSessionCfg Text
"spock" Int
20

-- | Stored JSON cannot be decoded as the application's session type. No stored
-- content is included in this exception.
data SessionDecodeError = SessionDecodeError deriving (SessionDecodeError -> SessionDecodeError -> Bool
(SessionDecodeError -> SessionDecodeError -> Bool)
-> (SessionDecodeError -> SessionDecodeError -> Bool)
-> Eq SessionDecodeError
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: SessionDecodeError -> SessionDecodeError -> Bool
== :: SessionDecodeError -> SessionDecodeError -> Bool
$c/= :: SessionDecodeError -> SessionDecodeError -> Bool
/= :: SessionDecodeError -> SessionDecodeError -> Bool
Eq, Int -> SessionDecodeError -> ShowS
[SessionDecodeError] -> ShowS
SessionDecodeError -> String
(Int -> SessionDecodeError -> ShowS)
-> (SessionDecodeError -> String)
-> ([SessionDecodeError] -> ShowS)
-> Show SessionDecodeError
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> SessionDecodeError -> ShowS
showsPrec :: Int -> SessionDecodeError -> ShowS
$cshow :: SessionDecodeError -> String
show :: SessionDecodeError -> String
$cshowList :: [SessionDecodeError] -> ShowS
showList :: [SessionDecodeError] -> ShowS
Show)
instance Exception SessionDecodeError

-- | Create the adapter's table. Run once during deployment/startup, before
-- concurrent workers. All ordinary operations are scoped to a namespace.
initializePostgresqlSessions :: Connection -> IO ()
initializePostgresqlSessions :: Connection -> IO ()
initializePostgresqlSessions Connection
connection = IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection -> Query -> IO Int64
execute_ Connection
connection
  Query
"CREATE TABLE IF NOT EXISTS spock_sessions (\
  \namespace TEXT NOT NULL, session_id TEXT NOT NULL, csrf_token TEXT NOT NULL,\
  \valid_until TIMESTAMPTZ NOT NULL, session_data JSONB NOT NULL,\
  \PRIMARY KEY (namespace, session_id))"

newPostgresqlSessionStore :: (ToJSON sess, FromJSON sess) =>
  PostgresqlSessionCfg -> Pool Connection -> IO (SessionStore (Session conn sess st) PostgresqlTx)
newPostgresqlSessionStore :: forall sess conn st.
(ToJSON sess, FromJSON sess) =>
PostgresqlSessionCfg
-> Pool Connection
-> IO (SessionStore (Session conn sess st) PostgresqlTx)
newPostgresqlSessionStore PostgresqlSessionCfg
cfg Pool Connection
pool = do
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (PostgresqlSessionCfg -> Int
psc_maxRetries PostgresqlSessionCfg
cfg Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IOError -> IO ()
forall a. HasCallStack => IOError -> IO a
ioError (IOError -> IO ()) -> IOError -> IO ()
forall a b. (a -> b) -> a -> b
$ String -> IOError
userError String
"psc_maxRetries must be nonnegative"
  let namespace :: Text
namespace = PostgresqlSessionCfg -> Text
psc_namespace PostgresqlSessionCfg
cfg
      store :: SessionStore (Session conn sess st) PostgresqlTx
store = SessionStore
        { ss_runTx :: forall a. PostgresqlTx a -> IO a
ss_runTx = \(PostgresqlTx ReaderT Connection IO a
action) -> Pool Connection -> (Connection -> IO a) -> IO a
forall a r. Pool a -> (a -> IO r) -> IO r
withResource Pool Connection
pool ((Connection -> IO a) -> IO a) -> (Connection -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Connection
connection ->
            Int -> Connection -> IO a -> IO a
forall a. Int -> Connection -> IO a -> IO a
retryTransaction (PostgresqlSessionCfg -> Int
psc_maxRetries PostgresqlSessionCfg
cfg) Connection
connection (ReaderT Connection IO a -> Connection -> IO a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT Connection IO a
action Connection
connection),
          ss_loadSession :: Text -> PostgresqlTx (Maybe (Session conn sess st))
ss_loadSession = \Text
sid -> (Connection -> IO (Maybe (Session conn sess st)))
-> PostgresqlTx (Maybe (Session conn sess st))
forall a. (Connection -> IO a) -> PostgresqlTx a
onConnection ((Connection -> IO (Maybe (Session conn sess st)))
 -> PostgresqlTx (Maybe (Session conn sess st)))
-> (Connection -> IO (Maybe (Session conn sess st)))
-> PostgresqlTx (Maybe (Session conn sess st))
forall a b. (a -> b) -> a -> b
$ \Connection
connection -> do
            rows <- Connection
-> Query -> (Text, Text) -> IO [(Text, Text, UTCTime, Text)]
forall q r.
(ToRow q, FromRow r) =>
Connection -> Query -> q -> IO [r]
query Connection
connection
              Query
"SELECT session_id, csrf_token, valid_until, session_data::text FROM spock_sessions WHERE namespace = ? AND session_id = ?"
              (Text
namespace, Text
sid)
            case rows of
              [] -> Maybe (Session conn sess st) -> IO (Maybe (Session conn sess st))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Session conn sess st)
forall a. Maybe a
Nothing
              (Text, Text, UTCTime, Text)
row : [(Text, Text, UTCTime, Text)]
_ -> Session conn sess st -> Maybe (Session conn sess st)
forall a. a -> Maybe a
Just (Session conn sess st -> Maybe (Session conn sess st))
-> IO (Session conn sess st) -> IO (Maybe (Session conn sess st))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Text, Text, UTCTime, Text) -> IO (Session conn sess st)
forall sess conn st.
FromJSON sess =>
(Text, Text, UTCTime, Text) -> IO (Session conn sess st)
decodeSession (Text, Text, UTCTime, Text)
row,
          ss_deleteSession :: Text -> PostgresqlTx ()
ss_deleteSession = \Text
sid -> (Connection -> IO ()) -> PostgresqlTx ()
forall a. (Connection -> IO a) -> PostgresqlTx a
onConnection ((Connection -> IO ()) -> PostgresqlTx ())
-> (Connection -> IO ()) -> PostgresqlTx ()
forall a b. (a -> b) -> a -> b
$ \Connection
connection -> IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection -> Query -> (Text, Text) -> IO Int64
forall q. ToRow q => Connection -> Query -> q -> IO Int64
execute Connection
connection
            Query
"DELETE FROM spock_sessions WHERE namespace = ? AND session_id = ?" (Text
namespace, Text
sid),
          ss_storeSession :: Session conn sess st -> PostgresqlTx ()
ss_storeSession = \Session conn sess st
session -> (Connection -> IO ()) -> PostgresqlTx ()
forall a. (Connection -> IO a) -> PostgresqlTx a
onConnection ((Connection -> IO ()) -> PostgresqlTx ())
-> (Connection -> IO ()) -> PostgresqlTx ()
forall a b. (a -> b) -> a -> b
$ \Connection
connection -> IO Int64 -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Int64 -> IO ()) -> IO Int64 -> IO ()
forall a b. (a -> b) -> a -> b
$ Connection
-> Query -> (Text, Text, Text, UTCTime, Text) -> IO Int64
forall q. ToRow q => Connection -> Query -> q -> IO Int64
execute Connection
connection
            Query
"INSERT INTO spock_sessions (namespace, session_id, csrf_token, valid_until, session_data) VALUES (?, ?, ?, ?, ?::jsonb) \
            \ON CONFLICT (namespace, session_id) DO UPDATE SET csrf_token = EXCLUDED.csrf_token, valid_until = EXCLUDED.valid_until, session_data = EXCLUDED.session_data"
            (Text
namespace, Session conn sess st -> Text
forall conn sess st. Session conn sess st -> Text
sess_id Session conn sess st
session, Session conn sess st -> Text
forall conn sess st. Session conn sess st -> Text
sess_csrfToken Session conn sess st
session, Session conn sess st -> UTCTime
forall conn sess st. Session conn sess st -> UTCTime
sess_validUntil Session conn sess st
session, ByteString -> Text
T.decodeUtf8 (ByteString -> Text) -> ByteString -> Text
forall a b. (a -> b) -> a -> b
$ LazyByteString -> ByteString
LBS.toStrict (LazyByteString -> ByteString) -> LazyByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ sess -> LazyByteString
forall a. ToJSON a => a -> LazyByteString
encode (sess -> LazyByteString) -> sess -> LazyByteString
forall a b. (a -> b) -> a -> b
$ Session conn sess st -> sess
forall conn sess st. Session conn sess st -> sess
sess_data Session conn sess st
session),
          ss_toList :: PostgresqlTx [Session conn sess st]
ss_toList = (Connection -> IO [Session conn sess st])
-> PostgresqlTx [Session conn sess st]
forall a. (Connection -> IO a) -> PostgresqlTx a
onConnection ((Connection -> IO [Session conn sess st])
 -> PostgresqlTx [Session conn sess st])
-> (Connection -> IO [Session conn sess st])
-> PostgresqlTx [Session conn sess st]
forall a b. (a -> b) -> a -> b
$ \Connection
connection -> do
            rows <- Connection
-> Query -> Only Text -> IO [(Text, Text, UTCTime, Text)]
forall q r.
(ToRow q, FromRow r) =>
Connection -> Query -> q -> IO [r]
query Connection
connection
              Query
"SELECT session_id, csrf_token, valid_until, session_data::text FROM spock_sessions WHERE namespace = ? ORDER BY session_id" (Text -> Only Text
forall a. a -> Only a
Only Text
namespace)
            mapM decodeSession rows,
          ss_filterSessions :: (Session conn sess st -> Bool) -> PostgresqlTx ()
ss_filterSessions = \Session conn sess st -> Bool
predicate -> do
            sessions <- SessionStore (Session conn sess st) PostgresqlTx
-> PostgresqlTx [Session conn sess st]
forall sess (tx :: * -> *). SessionStore sess tx -> tx [sess]
ss_toList SessionStore (Session conn sess st) PostgresqlTx
store
            forM_ sessions $ \Session conn sess st
session -> Bool -> PostgresqlTx () -> PostgresqlTx ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Session conn sess st -> Bool
predicate Session conn sess st
session) (PostgresqlTx () -> PostgresqlTx ())
-> PostgresqlTx () -> PostgresqlTx ()
forall a b. (a -> b) -> a -> b
$ SessionStore (Session conn sess st) PostgresqlTx
-> Text -> PostgresqlTx ()
forall sess (tx :: * -> *). SessionStore sess tx -> Text -> tx ()
ss_deleteSession SessionStore (Session conn sess st) PostgresqlTx
store (Session conn sess st -> Text
forall conn sess st. Session conn sess st -> Text
sess_id Session conn sess st
session),
          ss_mapSessions :: (Session conn sess st -> PostgresqlTx (Session conn sess st))
-> PostgresqlTx ()
ss_mapSessions = \Session conn sess st -> PostgresqlTx (Session conn sess st)
f -> do
            sessions <- SessionStore (Session conn sess st) PostgresqlTx
-> PostgresqlTx [Session conn sess st]
forall sess (tx :: * -> *). SessionStore sess tx -> tx [sess]
ss_toList SessionStore (Session conn sess st) PostgresqlTx
store
            forM_ sessions $ \Session conn sess st
session -> do
              updated <- Session conn sess st -> PostgresqlTx (Session conn sess st)
f Session conn sess st
session
              -- Like the STM store, mapping preserves the entry's key.
              ss_storeSession store (updated { sess_id = sess_id session })
        }
  SessionStore (Session conn sess st) PostgresqlTx
-> IO (SessionStore (Session conn sess st) PostgresqlTx)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SessionStore (Session conn sess st) PostgresqlTx
forall {conn} {st}.
SessionStore (Session conn sess st) PostgresqlTx
store

onConnection :: (Connection -> IO a) -> PostgresqlTx a
onConnection :: forall a. (Connection -> IO a) -> PostgresqlTx a
onConnection = ReaderT Connection IO a -> PostgresqlTx a
forall a. ReaderT Connection IO a -> PostgresqlTx a
PostgresqlTx (ReaderT Connection IO a -> PostgresqlTx a)
-> ((Connection -> IO a) -> ReaderT Connection IO a)
-> (Connection -> IO a)
-> PostgresqlTx a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Connection -> IO a) -> ReaderT Connection IO a
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT

decodeSession :: FromJSON sess => (Text, Text, UTCTime, Text) -> IO (Session conn sess st)
decodeSession :: forall sess conn st.
FromJSON sess =>
(Text, Text, UTCTime, Text) -> IO (Session conn sess st)
decodeSession (Text
sid, Text
csrf, UTCTime
expiry, Text
payload) =
  case ByteString -> Either String sess
forall a. FromJSON a => ByteString -> Either String a
eitherDecodeStrict' (Text -> ByteString
T.encodeUtf8 Text
payload) of
    Left String
_ -> SessionDecodeError -> IO (Session conn sess st)
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SessionDecodeError
SessionDecodeError
    Right sess
value -> Session conn sess st -> IO (Session conn sess st)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Session conn sess st -> IO (Session conn sess st))
-> Session conn sess st -> IO (Session conn sess st)
forall a b. (a -> b) -> a -> b
$ Text -> Text -> UTCTime -> sess -> Session conn sess st
forall conn sess st.
Text -> Text -> UTCTime -> sess -> Session conn sess st
Session Text
sid Text
csrf UTCTime
expiry sess
value

retryTransaction :: Int -> Connection -> IO a -> IO a
retryTransaction :: forall a. Int -> Connection -> IO a -> IO a
retryTransaction Int
limit Connection
connection IO a
action = Int -> IO a
go Int
0
  where
    go :: Int -> IO a
go Int
attempt = TransactionMode -> Connection -> IO a -> IO a
forall a. TransactionMode -> Connection -> IO a -> IO a
withTransactionMode (IsolationLevel -> ReadWriteMode -> TransactionMode
TransactionMode IsolationLevel
Serializable ReadWriteMode
ReadWrite) Connection
connection IO a
action
      IO a -> (SqlError -> IO a) -> IO a
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \(SqlError
err :: SqlError) ->
        if SqlError -> ByteString
sqlState SqlError
err ByteString -> [ByteString] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [ByteString
"40001", ByteString
"40P01"] Bool -> Bool -> Bool
&& Int
attempt Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
limit
          then do
            Int -> IO ()
threadDelay (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
100000 (Int
1000 Int -> Int -> Int
forall a. Num a => a -> a -> a
* (Int
attempt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)))
            Int -> IO a
go (Int
attempt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
          else SqlError -> IO a
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO SqlError
err