Source code for airflow.hooks.subprocess
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations
import contextlib
import os
import signal
from collections import namedtuple
from subprocess import PIPE, STDOUT, Popen
from tempfile import TemporaryDirectory, gettempdir
from airflow.hooks.base import BaseHook
[docs]SubprocessResult = namedtuple('SubprocessResult', ['exit_code', 'output'])
[docs]class SubprocessHook(BaseHook):
"""Hook for running processes with the ``subprocess`` module"""
def __init__(self) -> None:
self.sub_process: Popen[bytes] | None = None
super().__init__()
[docs] def run_command(
self,
command: list[str],
env: dict[str, str] | None = None,
output_encoding: str = 'utf-8',
cwd: str | None = None,
) -> SubprocessResult:
"""
Execute the command.
If ``cwd`` is None, execute the command in a temporary directory which will be cleaned afterwards.
If ``env`` is not supplied, ``os.environ`` is passed
:param command: the command to run
:param env: Optional dict containing environment variables to be made available to the shell
environment in which ``command`` will be executed. If omitted, ``os.environ`` will be used.
Note, that in case you have Sentry configured, original variables from the environment
will also be passed to the subprocess with ``SUBPROCESS_`` prefix. See
:doc:`/logging-monitoring/errors` for details.
:param output_encoding: encoding to use for decoding stdout
:param cwd: Working directory to run the command in.
If None (default), the command is run in a temporary directory.
:return: :class:`namedtuple` containing ``exit_code`` and ``output``, the last line from stderr
or stdout
"""
self.log.info('Tmp dir root location: \n %s', gettempdir())
with contextlib.ExitStack() as stack:
if cwd is None:
cwd = stack.enter_context(TemporaryDirectory(prefix='airflowtmp'))
def pre_exec():
# Restore default signal disposition and invoke setsid
for sig in ('SIGPIPE', 'SIGXFZ', 'SIGXFSZ'):
if hasattr(signal, sig):
signal.signal(getattr(signal, sig), signal.SIG_DFL)
os.setsid()
self.log.info('Running command: %s', command)
self.sub_process = Popen(
command,
stdout=PIPE,
stderr=STDOUT,
cwd=cwd,
env=env if env or env == {} else os.environ,
preexec_fn=pre_exec,
)
self.log.info('Output:')
line = ''
if self.sub_process is None:
raise RuntimeError("The subprocess should be created here and is None!")
if self.sub_process.stdout is not None:
for raw_line in iter(self.sub_process.stdout.readline, b''):
line = raw_line.decode(output_encoding, errors='backslashreplace').rstrip()
self.log.info("%s", line)
self.sub_process.wait()
self.log.info('Command exited with return code %s', self.sub_process.returncode)
return_code: int = self.sub_process.returncode
return SubprocessResult(exit_code=return_code, output=line)
[docs] def send_sigterm(self):
"""Sends SIGTERM signal to ``self.sub_process`` if one exists."""
self.log.info('Sending SIGTERM signal to process group')
if self.sub_process and hasattr(self.sub_process, 'pid'):
os.killpg(os.getpgid(self.sub_process.pid), signal.SIGTERM)