amoghrajesh commented on code in PR #64022:
URL: https://github.com/apache/airflow/pull/64022#discussion_r3206329608
##########
airflow-core/src/airflow/jobs/triggerer_job_runner.py:
##########
@@ -202,33 +202,47 @@ def _exit_gracefully(self, signum, frame) -> None:
sys.exit(os.EX_SOFTWARE)
def _execute(self) -> int | None:
- self.log.info("Starting the triggerer")
- self.register_signals()
- stats_factory = stats_utils.get_stats_factory(Stats)
- Stats.initialize(factory=stats_factory)
+ # Mark as server context for secrets backend detection when handling
GetConnection
+ # requests from the TriggerRunner subprocess (needs MetastoreBackend).
+ # The subprocess explicitly sets _AIRFLOW_PROCESS_CONTEXT=client to
prevent
+ # inheriting server privileges (runs user trigger/callback code).
+ # Similar to DagProcessorManager / DagProcessor child pattern.
+ _prev_ctx = os.environ.get("_AIRFLOW_PROCESS_CONTEXT")
+ os.environ["_AIRFLOW_PROCESS_CONTEXT"] = "server"
try:
- # Kick off runner sub-process without DB access
- self.trigger_runner = TriggerRunnerSupervisor.start(
- job=self.job,
- capacity=self.capacity,
- logger=log,
- queues=self.queues,
- )
+ self.log.info("Starting the triggerer")
+ self.register_signals()
+ stats_factory = stats_utils.get_stats_factory(Stats)
+ Stats.initialize(factory=stats_factory)
+ self.trigger_runner = None
+ try:
+ # Kick off runner sub-process without DB access
+ self.trigger_runner = TriggerRunnerSupervisor.start(
+ job=self.job,
+ capacity=self.capacity,
+ logger=log,
+ queues=self.queues,
+ )
- # Run the main DB comms loop in this process
- self.trigger_runner.run()
- return self.trigger_runner._exit_code
- except Exception:
- self.log.exception("Exception when executing
TriggerRunnerSupervisor.run")
- raise
+ # Run the main DB comms loop in this process
+ self.trigger_runner.run()
+ return self.trigger_runner._exit_code
+ except Exception:
+ self.log.exception("Exception when executing
TriggerRunnerSupervisor.run")
+ raise
+ finally:
+ self.log.info("Waiting for triggers to clean up")
+ # Tell the subprocess to stop and then wait for it.
+ # If the user interrupts/terms again, _graceful_exit will
allow them
+ # to force-kill here.
+ if self.trigger_runner is not None:
+ self.trigger_runner.kill(escalation_delay=10, force=True)
+ self.log.info("Exited trigger loop")
finally:
- self.log.info("Waiting for triggers to clean up")
- # Tell the subtproc to stop and then wait for it.
- # If the user interrupts/terms again, _graceful_exit will allow
them
- # to force-kill here.
- self.trigger_runner.kill(escalation_delay=10, force=True)
- self.log.info("Exited trigger loop")
- return None
+ if _prev_ctx is None:
+ os.environ.pop("_AIRFLOW_PROCESS_CONTEXT", None)
+ else:
+ os.environ["_AIRFLOW_PROCESS_CONTEXT"] = _prev_ctx
Review Comment:
Is there a need for the double `try`? The outer try/finally only exists to
restore `_prev_ctx`, but that restore can just go at the end of the existing
inner finally block? Something like:
It can just be:
```python
def _execute(self) -> int | None:
_prev_ctx = os.environ.get("_AIRFLOW_PROCESS_CONTEXT")
os.environ["_AIRFLOW_PROCESS_CONTEXT"] = "server"
self.log.info("Starting the triggerer")
self.register_signals()
stats_factory = stats_utils.get_stats_factory(Stats)
Stats.initialize(factory=stats_factory)
self.trigger_runner = None
try:
self.trigger_runner = TriggerRunnerSupervisor.start(
job=self.job,
capacity=self.capacity,
logger=log,
queues=self.queues,
)
self.trigger_runner.run()
return self.trigger_runner._exit_code
except Exception:
self.log.exception("Exception when executing
TriggerRunnerSupervisor.run")
raise
finally:
self.log.info("Waiting for triggers to clean up")
if self.trigger_runner is not None:
self.trigger_runner.kill(escalation_delay=10, force=True)
self.log.info("Exited trigger loop")
if _prev_ctx is None:
os.environ.pop("_AIRFLOW_PROCESS_CONTEXT", None)
else:
os.environ["_AIRFLOW_PROCESS_CONTEXT"] = _prev_ctx
```
##########
airflow-core/tests/unit/jobs/test_triggerer_job.py:
##########
@@ -1284,6 +1284,36 @@ def test_stats_initialize_called_on_execute(self,
mock_supervisor_start, stats_i
call_kwargs = stats_init_mock.call_args.kwargs
assert "factory" in call_kwargs
+ @patch.object(TriggerRunnerSupervisor, "start")
Review Comment:
Could you also add a test for the client half of the fix? The existing
test_execute_sets_server_process_context covers _execute() -> server, but
there's no symmetric test verifying that `TriggerRunner.run()` sets
`_AIRFLOW_PROCESS_CONTEXT=client`. Since thats the half which prevents user
code inheriting server privileges, it should have coverage too.
Something similar to `test_execute_sets_server_process_context`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
