perf: use pg_basetype on PG 17+ for schema cache

CHANGELOG.md

@@ -17,6 +17,7 @@ All notable changes to this project will be documented in this file. From versio
 - Support running the admin server on a unix socket by @wolfgangwalther in #5003
 - Add config `server-reuseport` to allow starting multiple PostgREST instances using the same port on supported platforms by @mkleczek in #4703, #4694
 - Make config `log-level` reloadable by @taimoorzaeem in #5113
+- Optimize schema cache domain type resolution by using `pg_basetype` on PostgreSQL 17+ by @joelonsql in #4567
 
 ### Fixed
 

src/library/PostgREST/App.hs

@@ -219,7 +219,8 @@ postgrestResponse appState conf@AppConfig{..} maybeSchemaCache jwtTime authResul
   liftIO $ when shouldShowWarnings $
     observer $ LegacyTargetNameWarningObs (legacyWarnMsg, legacyWarnHint) iMethod (iPath <> Wai.rawQueryString req) -- TODO maybe store rawQueryString in ApiRequest for consistency
 
-  let mainQ = Query.mainQuery plan conf apiReq authResult configDbPreRequest
+  pgVer <- liftIO $ AppState.getPgVersion appState
+  let mainQ = Query.mainQuery pgVer plan conf apiReq authResult configDbPreRequest
       tx = MainTx.mainTx mainQ conf authResult apiReq plan sCache
       obsQuery s = when configLogQuery $ observer $ QueryObs mainQ s
 

src/library/PostgREST/AppState/Reload.hs

@@ -92,8 +92,9 @@ retryingSchemaCacheLoad appState@AppState{stateObserver=observer} =
     qSchemaCache :: IO (Maybe SchemaCache)
     qSchemaCache = do
       conf@AppConfig{..} <- getConfig appState
+      pgVer <- getPgVersion appState
       (resultTime, result) <-
-        timeItT $ usePool appState (SQL.transactionNoRetry SQL.ReadCommitted SQL.Read $ querySchemaCache conf)
+        timeItT $ usePool appState (SQL.transactionNoRetry SQL.ReadCommitted SQL.Read $ querySchemaCache pgVer conf)
       case result of
         Left e -> do
           markSchemaCachePending appState

src/library/PostgREST/CLI.hs

@@ -63,8 +63,9 @@ runAppCommand conf@AppConfig{..} runCmd = do
 dumpSchema :: AppState -> IO LBS.ByteString
 dumpSchema appState = do
   conf@AppConfig{..} <- AppState.getConfig appState
+  pgVer <- AppState.getPgVersion appState
   result <-
-    AppState.usePool appState (SQL.transactionNoRetry SQL.ReadCommitted SQL.Read $ querySchemaCache conf)
+    AppState.usePool appState (SQL.transactionNoRetry SQL.ReadCommitted SQL.Read $ querySchemaCache pgVer conf)
   case result of
     Left e -> do
       let observer = AppState.getObserver appState

src/library/PostgREST/Config/PgVersion.hs

@@ -4,6 +4,7 @@ module PostgREST.Config.PgVersion
   ( PgVersion(..)
   , minimumPgVersion
   , pgVersion150
+  , pgVersion170
   , pgVersion180
   , pgVersion190
   ) where
@@ -33,6 +34,9 @@ pgVersion140 = PgVersion 140000 "14.0" "14.0"
 pgVersion150 :: PgVersion
 pgVersion150 = PgVersion 150000 "15.0" "15.0"
 
+pgVersion170 :: PgVersion
+pgVersion170 = PgVersion 170000 "17.0" "17.0"
+
 pgVersion180 :: PgVersion
 pgVersion180 = PgVersion 180000 "18.0" "18.0"
 

src/library/PostgREST/Query.hs

@@ -22,6 +22,7 @@ import PostgREST.ApiRequest              (ApiRequest (..))
 import PostgREST.ApiRequest.Preferences  (Preferences (..), shouldExplainCount)
 import PostgREST.Auth.Types              (AuthResult (..))
 import PostgREST.Config                  (AppConfig (..))
+import PostgREST.Config.PgVersion        (PgVersion)
 import PostgREST.Plan                    (ActionPlan (..), CrudPlan (..),
                                           DbActionPlan (..), InspectPlan (..))
 import PostgREST.SchemaCache.Identifiers (QualifiedIdentifier (..))
@@ -38,9 +39,9 @@ data MainQuery = MainQuery
   , mqExplain :: Maybe SQL.Snippet     -- ^ the explain query that gets generated for the "Prefer: count=estimated" case
   }
 
-mainQuery :: ActionPlan -> AppConfig -> ApiRequest -> AuthResult -> Maybe QualifiedIdentifier -> MainQuery
+mainQuery :: PgVersion -> ActionPlan -> AppConfig -> ApiRequest -> AuthResult -> Maybe QualifiedIdentifier -> MainQuery
-mainQuery (NoDb _) _ _ _ _ = MainQuery mempty Nothing mempty (mempty, mempty, mempty) mempty
+mainQuery _ (NoDb _) _ _ _ _ = MainQuery mempty Nothing mempty (mempty, mempty, mempty) mempty
-mainQuery (Db plan) conf@AppConfig{..} apiReq@ApiRequest{iTopLevelRange=range, iPreferences=Preferences{..}} authRes preReq =
+mainQuery pgVer (Db plan) conf@AppConfig{..} apiReq@ApiRequest{iTopLevelRange=range, iPreferences=Preferences{..}} authRes preReq =
   let genQ = MainQuery (PreQuery.txVarQuery plan conf authRes apiReq) (PreQuery.preReqQuery <$> preReq) in
   case plan of
     DbCrud _ WrappedReadPlan{..} ->
@@ -52,4 +53,4 @@ mainQuery (Db plan) conf@AppConfig{..} apiReq@ApiRequest{iTopLevelRange=range, i
     DbCrud _ CallReadPlan{..} ->
       genQ (Statements.mainCall crProc crCallPlan crReadPlan preferCount configDbMaxRows range pMedia crHandler) (mempty, mempty, mempty) mempty
     MayUseDb InspectPlan{ipSchema=tSchema} ->
-      genQ mempty (SqlFragment.accessibleTables tSchema, SqlFragment.accessibleFuncs tSchema, SqlFragment.schemaDescription tSchema) mempty
+      genQ mempty (SqlFragment.accessibleTables tSchema, SqlFragment.accessibleFuncs pgVer tSchema, SqlFragment.schemaDescription tSchema) mempty

src/library/PostgREST/Query/SqlFragment.hs

@@ -68,6 +68,7 @@ import PostgREST.ApiRequest.Types        (AggregateFunction (..), Alias, Cast,
                                           Operation (..), OrderDirection (..),
                                           OrderNulls (..), QuantOperator (..),
                                           SimpleOperator (..))
+import PostgREST.Config.PgVersion        (PgVersion, pgVersion170)
 import PostgREST.MediaType               (MTVndPlanFormat (..),
                                           MTVndPlanOption (..))
 import PostgREST.Plan.ReadPlan           (JoinCondition (..))
@@ -612,39 +613,60 @@ accessibleTables schema = SQL.sql (encodeUtf8 [trimming|
   where
     encodedSchema = SQL.encoderAndParam (HE.nonNullable HE.text) schema
 
-accessibleFuncs :: Text -> SQL.Snippet
+accessibleFuncs :: PgVersion -> Text -> SQL.Snippet
-accessibleFuncs schema = baseFuncSqlQuery <> "AND p.pronamespace = " <> encodedSchema <> "::regnamespace"
+accessibleFuncs pgVer schema = baseFuncSqlQuery pgVer <> "AND p.pronamespace = " <> encodedSchema <> "::regnamespace"
   where
     encodedSchema = SQL.encoderAndParam (HE.nonNullable HE.text) schema
 
-baseFuncSqlQuery :: SQL.Snippet
+baseTypesCte :: PgVersion -> Text
-baseFuncSqlQuery = SQL.sql $ encodeUtf8 [trimming|
+baseTypesCte pgVer
+  | pgVer >= pgVersion170 = [trimming|
+      -- Get base types using pg_basetype() (PG 17+)
+      base_types AS (
+        SELECT
+          t.oid,
+          bt.typnamespace AS base_namespace,
+          bt.oid AS base_type
+        FROM pg_type t
+        JOIN pg_type bt ON bt.oid = pg_basetype(t.oid)
+      )
+    |]
+  | otherwise = [trimming|
+      -- Recursively get the base types of domains (PG < 17)
+      base_types AS (
+        WITH RECURSIVE
+        recurse AS (
+          SELECT
+            oid,
+            typbasetype,
+            typnamespace AS base_namespace,
+            COALESCE(NULLIF(typbasetype, 0), oid) AS base_type
+          FROM pg_type
+          UNION
+          SELECT
+            t.oid,
+            b.typbasetype,
+            b.typnamespace AS base_namespace,
+            COALESCE(NULLIF(b.typbasetype, 0), b.oid) AS base_type
+          FROM recurse t
+          JOIN pg_type b ON t.typbasetype = b.oid
+        )
+        SELECT
+          oid,
+          base_namespace,
+          base_type
+        FROM recurse
+        WHERE typbasetype = 0
+      )
+    |]
+
+-- | SQL query to get accessible functions for OpenAPI.
+baseFuncSqlQuery :: PgVersion -> SQL.Snippet
+baseFuncSqlQuery pgVer =
+  let baseCte = baseTypesCte pgVer
+  in SQL.sql $ encodeUtf8 [trimming|
   WITH
-  base_types AS (
+  $baseCte,
-    WITH RECURSIVE
-    recurse AS (
-      SELECT
-        oid,
-        typbasetype,
-        typnamespace AS base_namespace,
-        COALESCE(NULLIF(typbasetype, 0), oid) AS base_type
-      FROM pg_type
-      UNION
-      SELECT
-        t.oid,
-        b.typbasetype,
-        b.typnamespace AS base_namespace,
-        COALESCE(NULLIF(b.typbasetype, 0), b.oid) AS base_type
-      FROM recurse t
-      JOIN pg_type b ON t.typbasetype = b.oid
-    )
-    SELECT
-      oid,
-      base_namespace,
-      base_type
-    FROM recurse
-    WHERE typbasetype = 0
-  ),
   arguments AS (
     SELECT
       oid,

src/library/PostgREST/SchemaCache.hs

@@ -46,6 +46,7 @@ import NeatInterpolation          (trimming)
 
 import PostgREST.Config                      (AppConfig (..), LogLevel (..))
 import PostgREST.Config.Database             (TimezoneNames, toIsolationLevel)
+import PostgREST.Config.PgVersion            (PgVersion, pgVersion170)
 import PostgREST.SchemaCache.Identifiers     (FieldName,
                                               QualifiedIdentifier (..),
                                               RelIdentifier (..), Schema,
@@ -146,16 +147,16 @@ type SqlQuery = ByteString
 maxDbTablesForFuzzySearch :: Int
 maxDbTablesForFuzzySearch = 500
 
-querySchemaCache :: AppConfig -> SQL.Transaction (SchemaCache, Maybe QueryTimings)
+querySchemaCache :: PgVersion -> AppConfig -> SQL.Transaction (SchemaCache, Maybe QueryTimings)
-querySchemaCache conf@AppConfig{..} = do
+querySchemaCache pgVer 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
 
   for_ configInternalSCQuerySleepFst (`SQL.statement` sleepCall) -- only used for testing
 
-  tabs    <- sqlTimedStmt gucTbls  conf   allTables
+  tabs    <- sqlTimedStmt gucTbls  conf   (allTables pgVer configDbPreparedStatements)
   keyDeps <- sqlTimedStmt gucKDeps conf   allViewsKeyDependencies
   m2oRels <- sqlTimedStmt gucRels  mempty allM2OandO2ORels
-  funcs   <- sqlTimedStmt gucFuncs conf   allFunctions
+  funcs   <- sqlTimedStmt gucFuncs conf   (allFunctions pgVer configDbPreparedStatements)
   cRels   <- sqlTimedStmt gucCRels mempty allComputedRels
   reps    <- sqlTimedStmt gucDReps conf   dataRepresentations
   mHdlers <- sqlTimedStmt gucMHdrs conf   mediaHandlers
@@ -374,47 +375,44 @@ dataRepresentations = SQL.Statement sql mempty decodeRepresentations True
        OR (dst_t.typtype = 'd' AND c.castsource IN ('json'::regtype::oid , 'text'::regtype::oid)))
     |]
 
-allFunctions :: SQL.Statement AppConfig RoutineMap
+allFunctions :: PgVersion -> Bool -> SQL.Statement AppConfig RoutineMap
-allFunctions = SQL.Statement funcsSqlQuery params decodeFuncs True
+allFunctions pgVer = SQL.Statement (funcsSqlQuery pgVer) params decodeFuncs
   where
     params =
       (map escapeIdent . toList . configDbSchemas >$< arrayParam HE.text) <>
       (configDbHoistedTxSettings >$< arrayParam HE.text)
 
-baseTypesCte :: Text
+baseTypesCte :: PgVersion -> Text
-baseTypesCte = [trimming|
+baseTypesCte pgVer
-  -- Recursively get the base types of domains
+  | pgVer >= pgVersion170 = [trimming|
-  base_types AS (
+      base_types AS (
-    WITH RECURSIVE
+        SELECT t.oid, bt.typnamespace AS base_namespace, bt.oid AS base_type
-    recurse AS (
+        FROM pg_type t
-      SELECT
+        JOIN pg_type bt ON bt.oid = pg_basetype(t.oid)
-        oid,
+      )
-        typbasetype,
+    |]
-        typnamespace AS base_namespace,
+  | otherwise = [trimming|
-        COALESCE(NULLIF(typbasetype, 0), oid) AS base_type
+      base_types AS (
-      FROM pg_type
+        WITH RECURSIVE recurse AS (
-      UNION
+          SELECT oid, typbasetype, typnamespace AS base_namespace,
-      SELECT
+            COALESCE(NULLIF(typbasetype, 0), oid) AS base_type
-        t.oid,
+          FROM pg_type
-        b.typbasetype,
+          UNION
-        b.typnamespace AS base_namespace,
+          SELECT t.oid, b.typbasetype, b.typnamespace AS base_namespace,
-        COALESCE(NULLIF(b.typbasetype, 0), b.oid) AS base_type
+            COALESCE(NULLIF(b.typbasetype, 0), b.oid) AS base_type
-      FROM recurse t
+          FROM recurse t
-      JOIN pg_type b ON t.typbasetype = b.oid
+          JOIN pg_type b ON t.typbasetype = b.oid
-    )
+        )
-    SELECT
+        SELECT oid, base_namespace, base_type FROM recurse WHERE typbasetype = 0
-      oid,
+      )
-      base_namespace,
+    |]
-      base_type
-    FROM recurse
-    WHERE typbasetype = 0
-  )
-|]
 
-funcsSqlQuery :: SqlQuery
+funcsSqlQuery :: PgVersion -> SqlQuery
-funcsSqlQuery = encodeUtf8 [trimming|
+funcsSqlQuery pgVer =
+  let baseCte = baseTypesCte pgVer
+  in encodeUtf8 [trimming|
   WITH
-  $baseTypesCte,
+  $baseCte,
   arguments AS (
     SELECT
       oid,
@@ -587,22 +585,23 @@ addViewPrimaryKeys tabs keyDeps =
     takeFirstPK = mapMaybe (head . snd)
     indexedDeps = HM.fromListWith (++) $ fmap ((keyDepType &&& keyDepView) &&& pure) keyDeps
 
-allTables :: SQL.Statement AppConfig TablesMap
+allTables :: PgVersion -> Bool -> SQL.Statement AppConfig TablesMap
-allTables = SQL.Statement tablesSqlQuery params decodeTables True
+allTables pgVer = SQL.Statement (tablesSqlQuery pgVer) params decodeTables
   where
     params = map escapeIdent . toList . configDbSchemas >$< arrayParam HE.text
 
 -- | Gets tables with their PK cols
-tablesSqlQuery :: SqlQuery
+tablesSqlQuery :: PgVersion -> SqlQuery
-tablesSqlQuery =
+tablesSqlQuery pgVer =
   -- the tbl_constraints/key_col_usage CTEs are based on the standard "information_schema.table_constraints"/"information_schema.key_column_usage" views,
   -- we cannot use those directly as they include the following privilege filter:
   -- (pg_has_role(ss.relowner, 'USAGE'::text) OR has_column_privilege(ss.roid, a.attnum, 'SELECT, INSERT, UPDATE, REFERENCES'::text));
   -- on the "columns" CTE, left joining on pg_depend and pg_class is used to obtain the sequence name as a column default in case there are GENERATED .. AS IDENTITY,
   -- generated columns are only available from pg >= 10 but the query is agnostic to versions. dep.deptype = 'i' is done because there are other 'a' dependencies on PKs
-  encodeUtf8 [trimming|
+  let baseCte = baseTypesCte pgVer
+  in encodeUtf8 [trimming|
   WITH
-  $baseTypesCte,
+  $baseCte,
   columns AS (
       SELECT
           c.oid AS relid,

test/io/test_io.py

@@ -783,7 +783,7 @@ def test_log_query(level, defaultenv):
         )
         infinite_recursion_5xx_regx = r'.+: WITH pgrst_source AS.+SELECT "public"\."infinite_recursion"\.\* FROM "public"\."infinite_recursion".+_postgrest_t'
         root_tables_regx = r".+: SELECT   n.nspname AS table_schema, .+ FROM pg_class c .+ ORDER BY table_schema, table_name"
-        root_procs_regx = r".+: WITH base_types AS \(.+\) SELECT   pn.nspname AS proc_schema, .+ FROM pg_proc p.+AND p.pronamespace = \$1::regnamespace"
+        root_procs_regx = r".+: WITH.+base_types AS.+pn\.nspname AS proc_schema.+FROM pg_proc p.+p\.pronamespace = \$1::regnamespace"
         root_descr_regx = r".+: SELECT pg_catalog\.obj_description\(\$1::regnamespace, 'pg_namespace'\)"
         set_config_regx = (
             r".+: select set_config\('search_path', \$1, true\), set_config\("

test/observability/Main.hs

@@ -48,7 +48,7 @@ main = do
   actualPgVersion <- either (panic . show) id <$> P.use pool queryPgVersion
 
   -- cached schema cache so most tests run fast
-  baseSchemaCache <- loadSCache pool testCfg
+  baseSchemaCache <- loadSCache pool actualPgVersion testCfg
 
   let
     initApp sCache config = do
@@ -72,5 +72,5 @@ main = do
       describe "Feature.SchemaCacheSpec" Observation.SchemaCacheSpec.spec
 
   where
-    loadSCache pool conf =
+    loadSCache pool pgVersion conf =
-      either (panic.show) fst <$> P.use pool (HT.transaction HT.ReadCommitted HT.Read $ querySchemaCache conf)
+      either (panic.show) fst <$> P.use pool (HT.transaction HT.ReadCommitted HT.Read $ querySchemaCache pgVersion conf)

test/spec/Main.hs

@@ -88,7 +88,7 @@ main = do
   actualPgVersion <- either (panic . show) id <$> P.use pool queryPgVersion
 
   -- cached schema cache so most tests run fast
-  baseSchemaCache <- loadSCache pool baseCfg
+  baseSchemaCache <- loadSCache pool actualPgVersion baseCfg
   metricsState <- Metrics.init (configDbPoolSize baseCfg)
 
   let
@@ -105,7 +105,7 @@ main = do
 
     -- For tests that run with a different SchemaCache (depends on configSchemas)
     appDbs config = do
-      customSchemaCache <- loadSCache pool config
+      customSchemaCache <- loadSCache pool actualPgVersion config
       initApp customSchemaCache config
 
     withConfig config = before (app config)
@@ -187,5 +187,5 @@ main = do
     describe "Feature.Query.PgSafeUpdateSpec.spec" $ Feature.Query.PgSafeUpdateSpec.spec withConfig
 
   where
-    loadSCache pool conf =
+    loadSCache pool pgVersion conf =
-      either (panic.show) fst <$> P.use pool (HT.transaction HT.ReadCommitted HT.Read $ querySchemaCache conf)
+      either (panic.show) fst <$> P.use pool (HT.transaction HT.ReadCommitted HT.Read $ querySchemaCache pgVersion conf)