Add psycopg db utils for dispatcherd use

AlanCoding · AlanCoding · commit e62b0490b583 · 2025-04-28T13:42:30.000-04:00
Add more docstrings

Test and docstring for dispatcherd factory function

Skip for sqlite3

Implement code style and commenting review comments

pick up CI fix established elsewhere
diff --git a/ansible_base/authentication/authenticator_plugins/ldap.py b/ansible_base/authentication/authenticator_plugins/ldap.py
@@ -363,7 +363,7 @@ def validate(self, attrs):
 class LDAPSettings(BaseLDAPSettings):
     def __init__(self, prefix: str = 'AUTH_LDAP_', defaults: dict = {}):
         # This init method double checks the passed defaults while initializing a settings objects
-        super(LDAPSettings, self).__init__(prefix, defaults)
+        super().__init__(prefix, defaults)
 
         # SERVER_URI needs to be a string, not an array
         setattr(self, 'SERVER_URI', ','.join(defaults['SERVER_URI']))
diff --git a/ansible_base/lib/utils/db.py b/ansible_base/lib/utils/db.py
@@ -1,8 +1,13 @@
 import logging
 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, OperationalError, connection, connections, transaction
+from django.db.backends.postgresql.base import DatabaseWrapper as PsycopgDatabaseWrapper
 from django.db.migrations.executor import MigrationExecutor
 
 logger = logging.getLogger(__name__)
@@ -164,3 +169,82 @@ 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 dispatcherd connection factory, just returns the Django connection
+
+    dispatcherd passes config info as kwargs, but in this case we just want to ignore then.
+    Because the point of this it to not reconnect, but rely on existing Django connection management.
+    """
+    if connection.connection is None:
+        connection.ensure_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:
+    """Returns a string that psycopg can take as conninfo for Connection class.
+
+    Example return value: "dbname=postgres user=postgres"
+    """
+    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 database settings dictionaries, combine them and return a new settings_dict"""
+    settings_dict = deepcopy(settings_dict1)
+
+    # Apply overrides specifically for the listener connection
+    for k, v in settings_dict2.items():
+        if k != 'OPTIONS':
+            settings_dict[k] = v
+
+    # Merge the database OPTIONS
+    # https://docs.djangoproject.com/en/5.2/ref/databases/#postgresql-connection-settings
+    # These are not expected to be nested, as they are psycopg params
+    settings_dict.setdefault('OPTIONS', {})
+    # extra_options are used by AWX to set application_name, which is generally a good idea
+    settings_dict['OPTIONS'].update(extra_options)
+    # Apply overrides from nested OPTIONS for the listener connection
+    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:
+    """Returns a dictionary that can be used as kwargs to create a psycopg.Connection
+
+    This should use the same connection parameters as Django does.
+    However, this also allows overrides specified by
+    - PG_NOTIFY_DATABASES, higher precedence, preferred setting
+    - LISTENER_DATABASES, lower precedence, deprecated AWX setting.
+    """
+    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,21 @@
 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_connection_from_django,
+    psycopg_kwargs_from_settings_dict,
+)
 
 
 @pytest.mark.django_db
@@ -21,6 +31,130 @@ def skip_if_sqlite(self):
             pytest.skip('Advisory lock is not written for sqlite')
 
 
+class TestPGNotifyConnection(SkipIfSqlite):
+    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)
+
+    @pytest.mark.django_db
+    def test_psycopg_connection_from_django_existing_conn(self):
+        if connection.vendor == 'sqlite':
+            pytest.skip('Advisory lock is not written for sqlite')
+        assert isinstance(psycopg_connection_from_django(), psycopg.Connection)
+
+    @pytest.mark.django_db(transaction=True)
+    def test_psycopg_connection_from_django_new_conn(self):
+        if connection.vendor == 'sqlite':
+            pytest.skip('Advisory lock is not written for sqlite')
+        connection.close()
+        assert isinstance(psycopg_connection_from_django(), psycopg.Connection)
+
+
+class TestAdvisoryLock:
+    @pytest.fixture(autouse=True)
+    def skip_if_sqlite(self):
+        if connection.vendor == 'sqlite':
+            pytest.skip('Advisory lock is not written for sqlite')
+
+
 class TestAdvisoryLock(SkipIfSqlite):
     THREAD_WAIT_TIME = 0.1
 
@@ -40,7 +174,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
