External-Mirrors/awx
2
0
Fork 0
mirror of https://github.com/ansible/awx.git synced 2026-05-23 08:37:48 -02:30
Code Issues Packages Projects Releases Wiki Activity

[wsrelay] Port back to psycopg3

Browse Source 
Signed-off-by: Rick Elrod <rick@elrod.me>
This commit is contained in:
Rick Elrod
2023-05-30 11:31:49 -05:00
committed by John Westcott IV
parent a665d96026
commit f9bd780d62
1 changed files with 52 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
awx/main/wsrelay.py
View File

@@ -12,7 +12,7 @@ from channels.layers import get_channel_layer
from django.conf import settings from django.conf import settings
from django.apps import apps from django.apps import apps
import asyncpg import psycopg
from awx.main.analytics.broadcast_websocket import ( from awx.main.analytics.broadcast_websocket import (
RelayWebsocketStats, RelayWebsocketStats,
@@ -209,13 +209,17 @@ class WebSocketRelayManager(object):
# hostname -> ip # hostname -> ip
self.known_hosts: Dict[str, str] = dict() self.known_hosts: Dict[str, str] = dict()
async def on_ws_heartbeat(self, conn, pid, channel, payload): async def on_ws_heartbeat(self, conn):
await conn.execute("LISTEN web_ws_heartbeat")
async for notif in conn.notifies():
if notif is None:
continue
try: try:
if not payload or channel != "web_ws_heartbeat": if not notif.payload or notif.channel != "web_ws_heartbeat":
return return
try: try:
payload = json.loads(payload) payload = json.loads(notif.payload)
except json.JSONDecodeError: except json.JSONDecodeError:
logmsg = "Failed to decode message from pg_notify channel `web_ws_heartbeat`" logmsg = "Failed to decode message from pg_notify channel `web_ws_heartbeat`"
if logger.isEnabledFor(logging.DEBUG): if logger.isEnabledFor(logging.DEBUG):
@@ -282,16 +286,16 @@ class WebSocketRelayManager(object):
# Set up a pg_notify consumer for allowing web nodes to "provision" and "deprovision" themselves gracefully. # Set up a pg_notify consumer for allowing web nodes to "provision" and "deprovision" themselves gracefully.
database_conf = settings.DATABASES['default'] database_conf = settings.DATABASES['default']
async_conn = await asyncpg.connect( async_conn = await psycopg.AsyncConnection.connect(
database=database_conf['NAME'], dbname=database_conf['NAME'],
host=database_conf['HOST'], host=database_conf['HOST'],
user=database_conf['USER'], user=database_conf['USER'],
password=database_conf['PASSWORD'], password=database_conf['PASSWORD'],
port=database_conf['PORT'], port=database_conf['PORT'],
# We cannot include these because asyncpg doesn't allow all the options that psycopg does. **database_conf.get("OPTIONS", {}),
# **database_conf.get("OPTIONS", {}),
) )
await async_conn.add_listener("web_ws_heartbeat", self.on_ws_heartbeat) await async_conn.set_autocommit(True)
event_loop.create_task(self.on_ws_heartbeat(async_conn))
# Establishes a websocket connection to /websocket/relay on all API servers # Establishes a websocket connection to /websocket/relay on all API servers
while True: while True: