diff --git a/psycodict/database.py b/psycodict/database.py index 9d108df..c787dd2 100644 --- a/psycodict/database.py +++ b/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) diff --git a/psycodict/notifications.py b/psycodict/notifications.py new file mode 100644 index 0000000..d5cff74 --- /dev/null +++ b/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 diff --git a/psycodict/table.py b/psycodict/table.py index de005fc..98114e8 100644 --- a/psycodict/table.py +++ b/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: diff --git a/tests/test_notifications.py b/tests/test_notifications.py new file mode 100644 index 0000000..8bde256 --- /dev/null +++ b/tests/test_notifications.py @@ -0,0 +1,327 @@ +# -*- coding: utf-8 -*- +""" +Tests for psycodict's LISTEN/NOTIFY support (``psycodict/notifications.py``). + +These exercise the two halves of the feature: + +* *emission* -- ``db.notify`` and the schema-change hooks on + ``create_table``/``drop_table``/``add_column``/``drop_column`` and the reload + swap, all of which send ``pg_notify`` transactionally on the main connection; +* *listening* -- a :class:`~psycodict.notifications.NotificationListener` with + its own ``autocommit`` connection, collected through ``poll`` /``listen``. + +Cross-connection delivery is the whole point (the sender is the database's main +connection, the receiver a dedicated one), so every test opens a listener, +performs an operation, and then polls. A notification is delivered only to a +``LISTEN`` issued *before* it was sent and only once its transaction commits, so +listeners are always created before the operation under test. + +Like ``test_devmirror``, this module builds its own ``Configuration`` from the +standard libpq environment variables rather than leaning on the shared session +fixtures, so it is self-contained. Poll windows are kept short (waits use +``stop_after`` semantics and so return as soon as the expected notification +arrives; the only full-length waits are the deliberately-empty polls) to keep +the added suite runtime to about a second. +""" +import os +import time +import uuid + +import pytest + +from psycodict.notifications import NotificationListener, SCHEMA_CHANNEL, validate_channel_name +from psycodict.utils import DelayCommit + + +# How long a poll may wait for a notification we expect to arrive. poll() +# returns as soon as it does, so this is only an upper bound guarding against a +# hang, not the time a passing test takes. +WAIT = 2.0 +# How long a poll waits when we expect *nothing*: long enough that a mistaken +# notification would have arrived, short enough to keep the suite quick. +QUIET = 0.4 + +COLUMNS = [("n", "integer"), ("label", "text"), ("data", "jsonb")] + + +def _connection_kwargs(): + # Mirror tests/conftest.py so the module reads the same libpq environment. + return { + "host": os.environ.get("PGHOST", "localhost"), + "port": int(os.environ.get("PGPORT", 5432)), + "user": os.environ.get("PGUSER", "postgres"), + "password": os.environ.get("PGPASSWORD", ""), + "dbname": os.environ.get("PGDATABASE", "psycodict_test"), + } + + +@pytest.fixture(scope="module") +def notif_config(tmp_path_factory): + """A ``Configuration`` from the libpq environment (see ``test_devmirror``).""" + from psycodict.config import Configuration + + tmp = tmp_path_factory.mktemp("notif") + conn = _connection_kwargs() + config_file = tmp / "config.ini" + with open(config_file, "w") as F: + F.write("[logging]\nslowcutoff = 0.1\nslowlogfile = %s\n" % (tmp / "slow.log")) + F.write("[postgresql]\n") + for key, val in conn.items(): + F.write("%s = %s\n" % (key, val)) + return Configuration( + defaults={"config_file": str(config_file), "secrets_file": str(tmp / "secrets.ini")}, + readargs=False, + ) + + +@pytest.fixture(scope="module") +def notif_db(notif_config): + """A ``PostgresDatabase`` on the private test database, meta tables ensured.""" + import psycopg + + from psycodict.database import PostgresDatabase + + try: + database = PostgresDatabase(config=notif_config, create=True) + except psycopg.OperationalError as err: + conn = _connection_kwargs() + message = "no PostgreSQL server at %s:%s/%s as %s (%s)" % ( + conn["host"], conn["port"], conn["dbname"], conn["user"], + str(err).strip().split("\n")[0], + ) + if os.environ.get("PSYCODICT_TEST_DB_REQUIRED"): + raise RuntimeError("PSYCODICT_TEST_DB_REQUIRED is set but there is %s" % message) + pytest.skip(message, allow_module_level=True) + yield database + database.conn.close() + + +@pytest.fixture +def make_table(notif_db): + """Factory for uniquely named tables on ``notif_db``, dropped afterwards.""" + created = [] + + def make(): + name = "test_notif_%s" % uuid.uuid4().hex[:12] + notif_db.create_table(name, COLUMNS, label_col="label", sort=["n"]) + created.append(name) + return notif_db[name] + + yield make + + for name in reversed(created): + try: + if name in notif_db.tablenames: + notif_db.drop_table(name, force=True) + except Exception: # pragma: no cover - cleanup must not mask failures + pass + + +def _unique_name(): + return "test_notif_%s" % uuid.uuid4().hex[:12] + + +# --------------------------------------------------------------------------- +# db.notify / general pub-sub +# --------------------------------------------------------------------------- + +def test_notify_is_received_by_listener(notif_db): + channel = "psycodict_test_chan" + with notif_db.listener(channels=channel) as listener: + notif_db.notify(channel, "hello world") + got = listener.poll(WAIT) + assert got == [(channel, "hello world")] + + +def test_notify_defaults_to_empty_payload(notif_db): + channel = "psycodict_test_chan" + with notif_db.listener(channels=channel) as listener: + notif_db.notify(channel) + got = listener.poll(WAIT) + assert got == [(channel, "")] + + +def test_notify_rejects_a_bad_channel(notif_db): + for bad in ["has space", "1leading", "has-dash", "", "sneaky;DROP"]: + with pytest.raises(ValueError): + notif_db.notify(bad, "x") + + +def test_validate_channel_name_accepts_and_rejects(): + assert validate_channel_name("psycodict_schema") == "psycodict_schema" + assert validate_channel_name("_ok9") == "_ok9" + for bad in ["9no", "a b", "a.b", "a-b", "", 3]: + with pytest.raises(ValueError): + validate_channel_name(bad) + + +# --------------------------------------------------------------------------- +# schema-change emission +# --------------------------------------------------------------------------- + +def test_create_table_and_drop_table_emit_on_schema_channel(notif_db): + name = _unique_name() + with notif_db.listener() as listener: # defaults to the schema channel + notif_db.create_table(name, COLUMNS, label_col="label", sort=["n"]) + try: + created = listener.poll(WAIT) + assert created == [(SCHEMA_CHANNEL, name)] + finally: + notif_db.drop_table(name, force=True) + dropped = listener.poll(WAIT) + assert dropped == [(SCHEMA_CHANNEL, name)] + + +def test_add_column_and_drop_column_emit(notif_db, make_table): + table = make_table() + name = table.search_table + with notif_db.listener() as listener: + table.add_column("newcol", "integer") + assert listener.poll(WAIT) == [(SCHEMA_CHANNEL, name)] + table.drop_column("newcol", force=True) + assert listener.poll(WAIT) == [(SCHEMA_CHANNEL, name)] + + +def test_reload_swap_emits_exactly_once(notif_db, make_table, tmp_path): + table = make_table() + name = table.search_table + table.insert_many([{"n": i, "label": "l%d" % i, "data": {"k": i}} for i in range(10)]) + searchfile = str(tmp_path / "search.txt") + table.copy_to(searchfile) + with notif_db.listener() as listener: + notif_db[name].reload(searchfile) + got = listener.poll(WAIT) + # The reload rewrites _tmp tables and swaps them in with internal DDL, + # but only the final swap announces the table -- exactly once. + extra = listener.poll(QUIET) + assert got == [(SCHEMA_CHANNEL, name)] + assert extra == [] + + +# --------------------------------------------------------------------------- +# transactional semantics +# --------------------------------------------------------------------------- + +def test_rolled_back_transaction_delivers_no_notification(notif_db): + channel = "psycodict_test_chan" + with notif_db.listener(channels=channel) as listener: + with pytest.raises(RuntimeError): + with DelayCommit(notif_db): + notif_db.notify(channel, "doomed") + raise RuntimeError("boom") # forces DelayCommit to roll back + # The NOTIFY rode the rolled-back transaction, so nothing is delivered. + assert listener.poll(QUIET) == [] + + +def test_committed_delaycommit_delivers_once(notif_db): + channel = "psycodict_test_chan" + with notif_db.listener(channels=channel) as listener: + with DelayCommit(notif_db): + notif_db.notify(channel, "kept") + # Not delivered yet: the transaction has not committed. + assert listener.poll(QUIET) == [] + # DelayCommit committed on exit; now it arrives. + assert listener.poll(WAIT) == [(channel, "kept")] + + +# --------------------------------------------------------------------------- +# channel subscription and listener lifecycle +# --------------------------------------------------------------------------- + +def test_listener_only_receives_subscribed_channels(notif_db): + subscribed = "psycodict_test_sub" + other = "psycodict_test_other" + with notif_db.listener(channels=subscribed) as listener: + notif_db.notify(other, "ignored") + notif_db.notify(subscribed, "wanted") + got = listener.poll(WAIT) + got += listener.poll(QUIET) # give the unsubscribed one time to (not) show + assert got == [(subscribed, "wanted")] + + +def test_listener_can_subscribe_to_several_channels(notif_db): + a, b = "psycodict_test_a", "psycodict_test_b" + with notif_db.listener(channels=(a, b)) as listener: + notif_db.notify(a, "1") + notif_db.notify(b, "2") + got = listener.poll(WAIT) + got += listener.poll(QUIET) + assert sorted(got) == [(a, "1"), (b, "2")] + + +def test_poll_timeout_returns_empty_quickly(notif_db): + with notif_db.listener() as listener: + start = time.time() + got = listener.poll(QUIET) + elapsed = time.time() - start + assert got == [] + # Waited about QUIET seconds, and certainly not a whole extra second. + assert elapsed < QUIET + 1.0 + + +def test_context_manager_closes_the_connection(notif_db): + listener = notif_db.listener() + conn = listener._conn + with listener as entered: + assert entered is listener + assert not conn.closed + assert listener._conn is None + assert conn.closed + # Using a closed listener is an error, not a silent no-op. + with pytest.raises(RuntimeError): + listener.poll(0) + + +def test_close_is_idempotent(notif_db): + listener = notif_db.listener() + listener.close() + listener.close() # second close must not raise + assert listener._conn is None + + +def test_listener_requires_a_channel(notif_db): + with pytest.raises(ValueError): + notif_db.listener(channels=()) + + +def test_listen_iterator_yields_until_timeout(notif_db): + channel = "psycodict_test_iter" + with notif_db.listener(channels=channel) as listener: + notif_db.notify(channel, "a") + notif_db.notify(channel, "b") + received = [] + # A bounded timeout makes the generator stop on its own; without it the + # iterator would block forever waiting for more. + for item in listener.listen(timeout=WAIT): + received.append(item) + if len(received) == 2: + break + assert received == [(channel, "a"), (channel, "b")] + + +def test_standalone_listener_matches_the_factory(notif_db): + # db.listener() is a thin convenience over the public class. + channel = "psycodict_test_direct" + with NotificationListener(notif_db.config, channels=channel) as listener: + notif_db.notify(channel, "direct") + assert listener.poll(WAIT) == [(channel, "direct")] + + +def test_listener_honors_connection_overrides(notif_config): + # Connection overrides passed to PostgresDatabase(...) must also reach the + # listener's dedicated connection; otherwise a caller who overrode e.g. the + # database or host would have the listener subscribe somewhere other than + # where the sender writes. application_name is an observable stand-in that + # needs no second database to test. + from psycodict.database import PostgresDatabase + + marker = "psycodict_listener_marker" + db = PostgresDatabase(config=notif_config, application_name=marker) + try: + with db.listener() as listener: + got = listener._conn.execute( + "SELECT current_setting('application_name')" + ).fetchone()[0] + assert got == marker + finally: + db.conn.close()