Psycopg2, properly terminate notification listener

Mr.D 0
Viewed 16

I've implemented simple notification listener with psycopg2, which will listen trigger events from PostgreSQL table changes:

def trigger_event(self, connection) -> List[dict]:
    cursor = connection.cursor()
    cursor.execute("LISTEN mychannel")
    while True:
        select.select([connection], [], [])
        connection.poll()
        while connection.notifies:
            notify = connection.notifies.pop()
            yield notify.payload

for payload in self.trigger_event(connection):
    print(payload)

When I start this code, it works properly. However, when I try to it, it does not listen keyboard interrupts or any signal, or termination takes very long time.

Is it possible to make Psycopg2 notification listening stop without any hard process killing?

psycopg2 postgresql python
1 Answers

The documentation for select.select states, regarding the timeout argument:

The optional timeout argument specifies a time-out as a floating point number in seconds. When the timeout argument is omitted the function blocks until at least one file descriptor is ready

so providing a value for timeout should reduce the worst case wait time to the value of the timeout.

select.select([connection], [], [], 5)

To stop the listening programmatically, you could listen for a particular value in the channel:

    def trigger_event(self, connection):
        cursor = connection.cursor()
        cursor.execute('LISTEN mychannel;')
        try:
            while True:
                if select.select([connection], [], [], 5) == ([], [], []):
                    print('Timeout')
                else:
                    connection.poll()
                    while connection.notifies:
                        notify = connection.notifies.pop()
                        if notify.payload == 'TERMINATE':
                            return
                        yield notify.payload
        finally:
            cursor.execute('UNLISTEN mychannel')
            cursor.close()

or close the generator (suing the same finally block as above):

gen = self.trigger_event(connection)
for payload in gen:
    print(payload)
    if some_condition:
        gen.close()
snakecharmerb0
Related