refactor: add internal-schema-cache-lock-id

Author: mkleczek

nix/tools/tests.nix

@@ -78,6 +78,7 @@ let
   ioTestPython =
     python3.withPackages (ps: [
       ps.pyjwt
+      ps.psycopg
       ps.pytest
       ps.pytest-xdist
       ps.pyyaml

src/PostgREST/Config.hs

@@ -127,6 +127,7 @@ data AppConfig = AppConfig
   , configRoleSettings             :: RoleSettings
   , configRoleIsoLvl               :: RoleIsolationLvl
   , configInternalSCQuerySleep     :: Maybe Int32
+  , configInternalSCLockId         :: Maybe Int32
   }
 
 data LogLevel = LogCrit | LogError | LogWarn | LogInfo | LogDebug
@@ -326,6 +327,7 @@ parser optPath env dbSettings roleSettings roleIsolationLvl =
     <*> pure roleSettings
     <*> pure roleIsolationLvl
     <*> optInt "internal-schema-cache-query-sleep"
+    <*> optInt "internal-schema-cache-lock-id"
   where
     parseErrorVerbosity :: C.Key -> C.Parser C.Config Verbosity
     parseErrorVerbosity k =

src/PostgREST/Logger.hs

@@ -32,7 +32,6 @@ import PostgREST.Debounce      (makeDebouncer)
 import PostgREST.Logger.Apache (apacheFormat)
 import PostgREST.Observation
 import PostgREST.Query         (MainQuery (..))
-import PostgREST.SchemaCache   (queryTimingsWLabels)
 
 import qualified Data.ByteString.Lazy       as LBS
 import qualified Data.Text                  as T
@@ -160,7 +159,7 @@ observationMessages = \case
       <> ". " <> jsonMessage usageErr
   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
+    let showTimings qt = [ T.intercalate ", " $ (\(l, v) -> T.decodeUtf8 l <> ": " <> v <> " ms") <$> qt ] in
     maybe mempty showTimings timings
   SchemaCacheLoadedObs resultTime summary ->
     [

src/PostgREST/Observation.hs

@@ -19,7 +19,7 @@ import           Network.HTTP.Types.Status  (Status)
 import qualified Network.Wai                as Wai
 import           PostgREST.Config.PgVersion
 import           PostgREST.Query            (MainQuery)
-import           PostgREST.SchemaCache      (QueryTimings)
+import           PostgREST.SqlTransaction   (QueryTimings)
 
 import Protolude hiding (toList)
 

src/PostgREST/SchemaCache.hs

@@ -8,67 +8,74 @@ The schema cache is necessary for resource embedding, foreign keys are used for
 
 These queries are executed once at startup or when PostgREST is reloaded.
 -}
+{-# LANGUAGE AllowAmbiguousTypes   #-}
+{-# LANGUAGE DataKinds             #-}
 {-# LANGUAGE DeriveAnyClass        #-}
 {-# LANGUAGE DeriveGeneric         #-}
 {-# LANGUAGE FlexibleContexts      #-}
+{-# LANGUAGE FlexibleInstances     #-}
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE NamedFieldPuns        #-}
+{-# LANGUAGE PolyKinds             #-}
 {-# LANGUAGE QuasiQuotes           #-}
 {-# LANGUAGE RecordWildCards       #-}
 {-# LANGUAGE ScopedTypeVariables   #-}
-{-# LANGUAGE TypeSynonymInstances  #-}
+{-# LANGUAGE TypeApplications      #-}
+{-# LANGUAGE TypeFamilies          #-}
 
 module PostgREST.SchemaCache
   ( SchemaCache(..)
   , TablesFuzzyIndex
   , querySchemaCache
   , showSummary
   , decodeFuncs
-  , QueryTimings(..)
-  , queryTimingsWLabels
   ) where
 
 import           Data.Aeson ((.=))
 import qualified Data.Aeson as JSON
 
-import qualified Data.ByteString.Char8      as BS
 import qualified Data.HashMap.Strict        as HM
 import qualified Data.HashMap.Strict.InsOrd as HMI
 import qualified Data.Set                   as S
 import qualified Data.Text                  as T
 import qualified Hasql.Decoders             as HD
 import qualified Hasql.Encoders             as HE
 import qualified Hasql.Statement            as SQL
-import qualified Hasql.Transaction          as SQL
+import qualified Hasql.Transaction          as SQL hiding (sql,
+                                                    statement)
 
 import Data.Functor.Contravariant ((>$<))
 import NeatInterpolation          (trimming)
 
-import PostgREST.Config                      (AppConfig (..),
-                                              LogLevel (..))
-import PostgREST.Config.Database             (TimezoneNames,
-                                              toIsolationLevel)
-import PostgREST.SchemaCache.Identifiers     (FieldName,
-                                              QualifiedIdentifier (..),
-                                              RelIdentifier (..),
-                                              Schema, escapeIdent,
-                                              isAnyElement)
-import PostgREST.SchemaCache.Relationship    (Cardinality (..),
-                                              Junction (..),
-                                              Relationship (..),
-                                              RelationshipsMap)
-import PostgREST.SchemaCache.Representations (DataRepresentation (..),
-                                              RepresentationsMap)
-import PostgREST.SchemaCache.Routine         (FuncVolatility (..),
-                                              MediaHandler (..),
-                                              MediaHandlerMap,
-                                              PgType (..),
-                                              RetType (..),
-                                              Routine (..),
-                                              RoutineMap,
-                                              RoutineParam (..))
-import PostgREST.SchemaCache.Table           (Column (..), ColumnMap,
-                                              Table (..), TablesMap)
+import           PostgREST.Config                      (AppConfig (..),
+                                                        LogLevel (..))
+import           PostgREST.Config.Database             (TimezoneNames,
+                                                        toIsolationLevel)
+import           PostgREST.SchemaCache.Identifiers     (FieldName,
+                                                        QualifiedIdentifier (..),
+                                                        RelIdentifier (..),
+                                                        Schema,
+                                                        escapeIdent,
+                                                        isAnyElement)
+import           PostgREST.SchemaCache.Relationship    (Cardinality (..),
+                                                        Junction (..),
+                                                        Relationship (..),
+                                                        RelationshipsMap)
+import           PostgREST.SchemaCache.Representations (DataRepresentation (..),
+                                                        RepresentationsMap)
+import           PostgREST.SchemaCache.Routine         (FuncVolatility (..),
+                                                        MediaHandler (..),
+                                                        MediaHandlerMap,
+                                                        PgType (..),
+                                                        RetType (..),
+                                                        Routine (..),
+                                                        RoutineMap,
+                                                        RoutineParam (..))
+import           PostgREST.SchemaCache.Table           (Column (..),
+                                                        ColumnMap,
+                                                        Table (..),
+                                                        TablesMap)
+import qualified PostgREST.SqlTransaction              as SQL
 
 import qualified PostgREST.MediaType as MediaType
 
@@ -153,48 +160,47 @@ type SqlQuery = ByteString
 maxDbTablesForFuzzySearch :: Int
 maxDbTablesForFuzzySearch = 500
 
-querySchemaCache :: AppConfig -> SQL.Transaction (SchemaCache, Maybe QueryTimings)
-querySchemaCache conf@AppConfig{..} = do
-  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
-  tabs    <- sqlTimedStmt gucTbls  conf   allTables
-  keyDeps <- sqlTimedStmt gucKDeps conf   allViewsKeyDependencies
-  m2oRels <- sqlTimedStmt gucRels  mempty allM2OandO2ORels
-  funcs   <- sqlTimedStmt gucFuncs conf   allFunctions
-  cRels   <- sqlTimedStmt gucCRels mempty allComputedRels
-  reps    <- sqlTimedStmt gucDReps conf   dataRepresentations
-  mHdlers <- sqlTimedStmt gucMHdrs conf   mediaHandlers
-  tzones  <- if configDbTimezoneEnabled
-    then sqlTimedStmt gucTzones mempty timezones
-    else pure S.empty
-  _       <-
-    let sleepCall = SQL.Statement "select pg_sleep($1 / 1000.0)" (param HE.int4) HD.noResult True in
-    for_ configInternalSCQuerySleep (`SQL.statement` sleepCall) -- only used for testing
-
-  qsTime <-
-    if isLogDebug
-      then Just <$> SQL.statement mempty (extractTimings configDbTimezoneEnabled)
-      else pure Nothing
-
-  let tabsWViewsPks = addViewPrimaryKeys tabs keyDeps
-      rels          = addInverseRels $ addM2MRels tabsWViewsPks $ addViewM2OAndO2ORels keyDeps m2oRels
-
-  return (removeInternal schemas $ SchemaCache {
-      dbTables = tabsWViewsPks
-    , dbRelationships = getOverrideRelationshipsMap rels cRels
-    , dbRoutines = funcs
-    , dbRepresentations = reps
-    , dbMediaHandlers = HM.union mHdlers initialMediaHandlers -- the custom handlers will override the initial ones
-    , dbTimezones = tzones
-
-    , 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))
-    }, qsTime)
+querySchemaCache :: AppConfig -> SQL.Transaction (SchemaCache, Maybe SQL.QueryTimings)
+querySchemaCache conf@AppConfig{..} =
+  -- if configInternalSCLockId is set run queries step-by-step waiting for lock release before each
+  SQL.runSteppedTransaction @(SQL.SqlTransaction SchemaCacheLabel) configInternalSCLockId $
+    -- if log level is debug then time queries
+    SQL.runTimed @(SQL.SqlTransaction SchemaCacheLabel) isLogDebug $ do
+      SQL.sql @NoStep "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
+      tabs    <- SQL.statement @Tables                conf   allTables
+      keyDeps <- SQL.statement @KeyDependencies       conf   allViewsKeyDependencies
+      m2oRels <- SQL.statement @Relationships         mempty allM2OandO2ORels
+      funcs   <- SQL.statement @Functions             conf   allFunctions
+      cRels   <- SQL.statement @ComputedRelationships mempty allComputedRels
+      reps    <- SQL.statement @DataRepresentations   conf   dataRepresentations
+      mHdlers <- SQL.statement @MediaHandlers         conf   mediaHandlers
+      tzones  <- if configDbTimezoneEnabled
+        then SQL.statement @Timezones mempty timezones
+        else pure S.empty
+
+      let tabsWViewsPks = addViewPrimaryKeys tabs keyDeps
+          rels          = addInverseRels $ addM2MRels tabsWViewsPks $ addViewM2OAndO2ORels keyDeps m2oRels
+
+      return $ removeInternal schemas $ SchemaCache {
+          dbTables = tabsWViewsPks
+        , dbRelationships = getOverrideRelationshipsMap rels cRels
+        , dbRoutines = funcs
+        , dbRepresentations = reps
+        , dbMediaHandlers = HM.union mHdlers initialMediaHandlers -- the custom handlers will override the initial ones
+        , dbTimezones = tzones
+
+        , 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))
+        }
+      -- only used for testing
+      -- TODO remove configInternalSCQuerySleep once all tests are migrated to stepped execution
+      <* for_ configInternalSCQuerySleep (\sleep -> SQL.statement @NoStep sleep sleepCall)
   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
 
 -- | overrides detected relationships with the computed relationships and gets the RelationshipsMap
 getOverrideRelationshipsMap :: [Relationship] -> [Relationship] -> RelationshipsMap
@@ -1155,71 +1161,19 @@ nullableColumn = HD.column . HD.nullable
 arrayColumn :: HD.Value a -> HD.Row [a]
 arrayColumn = column . HD.listArray . HD.nonNullable
 
-{-
- - Times a sql statement inside a transaction, for this:
- -
- - 1. We start a timer: select set_config('pgrst.tmp_x', clock_timestamp()::text, false);
- - 2. Run the statement: select ....
- - 3. End the timer:  select set_config('pgrst.tmp_x', (clock_timestamp() - current_setting('pgrst.tmp_x', false)::timestamptz)::text, false);
- -
- - We can do this for several statements inside the transaction. The timings are later captured at the end of the transaction with extractTimings.
- -}
-sqlTimedStatement :: Bool -> ByteString -> a -> SQL.Statement a b -> SQL.Transaction b
-sqlTimedStatement isLogDebug guc params stmt =
-  if isLogDebug then
-    SQL.sql sFrag >> SQL.statement params stmt <* SQL.sql eFrag
-  else
-    SQL.statement params stmt
-  where
-    sFrag = "select set_config('pgrst." <> guc <> "', clock_timestamp()::text, true)"
-    eFrag = "select set_config('pgrst." <> guc <> "', (clock_timestamp() - current_setting('pgrst." <> guc <> "', false)::timestamptz)::text, true)"
-
--- Extract all the generated timings (see sqlTimedStatement) converting the value to milliseconds.
-extractTimings :: Bool -> SQL.Statement () QueryTimings
-extractTimings hasTimezones = SQL.Statement sql HE.noParams decodeThem True
-  where
-    qFrag setting = "extract('milliseconds' from current_setting('pgrst." <> setting <> "', false)::interval)::text"
-    sql = "SELECT " <> BS.intercalate ","
-      [ qFrag gucTbls,  qFrag gucKDeps, qFrag gucRels
-      , qFrag gucFuncs, qFrag gucCRels, qFrag gucDReps
-      , qFrag gucMHdrs, if hasTimezones then qFrag gucTzones else "'0.0'"
-      ]
-    decodeThem :: HD.Result QueryTimings
-    decodeThem = HD.singleRow $
-      QueryTimings
-        <$> column HD.text <*> column HD.text <*> column HD.text
-        <*> column HD.text <*> column HD.text <*> column HD.text
-        <*> column HD.text <*> column HD.text
-
-data QueryTimings = QueryTimings
-  { qtTables  :: Text
-  , qtKeyDeps :: Text
-  , qtRels    :: Text
-  , qtFuncs   :: Text
-  , qtCRels   :: Text
-  , qtDReps   :: Text
-  , qtMHdrs   :: Text
-  , qtTzones  :: Text
-  } deriving (Show)
-
-queryTimingsWLabels :: QueryTimings -> [(ByteString, Text)]
-queryTimingsWLabels qt =
-  [ (gucTbls,   qtTables qt)
-  , (gucKDeps,  qtKeyDeps qt)
-  , (gucRels,   qtRels qt)
-  , (gucFuncs,  qtFuncs qt)
-  , (gucCRels,  qtCRels qt)
-  , (gucDReps,  qtDReps qt)
-  , (gucMHdrs,  qtMHdrs qt)
-  , (gucTzones, qtTzones qt)
-  ]
-
-gucTbls, gucKDeps, gucRels, gucFuncs, gucCRels, gucDReps, gucMHdrs, gucTzones :: ByteString
-gucTbls   = "tables"
-gucKDeps  = "keydeps"
-gucRels   = "rels"
-gucFuncs  = "funcs"
-gucCRels  = "comprels"
-gucDReps  = "dreps"
-gucMHdrs  = "mhandlers"
-gucTzones = "tzones"
+data SchemaCacheLabel = Step SQL.LockSpec SQL.TimingSpec
+
+instance SQL.LabelKind SchemaCacheLabel where
+  type LabelConstraint SchemaCacheLabel label = (SQL.SqlBreakpoint label, SQL.SqlTiming label)
+instance SQL.HasTimingsQueryLabel SchemaCacheLabel where
+  type TimingsQueryLabel SchemaCacheLabel = NoStep
+
+type Tables                = Step (SQL.Lock 0) (SQL.Timing "tables")
+type KeyDependencies       = Step (SQL.Lock 1) (SQL.Timing "keydeps")
+type Relationships         = Step (SQL.Lock 2) (SQL.Timing "rels")
+type Functions             = Step (SQL.Lock 3) (SQL.Timing "funcs")
+type ComputedRelationships = Step (SQL.Lock 4) (SQL.Timing "comprels")
+type DataRepresentations   = Step (SQL.Lock 5) (SQL.Timing "dreps")
+type MediaHandlers         = Step (SQL.Lock 6) (SQL.Timing "mhandlers")
+type Timezones             = Step (SQL.Lock 7) (SQL.Timing "tzones")
+type NoStep                = Step SQL.NoLock SQL.NoTiming

test/io/conftest.py

@@ -1,7 +1,7 @@
 import os
 import pytest
 from syrupy.extensions.json import SingleFileSnapshotExtension
-from postgrest import run
+from postgrest import SchemaCacheLocks, run
 
 
 @pytest.fixture
@@ -57,16 +57,13 @@ def replicaenv(defaultenv):
 
 
 @pytest.fixture
-def slow_schema_cache_env(defaultenv):
-    "Slow schema cache load environment PostgREST."
-    return {
-        **defaultenv,
-        "PGRST_INTERNAL_SCHEMA_CACHE_QUERY_SLEEP": "1000",  # this does a pg_sleep internally, it will cause the schema cache query to be slow
-        # the slow schema cache query will keep using one pool connection until it finishes
-        # to prevent requests waiting for PGRST_DB_POOL_ACQUISITION_TIMEOUT we'll increase the pool size (must be >= 2)
-        "PGRST_DB_POOL": "2",
-        "PGRST_DB_CHANNEL_ENABLED": "true",
-    }
+def schema_cache_locks(baseenv):
+    "Factory for controlling step-by-step querySchemaCache execution."
+
+    def factory(lock_id=None, max_step=15):
+        return SchemaCacheLocks(baseenv, lock_id=lock_id, max_step=max_step)
+
+    return factory
 
 
 @pytest.fixture