add: load schema cache dump from URI at startup

Author: mkleczek

CHANGELOG.md

@@ -15,6 +15,7 @@ All notable changes to this project will be documented in this file. From versio
 - Log schema cache queries timings on `log-level=debug` by @steve-chavez in #4805
 - Add GHC runtime metrics to the metrics endpoint by @mkleczek in #4862
 - Support running the admin server on a unix socket by @wolfgangwalther in #5003
+- Add `--schema-cache-uri` to asynchronously load an initial schema cache dump from `http(s)` or local `file` URIs by @mkleczek in #5031
 
 ### Fixed
 

docs/references/cli.rst

@@ -8,7 +8,7 @@ PostgREST provides a CLI with the options listed below:
 .. code:: text
 
   Usage: postgrest [-v|--version] [-e|--example] [--dump-config | --dump-schema | --ready]
-                 [FILENAME]
+                   [--schema-cache-uri URI] [FILENAME]
 
     PostgREST / create a REST API to an existing Postgres
     database
@@ -22,6 +22,8 @@ PostgREST provides a CLI with the options listed below:
                              output structure is unstable)
     --ready                  Checks the health of PostgREST by doing a request on
                              the admin server /ready endpoint
+    --schema-cache-uri URI
+                             Try pre-loading schema cache from provided URI
     FILENAME                 Path to configuration file
 
 FILENAME
@@ -74,6 +76,15 @@ Dump Schema
 
 Dumps the schema cache in JSON format.
 
+Schema Cache Load URI
+---------------------
+
+.. code:: bash
+
+  $ postgrest --schema-cache-uri http://localhost/schema-cache.json
+
+Loads the initial schema cache from the given URI when starting PostgREST.
+
 Ready Flag
 ----------
 

src/PostgREST/AppState.hs

@@ -1,7 +1,9 @@
-{-# LANGUAGE LambdaCase      #-}
-{-# LANGUAGE NamedFieldPuns  #-}
-{-# LANGUAGE RecordWildCards #-}
-{-# LANGUAGE RecursiveDo     #-}
+{-# LANGUAGE LambdaCase       #-}
+{-# LANGUAGE NamedFieldPuns   #-}
+{-# LANGUAGE RecordWildCards  #-}
+{-# LANGUAGE RecursiveDo      #-}
+{-# LANGUAGE TupleSections    #-}
+{-# LANGUAGE TypeApplications #-}
 
 module PostgREST.AppState
   ( AppState
@@ -48,13 +50,14 @@ import Control.AutoUpdate (defaultUpdateSettings, mkAutoUpdate,
 import Control.Retry      (RetryPolicy, RetryStatus (..), capDelay,
                            exponentialBackoff, retrying,
                            rsPreviousDelay)
-import Data.IORef         (IORef, atomicWriteIORef, newIORef,
-                           readIORef)
+import Data.IORef         (IORef, atomicModifyIORef, atomicWriteIORef,
+                           newIORef, readIORef)
 import Data.Time.Clock    (UTCTime, getCurrentTime)
 
 import Control.Concurrent.STM            (TMVar, newEmptyTMVarIO,
                                           putTMVar, readTMVar,
-                                          tryReadTMVar, tryTakeTMVar)
+                                          tryPutTMVar, tryReadTMVar,
+                                          tryTakeTMVar)
 import PostgREST.Auth.JwtCache           (JwtCacheState, update)
 import PostgREST.Config                  (AppConfig (..),
                                           readAppConfig,
@@ -70,7 +73,13 @@ import PostgREST.SchemaCache             (SchemaCache (..),
                                           showSummary)
 import PostgREST.SchemaCache.Identifiers (quoteQi)
 
-import Protolude
+import qualified Data.Aeson           as JSON
+import qualified Data.ByteString.Lazy as LBS
+import qualified Network.HTTP.Client  as HC
+import           Network.URI          (URI (..), URIAuth (..),
+                                       parseURI, unEscapeString)
+import           Protolude
+import           System.IO.Error      (IOError, userError)
 
 data AppState = AppState
   -- | Database connection pool
@@ -109,16 +118,56 @@ newtype SchemaCacheStatus = SchemaCacheStatus
   { getSCStatusTMVar :: TMVar Bool
   }
 
-init :: AppConfig -> IO () -> IO AppState
-init conf@AppConfig{configLogLevel, configDbPoolSize} appKiller = do
+init :: AppConfig -> IO () -> Maybe Text -> IO AppState
+init conf@AppConfig{configLogLevel, configDbPoolSize} appKiller schemaCacheLoadUri = do
   loggerState  <- Logger.init
   metricsState <- Metrics.init configDbPoolSize
   let observer = liftA2 (>>) (Logger.observationLogger loggerState configLogLevel) (Metrics.observationMetrics metricsState)
 
   observer $ AppStartObs prettyVersion
 
   pool <- initPool conf observer
-  initWithPool pool conf loggerState metricsState observer appKiller
+  appState <- initWithPool pool conf loggerState metricsState observer appKiller
+
+  runInitialSchemaCacheLoader observer schemaCacheLoadUri appState
+
+  pure appState
+
+runInitialSchemaCacheLoader :: ObservationHandler -> Maybe Text -> AppState -> IO ()
+runInitialSchemaCacheLoader observer schemaCacheLoadUri AppState{stateSchemaCache, stateSCacheStatus=SchemaCacheStatus{getSCStatusTMVar}} = do
+  void $ forkIO $
+    traverse (fetchInitialSchemaCache observer) schemaCacheLoadUri >>= foldMap setInitialSchemaCache . join
+  where
+    setInitialSchemaCache sc =
+      whenM (atomically $ tryPutTMVar getSCStatusTMVar True) $
+        atomicModifyIORef stateSchemaCache $ (, ()) . maybe (Just sc) Just
+
+fetchInitialSchemaCache :: ObservationHandler -> Text -> IO (Maybe SchemaCache)
+fetchInitialSchemaCache observer uri = flip catches [
+  Handler (handleError @IOError),
+  Handler (handleError @HC.HttpException),
+  Handler (handleError @JSON.AesonException)
+  ] $ maybe (throwIO $ userError "Invalid schema cache dump URI") pure =<< traverse fetchURI (parseURI $ toS uri)
+  where
+    handleError :: Show e => e -> IO (Maybe a)
+    handleError = (Nothing <$) . observer . SchemaCacheDumpFailureObs uri . show
+
+    fetchURI URI{uriScheme, ..}
+      | uriScheme == "file:" = do
+          path <- fileURIPath uriAuthority uriPath
+          Just <$> (JSON.throwDecode =<< LBS.readFile path)
+      | uriScheme `elem` ["http:", "https:"] = do
+          request <- HC.parseUrlThrow $ toS uri
+          manager <- HC.newManager HC.defaultManagerSettings
+          HC.withResponse request manager $ \response ->
+            Just <$> (JSON.throwDecode . LBS.fromChunks =<< HC.brConsume (HC.responseBody response))
+      | otherwise =
+          throwIO $ userError $ "Unsupported schema cache dump URI scheme: " <> uriScheme
+
+    fileURIPath Nothing path = pure $ unEscapeString path
+    fileURIPath (Just URIAuth{uriRegName=""}) path = pure $ unEscapeString path
+    fileURIPath (Just URIAuth{uriRegName="localhost"}) path = pure $ unEscapeString path
+    fileURIPath _ _ = throwIO $ userError "Only local file URIs are supported"
 
 initWithPool :: SQL.Pool -> AppConfig -> Logger.LoggerState -> Metrics.MetricsState -> ObservationHandler -> IO () -> IO AppState
 initWithPool pool conf loggerState metricsState observer appKiller = mdo

src/PostgREST/CLI.hs

@@ -1,5 +1,6 @@
-{-# LANGUAGE NamedFieldPuns  #-}
-{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE MonadComprehensions #-}
+{-# LANGUAGE NamedFieldPuns      #-}
+{-# LANGUAGE RecordWildCards     #-}
 module PostgREST.CLI
   ( main
   , CLI (..)
@@ -47,7 +48,7 @@ runAppCommand conf@AppConfig{..} runCmd = do
   -- explicitly close the connections to PostgreSQL on shutdown.
   -- 'AppState.destroy' takes care of that.
   bracket
-    (AppState.init conf (killThread mainThreadId))
+    (AppState.init conf (killThread mainThreadId) [schemaCacheLoadUri | CmdRun (Just schemaCacheLoadUri) <- Just runCmd])
     AppState.destroy
     (\appState -> case runCmd of
       CmdDumpConfig -> do
@@ -56,7 +57,7 @@ runAppCommand conf@AppConfig{..} runCmd = do
       CmdDumpSchema -> do
         when configDbConfig $ AppState.readInDbConfig True appState
         putStrLn =<< dumpSchema appState
-      CmdRun -> App.run appState)
+      CmdRun _ -> App.run appState)
 
 -- | Dump SchemaCache schema to JSON
 dumpSchema :: AppState -> IO LBS.ByteString
@@ -85,7 +86,7 @@ data ClientCommand
   = CmdReady
 
 data RunCommand
-  = CmdRun
+  = CmdRun (Maybe Text)
   | CmdDumpConfig
   | CmdDumpSchema
 
@@ -119,25 +120,34 @@ readCLIShowHelp =
     cliParser :: O.Parser CLI
     cliParser =
       CLI
-        <$> (dumpConfigFlag <|> dumpSchemaFlag <|> readyFlag)
+        <$> (dumpConfigFlag <|> dumpSchemaFlag <|> readyFlag <|> runCommand)
         <*> O.optional configFileOption
 
+    runCommand =
+      Run . CmdRun <$> O.optional schemaCacheUriOption
+
+    schemaCacheUriOption =
+      O.strOption $
+        O.long "schema-cache-uri"
+        <> O.metavar "URI"
+        <> O.help "Try pre-loading schema cache from provided URI"
+
     configFileOption =
       O.strArgument $
         O.metavar "FILENAME"
         <> O.help "Path to configuration file"
 
     dumpConfigFlag =
-      O.flag (Run CmdRun) (Run CmdDumpConfig) $
+      O.flag' (Run CmdDumpConfig) $
         O.long "dump-config"
         <> O.help "Dump loaded configuration and exit"
 
     dumpSchemaFlag =
-      O.flag (Run CmdRun) (Run CmdDumpSchema) $
+      O.flag' (Run CmdDumpSchema) $
         O.long "dump-schema"
         <> O.help "Dump loaded schema as JSON and exit (for debugging, output structure is unstable)"
 
     readyFlag =
-      O.flag (Run CmdRun) (Client CmdReady) $
+      O.flag' (Client CmdReady) $
         O.long "ready"
         <> O.help "Checks the health of PostgREST by doing a request on the admin server /ready endpoint"

src/PostgREST/Logger.hs

@@ -160,6 +160,8 @@ observationMessages = \case
       <> " and "
       <> "db-extra-search-path=" <> T.intercalate "," extraPaths
       <> ". " <> jsonMessage usageErr
+  SchemaCacheDumpFailureObs uri err ->
+    pure $ "Failed to load schema cache dump from " <> uri <> ". " <> err
   SchemaCacheQueriedObs resultTime timings ->
     [ "Schema cache queried in " <> showMillis resultTime  <> " milliseconds " ] <>
     let showTimings qt = [ T.intercalate ", " $ (\(l, v) -> T.decodeUtf8 l <> ": " <> v <> " ms") <$> queryTimingsWLabels qt ] in

src/PostgREST/MediaType.hs

@@ -39,17 +39,17 @@ data MediaType
   | MTVndSingularJSON Bool
   -- TODO MTVndPlan should only have its options as [Text]. Its ResultAggregate should have the typed attributes.
   | MTVndPlan MediaType MTVndPlanFormat [MTVndPlanOption]
-  deriving (Eq, Show, Generic, JSON.ToJSON)
+  deriving (Eq, Show, Generic, JSON.FromJSON, JSON.ToJSON)
 instance Hashable MediaType
 
 data MTVndPlanOption
   = PlanAnalyze | PlanVerbose | PlanSettings | PlanBuffers | PlanWAL
-  deriving (Eq, Show, Generic, JSON.ToJSON)
+  deriving (Eq, Show, Generic, JSON.FromJSON, JSON.ToJSON)
 instance Hashable MTVndPlanOption
 
 data MTVndPlanFormat
   = PlanJSON | PlanText
-  deriving (Eq, Show, Generic, JSON.ToJSON)
+  deriving (Eq, Show, Generic, JSON.FromJSON, JSON.ToJSON)
 instance Hashable MTVndPlanFormat
 
 -- | Convert MediaType to a Content-Type HTTP Header

src/PostgREST/Observation.hs

@@ -34,6 +34,7 @@ data Observation
   | DBConnectedObs Text
   | SchemaCacheEmptyObs
   | SchemaCacheErrorObs (NonEmpty Text) [Text] SQL.UsageError
+  | SchemaCacheDumpFailureObs Text Text
   | SchemaCacheQueriedObs Double (Maybe QueryTimings)
   | SchemaCacheLoadedObs Double Text
   | ConnectionRetryObs Int

src/PostgREST/SchemaCache.hs

@@ -28,7 +28,7 @@ module PostgREST.SchemaCache
   , queryTimingsWLabels
   ) where
 
-import           Data.Aeson ((.=))
+import           Data.Aeson ((.:), (.=))
 import qualified Data.Aeson as JSON
 
 import qualified Data.ByteString.Char8      as BS
@@ -101,6 +101,17 @@ instance JSON.ToJSON SchemaCache where
     , "dbTimezones"       .= JSON.toJSON tzs
     ]
 
+instance JSON.FromJSON SchemaCache where
+  parseJSON = JSON.withObject "SchemaCache" $ \o -> do
+    tabs <- o .: "dbTables"
+    SchemaCache tabs
+      <$> o .: "dbRelationships"
+      <*> o .: "dbRoutines"
+      <*> o .: "dbRepresentations"
+      <*> o .: "dbMediaHandlers"
+      <*> o .: "dbTimezones"
+      <*> pure (tablesFuzzyIndex tabs)
+
 showSummary :: SchemaCache -> Text
 showSummary (SchemaCache tbls rels routs reps mediaHdlrs tzs _) =
   T.intercalate ", "
@@ -191,14 +202,20 @@ querySchemaCache conf@AppConfig{..} = do
     , dbTablesFuzzyIndex =
         -- Only build fuzzy index for schemas with a reasonable number of tables
         -- Fuzzy.FuzzySet is memory heavy we just don't use it for large schemas
-        Fuzzy.fromList <$> HM.filter ((< maxDbTablesForFuzzySearch) . length) (HM.fromListWith (<>) ((qiSchema &&& pure . qiName) <$> HM.keys tabsWViewsPks))
+        tablesFuzzyIndex tabsWViewsPks
     }, qsTime)
   where
     schemas = toList configDbSchemas
     isLogDebug = configLogLevel == LogDebug
     sqlTimedStmt = sqlTimedStatement isLogDebug
     sleepCall = SQL.Statement "select pg_sleep($1 / 1000.0)" (param HE.int4) HD.noResult True
 
+tablesFuzzyIndex :: TablesMap -> TablesFuzzyIndex
+tablesFuzzyIndex tabs =
+  Fuzzy.fromList <$>
+    HM.filter ((< maxDbTablesForFuzzySearch) . length)
+      (HM.fromListWith (<>) ((qiSchema &&& pure . qiName) <$> HM.keys tabs))
+
 -- | overrides detected relationships with the computed relationships and gets the RelationshipsMap
 getOverrideRelationshipsMap :: [Relationship] -> [Relationship] -> RelationshipsMap
 getOverrideRelationshipsMap rels cRels =

src/PostgREST/SchemaCache/Identifiers.hs

@@ -20,7 +20,7 @@ import qualified Data.Text  as T
 import Protolude
 
 data RelIdentifier = RelId QualifiedIdentifier | RelAnyElement
-  deriving (Eq, Ord, Generic, JSON.ToJSON, JSON.ToJSONKey, Show)
+  deriving (Eq, Ord, Generic, JSON.FromJSON, JSON.FromJSONKey, JSON.ToJSON, JSON.ToJSONKey, Show)
 instance Hashable RelIdentifier
 
 -- | Represents a pg identifier with a prepended schema name "schema.table".
@@ -30,7 +30,7 @@ data QualifiedIdentifier = QualifiedIdentifier
   { qiSchema :: Schema
   , qiName   :: TableName
   }
-  deriving (Eq, Show, Ord, Generic, JSON.ToJSON, JSON.ToJSONKey)
+  deriving (Eq, Show, Ord, Generic, JSON.FromJSON, JSON.FromJSONKey, JSON.ToJSON, JSON.ToJSONKey)
 
 instance Hashable QualifiedIdentifier
 

src/PostgREST/SchemaCache/Relationship.hs

@@ -35,7 +35,7 @@ data Relationship = Relationship
   , relToOne        :: Bool
   , relIsSelf       :: Bool
   }
-  deriving (Eq, Show, Ord, Generic, JSON.ToJSON)
+  deriving (Eq, Show, Ord, Generic, JSON.FromJSON, JSON.ToJSON)
 
 -- | The relationship cardinality
 -- | https://en.wikipedia.org/wiki/Cardinality_(data_modeling)
@@ -48,7 +48,7 @@ data Cardinality
   -- ^ one-to-one, this is a refinement over M2O, operating on it is pretty much the same as M2O when isParent == False
   | M2M Junction
   -- ^ many-to-many
-  deriving (Eq, Show, Ord, Generic, JSON.ToJSON)
+  deriving (Eq, Show, Ord, Generic, JSON.FromJSON, JSON.ToJSON)
 
 type FKConstraint = Text
 
@@ -60,7 +60,7 @@ data Junction = Junction
   , junColsSource  :: [(FieldName, FieldName)]
   , junColsTarget  :: [(FieldName, FieldName)]
   }
-  deriving (Eq, Show, Ord, Generic, JSON.ToJSON)
+  deriving (Eq, Show, Ord, Generic, JSON.FromJSON, JSON.ToJSON)
 
 -- | Key based on the source table and the foreign table schema
 type RelationshipsMap = HM.HashMap (QualifiedIdentifier, Schema)  [Relationship]