Add psycopg db utils for dispatcherd use

AlanCoding · AlanCoding · commit 93371e4cb82b · 2025-04-18T08:00:39.000-04:00
diff --git a/ansible_base/lib/utils/db.py b/ansible_base/lib/utils/db.py
@@ -1,7 +1,12 @@
 from contextlib import contextmanager
+from copy import deepcopy
+from typing import Union
 from zlib import crc32
 
+import psycopg
+from django.conf import settings
 from django.db import DEFAULT_DB_ALIAS, connection, connections, transaction
+from django.db.backends.postgresql.base import DatabaseWrapper as PsycopgDatabaseWrapper
 from django.db.migrations.executor import MigrationExecutor
 
 
@@ -149,3 +154,62 @@ def advisory_lock(*args, lock_session_timeout_milliseconds=0, **kwargs):
         yield True
     else:
         raise RuntimeError(f'Advisory lock not implemented for database type {connection.vendor}')
+
+
+# Django settings.DATABASES['alias'] dictionary type
+dj_db_dict = dict[str, Union[str, int]]
+
+
+def psycopg_connection_from_django(**kwargs) -> psycopg.Connection:
+    "Compatibility with dispatcher connection factory, just returns the Django connection"
+    return connection.connection
+
+
+def psycopg_kwargs_from_settings_dict(settings_dict: dj_db_dict) -> dict:
+    """Return psycopg connection creation kwargs given Django db settings info
+
+    :param dict setting_dict: DATABASES in Django settings
+    :return: kwargs that can be passed to psycopg.connect, or connection classes"""
+    psycopg_params = PsycopgDatabaseWrapper(settings_dict).get_connection_params().copy()
+    psycopg_params.pop('cursor_factory', None)
+    psycopg_params.pop('context', None)
+    return psycopg_params
+
+
+def psycopg_conn_string_from_settings_dict(settings_dict: dj_db_dict) -> str:
+    conn_params = psycopg_kwargs_from_settings_dict(settings_dict)
+    return psycopg.conninfo.make_conninfo(**conn_params)
+
+
+def combine_settings_dict(settings_dict1: dj_db_dict, settings_dict2: dj_db_dict, **extra_options) -> dj_db_dict:
+    """Given two Django settings dictionaries, combine them and return a new settings_dict"""
+    settings_dict = deepcopy(settings_dict1)
+    settings_dict['OPTIONS'] = deepcopy(settings_dict.get('OPTIONS', {}))
+
+    # These extra options are used by AWX to set application_name
+    settings_dict['OPTIONS'].update(extra_options)
+
+    # Apply overrides specifically for the listener connection
+    for k, v in settings_dict2.items():
+        if k != 'OPTIONS':
+            settings_dict[k] = v
+
+    for k, v in settings_dict2.get('OPTIONS', {}).items():
+        settings_dict['OPTIONS'][k] = v
+
+    return settings_dict
+
+
+def get_pg_notify_params(alias: str = DEFAULT_DB_ALIAS, **extra_options) -> dict:
+    pg_notify_overrides = {}
+    if hasattr(settings, 'PG_NOTIFY_DATABASES'):
+        pg_notify_overrides = settings.PG_NOTIFY_DATABASES.get(alias, {})
+    elif hasattr(settings, 'LISTENER_DATABASES'):
+        pg_notify_overrides = settings.LISTENER_DATABASES.get(alias, {})
+
+    settings_dict = combine_settings_dict(settings.DATABASES[alias], pg_notify_overrides, **extra_options)
+
+    # Reuse the Django postgres DB backend to create params for the psycopg library
+    psycopg_params = psycopg_kwargs_from_settings_dict(settings_dict)
+
+    return psycopg_params
diff --git a/test_app/tests/lib/utils/test_db.py b/test_app/tests/lib/utils/test_db.py
@@ -1,11 +1,20 @@
 import threading
 import time
 
+import psycopg
 import pytest
+from django.conf import settings
 from django.db import connection
 from django.db.utils import OperationalError
+from django.test import override_settings
 
-from ansible_base.lib.utils.db import advisory_lock, migrations_are_complete
+from ansible_base.lib.utils.db import (
+    advisory_lock,
+    get_pg_notify_params,
+    migrations_are_complete,
+    psycopg_conn_string_from_settings_dict,
+    psycopg_kwargs_from_settings_dict,
+)
 
 
 @pytest.mark.django_db
@@ -14,6 +23,110 @@ def test_migrations_are_complete():
     assert migrations_are_complete()
 
 
+class TestPGNotifyConnection:
+    TEST_DATABASE_DICT = {
+        "default": {
+            "ENGINE": "django.db.backends.postgresql",
+            "HOST": "https://foo.invalid",
+            "PORT": 55434,
+            "USER": "dab_user",
+            "PASSWORD": "dabbing",
+            "NAME": "dab_db",
+        }
+    }
+    PSYCOPG_KWARGS = {
+        'dbname': 'dab_db',
+        'client_encoding': 'UTF8',
+        # kwargs containing objects can not be compared so they will be ignored
+        # 'cursor_factory': <class 'django.db.backends.postgresql.base.Cursor'>,
+        'user': 'dab_user',
+        'password': 'dabbing',
+        'host': 'https://foo.invalid',
+        'port': 55434,
+        # 'context': <psycopg.adapt.AdaptersMap object at 0x7f537f2d9f70>,
+        'prepare_threshold': None,
+    }
+
+    @pytest.fixture
+    def mock_settings(self):
+        with override_settings(DATABASES=self.TEST_DATABASE_DICT, USE_TZ=False):
+            yield
+
+    def test_default_behavior(self, mock_settings):
+        params = get_pg_notify_params()
+        assert params == self.PSYCOPG_KWARGS
+
+    def test_pg_notify_extra_options(self, mock_settings):
+        params = get_pg_notify_params(application_name='joe_connection')
+        expected = self.PSYCOPG_KWARGS.copy()
+        expected['application_name'] = 'joe_connection'
+        assert params == expected
+
+    def test_lister_databases(self, mock_settings):
+        LISTENER_DATABASES = {"default": {"HOST": "https://foo.anotherhost.invalid"}}
+        with override_settings(LISTENER_DATABASES=LISTENER_DATABASES):
+            params = get_pg_notify_params()
+            assert params['host'] == "https://foo.anotherhost.invalid"
+
+    def test_pg_notify_databases(self, mock_settings):
+        PG_NOTIFY_DATABASES = {"default": {"HOST": "https://foo.anotherhost2.invalid"}}
+        with override_settings(PG_NOTIFY_DATABASES=PG_NOTIFY_DATABASES):
+            params = get_pg_notify_params()
+            assert params['host'] == "https://foo.anotherhost2.invalid"
+
+    def test_psycopg_kwargs_from_settings_dict(self):
+        "More of a unit test, doing the same thing"
+        test_dict = self.TEST_DATABASE_DICT["default"].copy()
+        test_dict['OPTIONS'] = {'autocommit': True}
+        with override_settings(USE_TZ=False):
+            psycopg_params = psycopg_kwargs_from_settings_dict(test_dict)
+            expected_kwargs = self.PSYCOPG_KWARGS.copy()
+            expected_kwargs['autocommit'] = True
+            assert psycopg_params == expected_kwargs
+
+    def test_psycopg_kwargs_use(self):
+        "This assures that the data we get for the kwargs are usable, and demos how to use"
+        if connection.vendor == 'sqlite':
+            pytest.skip('Test needs to connect to postgres which is not running')
+
+        test_dict = settings.DATABASES['default'].copy()
+        test_dict['OPTIONS'] = {'autocommit': True}
+        with override_settings(USE_TZ=False):
+            psycopg_params = psycopg_kwargs_from_settings_dict(test_dict)
+
+        psycopg.connect(**psycopg_params)
+
+    def test_listener_string_production(self):
+        "This is a test to correspond to PG_NOTIFY_DSN_SERVER type settings in eda-server"
+        with override_settings(USE_TZ=False):
+            args = psycopg_conn_string_from_settings_dict(
+                {
+                    "ENGINE": "django.db.backends.postgresql",
+                    "HOST": "127.0.0.1",
+                    "PORT": 5432,
+                    "USER": "postgres",
+                    "PASSWORD": "DB_PASSWORD",
+                    "NAME": "eda",
+                    "OPTIONS": {
+                        "sslmode": "allow",
+                        "sslcert": "",
+                        "sslkey": "",
+                        "sslrootcert": "",
+                    },
+                }
+            )
+        assert args == (
+            "dbname=eda sslmode=allow sslcert='' sslkey='' sslrootcert='' client_encoding=UTF8 user=postgres password=DB_PASSWORD host=127.0.0.1 port=5432"
+        )
+
+    def test_listener_string_production_use(self):
+        "This assures that the data we get for the connection string is usable, and demos how to use"
+        if connection.vendor == 'sqlite':
+            pytest.skip('Test needs to connect to postgres which is not running')
+        args = psycopg_conn_string_from_settings_dict(settings.DATABASES['default'])
+        psycopg.connect(conninfo=args)
+
+
 class TestAdvisoryLock:
     @pytest.fixture(autouse=True)
     def skip_if_sqlite(self):
@@ -36,7 +149,7 @@ def background_task(django_db_blocker):
     def test_determine_lock_is_held(self, django_db_blocker):
         thread = threading.Thread(target=TestAdvisoryLock.background_task, args=(django_db_blocker,))
         thread.start()
-        for _ in range(5):
+        for _ in range(20):
             with advisory_lock('background_task_lock', wait=False) as held:
                 if held is False:
                     break
