mirror of
https://github.com/ansible/awx.git
synced 2026-02-12 07:04:45 -03:30
56 lines
1.5 KiB
Python
56 lines
1.5 KiB
Python
import psycopg2
|
|
import select
|
|
|
|
from contextlib import contextmanager
|
|
|
|
from django.conf import settings
|
|
|
|
|
|
NOT_READY = ([], [], [])
|
|
|
|
|
|
def get_local_queuename():
|
|
return settings.CLUSTER_HOST_ID
|
|
|
|
|
|
class PubSub(object):
|
|
def __init__(self, conn):
|
|
assert conn.autocommit, "Connection must be in autocommit mode."
|
|
self.conn = conn
|
|
|
|
def listen(self, channel):
|
|
with self.conn.cursor() as cur:
|
|
cur.execute('LISTEN "%s";' % channel)
|
|
|
|
def unlisten(self, channel):
|
|
with self.conn.cursor() as cur:
|
|
cur.execute('UNLISTEN "%s";' % channel)
|
|
|
|
def notify(self, channel, payload):
|
|
with self.conn.cursor() as cur:
|
|
cur.execute('SELECT pg_notify(%s, %s);', (channel, payload))
|
|
|
|
def events(self, select_timeout=5, yield_timeouts=False):
|
|
while True:
|
|
if select.select([self.conn], [], [], select_timeout) == NOT_READY:
|
|
if yield_timeouts:
|
|
yield None
|
|
else:
|
|
self.conn.poll()
|
|
while self.conn.notifies:
|
|
yield self.conn.notifies.pop(0)
|
|
|
|
def close(self):
|
|
self.conn.close()
|
|
|
|
|
|
@contextmanager
|
|
def pg_bus_conn():
|
|
conf = settings.DATABASES['default']
|
|
conn = psycopg2.connect(dbname=conf['NAME'], host=conf['HOST'], user=conf['USER'], password=conf['PASSWORD'], port=conf['PORT'], **conf.get("OPTIONS", {}))
|
|
# Django connection.cursor().connection doesn't have autocommit=True on
|
|
conn.set_session(autocommit=True)
|
|
pubsub = PubSub(conn)
|
|
yield pubsub
|
|
conn.close()
|