Fixed HookConfig.atomic flag, which was ignored in fire_hook method. (#242)

Author: droserasprout

CHANGELOG.md

@@ -6,7 +6,8 @@ Please use [this](https://docs.gitlab.com/ee/development/changelog.html) documen
 
 ### Fixed
 
-* context: Fixed CallbackManager pending_hooks loop when queue becomes empty
+* context: Fixed crash when calling `fire_hook` method.
+* context: Fixed `HookConfig.atomic` flag, which was ignored in `fire_hook` method.
 * index: Fixed loading handler callbacks from nested packages (@veqtor)
 
 ### Other

src/dipdup/context.py

@@ -3,6 +3,7 @@
 import sys
 import time
 from collections import deque
+from contextlib import AsyncExitStack
 from contextlib import contextmanager
 from contextlib import suppress
 from os.path import exists
@@ -58,6 +59,7 @@
 from dipdup.utils import FormattedLogger
 from dipdup.utils import slowdown
 from dipdup.utils.database import execute_sql_scripts
+from dipdup.utils.database import in_global_transaction
 from dipdup.utils.database import wipe_schema
 
 DatasourceT = TypeVar('DatasourceT', bound=Datasource)
@@ -331,6 +333,7 @@ async def fire_handler(
             handler_config=handler_config,
             datasource=datasource,
         )
+        # NOTE: Handlers are not atomic, levels are. Do not open transaction here.
         with self._callback_wrapper('handler', name):
             await handler_config.callback_fn(new_ctx, *args, **kwargs)
 
@@ -355,7 +358,12 @@ async def fire_hook(
         self._verify_arguments(new_ctx, *args, **kwargs)
 
         async def _wrapper():
-            with self._callback_wrapper('hook', name):
+            async with AsyncExitStack() as stack:
+
+                stack.enter_context(self._callback_wrapper('hook', name))
+                if hook_config.atomic:
+                    await stack.enter_async_context(in_global_transaction())
+
                 await hook_config.callback_fn(ctx, *args, **kwargs)
 
         if wait:

src/dipdup/scheduler.py

@@ -1,5 +1,4 @@
 import json
-from contextlib import AsyncExitStack
 from datetime import datetime
 from functools import partial
 from typing import Any
@@ -16,7 +15,6 @@
 from dipdup.context import HookContext
 from dipdup.exceptions import ConfigurationError
 from dipdup.utils import FormattedLogger
-from dipdup.utils.database import in_global_transaction
 
 DEFAULT_CONFIG = {
     'apscheduler.jobstores.default.class': 'apscheduler.jobstores.memory:MemoryJobStore',
@@ -51,14 +49,10 @@ async def _job_wrapper(ctx: DipDupContext, *args, **kwargs) -> None:
             hook_config=hook_config,
         )
 
-        async with AsyncExitStack() as stack:
-            if hook_config.atomic:
-                await stack.enter_async_context(in_global_transaction())
+        await job_ctx.fire_hook(hook_config.callback, *args, **kwargs)
 
-            await job_ctx.fire_hook(hook_config.callback, *args, **kwargs)
-
-            if job_config.daemon:
-                raise ConfigurationError('Daemon jobs are intended to run forever')
+        if job_config.daemon:
+            raise ConfigurationError('Daemon jobs are intended to run forever')
 
     logger = FormattedLogger(
         name=f'dipdup.hooks.{hook_config.callback}',