refactor: move connection pool functions to Pool.hs

Signed-off-by: Taimoor Zaeem <taimoorzaeem@gmail.com>
2026-07-25 14:39:48 +05:00
parent a108d66968
commit e32e2b07b5
3 changed files with 113 additions and 95 deletions
@@ -46,6 +46,7 @@ library
  exposed-modules:    PostgREST.Admin
                      PostgREST.App
                      PostgREST.AppState
+                      PostgREST.AppState.Pool
                      PostgREST.AppState.Types
                      PostgREST.Auth
                      PostgREST.Auth.Jwt
@@ -29,15 +29,10 @@ module PostgREST.AppState
  , waitForSchemaCacheLoaded
  ) where

-import qualified Data.ByteString.Char8      as BS
-import           Data.Either.Combinators    (whenLeft)
 import qualified Hasql.Pool                 as SQL
-import qualified Hasql.Pool.Config          as SQL
 import qualified Hasql.Session              as SQL
 import qualified Hasql.Transaction.Sessions as SQL
-import qualified Network.HTTP.Types.Status  as HTTP
 import qualified PostgREST.Auth.JwtCache    as JwtCache
-import qualified PostgREST.Error            as Error
 import qualified PostgREST.Logger           as Logger
 import qualified PostgREST.Metrics          as Metrics
 import           PostgREST.Observation
@@ -52,9 +47,9 @@ import Data.Time.Clock    (getCurrentTime)

 import Control.Concurrent.STM            (newEmptyTMVarIO, putTMVar, readTMVar,
                                          tryReadTMVar, tryTakeTMVar)
+import PostgREST.AppState.Pool           (destroy, flushPool, initPool, usePool)
 import PostgREST.Auth.JwtCache           (update)
-import PostgREST.Config                  (AppConfig (..), readAppConfig,
-                                          toConnectionSettings)
+import PostgREST.Config                  (AppConfig (..), readAppConfig)
 import PostgREST.Config.Database         (queryDbSettings, queryPgVersion,
                                          queryRoleSettings)
 import PostgREST.Config.PgVersion        (PgVersion (..), minimumPgVersion)
@@ -98,94 +93,6 @@ initWithPool pool conf loggerState metricsState observer appKiller = mdo

  return appState

-- | Destroy the pool on shutdown.
-- | Differs from flushPool in not emiting PoolFlushed observation.
-destroy :: AppState -> IO ()
-destroy AppState{..} = SQL.release statePool
-
-initPool :: AppConfig -> ObservationHandler -> IO SQL.Pool
-initPool cfg@AppConfig{..} observer = do
-  SQL.acquire $ SQL.settings
-    [ SQL.size configDbPoolSize
-    , SQL.acquisitionTimeout $ fromIntegral configDbPoolAcquisitionTimeout
-    , SQL.agingTimeout $ fromIntegral configDbPoolMaxLifetime
-    , SQL.idlenessTimeout $ fromIntegral configDbPoolMaxIdletime
-    , SQL.staticConnectionSettings $ toConnectionSettings identity cfg
-    , SQL.observationHandler $ observer . HasqlPoolObs
-    ]
-
-- | Run an action with a database connection.
-usePool :: AppState -> SQL.Session a -> IO (Either SQL.UsageError a)
-usePool appState@AppState{stateObserver=observer, ..} sess = do
-    observer PoolRequest
-
-    res <- SQL.use statePool sess
-
-    observer PoolRequestFullfilled
-
-    whenLeft res (\case
-      SQL.AcquisitionTimeoutUsageError ->
-        observer PoolAcqTimeoutObs
-      err@(SQL.ConnectionUsageError e) ->
-        let failureMessage = BS.unpack $ fromMaybe mempty e in
-        when (("FATAL:  password authentication failed" `isInfixOf` failureMessage) || ("no password supplied" `isInfixOf` failureMessage)) $ do
-          observer $ ExitDBFatalError ServerAuthError err
-          killApp appState
-      err@(SQL.SessionUsageError (SQL.QueryError tpl _ (SQL.ResultError resultErr))) ->
-        handleResultError err tpl resultErr
-      err@(SQL.SessionUsageError (SQL.PipelineError (SQL.ResultError resultErr))) ->
-        -- Passing the empty template will not work for schema cache queries, see TODO further below.
-        handleResultError err mempty resultErr
-      err@(SQL.SessionUsageError (SQL.QueryError _ _ (SQL.ClientError _))) ->
-        -- An error on the client-side, usually indicates problems with connection
-        observer $ QueryErrorCodeHighObs err
-      SQL.SessionUsageError (SQL.PipelineError (SQL.ClientError _))  -> pure ()
-      )
-
-    return res
-  where
-    handleResultError err tpl resultErr = do
-      case resultErr of
-        SQL.UnexpectedResult{} -> do
-          observer $ ExitDBFatalError ServerPgrstBug err
-          killApp appState
-        SQL.RowError{} -> do
-          observer $ ExitDBFatalError ServerPgrstBug err
-          killApp appState
-        SQL.UnexpectedAmountOfRows{} -> do
-          observer $ ExitDBFatalError ServerPgrstBug err
-          killApp appState
-        -- Check for a syntax error (42601 is the pg code) only for queries that don't have `WITH pgrst_source` as prefix.
-        -- This would mean the error is on our schema cache queries, so we treat it as fatal.
-        -- TODO have a better way to mark this as a schema cache query
-        SQL.ServerError "42601" _ _ _ _ ->
-          unless ("WITH pgrst_source" `BS.isPrefixOf` tpl) $ do
-            observer $ ExitDBFatalError ServerPgrstBug err
-            killApp appState
-        -- Check for a "prepared statement <name> already exists" error (Code 42P05: duplicate_prepared_statement).
-        -- This would mean that a connection pooler in transaction mode is being used
-        -- while prepared statements are enabled in the PostgREST configuration,
-        -- both of which are incompatible with each other.
-        SQL.ServerError "42P05" _ _ _ _ -> do
-          observer $ ExitDBFatalError ServerError42P05 err
-          killApp appState
-        -- Check for a "transaction blocks not allowed in statement pooling mode" error (Code 08P01: protocol_violation).
-        -- This would mean that a connection pooler in statement mode is being used which is not supported in PostgREST.
-        SQL.ServerError "08P01" "transaction blocks not allowed in statement pooling mode" _ _ _ -> do
-          observer $ ExitDBFatalError ServerError08P01 err
-          killApp appState
-        SQL.ServerError{} ->
-          when (Error.status (Error.PgError False err) >= HTTP.status500) $
-            observer $ QueryErrorCodeHighObs err
-
-- | Flush the connection pool so that any future use of the pool will
-- use connections freshly established after this call.
-- | Emits PoolFlushed observation
-flushPool :: AppState -> IO ()
-flushPool AppState{..} = do
-  SQL.release statePool
-  stateObserver PoolFlushed
-
 isConnEstablished :: AppState -> IO Bool
 isConnEstablished appState = do
  AppConfig{..} <- getConfig appState
@@ -0,0 +1,110 @@
+{-# LANGUAGE LambdaCase      #-}
+{-# LANGUAGE RecordWildCards #-}
+module PostgREST.AppState.Pool
+  ( destroy
+  , initPool
+  , flushPool
+  , usePool
+  ) where
+
+import qualified Data.ByteString.Char8     as BS
+import           Data.Either.Combinators   (whenLeft)
+import qualified Hasql.Pool                as SQL
+import qualified Hasql.Pool.Config         as SQL
+import qualified Hasql.Session             as SQL
+import qualified Network.HTTP.Types.Status as HTTP
+import qualified PostgREST.Error           as Error
+import           PostgREST.Observation
+
+import PostgREST.Config (AppConfig (..), toConnectionSettings)
+
+import PostgREST.AppState.Types
+import Protolude
+
+-- | Destroy the pool on shutdown.
+-- | Differs from flushPool in not emiting PoolFlushed observation.
+destroy :: AppState -> IO ()
+destroy AppState{..} = SQL.release statePool
+
+initPool :: AppConfig -> ObservationHandler -> IO SQL.Pool
+initPool cfg@AppConfig{..} observer = do
+  SQL.acquire $ SQL.settings
+    [ SQL.size configDbPoolSize
+    , SQL.acquisitionTimeout $ fromIntegral configDbPoolAcquisitionTimeout
+    , SQL.agingTimeout $ fromIntegral configDbPoolMaxLifetime
+    , SQL.idlenessTimeout $ fromIntegral configDbPoolMaxIdletime
+    , SQL.staticConnectionSettings $ toConnectionSettings identity cfg
+    , SQL.observationHandler $ observer . HasqlPoolObs
+    ]
+
+-- | Run an action with a database connection.
+usePool :: AppState -> SQL.Session a -> IO (Either SQL.UsageError a)
+usePool appState@AppState{stateObserver=observer, ..} sess = do
+    observer PoolRequest
+
+    res <- SQL.use statePool sess
+
+    observer PoolRequestFullfilled
+
+    whenLeft res (\case
+      SQL.AcquisitionTimeoutUsageError ->
+        observer PoolAcqTimeoutObs
+      err@(SQL.ConnectionUsageError e) ->
+        let failureMessage = BS.unpack $ fromMaybe mempty e in
+        when (("FATAL:  password authentication failed" `isInfixOf` failureMessage) || ("no password supplied" `isInfixOf` failureMessage)) $ do
+          observer $ ExitDBFatalError ServerAuthError err
+          killApp appState
+      err@(SQL.SessionUsageError (SQL.QueryError tpl _ (SQL.ResultError resultErr))) ->
+        handleResultError err tpl resultErr
+      err@(SQL.SessionUsageError (SQL.PipelineError (SQL.ResultError resultErr))) ->
+        -- Passing the empty template will not work for schema cache queries, see TODO further below.
+        handleResultError err mempty resultErr
+      err@(SQL.SessionUsageError (SQL.QueryError _ _ (SQL.ClientError _))) ->
+        -- An error on the client-side, usually indicates problems with connection
+        observer $ QueryErrorCodeHighObs err
+      SQL.SessionUsageError (SQL.PipelineError (SQL.ClientError _))  -> pure ()
+      )
+
+    return res
+  where
+    handleResultError err tpl resultErr = do
+      case resultErr of
+        SQL.UnexpectedResult{} -> do
+          observer $ ExitDBFatalError ServerPgrstBug err
+          killApp appState
+        SQL.RowError{} -> do
+          observer $ ExitDBFatalError ServerPgrstBug err
+          killApp appState
+        SQL.UnexpectedAmountOfRows{} -> do
+          observer $ ExitDBFatalError ServerPgrstBug err
+          killApp appState
+        -- Check for a syntax error (42601 is the pg code) only for queries that don't have `WITH pgrst_source` as prefix.
+        -- This would mean the error is on our schema cache queries, so we treat it as fatal.
+        -- TODO have a better way to mark this as a schema cache query
+        SQL.ServerError "42601" _ _ _ _ ->
+          unless ("WITH pgrst_source" `BS.isPrefixOf` tpl) $ do
+            observer $ ExitDBFatalError ServerPgrstBug err
+            killApp appState
+        -- Check for a "prepared statement <name> already exists" error (Code 42P05: duplicate_prepared_statement).
+        -- This would mean that a connection pooler in transaction mode is being used
+        -- while prepared statements are enabled in the PostgREST configuration,
+        -- both of which are incompatible with each other.
+        SQL.ServerError "42P05" _ _ _ _ -> do
+          observer $ ExitDBFatalError ServerError42P05 err
+          killApp appState
+        -- Check for a "transaction blocks not allowed in statement pooling mode" error (Code 08P01: protocol_violation).
+        -- This would mean that a connection pooler in statement mode is being used which is not supported in PostgREST.
+        SQL.ServerError "08P01" "transaction blocks not allowed in statement pooling mode" _ _ _ -> do
+          observer $ ExitDBFatalError ServerError08P01 err
+          killApp appState
+        SQL.ServerError{} ->
+          when (Error.status (Error.PgError False err) >= HTTP.status500) $
+            observer $ QueryErrorCodeHighObs err
+
+-- | Flush the connection pool so that any future use of the pool will
+-- use connections freshly established after this call.
+-- | Emits PoolFlushed observation
+flushPool :: AppState -> IO ()
+flushPool AppState{..} = do
+  SQL.release statePool
+  stateObserver PoolFlushed