Skip to content

Commit 32a2995

Browse files
mkleczekmkleczek
committed
refactor: add internal-schema-cache-lock-id
1 parent 1a6ba20 commit 32a2995

11 files changed

Lines changed: 295 additions & 119 deletions

File tree

‎nix/tools/tests.nix‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@ let
7878
ioTestPython =
7979
python3.withPackages (ps: [
8080
ps.pyjwt
81+
ps.psycopg
8182
ps.pytest
8283
ps.pytest-xdist
8384
ps.pyyaml

‎postgrest.cabal‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ library
6565
PostgREST.SchemaCache.Relationship
6666
PostgREST.SchemaCache.Representations
6767
PostgREST.SchemaCache.Table
68+
PostgREST.SqlTransaction
6869
PostgREST.Error
6970
PostgREST.Error.Types
7071
PostgREST.Listener

‎src/PostgREST/Config.hs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,7 @@ data AppConfig = AppConfig
127127
, configRoleSettings :: RoleSettings
128128
, configRoleIsoLvl :: RoleIsolationLvl
129129
, configInternalSCQuerySleep :: Maybe Int32
130+
, configInternalSCLockId :: Maybe Int32
130131
}
131132

132133
data LogLevel = LogCrit | LogError | LogWarn | LogInfo | LogDebug
@@ -326,6 +327,7 @@ parser optPath env dbSettings roleSettings roleIsolationLvl =
326327
<*> pure roleSettings
327328
<*> pure roleIsolationLvl
328329
<*> optInt "internal-schema-cache-query-sleep"
330+
<*> optInt "internal-schema-cache-lock-id"
329331
where
330332
parseErrorVerbosity :: C.Key -> C.Parser C.Config Verbosity
331333
parseErrorVerbosity k =

‎src/PostgREST/SchemaCache.hs‎

Lines changed: 87 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -8,15 +8,13 @@ The schema cache is necessary for resource embedding, foreign keys are used for
88
99
These queries are executed once at startup or when PostgREST is reloaded.
1010
-}
11-
{-# LANGUAGE DeriveAnyClass #-}
12-
{-# LANGUAGE DeriveGeneric #-}
13-
{-# LANGUAGE FlexibleContexts #-}
14-
{-# LANGUAGE MultiParamTypeClasses #-}
15-
{-# LANGUAGE NamedFieldPuns #-}
16-
{-# LANGUAGE QuasiQuotes #-}
17-
{-# LANGUAGE RecordWildCards #-}
18-
{-# LANGUAGE ScopedTypeVariables #-}
19-
{-# LANGUAGE TypeSynonymInstances #-}
11+
{-# LANGUAGE DeriveAnyClass #-}
12+
{-# LANGUAGE DeriveGeneric #-}
13+
{-# LANGUAGE FlexibleContexts #-}
14+
{-# LANGUAGE NamedFieldPuns #-}
15+
{-# LANGUAGE QuasiQuotes #-}
16+
{-# LANGUAGE RecordWildCards #-}
17+
{-# LANGUAGE ScopedTypeVariables #-}
2018

2119
module PostgREST.SchemaCache
2220
( SchemaCache(..)
@@ -44,31 +42,36 @@ import qualified Hasql.Transaction as SQL
4442
import Data.Functor.Contravariant ((>$<))
4543
import NeatInterpolation (trimming)
4644

47-
import PostgREST.Config (AppConfig (..),
48-
LogLevel (..))
49-
import PostgREST.Config.Database (TimezoneNames,
50-
toIsolationLevel)
51-
import PostgREST.SchemaCache.Identifiers (FieldName,
52-
QualifiedIdentifier (..),
53-
RelIdentifier (..),
54-
Schema, escapeIdent,
55-
isAnyElement)
56-
import PostgREST.SchemaCache.Relationship (Cardinality (..),
57-
Junction (..),
58-
Relationship (..),
59-
RelationshipsMap)
60-
import PostgREST.SchemaCache.Representations (DataRepresentation (..),
61-
RepresentationsMap)
62-
import PostgREST.SchemaCache.Routine (FuncVolatility (..),
63-
MediaHandler (..),
64-
MediaHandlerMap,
65-
PgType (..),
66-
RetType (..),
67-
Routine (..),
68-
RoutineMap,
69-
RoutineParam (..))
70-
import PostgREST.SchemaCache.Table (Column (..), ColumnMap,
71-
Table (..), TablesMap)
45+
import PostgREST.Config (AppConfig (..),
46+
LogLevel (..))
47+
import PostgREST.Config.Database (TimezoneNames,
48+
toIsolationLevel)
49+
import PostgREST.SchemaCache.Identifiers (FieldName,
50+
QualifiedIdentifier (..),
51+
RelIdentifier (..),
52+
Schema,
53+
escapeIdent,
54+
isAnyElement)
55+
import PostgREST.SchemaCache.Relationship (Cardinality (..),
56+
Junction (..),
57+
Relationship (..),
58+
RelationshipsMap)
59+
import PostgREST.SchemaCache.Representations (DataRepresentation (..),
60+
RepresentationsMap)
61+
import PostgREST.SchemaCache.Routine (FuncVolatility (..),
62+
MediaHandler (..),
63+
MediaHandlerMap,
64+
PgType (..),
65+
RetType (..),
66+
Routine (..),
67+
RoutineMap,
68+
RoutineParam (..))
69+
import PostgREST.SchemaCache.Table (Column (..),
70+
ColumnMap,
71+
Table (..),
72+
TablesMap)
73+
import PostgREST.SqlTransaction (SqlTransaction)
74+
import qualified PostgREST.SqlTransaction as SQLT
7275

7376
import qualified PostgREST.MediaType as MediaType
7477

@@ -157,46 +160,55 @@ maxDbTablesForFuzzySearch = 500
157160
querySchemaCache :: AppConfig -> SQL.Transaction SchemaCache
158161
querySchemaCache conf@AppConfig{..} = do
159162
SQL.sql "set local schema ''" -- This voids the search path. The following queries need this for getting the fully qualified name(schema.name) of every db object
160-
tabs <- sqlTimedStmt gucTbls conf allTables
161-
keyDeps <- sqlTimedStmt gucKDeps conf allViewsKeyDependencies
162-
m2oRels <- sqlTimedStmt gucRels mempty allM2OandO2ORels
163-
funcs <- sqlTimedStmt gucFuncs conf allFunctions
164-
cRels <- sqlTimedStmt gucCRels mempty allComputedRels
165-
reps <- sqlTimedStmt gucDReps conf dataRepresentations
166-
mHdlers <- sqlTimedStmt gucMHdrs conf mediaHandlers
167-
tzones <- if configDbTimezoneEnabled
168-
then sqlTimedStmt gucTzones mempty timezones
169-
else pure S.empty
170-
_ <-
171-
let sleepCall = SQL.Statement "select pg_sleep($1 / 1000.0)" (param HE.int4) HD.noResult True in
172-
for_ configInternalSCQuerySleep (`SQL.statement` sleepCall) -- only used for testing
173-
174-
qsTime <-
175-
if isLogDebug
176-
then Just <$> SQL.statement mempty (extractTimings configDbTimezoneEnabled)
177-
else pure Nothing
178-
179-
let tabsWViewsPks = addViewPrimaryKeys tabs keyDeps
180-
rels = addInverseRels $ addM2MRels tabsWViewsPks $ addViewM2OAndO2ORels keyDeps m2oRels
181-
182-
return $ removeInternal schemas $ SchemaCache {
183-
dbTables = tabsWViewsPks
184-
, dbRelationships = getOverrideRelationshipsMap rels cRels
185-
, dbRoutines = funcs
186-
, dbRepresentations = reps
187-
, dbMediaHandlers = HM.union mHdlers initialMediaHandlers -- the custom handlers will override the initial ones
188-
, dbTimezones = tzones
189-
190-
, dbTablesFuzzyIndex =
191-
-- Only build fuzzy index for schemas with a reasonable number of tables
192-
-- Fuzzy.FuzzySet is memory heavy we just don't use it for large schemas
193-
Fuzzy.fromList <$> HM.filter ((< maxDbTablesForFuzzySearch) . length) (HM.fromListWith (<>) ((qiSchema &&& pure . qiName) <$> HM.keys tabsWViewsPks))
194-
, dbQueryTimings = qsTime
195-
}
163+
-- if configInternalSCLockId is set run queries step-by-step waiting for lock release before each
164+
maybe
165+
querySchemaCacheSqlTransaction
166+
(SQLT.runSteppedTransaction querySchemaCacheSqlTransaction)
167+
configInternalSCLockId
168+
-- TODO remove configInternalSCQuerySleep once all tests are migrated to stepped execution
169+
<* for_ configInternalSCQuerySleep (`SQL.statement` sleepCall) -- only used for testing
196170
where
171+
querySchemaCacheSqlTransaction :: forall m. SqlTransaction m => m SchemaCache
172+
querySchemaCacheSqlTransaction = do
173+
tabs <- sqlTimedStmt gucTbls conf allTables
174+
keyDeps <- sqlTimedStmt gucKDeps conf allViewsKeyDependencies
175+
m2oRels <- sqlTimedStmt gucRels mempty allM2OandO2ORels
176+
funcs <- sqlTimedStmt gucFuncs conf allFunctions
177+
cRels <- sqlTimedStmt gucCRels mempty allComputedRels
178+
reps <- sqlTimedStmt gucDReps conf dataRepresentations
179+
mHdlers <- sqlTimedStmt gucMHdrs conf mediaHandlers
180+
tzones <- if configDbTimezoneEnabled
181+
then sqlTimedStmt gucTzones mempty timezones
182+
else pure S.empty
183+
184+
qsTime <-
185+
if isLogDebug
186+
then Just <$> SQLT.statement mempty (extractTimings configDbTimezoneEnabled)
187+
else pure Nothing
188+
189+
let tabsWViewsPks = addViewPrimaryKeys tabs keyDeps
190+
rels = addInverseRels $ addM2MRels tabsWViewsPks $ addViewM2OAndO2ORels keyDeps m2oRels
191+
192+
return $ removeInternal schemas $ SchemaCache {
193+
dbTables = tabsWViewsPks
194+
, dbRelationships = getOverrideRelationshipsMap rels cRels
195+
, dbRoutines = funcs
196+
, dbRepresentations = reps
197+
, dbMediaHandlers = HM.union mHdlers initialMediaHandlers -- the custom handlers will override the initial ones
198+
, dbTimezones = tzones
199+
200+
, dbTablesFuzzyIndex =
201+
-- Only build fuzzy index for schemas with a reasonable number of tables
202+
-- Fuzzy.FuzzySet is memory heavy we just don't use it for large schemas
203+
Fuzzy.fromList <$> HM.filter ((< maxDbTablesForFuzzySearch) . length) (HM.fromListWith (<>) ((qiSchema &&& pure . qiName) <$> HM.keys tabsWViewsPks))
204+
, dbQueryTimings = qsTime
205+
}
206+
197207
schemas = toList configDbSchemas
198208
isLogDebug = configLogLevel == LogDebug
209+
sqlTimedStmt :: SqlTransaction m => ByteString -> a -> SQL.Statement a b -> m b
199210
sqlTimedStmt = sqlTimedStatement isLogDebug
211+
sleepCall = SQL.Statement "select pg_sleep($1 / 1000.0)" (param HE.int4) HD.noResult True
200212

201213
-- | overrides detected relationships with the computed relationships and gets the RelationshipsMap
202214
getOverrideRelationshipsMap :: [Relationship] -> [Relationship] -> RelationshipsMap
@@ -1167,22 +1179,22 @@ arrayColumn = column . HD.listArray . HD.nonNullable
11671179
-
11681180
- We can do this for several statements inside the transaction. The timings are later captured at the end of the transaction with extractTimings.
11691181
-}
1170-
sqlTimedStatement :: Bool -> ByteString -> a -> SQL.Statement a b -> SQL.Transaction b
1182+
sqlTimedStatement :: SqlTransaction m => Bool -> ByteString -> a -> SQL.Statement a b -> m b
11711183
sqlTimedStatement isLogDebug guc params stmt =
11721184
if isLogDebug then
1173-
SQL.sql sFrag >> SQL.statement params stmt <* SQL.sql eFrag
1185+
SQLT.sql sFrag >> SQLT.statement params stmt <* SQLT.sql eFrag
11741186
else
1175-
SQL.statement params stmt
1187+
SQLT.statement params stmt
11761188
where
11771189
sFrag = "select set_config('pgrst." <> guc <> "', clock_timestamp()::text, true)"
11781190
eFrag = "select set_config('pgrst." <> guc <> "', (clock_timestamp() - current_setting('pgrst." <> guc <> "', false)::timestamptz)::text, true)"
11791191

11801192
-- Extract all the generated timings (see sqlTimedStatement) converting the value to milliseconds.
11811193
extractTimings :: Bool -> SQL.Statement () QueryTimings
1182-
extractTimings hasTimezones = SQL.Statement sql HE.noParams decodeThem True
1194+
extractTimings hasTimezones = SQL.Statement query HE.noParams decodeThem True
11831195
where
11841196
qFrag setting = "extract('milliseconds' from current_setting('pgrst." <> setting <> "', false)::interval)::text"
1185-
sql = "SELECT " <> BS.intercalate ","
1197+
query = "SELECT " <> BS.intercalate ","
11861198
[ qFrag gucTbls, qFrag gucKDeps, qFrag gucRels
11871199
, qFrag gucFuncs, qFrag gucCRels, qFrag gucDReps
11881200
, qFrag gucMHdrs, if hasTimezones then qFrag gucTzones else "'0.0'"

‎src/PostgREST/SqlTransaction.hs‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
{-# LANGUAGE DerivingStrategies #-}
2+
{-# LANGUAGE GeneralisedNewtypeDeriving #-}
3+
{-# LANGUAGE KindSignatures #-}
4+
{-# LANGUAGE ScopedTypeVariables #-}
5+
{-# LANGUAGE TypeApplications #-}
6+
7+
module PostgREST.SqlTransaction
8+
( SqlTransaction(..)
9+
, WithLock
10+
, runSteppedTransaction
11+
) where
12+
13+
import Protolude
14+
15+
import Data.Functor.Contravariant ((>$<))
16+
import qualified Hasql.Decoders as HD
17+
import qualified Hasql.Encoders as HE
18+
import qualified Hasql.Statement as SQL
19+
import qualified Hasql.Transaction as SQL
20+
21+
class Monad m => SqlTransaction (m :: Type -> Type) where
22+
sql :: ByteString -> m ()
23+
statement :: a -> SQL.Statement a b -> m b
24+
25+
instance SqlTransaction SQL.Transaction where
26+
sql = SQL.sql
27+
statement = SQL.statement
28+
29+
newtype WithLock m a = WithLock (ReaderT Int32 (StateT Int16 m) a)
30+
deriving newtype (Functor, Applicative, Monad)
31+
32+
instance SqlTransaction m => SqlTransaction (WithLock m) where
33+
sql = lockNext . sql @m
34+
statement = (lockNext .) . statement @m
35+
36+
runSteppedTransaction :: SqlTransaction m => WithLock m a -> Int32 -> m a
37+
runSteppedTransaction (WithLock m) lockId =
38+
evalStateT (runReaderT m lockId) 0
39+
40+
lockNext :: forall m a. SqlTransaction m => m a -> WithLock m a
41+
lockNext tx = WithLock $ do
42+
lockId <- ask
43+
nextLock <- get
44+
put $ succ nextLock
45+
lift . lift $ do
46+
statement @m (lockId, nextLock) advisoryLockStatement
47+
tx
48+
where
49+
advisoryLockStatement =
50+
SQL.Statement
51+
"SELECT pg_advisory_xact_lock($1, $2)"
52+
((fst >$< param HE.int4) <> (snd >$< param HE.int2))
53+
HD.noResult
54+
False
55+
param = HE.param . HE.nonNullable

‎test/io/conftest.py‎

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import os
22
import pytest
33
from syrupy.extensions.json import SingleFileSnapshotExtension
4-
from postgrest import run
4+
from postgrest import SchemaCacheLocks, run
55

66

77
@pytest.fixture
@@ -57,16 +57,13 @@ def replicaenv(defaultenv):
5757

5858

5959
@pytest.fixture
60-
def slow_schema_cache_env(defaultenv):
61-
"Slow schema cache load environment PostgREST."
62-
return {
63-
**defaultenv,
64-
"PGRST_INTERNAL_SCHEMA_CACHE_QUERY_SLEEP": "1000", # this does a pg_sleep internally, it will cause the schema cache query to be slow
65-
# the slow schema cache query will keep using one pool connection until it finishes
66-
# to prevent requests waiting for PGRST_DB_POOL_ACQUISITION_TIMEOUT we'll increase the pool size (must be >= 2)
67-
"PGRST_DB_POOL": "2",
68-
"PGRST_DB_CHANNEL_ENABLED": "true",
69-
}
60+
def schema_cache_locks(baseenv):
61+
"Factory for controlling step-by-step querySchemaCache execution."
62+
63+
def factory(lock_id=None, max_step=15):
64+
return SchemaCacheLocks(baseenv, lock_id=lock_id, max_step=max_step)
65+
66+
return factory
7067

7168

7269
@pytest.fixture

0 commit comments

Comments
 (0)