Skip to content

Commit ecd794f

Browse files
mkleczekmkleczek
committed
feat(cli): load schema cache dump at startup
1 parent 93c2cbe commit ecd794f

9 files changed

Lines changed: 96 additions & 32 deletions

File tree

‎src/PostgREST/AppState.hs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ module PostgREST.AppState
1414
, getTime
1515
, getJwtCacheState
1616
, init
17+
, initWithSchemaCache
1718
, initWithPool
1819
, putConfig -- For tests TODO refactoring
1920
, putSchemaCache
@@ -120,6 +121,13 @@ init conf@AppConfig{configLogLevel, configDbPoolSize} = do
120121
pool <- initPool conf observer
121122
initWithPool pool conf loggerState metricsState observer
122123

124+
initWithSchemaCache :: AppConfig -> SchemaCache -> IO AppState
125+
initWithSchemaCache conf sCache = do
126+
appState <- init conf
127+
putSchemaCache appState $ Just sCache
128+
markSchemaCacheLoaded appState
129+
pure appState
130+
123131
initWithPool :: SQL.Pool -> AppConfig -> Logger.LoggerState -> Metrics.MetricsState -> ObservationHandler -> IO AppState
124132
initWithPool pool conf loggerState metricsState observer = mdo
125133

‎src/PostgREST/CLI.hs‎

Lines changed: 38 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,12 @@ import qualified Data.ByteString.Char8 as BS
1212
import qualified Data.ByteString.Lazy as LBS
1313
import qualified Hasql.Transaction.Sessions as SQL
1414
import qualified Options.Applicative as O
15+
import System.IO.Error (isDoesNotExistError)
1516

1617
import PostgREST.AppState (AppState)
1718
import PostgREST.Config (AppConfig (..))
1819
import PostgREST.Observation (Observation (..))
19-
import PostgREST.SchemaCache (querySchemaCache)
20+
import PostgREST.SchemaCache (SchemaCache, querySchemaCache)
2021
import PostgREST.Version (prettyVersion)
2122

2223
import qualified PostgREST.App as App
@@ -28,25 +29,32 @@ import Protolude
2829

2930

3031
main :: CLI -> IO ()
31-
main CLI{cliCommand, cliPath} = do
32+
main CLI{cliCommand, cliPath, cliSchemaCachePath} = do
3233
conf <-
3334
either panic identity <$> Config.readAppConfig mempty cliPath Nothing mempty mempty
3435
case cliCommand of
3536
Client adminCmd -> runClientCommand conf adminCmd
36-
Run runCmd -> runAppCommand conf runCmd
37+
Run runCmd -> runAppCommand conf cliSchemaCachePath runCmd
3738

3839
-- | Run command using http-client to communicate with an already running postgrest
3940
runClientCommand :: AppConfig -> ClientCommand -> IO ()
4041
runClientCommand conf CmdReady = Client.ready conf
4142

4243
-- | Run postgrest with command
43-
runAppCommand :: AppConfig -> RunCommand -> IO ()
44-
runAppCommand conf@AppConfig{..} runCmd = do
44+
runAppCommand :: AppConfig -> Maybe FilePath -> RunCommand -> IO ()
45+
runAppCommand conf@AppConfig{..} schemaCachePath runCmd = do
46+
initAppState <- case runCmd of
47+
CmdRun -> do
48+
initialSchemaCache <- join <$> traverse readSchemaCacheDump schemaCachePath
49+
pure $ maybe (AppState.init conf) (AppState.initWithSchemaCache conf) initialSchemaCache
50+
_ ->
51+
pure $ AppState.init conf
52+
4553
-- Per https://github.com/PostgREST/postgrest/issues/268, we want to
4654
-- explicitly close the connections to PostgreSQL on shutdown.
4755
-- 'AppState.destroy' takes care of that.
4856
bracket
49-
(AppState.init conf)
57+
initAppState
5058
AppState.destroy
5159
(\appState -> case runCmd of
5260
CmdDumpConfig -> do
@@ -55,7 +63,8 @@ runAppCommand conf@AppConfig{..} runCmd = do
5563
CmdDumpSchema -> do
5664
when configDbConfig $ AppState.readInDbConfig True appState
5765
putStrLn =<< dumpSchema appState
58-
CmdRun -> App.run appState)
66+
CmdRun ->
67+
App.run appState)
5968

6069
-- | Dump SchemaCache schema to JSON
6170
dumpSchema :: AppState -> IO LBS.ByteString
@@ -70,10 +79,23 @@ dumpSchema appState = do
7079
exitFailure
7180
Right (sCache, _) -> return $ JSON.encode sCache
7281

82+
readSchemaCacheDump :: FilePath -> IO (Maybe SchemaCache)
83+
readSchemaCacheDump path =
84+
fmap decodeSchemaCache <$> catch (Just <$> LBS.readFile path) handleReadError
85+
where
86+
decodeSchemaCache =
87+
fromRight (panic $ "Error loading schema cache dump from " <> toS path) . JSON.eitherDecode
88+
89+
handleReadError :: IOException -> IO (Maybe LBS.ByteString)
90+
handleReadError e
91+
| isDoesNotExistError e = pure Nothing
92+
| otherwise = throwIO e
93+
7394
-- | Command line interface options
7495
data CLI = CLI
75-
{ cliCommand :: Command
76-
, cliPath :: Maybe FilePath
96+
{ cliCommand :: Command
97+
, cliPath :: Maybe FilePath
98+
, cliSchemaCachePath :: Maybe FilePath
7799
}
78100

79101
data Command
@@ -120,12 +142,19 @@ readCLIShowHelp =
120142
CLI
121143
<$> (dumpConfigFlag <|> dumpSchemaFlag <|> readyFlag)
122144
<*> O.optional configFileOption
145+
<*> O.optional loadSchemaCacheOption
123146

124147
configFileOption =
125148
O.strArgument $
126149
O.metavar "FILENAME"
127150
<> O.help "Path to configuration file"
128151

152+
loadSchemaCacheOption =
153+
O.strOption $
154+
O.long "load-schema-cache"
155+
<> O.metavar "FILENAME"
156+
<> O.help "Load schema cache from JSON dump at startup if the file exists"
157+
129158
dumpConfigFlag =
130159
O.flag (Run CmdRun) (Run CmdDumpConfig) $
131160
O.long "dump-config"

‎src/PostgREST/MediaType.hs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -39,17 +39,17 @@ data MediaType
3939
| MTVndSingularJSON Bool
4040
-- TODO MTVndPlan should only have its options as [Text]. Its ResultAggregate should have the typed attributes.
4141
| MTVndPlan MediaType MTVndPlanFormat [MTVndPlanOption]
42-
deriving (Eq, Show, Generic, JSON.ToJSON)
42+
deriving (Eq, Show, Generic, JSON.FromJSON, JSON.ToJSON)
4343
instance Hashable MediaType
4444

4545
data MTVndPlanOption
4646
= PlanAnalyze | PlanVerbose | PlanSettings | PlanBuffers | PlanWAL
47-
deriving (Eq, Show, Generic, JSON.ToJSON)
47+
deriving (Eq, Show, Generic, JSON.FromJSON, JSON.ToJSON)
4848
instance Hashable MTVndPlanOption
4949

5050
data MTVndPlanFormat
5151
= PlanJSON | PlanText
52-
deriving (Eq, Show, Generic, JSON.ToJSON)
52+
deriving (Eq, Show, Generic, JSON.FromJSON, JSON.ToJSON)
5353
instance Hashable MTVndPlanFormat
5454

5555
-- | Convert MediaType to a Content-Type HTTP Header

‎src/PostgREST/SchemaCache.hs‎

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ module PostgREST.SchemaCache
2828
, queryTimingsWLabels
2929
) where
3030

31-
import Data.Aeson ((.=))
31+
import Data.Aeson ((.:), (.=))
3232
import qualified Data.Aeson as JSON
3333

3434
import qualified Data.ByteString.Char8 as BS
@@ -101,6 +101,17 @@ instance JSON.ToJSON SchemaCache where
101101
, "dbTimezones" .= JSON.toJSON tzs
102102
]
103103

104+
instance JSON.FromJSON SchemaCache where
105+
parseJSON = JSON.withObject "SchemaCache" $ \o -> do
106+
tabs <- o .: "dbTables"
107+
SchemaCache tabs
108+
<$> o .: "dbRelationships"
109+
<*> o .: "dbRoutines"
110+
<*> o .: "dbRepresentations"
111+
<*> o .: "dbMediaHandlers"
112+
<*> o .: "dbTimezones"
113+
<*> pure (tablesFuzzyIndex tabs)
114+
104115
showSummary :: SchemaCache -> Text
105116
showSummary (SchemaCache tbls rels routs reps mediaHdlrs tzs _) =
106117
T.intercalate ", "
@@ -188,17 +199,22 @@ querySchemaCache conf@AppConfig{..} = do
188199
, dbMediaHandlers = HM.union mHdlers initialMediaHandlers -- the custom handlers will override the initial ones
189200
, dbTimezones = tzones
190201

191-
, dbTablesFuzzyIndex =
192-
-- Only build fuzzy index for schemas with a reasonable number of tables
193-
-- Fuzzy.FuzzySet is memory heavy we just don't use it for large schemas
194-
Fuzzy.fromList <$> HM.filter ((< maxDbTablesForFuzzySearch) . length) (HM.fromListWith (<>) ((qiSchema &&& pure . qiName) <$> HM.keys tabsWViewsPks))
202+
, dbTablesFuzzyIndex = tablesFuzzyIndex tabsWViewsPks
195203
}, qsTime)
196204
where
197205
schemas = toList configDbSchemas
198206
isLogDebug = configLogLevel == LogDebug
199207
sqlTimedStmt = sqlTimedStatement isLogDebug
200208
sleepCall = SQL.Statement "select pg_sleep($1 / 1000.0)" (param HE.int4) HD.noResult True
201209

210+
tablesFuzzyIndex :: TablesMap -> TablesFuzzyIndex
211+
tablesFuzzyIndex tabs =
212+
-- Only build fuzzy index for schemas with a reasonable number of tables.
213+
-- Fuzzy.FuzzySet is memory heavy so we do not use it for large schemas.
214+
Fuzzy.fromList <$>
215+
HM.filter ((< maxDbTablesForFuzzySearch) . length)
216+
(HM.fromListWith (<>) ((qiSchema &&& pure . qiName) <$> HM.keys tabs))
217+
202218
-- | overrides detected relationships with the computed relationships and gets the RelationshipsMap
203219
getOverrideRelationshipsMap :: [Relationship] -> [Relationship] -> RelationshipsMap
204220
getOverrideRelationshipsMap rels cRels =

‎src/PostgREST/SchemaCache/Identifiers.hs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ import qualified Data.Text as T
2020
import Protolude
2121

2222
data RelIdentifier = RelId QualifiedIdentifier | RelAnyElement
23-
deriving (Eq, Ord, Generic, JSON.ToJSON, JSON.ToJSONKey, Show)
23+
deriving (Eq, Ord, Generic, JSON.FromJSON, JSON.FromJSONKey, JSON.ToJSON, JSON.ToJSONKey, Show)
2424
instance Hashable RelIdentifier
2525

2626
-- | Represents a pg identifier with a prepended schema name "schema.table".
@@ -30,7 +30,7 @@ data QualifiedIdentifier = QualifiedIdentifier
3030
{ qiSchema :: Schema
3131
, qiName :: TableName
3232
}
33-
deriving (Eq, Show, Ord, Generic, JSON.ToJSON, JSON.ToJSONKey)
33+
deriving (Eq, Show, Ord, Generic, JSON.FromJSON, JSON.FromJSONKey, JSON.ToJSON, JSON.ToJSONKey)
3434

3535
instance Hashable QualifiedIdentifier
3636

‎src/PostgREST/SchemaCache/Relationship.hs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ data Relationship = Relationship
3535
, relToOne :: Bool
3636
, relIsSelf :: Bool
3737
}
38-
deriving (Eq, Show, Ord, Generic, JSON.ToJSON)
38+
deriving (Eq, Show, Ord, Generic, JSON.FromJSON, JSON.ToJSON)
3939

4040
-- | The relationship cardinality
4141
-- | https://en.wikipedia.org/wiki/Cardinality_(data_modeling)
@@ -48,7 +48,7 @@ data Cardinality
4848
-- ^ one-to-one, this is a refinement over M2O, operating on it is pretty much the same as M2O when isParent == False
4949
| M2M Junction
5050
-- ^ many-to-many
51-
deriving (Eq, Show, Ord, Generic, JSON.ToJSON)
51+
deriving (Eq, Show, Ord, Generic, JSON.FromJSON, JSON.ToJSON)
5252

5353
type FKConstraint = Text
5454

@@ -60,7 +60,7 @@ data Junction = Junction
6060
, junColsSource :: [(FieldName, FieldName)]
6161
, junColsTarget :: [(FieldName, FieldName)]
6262
}
63-
deriving (Eq, Show, Ord, Generic, JSON.ToJSON)
63+
deriving (Eq, Show, Ord, Generic, JSON.FromJSON, JSON.ToJSON)
6464

6565
-- | Key based on the source table and the foreign table schema
6666
type RelationshipsMap = HM.HashMap (QualifiedIdentifier, Schema) [Relationship]

‎src/PostgREST/SchemaCache/Routine.hs‎

Lines changed: 19 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ module PostgREST.SchemaCache.Routine
2020
, MediaHandler(..)
2121
) where
2222

23-
import Data.Aeson ((.=))
23+
import Data.Aeson ((.:), (.=))
2424
import qualified Data.Aeson as JSON
2525
import qualified Data.HashMap.Strict as HM
2626
import qualified Hasql.Transaction.Sessions as SQL
@@ -36,18 +36,18 @@ import Protolude
3636
data PgType
3737
= Scalar QualifiedIdentifier
3838
| Composite QualifiedIdentifier Bool -- True if the composite is a domain alias(used to work around a bug in pg 11 and 12, see QueryBuilder.hs)
39-
deriving (Eq, Show, Ord, Generic, JSON.ToJSON)
39+
deriving (Eq, Show, Ord, Generic, JSON.FromJSON, JSON.ToJSON)
4040

4141
data RetType
4242
= Single PgType
4343
| SetOf PgType
44-
deriving (Eq, Show, Ord, Generic, JSON.ToJSON)
44+
deriving (Eq, Show, Ord, Generic, JSON.FromJSON, JSON.ToJSON)
4545

4646
data FuncVolatility
4747
= Volatile
4848
| Stable
4949
| Immutable
50-
deriving (Eq, Show, Ord, Generic, JSON.ToJSON)
50+
deriving (Eq, Show, Ord, Generic, JSON.FromJSON, JSON.ToJSON)
5151

5252
type FuncSettings = [(Text,Text)]
5353

@@ -77,14 +77,27 @@ instance JSON.ToJSON Routine where
7777
, "pdFuncSettings" .= JSON.toJSON sets
7878
]
7979

80+
instance JSON.FromJSON Routine where
81+
parseJSON = JSON.withObject "Routine" $ \o ->
82+
Function
83+
<$> o .: "pdSchema"
84+
<*> o .: "pdName"
85+
<*> o .: "pdDescription"
86+
<*> o .: "pdParams"
87+
<*> o .: "pdReturnType"
88+
<*> o .: "pdVolatility"
89+
<*> o .: "pdHasVariadic"
90+
<*> pure Nothing
91+
<*> o .: "pdFuncSettings"
92+
8093
data RoutineParam = RoutineParam
8194
{ ppName :: Text
8295
, ppType :: Text
8396
, ppTypeMaxLength :: Text
8497
, ppReq :: Bool
8598
, ppVar :: Bool
8699
}
87-
deriving (Eq, Show, Ord, Generic, JSON.ToJSON)
100+
deriving (Eq, Show, Ord, Generic, JSON.FromJSON, JSON.ToJSON)
88101

89102
-- Order by least number of params in the case of overloaded functions
90103
instance Ord Routine where
@@ -109,7 +122,7 @@ data MediaHandler
109122
-- custom
110123
| CustomFunc QualifiedIdentifier RelIdentifier
111124
| NoAgg
112-
deriving (Eq, Show, Generic, JSON.ToJSON)
125+
deriving (Eq, Show, Generic, JSON.FromJSON, JSON.ToJSON)
113126

114127
funcReturnsSingle :: Routine -> Bool
115128
funcReturnsSingle proc = case proc of

‎src/PostgREST/SchemaCache/Table.hs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ data Table = Table
3434
, tablePKCols :: [FieldName]
3535
, tableColumns :: ColumnMap
3636
}
37-
deriving (Show, Generic, JSON.ToJSON)
37+
deriving (Show, Generic, JSON.FromJSON, JSON.ToJSON)
3838

3939
tableColumnsList :: Table -> [Column]
4040
tableColumnsList = HMI.elems . tableColumns
@@ -52,7 +52,7 @@ data Column = Column
5252
, colDefault :: Maybe Text
5353
, colEnum :: [Text]
5454
}
55-
deriving (Eq, Show, Ord, Generic, JSON.ToJSON)
55+
deriving (Eq, Show, Ord, Generic, JSON.FromJSON, JSON.ToJSON)
5656

5757
type TablesMap = HM.HashMap QualifiedIdentifier Table
5858
type ColumnMap = HMI.InsOrdHashMap FieldName Column

‎test/io/test_cli.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -259,7 +259,6 @@ def test_schema_cache_snapshot(baseenv, key, snapshot_yaml):
259259
assert formatted == snapshot_yaml
260260

261261

262-
@pytest.mark.xfail(reason="--load-schema-cache is not implemented", strict=True)
263262
def test_load_schema_cache_dump_at_startup(tmp_path, defaultenv, metapostgrest):
264263
"A valid schema cache dump should be loaded at startup."
265264

@@ -289,7 +288,6 @@ def test_load_schema_cache_dump_at_startup(tmp_path, defaultenv, metapostgrest):
289288
reset_statement_timeout(metapostgrest, role)
290289

291290

292-
@pytest.mark.xfail(reason="--load-schema-cache is not implemented", strict=True)
293291
def test_load_schema_cache_dump_missing_file_falls_back_to_database(
294292
tmp_path, defaultenv
295293
):

0 commit comments

Comments
 (0)