1- {-# LANGUAGE LambdaCase #-}
2- {-# LANGUAGE NamedFieldPuns #-}
3- {-# LANGUAGE RecordWildCards #-}
4- {-# LANGUAGE RecursiveDo #-}
1+ {-# LANGUAGE LambdaCase #-}
2+ {-# LANGUAGE NamedFieldPuns #-}
3+ {-# LANGUAGE RecordWildCards #-}
4+ {-# LANGUAGE RecursiveDo #-}
5+ {-# LANGUAGE TupleSections #-}
6+ {-# LANGUAGE TypeApplications #-}
57
68module PostgREST.AppState
79 ( AppState
@@ -49,13 +51,14 @@ import Control.AutoUpdate (defaultUpdateSettings, mkAutoUpdate,
4951import Control.Retry (RetryPolicy , RetryStatus (.. ), capDelay ,
5052 exponentialBackoff , retrying ,
5153 rsPreviousDelay )
52- import Data.IORef (IORef , atomicWriteIORef , newIORef ,
53- readIORef )
54+ import Data.IORef (IORef , atomicModifyIORef , atomicWriteIORef ,
55+ newIORef , readIORef )
5456import Data.Time.Clock (UTCTime , getCurrentTime )
5557
5658import Control.Concurrent.STM (TMVar , newEmptyTMVarIO ,
5759 putTMVar , readTMVar ,
58- tryReadTMVar , tryTakeTMVar )
60+ tryPutTMVar , tryReadTMVar ,
61+ tryTakeTMVar )
5962import PostgREST.Auth.JwtCache (JwtCacheState , update )
6063import PostgREST.Config (AppConfig (.. ),
6164 readAppConfig ,
@@ -71,6 +74,13 @@ import PostgREST.SchemaCache (SchemaCache (..),
7174 showSummary )
7275import PostgREST.SchemaCache.Identifiers (quoteQi )
7376
77+ import qualified Data.Aeson as JSON
78+ import qualified Data.ByteString.Lazy as LBS
79+ import qualified Network.HTTP.Client as HC
80+ import Network.URI (URI (.. ), URIAuth (.. ),
81+ parseURI , unEscapeString )
82+ import System.IO.Error (userError )
83+
7484import Protolude
7585
7686data AppState = AppState
@@ -110,16 +120,56 @@ newtype SchemaCacheStatus = SchemaCacheStatus
110120 { getSCStatusTMVar :: TMVar Bool
111121 }
112122
113- init :: AppConfig -> IO () -> IO AppState
114- init conf@ AppConfig {configLogLevel, configDbPoolSize} appKiller = do
123+ init :: AppConfig -> IO () -> Maybe Text -> IO AppState
124+ init conf@ AppConfig {configLogLevel, configDbPoolSize} appKiller schemaCacheLoadUri = do
115125 loggerState <- Logger. init
116126 metricsState <- Metrics. init configDbPoolSize
117127 let observer = liftA2 (>>) (Logger. observationLogger loggerState configLogLevel) (Metrics. observationMetrics metricsState)
118128
119129 observer $ AppStartObs prettyVersion
120130
121131 pool <- initPool conf observer
122- initWithPool pool conf loggerState metricsState observer appKiller
132+ appState <- initWithPool pool conf loggerState metricsState observer appKiller
133+
134+ runInitialSchemaCacheLoader observer schemaCacheLoadUri appState
135+
136+ pure appState
137+
138+ runInitialSchemaCacheLoader :: ObservationHandler -> Maybe Text -> AppState -> IO ()
139+ runInitialSchemaCacheLoader observer schemaCacheLoadUri AppState {stateSchemaCache, stateSCacheStatus= SchemaCacheStatus {getSCStatusTMVar}} = do
140+ void $ forkIO $
141+ traverse (fetchInitialSchemaCache observer) schemaCacheLoadUri >>= foldMap setInitialSchemaCache . join
142+ where
143+ setInitialSchemaCache sc =
144+ whenM (atomically $ tryPutTMVar getSCStatusTMVar True ) $
145+ atomicModifyIORef stateSchemaCache $ (, () ) . maybe (Just sc) Just
146+
147+ fetchInitialSchemaCache :: ObservationHandler -> Text -> IO (Maybe SchemaCache )
148+ fetchInitialSchemaCache observer uri = flip catches [
149+ Handler (handleError @ IOException ),
150+ Handler (handleError @ HC. HttpException ),
151+ Handler (handleError @ JSON. AesonException )
152+ ] $ maybe (throwIO $ userError " Invalid schema cache dump URI" ) pure =<< traverse fetchURI (parseURI $ toS uri)
153+ where
154+ handleError :: Show e => e -> IO (Maybe a )
155+ handleError = (Nothing <$ ) . observer . SchemaCacheInitialLoadFailureObs uri . show
156+
157+ fetchURI URI {uriScheme, .. }
158+ | uriScheme == " file:" = do
159+ path <- fileURIPath uriAuthority uriPath
160+ Just <$> (JSON. throwDecode =<< LBS. readFile path)
161+ | uriScheme `elem` [" http:" , " https:" ] = do
162+ request <- HC. parseUrlThrow $ toS uri
163+ manager <- HC. newManager HC. defaultManagerSettings
164+ HC. withResponse request manager $ \ response ->
165+ Just <$> (JSON. throwDecode . LBS. fromChunks =<< HC. brConsume (HC. responseBody response))
166+ | otherwise =
167+ throwIO $ userError $ " Unsupported schema cache dump URI scheme: " <> uriScheme
168+
169+ fileURIPath Nothing path = pure $ unEscapeString path
170+ fileURIPath (Just URIAuth {uriRegName= " " }) path = pure $ unEscapeString path
171+ fileURIPath (Just URIAuth {uriRegName= " localhost" }) path = pure $ unEscapeString path
172+ fileURIPath _ _ = throwIO $ userError " Only local file URIs are supported"
123173
124174initWithPool :: SQL. Pool -> AppConfig -> Logger. LoggerState -> Metrics. MetricsState -> ObservationHandler -> IO () -> IO AppState
125175initWithPool pool conf loggerState metricsState observer appKiller = mdo
0 commit comments