Add LISTEN/NOTIFY: schema-change notifications and pub/sub

psycodict/database.py

@@ -170,6 +170,11 @@ def __init__(self, config=None, secretsfile=None, create=False, **kwargs):
         self._nocommit_stack = 0
         self._silenced = False
         self._objects = []
+        # The connection overrides passed here take precedence over config for
+        # every connection this database opens, including the separate one a
+        # listener() opens (otherwise it could subscribe on a different server
+        # or database than the sender writes to).
+        self._connect_kwargs = dict(kwargs)
         self.conn = self._new_connection(**kwargs)
         PostgresBase.__init__(self, "db_all", self)
         if self._user == "webserver":
@@ -956,6 +961,7 @@ def create_table(
                     include_nones,
                 ],
             )
+            self._notify_schema_change(name)  # rides this transaction
         new_table = self._search_table_class_(
             self,
             name,
@@ -1028,6 +1034,7 @@ def drop_table(self, name, force=False):
                 print("Dropped {0}".format(tbl))
             self.tablenames.remove(name)
             delattr(self, name)
+            self._notify_schema_change(name)  # rides this transaction
 
     def rename_table(self, old_name, new_name):
         """
@@ -1131,6 +1138,11 @@ def rename_table(self, old_name, new_name):
             self.tablenames.append(new_name)
             self.tablenames.remove(old_name)
             self.tablenames.sort()
+            # A rename touches both names: the old one is gone, the new one
+            # appeared.  Announce both so a listener can drop the stale
+            # metadata and pick up the new table (rides this transaction).
+            self._notify_schema_change(old_name)
+            self._notify_schema_change(new_name)
 
     def copy_to(self, search_tables, data_folder, fail_on_error=True, **kwds):
         """
@@ -1629,3 +1641,80 @@ def show_slow_report(self, logfile, top=20, cutoff=None):
         """
         from .slowlog import show_slow_report
         show_slow_report(logfile, top=top, cutoff=cutoff, db=self)
+
+    # ---------------------------------------------------------------------
+    # LISTEN/NOTIFY support (see psycodict/notifications.py for the design).
+    # ---------------------------------------------------------------------
+
+    def notify(self, channel, payload=""):
+        """
+        Send a PostgreSQL notification on ``channel`` with the given payload.
+
+        The notification is sent with ``pg_notify`` on the main connection,
+        through ``_execute``, so it is *transactional*: PostgreSQL delivers it
+        when the surrounding transaction commits and drops it on rollback.
+        Called on its own (outside a ``DelayCommit``) it commits immediately and
+        so is delivered at once; called inside a ``DelayCommit`` it rides that
+        transaction and is delivered (or dropped) with it.
+
+        INPUT:
+
+        - ``channel`` -- the channel name; must be a plain identifier (letters,
+          digits and underscores, not starting with a digit)
+        - ``payload`` -- a string payload (default ``""``); received verbatim by
+          listeners
+
+        Subscribe with :meth:`listener` (or the standalone
+        :class:`~psycodict.notifications.NotificationListener`).
+        """
+        from .notifications import validate_channel_name
+        validate_channel_name(channel)
+        self._execute(SQL("SELECT pg_notify(%s, %s)"), [channel, payload])
+
+    def _notify_schema_change(self, tablename):
+        """
+        Announce that ``tablename``'s schema changed, on the schema channel.
+
+        Used by the schema-changing operations (create/drop/rename table,
+        add/drop column, reload swap).  Because it goes through :meth:`notify`
+        on the main connection, the announcement is part of the same
+        transaction as the change itself.
+        """
+        from .notifications import SCHEMA_CHANNEL
+        self.notify(SCHEMA_CHANNEL, tablename)
+
+    def listener(self, channels=None):
+        """
+        Return a :class:`~psycodict.notifications.NotificationListener`.
+
+        The listener opens its *own* ``autocommit`` connection from this
+        database's configuration and ``LISTEN``s on ``channels`` (default: the
+        schema channel ``"psycodict_schema"`` alone).  It is pull-based: call
+        ``poll(timeout)`` for a bounded batch, or iterate ``listen()``; use it
+        as a context manager to close the connection when done.
+
+        The intended follow-up use is a long-running website process that keeps
+        a listener on ``"psycodict_schema"`` and, whenever a table name arrives,
+        refreshes that table's cached metadata so newly created columns and
+        reloaded tables become visible without a restart.  That refresh
+        mechanism is proposed in a separate PR; psycodict ships the notification
+        plumbing here without depending on it, so the two can land in either
+        order.
+
+        Under a pre-forking web server each worker must build its own listener
+        after the fork, and a server in recovery (a hot standby) refuses
+        ``LISTEN`` outright; see the *Forking* and *Hot standbys* sections of
+        :mod:`psycodict.notifications`.
+
+        INPUT:
+
+        - ``channels`` -- a channel name or iterable of them; ``None`` (default)
+          means the schema channel only
+        """
+        from .notifications import NotificationListener, SCHEMA_CHANNEL
+        if channels is None:
+            channels = (SCHEMA_CHANNEL,)
+        # Pass the same connection overrides the database itself was opened
+        # with, so the listener's dedicated connection reaches the same server
+        # and database as the sender.
+        return NotificationListener(self.config, channels, **self._connect_kwargs)

psycodict/notifications.py

@@ -0,0 +1,246 @@
+# -*- coding: utf-8 -*-
+"""
+LISTEN/NOTIFY support for psycodict: schema-change notifications and a small
+general-purpose publish/subscribe primitive built on PostgreSQL's asynchronous
+notification mechanism (``LISTEN``/``NOTIFY``/``pg_notify``).
+
+Design
+------
+
+**Emission is transactional.**  Notifications are sent with
+``SELECT pg_notify(channel, payload)`` executed on the database's *main*
+connection, through the same ``_execute`` path (and therefore the same
+transaction and ``DelayCommit`` bookkeeping) as every other statement.  This is
+deliberate: PostgreSQL delivers a notification only when the transaction that
+sent it commits, and drops it if that transaction rolls back.  A schema change
+and its notification are thus atomic -- a listener never hears about a
+``create_table`` that was rolled back, and always hears about one that
+committed.  ``PostgresDatabase.notify`` and the internal schema hooks both rely
+on this; when they run inside a ``DelayCommit`` block the notification rides the
+surrounding transaction, and when they run standalone ``_execute`` commits
+immediately (delivering at once).
+
+**Listening uses a dedicated connection.**  A listener must *not* reuse the main
+connection: that connection is busy running the application's transactions and
+buffered server-side cursors, and a connection sitting in a transaction does not
+see notifications committed by others until it ends that transaction.
+:class:`NotificationListener` therefore opens its *own* psycopg connection from
+the same :class:`~psycodict.config.Configuration` options, in ``autocommit``
+mode, and issues ``LISTEN`` on it.  It is pull-based (synchronous): call
+:meth:`~NotificationListener.poll` to collect whatever has arrived within a
+bounded window, or iterate :meth:`~NotificationListener.listen`.  There are no
+background threads and no thread-delivered callbacks by design -- a caller (such
+as the LMFDB website) can wrap the pull loop in whatever concurrency model it
+prefers.
+
+**The schema channel contract.**  psycodict's schema-changing operations
+(:meth:`~psycodict.database.PostgresDatabase.create_table`, ``drop_table``,
+``rename_table``, :meth:`~psycodict.table.PostgresTable.add_column`,
+``drop_column`` and the reload swap ``reload_final_swap``) each emit on the
+single channel named by :data:`SCHEMA_CHANNEL` (``"psycodict_schema"``).  The
+payload is *just the affected table's name*, as a plain string -- nothing else.
+Keeping the payload to a bare table name keeps the contract simple and easy to
+consume; a richer (e.g. JSON) payload describing exactly what changed is a
+possible future extension, and is intentionally not attempted here.  A rename
+emits twice, once for the old name and once for the new one, so that a listener
+can drop the stale metadata and pick up the new table.
+
+Reconnection
+------------
+
+Reconnection is intentionally *not* handled automatically.  If the listening
+connection drops, :meth:`~NotificationListener.poll` / ``listen`` raise; a
+long-lived listener should catch the error and build a fresh listener (any
+notifications sent while it was disconnected are, per PostgreSQL semantics,
+lost -- ``LISTEN`` only receives notifications sent after it was issued).
+Keeping v1 free of silent auto-reconnect magic makes that loss visible to the
+caller rather than hiding it.
+
+Forking
+-------
+
+A listener, like any libpq connection, must not be used across ``fork()``:
+parent and child would share one socket, so each would receive an
+unpredictable subset of the notification stream, and an explicit ``close()``
+in either process sends a protocol Terminate message over the shared socket,
+killing the other process's subscription as well.  Create each listener in
+the process that will poll it -- under a pre-forking web server (e.g.
+``gunicorn --preload``) that means each worker builds its own listener after
+the fork, for example lazily on first use, with ``os.getpid()`` recorded at
+creation time to detect an inherited one.  A process that does find itself
+holding a listener from before a fork should simply abandon the object: drop
+the reference *without* calling ``close``.  That is safe -- psycopg
+deliberately skips the protocol shutdown when a connection object is
+garbage-collected in a process other than the one that created it, precisely
+to protect the parent's copy.
+
+Hot standbys
+------------
+
+Notifications do not traverse replication.  ``NOTIFY`` is not WAL-logged, so
+nothing is delivered on physical replicas (nor by logical replication, which
+publishes only data changes); moreover a server in recovery refuses the
+subscription itself -- ``LISTEN`` raises ``cannot execute LISTEN during
+recovery`` (SQLSTATE ``25006``), so building a listener against a hot
+standby fails outright.  Treat that error as permanent for the server rather
+than retrying: a process whose queries go to a standby must fall back to
+refreshing on a schedule or on error, or subscribe on the primary (bearing
+in mind that a notification can then arrive before the corresponding change
+has replayed on the standby).
+"""
+import re
+
+import psycopg
+from psycopg.sql import SQL, Identifier
+
+# The single channel on which psycodict announces schema changes.  The payload
+# is the affected table's name.  See the module docstring for the contract.
+SCHEMA_CHANNEL = "psycodict_schema"
+
+# A NOTIFY/LISTEN channel is an SQL identifier.  We restrict the channels that
+# psycodict will send or listen on to plain unquoted-identifier spellings
+# (a leading letter or underscore, then letters, digits or underscores) so that
+# a channel name can never smuggle in quoting or other surprises, and so that
+# the name a sender uses always matches the name a listener uses.
+_CHANNEL_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_]*\Z")
+
+
+def validate_channel_name(channel):
+    """
+    Check that ``channel`` is a plain identifier and return it, else raise.
+
+    INPUT:
+
+    - ``channel`` -- the candidate channel name
+
+    OUTPUT: ``channel`` unchanged, if it is a non-empty string of letters,
+    digits and underscores that does not start with a digit.
+    """
+    if not isinstance(channel, str) or not _CHANNEL_RE.match(channel):
+        raise ValueError(
+            "invalid notification channel %r: a channel must be a plain "
+            "identifier (letters, digits and underscores, not starting with a "
+            "digit)" % (channel,)
+        )
+    return channel
+
+
+class NotificationListener:
+    """
+    A pull-based subscriber to one or more PostgreSQL notification channels.
+
+    The listener owns a dedicated ``autocommit`` connection (separate from the
+    database's main connection) built from the same configuration options, and
+    issues ``LISTEN`` on each requested channel.  Retrieve notifications with
+    :meth:`poll` (bounded, returns a list) or by iterating :meth:`listen`.
+
+    It is a context manager; leaving the ``with`` block closes the connection::
+
+        with db.listener() as listener:
+            for channel, payload in listener.listen(timeout=60):
+                ...
+
+    A listener is bound to the process that created it and cannot subscribe
+    on a server in recovery; see *Forking* and *Hot standbys* in the module
+    docstring before using one in a forking application or against a replica.
+
+    INPUT:
+
+    - ``config`` -- a :class:`~psycodict.config.Configuration`; the
+      ``postgresql`` options are used to open the dedicated connection
+    - ``channels`` -- a channel name, or an iterable of them (default:
+      ``("psycodict_schema",)``); each is validated and ``LISTEN``ed on
+    - ``**connect_kwargs`` -- extra keyword arguments passed on to
+      ``psycopg.connect`` (e.g. keepalive settings), overriding the
+      configuration where they overlap
+    """
+
+    def __init__(self, config, channels=(SCHEMA_CHANNEL,), **connect_kwargs):
+        if isinstance(channels, str):
+            channels = (channels,)
+        self.channels = tuple(validate_channel_name(c) for c in channels)
+        if not self.channels:
+            raise ValueError("a NotificationListener needs at least one channel")
+
+        # Mirror PostgresDatabase._new_connection: start from the configured
+        # postgresql options, let explicit keyword arguments override them.
+        options = dict(config.options["postgresql"])
+        options.update(connect_kwargs)
+        # autocommit is essential: a connection sitting in an open transaction
+        # does not receive notifications committed by other sessions.
+        self._conn = psycopg.connect(**options, autocommit=True)
+        try:
+            for channel in self.channels:
+                # The channel has been validated as a plain identifier;
+                # Identifier still quotes it, so LISTEN is injection-safe.
+                self._conn.execute(SQL("LISTEN {0}").format(Identifier(channel)))
+        except Exception:
+            self._conn.close()
+            self._conn = None
+            raise
+
+    def _require_open(self):
+        if self._conn is None or self._conn.closed:
+            raise RuntimeError(
+                "this NotificationListener's connection is closed; create a "
+                "new listener (auto-reconnect is intentionally not provided)"
+            )
+        return self._conn
+
+    def poll(self, timeout=0.0):
+        """
+        Return the notifications available within ``timeout`` seconds.
+
+        Waits up to ``timeout`` seconds for the *first* notification, then
+        returns it together with any others already buffered, without blocking
+        further.  With ``timeout <= 0`` (the default) it does not wait at all,
+        returning only what is already buffered.  A timed-out poll that saw
+        nothing returns an empty list.
+
+        OUTPUT: a list of ``(channel, payload)`` pairs, in arrival order.
+        """
+        conn = self._require_open()
+        # Drain what is already buffered without blocking.
+        result = [(n.channel, n.payload) for n in conn.notifies(timeout=0)]
+        if not result and timeout and timeout > 0:
+            # Nothing buffered yet: block up to `timeout` for the first one
+            # (stop_after=1 returns as soon as it arrives rather than always
+            # waiting the whole window), then drain any that came with it.
+            result = [
+                (n.channel, n.payload)
+                for n in conn.notifies(timeout=timeout, stop_after=1)
+            ]
+            if result:
+                result += [(n.channel, n.payload) for n in conn.notifies(timeout=0)]
+        return result
+
+    def listen(self, timeout=None):
+        """
+        Yield ``(channel, payload)`` pairs as notifications arrive.
+
+        With ``timeout=None`` (the default) this blocks indefinitely, yielding
+        each notification as it is received (until the connection is closed).
+        With a numeric ``timeout`` the iterator stops after that many seconds,
+        whether or not anything arrived.
+
+        This is a thin wrapper over psycopg's own notification generator; break
+        out of the loop (or close the listener) to stop early.
+        """
+        conn = self._require_open()
+        for n in conn.notifies(timeout=timeout):
+            yield (n.channel, n.payload)
+
+    def close(self):
+        """
+        Close the dedicated connection.  Idempotent.
+        """
+        if self._conn is not None and not self._conn.closed:
+            self._conn.close()
+        self._conn = None
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+        return False

psycodict/table.py

@@ -2023,6 +2023,7 @@ def reload_final_swap(self, tables=None, metafile=None, ordered=False, sep="|"):
             # total, so recount and store it before the reinitialization
             # below reads meta_tables.
             self.stats._set_total(self.stats._slow_count({}, record=False))
+            self._db._notify_schema_change(self.search_table)  # rides this transaction
 
         # Reinitialize object
         tabledata = self._execute(
@@ -2458,6 +2459,7 @@ def add_column(self, name, datatype, description=None, label=False, force_descri
                 if label:
                     self.set_label(name)
                 self.column_description(name, description)
+                self._db._notify_schema_change(self.search_table)  # rides this transaction
             aborted = False
         finally:
             self._log_db_change("add_column", logid=logid, aborted=aborted, name=name, datatype=datatype)
@@ -2506,6 +2508,7 @@ def drop_column(self, name, force=False):
                 modifier = SQL("ALTER TABLE {0} DROP COLUMN {1}").format(Identifier(table), Identifier(name))
                 self._execute(modifier)
                 self.col_type.pop(name, None)
+                self._db._notify_schema_change(self.search_table)  # rides this transaction
             print("Column %s dropped" % (name))
             aborted = False
         finally: