Fix subscribe reconnect and add onsubscribe method - fixes porsager#315

porsager · porsager · commit 5097345b8f43 · 2022-04-22T12:39:15.000+02:00
diff --git a/README.md b/README.md
@@ -607,8 +607,15 @@ CREATE PUBLICATION alltables FOR ALL TABLES
 ```js
 const sql = postgres({ publications: 'alltables' })
 
-const { unsubscribe } = await sql.subscribe('insert:events', (row, { command, relation, key, old }) =>
-  // tell about new event row over eg. websockets or do something else
+const { unsubscribe } = await sql.subscribe(
+  'insert:events', 
+  function(row, { command, relation, key, old }) => {
+    // Callback function for each row change
+    // tell about new event row over eg. websockets or do something else
+  },
+  function onsubscribe() => {
+    // Callback on initial connect and potential reconnects
+  }
 )
 ```
 
diff --git a/src/index.js b/src/index.js
@@ -42,7 +42,7 @@ export default Postgres
 
 function Postgres(a, b) {
   const options = parseOptions(a, b)
-      , subscribe = Subscribe(Postgres, { ...options })
+      , subscribe = options.no_subscribe || Subscribe(Postgres, { ...options })
 
   let ending = false
 
diff --git a/src/subscribe.js b/src/subscribe.js
@@ -1,59 +1,90 @@
+const noop = () => { /* noop */ }
+
 export default function Subscribe(postgres, options) {
-  const listeners = new Map()
+  const subscribers = new Map()
+      , slot = 'postgresjs_' + Math.random().toString(36).slice(2)
+      , state = {}
 
   let connection
-
-  return async function subscribe(event, fn) {
-    event = parseEvent(event)
-
-    options.max = 1
-    options.onclose = onclose
-    options.fetch_types = false
-    options.connection = {
+    , stream
+    , ended = false
+
+  const sql = subscribe.sql = postgres({
+    ...options,
+    max: 1,
+    fetch_types: false,
+    idle_timeout: null,
+    max_lifetime: null,
+    connection: {
       ...options.connection,
       replication: 'database'
-    }
+    },
+    onclose: async function() {
+      if (ended)
+        return
+      stream = null
+      state.pid = state.secret = undefined
+      !ended && connected(await init(sql, slot, options.publications))
+      subscribers.forEach(event => event.forEach(({ onsubscribe }) => onsubscribe()))
+    },
+    no_subscribe: true
+  })
 
-    let stream
-      , ended = false
+  const end = sql.end
+      , close = sql.close
 
-    const sql = postgres(options)
-        , slot = 'postgresjs_' + Math.random().toString(36).slice(2)
-        , end = sql.end
+  sql.end = async() => {
+    ended = true
+    stream && (await new Promise(r => (stream.once('end', r), stream.end())))
+    return end()
+  }
 
-    sql.end = async() => {
-      ended = true
-      stream && (await new Promise(r => (stream.once('end', r), stream.end())))
-      return end()
-    }
+  sql.close = async() => {
+    stream && (await new Promise(r => (stream.once('end', r), stream.end())))
+    return close()
+  }
+
+  return subscribe
 
-    !connection && (subscribe.sql = sql, connection = init(sql, slot, options.publications))
+  async function subscribe(event, fn, onsubscribe = noop) {
+    event = parseEvent(event)
 
-    const fns = listeners.has(event)
-      ? listeners.get(event).add(fn)
-      : listeners.set(event, new Set([fn])).get(event)
+    if (!connection)
+      connection = init(sql, slot, options.publications)
+
+    const subscriber = { fn, onsubscribe }
+    const fns = subscribers.has(event)
+      ? subscribers.get(event).add(subscriber)
+      : subscribers.set(event, new Set([subscriber])).get(event)
 
     const unsubscribe = () => {
-      fns.delete(fn)
-      fns.size === 0 && listeners.delete(event)
+      fns.delete(subscriber)
+      fns.size === 0 && subscribers.delete(event)
     }
 
-    return connection.then(x => (stream = x, { unsubscribe }))
+    return connection.then(x => {
+      connected(x)
+      onsubscribe()
+      return { unsubscribe, state, sql }
+    })
+  }
 
-    async function onclose() {
-      stream = null
-      !ended && (stream = await init(sql, slot, options.publications))
-    }
+  function connected(x) {
+    stream = x.stream
+    state.pid = x.state.pid
+    state.secret = x.state.secret
   }
 
   async function init(sql, slot, publications) {
     if (!publications)
       throw new Error('Missing publication names')
 
-    const [x] = await sql.unsafe(
+    const xs = await sql.unsafe(
       `CREATE_REPLICATION_SLOT ${ slot } TEMPORARY LOGICAL pgoutput NOEXPORT_SNAPSHOT`
     )
 
+    const [x] = xs
+
     const stream = await sql.unsafe(
       `START_REPLICATION SLOT ${ slot } LOGICAL ${
         x.consistent_point
@@ -65,12 +96,10 @@ export default function Subscribe(postgres, options) {
     }
 
     stream.on('data', data)
-    stream.on('error', (error) => {
-      console.error('Logical Replication Error - Reconnecting', error) // eslint-disable-line
-      sql.end()
-    })
+    stream.on('error', sql.close)
+    stream.on('close', sql.close)
 
-    return stream
+    return { stream, state: xs.state }
 
     function data(x) {
       if (x[0] === 0x77)
@@ -99,7 +128,7 @@ export default function Subscribe(postgres, options) {
   }
 
   function call(x, a, b) {
-    listeners.has(x) && listeners.get(x).forEach(fn => fn(a, b, x))
+    subscribers.has(x) && subscribers.get(x).forEach(({ fn }) => fn(a, b, x))
   }
 }
 
diff --git a/tests/index.js b/tests/index.js
@@ -320,12 +320,12 @@ t('Undefined values throws', async() => {
 })
 
 t('Transform undefined', async() => {
-  const sql = postgres({ transform: { undefined: null } })
+  const sql = postgres({ ...options, transform: { undefined: null } })
   return [null, (await sql`select ${ undefined } as x`)[0].x]
 })
 
 t('Transform undefined in array', async() => {
-  const sql = postgres({ transform: { undefined: null } })
+  const sql = postgres({ ...options, transform: { undefined: null } })
   return [null, (await sql`select * from (values ${ sql([undefined, undefined]) }) as x(x, y)`)[0].y]
 })
 
@@ -1777,6 +1777,47 @@ t('subscribe', { timeout: 2 }, async() => {
   ]
 })
 
+t('subscribe reconnects and calls onsubscribe', { timeout: 4 }, async() => {
+  const sql = postgres({
+    database: 'postgres_js_test',
+    publications: 'alltables',
+    fetch_types: false
+  })
+
+  await sql.unsafe('create publication alltables for all tables')
+
+  const result = []
+  let onsubscribes = 0
+
+  const { unsubscribe, sql: subscribeSql } = await sql.subscribe(
+    '*',
+    (row, { command, old }) => result.push(command, row.name || row.id, old && old.name),
+    () => onsubscribes++
+  )
+
+  await sql`
+    create table test (
+      id serial primary key,
+      name text
+    )
+  `
+
+  await sql`insert into test (name) values ('Murray')`
+  await delay(10)
+  await subscribeSql.close()
+  await delay(500)
+  await sql`delete from test`
+  await delay(10)
+  await unsubscribe()
+  return [
+    '2insert,Murray,,delete,1,',
+    onsubscribes + result.join(','),
+    await sql`drop table test`,
+    await sql`drop publication alltables`,
+    await sql.end()
+  ]
+})
+
 t('Execute', async() => {
   const result = await new Promise((resolve) => {
     const sql = postgres({ ...options, fetch_types: false, debug:(id, query) => resolve(query) })
