Move psycopg param generator to DAB utility

Author: AlanCoding

ansible_base/lib/utils/db.py

@@ -1,7 +1,11 @@
 from contextlib import contextmanager
+from copy import deepcopy
+from typing import Union
 from zlib import crc32
 
+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 +153,51 @@ 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}')
+
+
+settings_dict_type = dict[str, Union[str, int]]
+
+
+def psycopg_kwargs_from_settings_dict(settings_dict: settings_dict_type) -> 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"""
+    return PsycopgDatabaseWrapper(settings_dict).get_connection_params()
+
+
+def combine_settings_dict(
+    settings_dict1: dict[str, Union[str, int]], settings_dict2: dict[str, Union[str, int]], **extra_options
+) -> dict[str, Union[str, int]]:
+    """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)
+    psycopg_params['autocommit'] = True
+
+    return psycopg_params

test_app/tests/lib/utils/test_db.py

@@ -4,8 +4,9 @@
 import pytest
 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_kwargs_from_settings_dict
 
 
 @pytest.mark.django_db
@@ -14,6 +15,73 @@ 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,
+        'autocommit': True,
+    }
+
+    @pytest.fixture
+    def mock_settings(self):
+        with override_settings(DATABASES=self.TEST_DATABASE_DICT, USE_TZ=False):
+            yield
+
+    def _trim_python_objects(self, psycopg_params):
+        # These remove those commented-out kwargs in PSYCOPG_KWARGS
+        psycopg_params.pop('cursor_factory')
+        psycopg_params.pop('context')
+        return psycopg_params
+
+    def test_default_behavior(self, mock_settings):
+        params = self._trim_python_objects(get_pg_notify_params())
+        assert params == self.PSYCOPG_KWARGS
+
+    def test_pg_notify_extra_options(self, mock_settings):
+        params = self._trim_python_objects(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 = self._trim_python_objects(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 = self._trim_python_objects(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 = self._trim_python_objects(psycopg_kwargs_from_settings_dict(test_dict))
+            assert psycopg_params == self.PSYCOPG_KWARGS
+
+
 class TestAdvisoryLock:
     @pytest.fixture(autouse=True)
     def skip_if_sqlite(self):