用pg_notify和Python监听PostgreSQL事件的教程

621 阅读1分钟

通过PostgreSQL,我们可以很容易地从一个连接发布和监听事件到另一个连接。这很酷,因为这些通知属于一个事务。在这个例子中,我将创建一个包装器来帮助我用Python监听事件。

为了通知事件,我只需要使用pg_notify函数。比如说。

select pg_notify('channel', 'xxx')

听取事件

import psycopg2

from pg_listener import on_db_event

dsn = f"dbname='gonzalo123' user='username' host='localhost' password='password'"

conn = psycopg2.connect(dsn)
conn.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)

for payload in on_db_event(conn, 'channel'):
    print(payload)

神奇的地方在于on_db_event。我们需要传递一个psycopg2连接,以及通道名称。当有人在该通道上触发事件时,我们可以迭代该函数并检索有效载荷。

def on_db_event(conn: connection, channel: str):
    with conn:
        with conn.cursor() as cursor:
            cursor.execute(f"LISTEN {channel};")
            logger.info(f"Waiting for notifications on channel '{channel}'.")

            while True:
                if select.select([conn], [], [], 5) != ([], [], []):
                    conn.poll()
                    while conn.notifies:
                        notify = conn.notifies.pop(0)
                        yield notify.payload

由于我经常使用Django,而Django使用一个连接封装器,我需要创建一个本地psycopg2连接。也许有可能从Django的连接中获取(如果你知道怎么做,请告诉我)。

def conn_from_django(django_connection):
    db_settings = django_connection.settings_dict
    dsn = f"dbname='{db_settings['NAME']}' " \
          f"user='{db_settings['USER']}' " \
          f"host='{db_settings['HOST']}' " \
          f"password='{db_settings['PASSWORD']}'"

    conn = psycopg2.connect(dsn)
    conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)

    return conn

你可以用pip来安装这个库

pip install pglistener-gonzalo123

源代码可在我的github上找到