You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
110 lines
4.0 KiB
110 lines
4.0 KiB
import getpass |
|
import logging |
|
import os |
|
from builtins import bytes |
|
from subprocess import PIPE, STDOUT, Popen |
|
from tempfile import NamedTemporaryFile, gettempdir |
|
|
|
from airflow.exceptions import AirflowException |
|
from airflow.models import BaseOperator |
|
from airflow.utils.decorators import apply_defaults |
|
from airflow.utils.file import TemporaryDirectory |
|
|
|
|
|
class SudoBashOperator(BaseOperator): |
|
""" |
|
Execute a Bash script, command or set of commands but sudo's as another user to execute them. |
|
|
|
:param bash_command: The command, set of commands or reference to a |
|
bash script (must be '.sh') to be executed. |
|
:type bash_command: string |
|
:param user: The user to run the command as. The Airflow worker user |
|
must have permission to sudo as that user |
|
:type user: string |
|
:param env: If env is not None, it must be a mapping that defines the |
|
environment variables for the new process; these are used instead |
|
of inheriting the current process environment, which is the default |
|
behavior. |
|
:type env: dict |
|
:type output_encoding: output encoding of bash command |
|
""" |
|
|
|
template_fields = ("bash_command", "user", "env", "output_encoding") |
|
template_ext = ( |
|
".sh", |
|
".bash", |
|
) |
|
ui_color = "#f0ede4" |
|
|
|
@apply_defaults |
|
def __init__( |
|
self, |
|
bash_command, |
|
user, |
|
xcom_push=False, |
|
env=None, |
|
output_encoding="utf-8", |
|
*args, |
|
**kwargs, |
|
): |
|
|
|
super(SudoBashOperator, self).__init__(*args, **kwargs) |
|
self.bash_command = bash_command |
|
self.user = user |
|
self.env = env |
|
self.xcom_push_flag = xcom_push |
|
self.output_encoding = output_encoding |
|
|
|
def execute(self, context): |
|
""" |
|
Execute the bash command in a temporary directory which will be cleaned afterwards. |
|
""" |
|
logging.info("tmp dir root location: \n" + gettempdir()) |
|
with TemporaryDirectory(prefix="airflowtmp") as tmp_dir: |
|
os.chmod(tmp_dir, 777) |
|
# Ensure the sudo user has perms to their current working directory for making tempfiles |
|
# This is not really a security flaw because the only thing in that dir is the |
|
# temp script, owned by the airflow user and any temp files made by the sudo user |
|
# and all of those will be created with the owning user's umask |
|
# If a process needs finer control over the tempfiles it creates, that process can chmod |
|
# them as they are created. |
|
with NamedTemporaryFile(dir=tmp_dir, prefix=self.task_id) as f: |
|
|
|
if self.user == getpass.getuser(): # don't try to sudo as yourself |
|
f.write(bytes(self.bash_command, "utf_8")) |
|
else: |
|
sudo_cmd = "sudo -u {} sh -c '{}'".format( |
|
self.user, self.bash_command |
|
) |
|
f.write(bytes(sudo_cmd, "utf_8")) |
|
f.flush() |
|
|
|
logging.info("Temporary script location: {0}".format(f.name)) |
|
logging.info("Running command: {}".format(self.bash_command)) |
|
self.sp = Popen( |
|
["bash", f.name], |
|
stdout=PIPE, |
|
stderr=STDOUT, |
|
cwd=tmp_dir, |
|
env=self.env, |
|
) |
|
|
|
logging.info("Output:") |
|
line = "" |
|
for line in iter(self.sp.stdout.readline, b""): |
|
line = line.decode(self.output_encoding).strip() |
|
logging.info(line) |
|
self.sp.wait() |
|
logging.info( |
|
"Command exited with return code {0}".format(self.sp.returncode) |
|
) |
|
|
|
if self.sp.returncode: |
|
raise AirflowException("Bash command failed") |
|
|
|
if self.xcom_push_flag: |
|
return line |
|
|
|
def on_kill(self): |
|
logging.warn("Sending SIGTERM signal to bash subprocess") |
|
self.sp.terminate()
|
|
|