Move config into its own module, pop inside utils method

Author: AlanCoding

ansible_base/lib/dynamic_config/settings_logic.py

@@ -331,4 +331,9 @@ def get_dab_settings(
                 }
             )
 
+    if 'ansible_base.task' in installed_apps:
+        # Queue used for dab_task tasks to re-schedule themselves if there is more work to do
+        # if you want to use a local queue instead, that would be more performant
+        dab_data['DAB_TASK_ADMIN_QUEUE'] = 'dab_broadcast'
+
     return dab_data

ansible_base/lib/utils/db.py

@@ -165,13 +165,14 @@ def psycopg_kwargs_from_settings_dict(settings_dict: dj_db_dict) -> dict:
 
     :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()
+    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)
-    conn_params.pop('cursor_factory')
-    conn_params.pop('context')
     return psycopg.conninfo.make_conninfo(**conn_params)
 
 
@@ -205,6 +206,5 @@ def get_pg_notify_params(alias: str = DEFAULT_DB_ALIAS, **extra_options) -> dict
 
     # 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

ansible_base/task/__init__.py

@@ -1,27 +0,0 @@
-from ansible_base.lib.utils.db import get_pg_notify_params
-
-
-def get_config():
-    psycopg_params = get_pg_notify_params()
-    psycopg_params.pop('autocommit')  # dispatcher automatically adds this, causes error, TODO: need pre-check
-    psycopg_params.pop('cursor_factory')
-    psycopg_params.pop('context')  # TODO: remove in inner method, makes non-async, not good
-
-    return {
-        "version": 2,
-        "brokers": {
-            "pg_notify": {
-                "config": psycopg_params
-                # TODO: hook in synchronous connection factory
-            }
-        },
-        "producers": {
-            "ScheduledProducer": {
-                "task_schedule": {
-                    "ansible_base.task.tasks.run_task_from_queue": {"schedule": 60},
-                    "ansible_base.task.tasks.manage_lost_tasks": {"schedule": 60*10}
-                }
-            },
-        },
-        "pool": {"max_workers": 4},  # TODO: to settings
-    }

ansible_base/task/apps.py

@@ -2,7 +2,7 @@
 
 from dispatcher.config import setup
 
-from ansible_base.task import get_config
+from ansible_base.task.config import get_config
 
 
 class TaskConfig(AppConfig):

ansible_base/task/config.py

@@ -0,0 +1,27 @@
+from django.conf import settings
+
+from ansible_base.lib.utils.db import get_pg_notify_params
+
+
+def get_config() -> dict:
+    "Returns dispatcher config from what is in Django settings"
+    psycopg_params = get_pg_notify_params()
+
+    return {
+        "version": 2,
+        "brokers": {
+            "pg_notify": {
+                "config": psycopg_params
+                # TODO: hook in synchronous connection factory
+            }
+        },
+        "producers": {
+            "ScheduledProducer": {
+                "task_schedule": {
+                    "ansible_base.task.tasks.run_task_from_queue": {"schedule": 60},
+                    "ansible_base.task.tasks.manage_lost_tasks": {"schedule": 60*10}
+                }
+            },
+        },
+        "pool": {"max_workers": 4},  # TODO: to settings
+    }

ansible_base/task/tasks.py

@@ -5,6 +5,7 @@
 from dispatcher.control import Control
 
 from django.db import transaction
+from django.conf import settings
 from django.utils.timezone import now, timedelta
 
 from ansible_base.lib.utils.db import get_pg_notify_params
@@ -14,7 +15,7 @@
 logger = logging.getLogger(__name__)
 
 
-@task(queue='dab_broadcast', bind=True)
+@task(queue=settings.DAB_TASK_ADMIN_QUEUE, bind=True)
 def run_task_from_queue(dispatcher):
     with transaction.atomic():
         task = Task.objects.filter(state=TASK_STATES.WAITING).select_for_update().first()
@@ -39,7 +40,7 @@ def run_task_from_queue(dispatcher):
     task.delete()
 
 
-@task(queue='dab_broadcast')
+@task(queue=settings.DAB_TASK_ADMIN_QUEUE)
 def manage_lost_tasks(grace_period: int = 10):
     cutoff_time = now() - timedelta(minutes=grace_period)
     for task in Task.objects.filter(state=TASK_STATES.RUNNING, started_at__lt=cutoff_time).iterator():
@@ -48,9 +49,7 @@ def manage_lost_tasks(grace_period: int = 10):
         psycopg_params.pop('cursor_factory')
         psycopg_params.pop('context')  # TODO: remove in inner method, makes non-async, not good
 
-        ctl = Control('dab_broadcast', config=psycopg_params)
-
-        # TODO: row-level lock for the task
+        ctl = Control(settings.DAB_TASK_ADMIN_QUEUE, config=psycopg_params)
 
         running_tasks = ctl.control_with_reply('running', data={'uuid': str(task.wrapper_uuid)})
 
@@ -65,8 +64,13 @@ def manage_lost_tasks(grace_period: int = 10):
 
         if not found:
             # TODO: feature of retry policy
-            logger.warning(f'Could not find task {task.name} {task.wrapper_uuid}, deleting entry')
-            task.delete()
+            try:
+                with transaction.atomic():
+                    task = Task.objects.get(uuid=task.uuid).select_for_update()
+                    logger.warning(f'Could not find task {task.name} {task.wrapper_uuid}, deleting entry')
+                    task.delete()
+            except Task.DoesNotExist:
+                logger.debug(f'task {task.name} uuid={task.uuid} already deleted, doing nothing')
         else:
             delta = now() - task.started_at
             logger.info(f'Noticed {task.name} {task.wrapper_uuid} running for {delta} seconds, seems to be fine')

test_app/tests/lib/utils/test_db.py

@@ -45,49 +45,44 @@ class TestPGNotifyConnection:
         '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())
+        params = 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'))
+        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 = self._trim_python_objects(get_pg_notify_params())
+            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 = self._trim_python_objects(get_pg_notify_params())
+            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 = self._trim_python_objects(psycopg_kwargs_from_settings_dict(test_dict))
-            assert psycopg_params == self.PSYCOPG_KWARGS
+            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"
@@ -97,7 +92,7 @@ def test_psycopg_kwargs_use(self):
         test_dict = settings.DATABASES['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))
+            psycopg_params = psycopg_kwargs_from_settings_dict(test_dict)
 
         psycopg.connect(**psycopg_params)
 

test_app/tests/task/test_basic.py

@@ -3,7 +3,7 @@
 
 import pytest
 
-from ansible_base.task import get_config
+from ansible_base.task.config import get_config
 
 from test_app.tasks import create_uuid_entry
 from test_app.models import UUIDModel

test_app/tests/task/test_multi_server.py

@@ -3,7 +3,7 @@
 
 import pytest
 
-from ansible_base.task import get_config
+from ansible_base.task.config import get_config
 
 from test_app.tasks import create_uuid_entry
 from test_app.models import UUIDModel