Merge pull request #162 from supabase/feat/pg-meta-cache

feat(server): cache pg-meta connections

Author: soedirgo

package-lock.json

@@ -57,6 +57,7 @@
   },
   "dependencies": {
     "@sinclair/typebox": "^0.25.1",
+    "lru-cache": "^7.18.1",
     "pg": "^8.7.1",
     "pg-format": "^1.0.4",
     "pgsql-parser": "^13.3.0",
@@ -70,6 +71,7 @@
     "@fastify/type-provider-typebox": "^2.4.0",
     "@types/crypto-js": "^4.1.1",
     "@types/jest": "^29.2.4",
+    "@types/lru-cache": "^7.10.10",
     "@types/node": "^16.18.3",
     "@types/pg": "^8.6.5",
     "@types/pg-format": "^1.0.1",

package.json

@@ -25,5 +25,6 @@ export const GENERATE_TYPES =
 export const GENERATE_TYPES_INCLUDED_SCHEMAS =
   GENERATE_TYPES && process.argv[5] === '--include-schemas' ? process.argv[6]?.split(',') ?? [] : []
 
+export const CONNECTIONS_CACHE_SIZE = Number(process.env.PG_META_CONNECTIONS_CACHE_SIZE || '100')
 export const DEFAULT_POOL_CONFIG = { max: 1, connectionTimeoutMillis: PG_CONN_TIMEOUT_SECS * 1000 }
 export const PG_META_REQ_HEADER = process.env.PG_META_REQ_HEADER || 'request-id'

src/server/constants.ts

@@ -0,0 +1,29 @@
+import LruCache from 'lru-cache'
+import PostgresMeta from '../lib/PostgresMeta.js'
+import { CONNECTIONS_CACHE_SIZE, DEFAULT_POOL_CONFIG } from './constants.js'
+
+const cache = new LruCache<string, PostgresMeta>({
+  max: CONNECTIONS_CACHE_SIZE,
+  dispose: async (value) => {
+    await value.end()
+  },
+})
+
+/**
+ * Retrieve a PostgresMeta handle from the cache. If it doesn't exist, create
+ * the handle and insert it into the cache.
+ */
+const get = (connectionString: string): PostgresMeta => {
+  let pgMeta: PostgresMeta | undefined = cache.get(connectionString)
+
+  if (pgMeta === undefined) {
+    pgMeta = new PostgresMeta({ ...DEFAULT_POOL_CONFIG, connectionString })
+    cache.set(connectionString, pgMeta)
+  }
+
+  return pgMeta
+}
+
+export default {
+  get,
+}

src/server/pgMetaCache.ts

@@ -1,6 +1,5 @@
 import { FastifyInstance } from 'fastify'
-import { PostgresMeta } from '../../lib/index.js'
-import { DEFAULT_POOL_CONFIG } from '../constants.js'
+import PgMetaCache from '../pgMetaCache.js'
 import { extractRequestForLogging } from '../utils.js'
 
 export default async (fastify: FastifyInstance) => {
@@ -22,15 +21,14 @@ export default async (fastify: FastifyInstance) => {
     const limit = request.query.limit
     const offset = request.query.offset
 
-    const pgMeta = new PostgresMeta({ ...DEFAULT_POOL_CONFIG, connectionString })
+    const pgMeta = PgMetaCache.get(connectionString)
     const { data, error } = await pgMeta.columns.list({
       includeSystemSchemas,
       includedSchemas,
       excludedSchemas,
       limit,
       offset,
     })
-    await pgMeta.end()
     if (error) {
       request.log.error({ error, request: extractRequestForLogging(request) })
       reply.code(500)
@@ -60,14 +58,13 @@ export default async (fastify: FastifyInstance) => {
       } = request
       const includeSystemSchemas = request.query.include_system_schemas === 'true'
 
-      const pgMeta: PostgresMeta = new PostgresMeta({ ...DEFAULT_POOL_CONFIG, connectionString })
+      const pgMeta = PgMetaCache.get(connectionString)
       const { data, error } = await pgMeta.columns.list({
         tableId: Number(tableId),
         includeSystemSchemas,
         limit: Number(limit),
         offset: Number(offset),
       })
-      await pgMeta.end()
       if (error) {
         request.log.error({ error, request: extractRequestForLogging(request) })
         reply.code(500)
@@ -82,9 +79,8 @@ export default async (fastify: FastifyInstance) => {
       } = request
       const ordinalPosition = ordinalPositionWithDot.slice(1)
 
-      const pgMeta = new PostgresMeta({ ...DEFAULT_POOL_CONFIG, connectionString })
+      const pgMeta = PgMetaCache.get(connectionString)
       const { data, error } = await pgMeta.columns.retrieve({ id: `${tableId}.${ordinalPosition}` })
-      await pgMeta.end()
       if (error) {
         request.log.error({ error, request: extractRequestForLogging(request) })
         reply.code(400)
@@ -104,9 +100,8 @@ export default async (fastify: FastifyInstance) => {
   }>('/', async (request, reply) => {
     const connectionString = request.headers.pg
 
-    const pgMeta = new PostgresMeta({ ...DEFAULT_POOL_CONFIG, connectionString })
+    const pgMeta = PgMetaCache.get(connectionString)
     const { data, error } = await pgMeta.columns.create(request.body as any)
-    await pgMeta.end()
     if (error) {
       request.log.error({ error, request: extractRequestForLogging(request) })
       reply.code(400)
@@ -126,9 +121,8 @@ export default async (fastify: FastifyInstance) => {
   }>('/:id(\\d+\\.\\d+)', async (request, reply) => {
     const connectionString = request.headers.pg
 
-    const pgMeta = new PostgresMeta({ ...DEFAULT_POOL_CONFIG, connectionString })
+    const pgMeta = PgMetaCache.get(connectionString)
     const { data, error } = await pgMeta.columns.update(request.params.id, request.body as any)
-    await pgMeta.end()
     if (error) {
       request.log.error({ error, request: extractRequestForLogging(request) })
       reply.code(400)
@@ -151,9 +145,8 @@ export default async (fastify: FastifyInstance) => {
     const connectionString = request.headers.pg
     const cascade = request.query.cascade === 'true'
 
-    const pgMeta = new PostgresMeta({ ...DEFAULT_POOL_CONFIG, connectionString })
+    const pgMeta = PgMetaCache.get(connectionString)
     const { data, error } = await pgMeta.columns.remove(request.params.id, { cascade })
-    await pgMeta.end()
     if (error) {
       request.log.error({ error, request: extractRequestForLogging(request) })
       reply.code(400)