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

Author: mkleczek

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/SchemaCache.hs

@@ -157,46 +157,56 @@ maxDbTablesForFuzzySearch = 500
 querySchemaCache :: AppConfig -> SQL.Transaction SchemaCache
 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))
-    , dbQueryTimings = qsTime
-    }
+  (`evalStateT` 0) $ do
+    tabs    <- sqlWithLock $ sqlTimedStmt gucTbls  conf   allTables
+    keyDeps <- sqlWithLock $ sqlTimedStmt gucKDeps conf   allViewsKeyDependencies
+    m2oRels <- sqlWithLock $ sqlTimedStmt gucRels  mempty allM2OandO2ORels
+    funcs   <- sqlWithLock $ sqlTimedStmt gucFuncs conf   allFunctions
+    cRels   <- sqlWithLock $ sqlTimedStmt gucCRels mempty allComputedRels
+    reps    <- sqlWithLock $ sqlTimedStmt gucDReps conf   dataRepresentations
+    mHdlers <- sqlWithLock $ sqlTimedStmt gucMHdrs conf   mediaHandlers
+    tzones  <- if configDbTimezoneEnabled
+      then sqlWithLock $ 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
+      lift $ for_ configInternalSCQuerySleep (`SQL.statement` sleepCall) -- only used for testing
+
+    qsTime <-
+      if isLogDebug
+        then Just <$> sqlWithLock (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))
+      , dbQueryTimings = qsTime
+      }
   where
     schemas = toList configDbSchemas
     isLogDebug = configLogLevel == LogDebug
     sqlTimedStmt = sqlTimedStatement isLogDebug
+    sqlWithLock stmt = maybe (lift stmt) (lockNext stmt) configInternalSCLockId
+    lockNext stmt lockId = do
+      nextLock <- get
+      put $ succ nextLock
+      lift $ do
+        SQL.statement (lockId, nextLock) $ SQL.Statement "SELECT pg_advisory_xact_lock($1, $2)" ((fst >$< param HE.int4) <> (snd >$< param HE.int2)) HD.noResult False
+        stmt
+
+
 
 -- | overrides detected relationships with the computed relationships and gets the RelationshipsMap
 getOverrideRelationshipsMap :: [Relationship] -> [Relationship] -> RelationshipsMap

test/io/conftest.py

@@ -61,6 +61,7 @@ def slow_schema_cache_env(defaultenv):
     "Slow schema cache load environment PostgREST."
     return {
         **defaultenv,
+        "PGRST_INTERNAL_SCHEMA_CACHE_LOCK_ID": "1111",
         "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)

test/io/test_io.py

@@ -1394,6 +1394,7 @@ def test_schema_cache_concurrent_notifications(slow_schema_cache_env):
     internal_sleep = (
         int(slow_schema_cache_env["PGRST_INTERNAL_SCHEMA_CACHE_QUERY_SLEEP"]) / 1000
     )
+    # TODO change to usage of PGRST_INTERNAL_SCHEMA_CACHE_LOCK_ID
 
     with run(env=slow_schema_cache_env, wait_for=None) as postgrest:
         time.sleep(2 * internal_sleep + 0.1)  # wait for readiness manually