def prepare_virtualenv(venv_directory: str, python_bin: str,
system_site_packages: bool,
requirements: List[str]) -> str:
"""
Creates a virtual environment and installs the additional python packages
:param venv_directory: The path for directory where the environment will be created
:type venv_directory: str
:param python_bin: Path for python binary
:type python_bin: str
:param system_site_packages: Whether to include system_site_packages in your virtualenv.
See virtualenv documentation for more information.
:type system_site_packages: bool
:param requirements: List of additional python packages
:type requirements: List[str]
:return: Path to a binary file with Python in a virtual environment.
:rtype: str
"""
virtualenv_cmd = _generate_virtualenv_cmd(venv_directory, python_bin,
system_site_packages)
execute_in_subprocess(virtualenv_cmd)
pip_cmd = _generate_pip_install_cmd(venv_directory, requirements)
if pip_cmd:
execute_in_subprocess(pip_cmd)
return '{}/bin/python'.format(venv_directory)
def execute_callable(self):
with TemporaryDirectory(prefix='venv') as tmp_dir:
if self.templates_dict:
self.op_kwargs['templates_dict'] = self.templates_dict
input_filename = os.path.join(tmp_dir, 'script.in')
output_filename = os.path.join(tmp_dir, 'script.out')
string_args_filename = os.path.join(tmp_dir, 'string_args.txt')
script_filename = os.path.join(tmp_dir, 'script.py')
prepare_virtualenv(venv_directory=tmp_dir,
python_bin=f'python{self.python_version}'
if self.python_version else None,
system_site_packages=self.system_site_packages,
requirements=self.requirements)
self._write_args(input_filename)
self._write_string_args(string_args_filename)
write_python_script(jinja_context=dict(
op_args=self.op_args,
op_kwargs=self.op_kwargs,
pickling_library=self.pickling_library.__name__,
python_callable=self.python_callable.__name__,
python_callable_source=dedent(
inspect.getsource(self.python_callable))),
filename=script_filename)
execute_in_subprocess(cmd=[
f'{tmp_dir}/bin/python', script_filename, input_filename,
output_filename, string_args_filename
])
return self._read_result(output_filename)
def test_should_print_all_messages1(self):
with self.assertLogs(log) as logs:
execute_in_subprocess(["bash", "-c", "echo CAT; echo KITTY;"])
msgs = [record.getMessage() for record in logs.records]
assert ["Executing cmd: bash -c 'echo CAT; echo KITTY;'", 'Output:', 'CAT', 'KITTY'] == msgs
def execute_callable(self):
with TemporaryDirectory(prefix='venv') as tmp_dir:
if self.templates_dict:
self.op_kwargs['templates_dict'] = self.templates_dict
# generate filenames
input_filename = os.path.join(tmp_dir, 'script.in')
output_filename = os.path.join(tmp_dir, 'script.out')
string_args_filename = os.path.join(tmp_dir, 'string_args.txt')
script_filename = os.path.join(tmp_dir, 'script.py')
# set up virtualenv
python_bin = 'python' + str(
self.python_version) if self.python_version else None
prepare_virtualenv(
venv_directory=tmp_dir,
python_bin=python_bin,
system_site_packages=self.system_site_packages,
requirements=self.requirements,
)
self._write_args(input_filename)
self._write_script(script_filename)
self._write_string_args(string_args_filename)
# execute command in virtualenv
execute_in_subprocess(
self._generate_python_cmd(tmp_dir, script_filename,
input_filename, output_filename,
string_args_filename))
return self._read_result(output_filename)
def execute(self, context):
hook = GoogleBaseHook(gcp_conn_id=self.gcp_conn_id)
self.project_id = self.project_id or hook.project_id
if not self.project_id:
raise AirflowException(
"The project id must be passed either as "
"keyword project_id parameter or as project_id extra "
"in GCP connection definition. Both are not set!")
# Write config to a temp file and set the environment variable to point to it.
# This is to avoid race conditions of reading/writing a single file
with tempfile.NamedTemporaryFile() as conf_file,\
patch_environ({KUBE_CONFIG_ENV_VAR: conf_file.name}), \
hook.provide_authorized_gcloud():
# Attempt to get/update credentials
# We call gcloud directly instead of using google-cloud-python api
# because there is no way to write kubernetes config to a file, which is
# required by KubernetesPodOperator.
# The gcloud command looks at the env variable `KUBECONFIG` for where to save
# the kubernetes config file.
cmd = [
"gcloud", "container", "clusters", "get-credentials",
self.cluster_name, "--zone", self.location, "--project",
self.project_id
]
if self.use_internal_ip:
cmd.append('--internal-ip')
execute_in_subprocess(cmd)
# Tell `KubernetesPodOperator` where the config file is located
self.config_file = os.environ[KUBE_CONFIG_ENV_VAR]
return super().execute(context)
def install_dependencies(go_module_path: str) -> None:
"""Install dependencies for a Go module.
:param go_module_path: The path to the directory containing the Go module.
:return:
"""
go_mod_tidy = ["go", "mod", "tidy"]
execute_in_subprocess(go_mod_tidy, cwd=go_module_path)
def execute_callable(self):
with TemporaryDirectory(prefix='venv') as tmp_dir:
requirements_file_name = f'{tmp_dir}/requirements.txt'
if not isinstance(self.requirements, str):
requirements_file_contents = "\n".join(
str(dependency) for dependency in self.requirements)
else:
requirements_file_contents = self.requirements
if not self.system_site_packages and self.use_dill:
requirements_file_contents += '\ndill'
with open(requirements_file_name, 'w') as file:
file.write(requirements_file_contents)
if self.templates_dict:
self.op_kwargs['templates_dict'] = self.templates_dict
input_filename = os.path.join(tmp_dir, 'script.in')
output_filename = os.path.join(tmp_dir, 'script.out')
string_args_filename = os.path.join(tmp_dir, 'string_args.txt')
script_filename = os.path.join(tmp_dir, 'script.py')
prepare_virtualenv(
venv_directory=tmp_dir,
python_bin=f'python{self.python_version}'
if self.python_version else None,
system_site_packages=self.system_site_packages,
requirements_file_path=requirements_file_name,
pip_install_options=self.pip_install_options,
)
self._write_args(input_filename)
self._write_string_args(string_args_filename)
write_python_script(
jinja_context=dict(
op_args=self.op_args,
op_kwargs=self.op_kwargs,
pickling_library=self.pickling_library.__name__,
python_callable=self.python_callable.__name__,
python_callable_source=self.get_python_source(),
),
filename=script_filename,
render_template_as_native_obj=self.dag.
render_template_as_native_obj,
)
execute_in_subprocess(cmd=[
f'{tmp_dir}/bin/python',
script_filename,
input_filename,
output_filename,
string_args_filename,
])
return self._read_result(output_filename)
def execute(self, context: 'Context') -> Optional[str]:
hook = GoogleBaseHook(gcp_conn_id=self.gcp_conn_id)
self.project_id = self.project_id or hook.project_id
if not self.project_id:
raise AirflowException(
"The project id must be passed either as "
"keyword project_id parameter or as project_id extra "
"in Google Cloud connection definition. Both are not set!")
# Write config to a temp file and set the environment variable to point to it.
# This is to avoid race conditions of reading/writing a single file
with tempfile.NamedTemporaryFile() as conf_file, patch_environ(
{KUBE_CONFIG_ENV_VAR:
conf_file.name}), hook.provide_authorized_gcloud():
# Attempt to get/update credentials
# We call gcloud directly instead of using google-cloud-python api
# because there is no way to write kubernetes config to a file, which is
# required by KubernetesPodOperator.
# The gcloud command looks at the env variable `KUBECONFIG` for where to save
# the kubernetes config file.
cmd = [
"gcloud",
"container",
"clusters",
"get-credentials",
self.cluster_name,
"--project",
self.project_id,
]
if self.impersonation_chain:
if isinstance(self.impersonation_chain, str):
impersonation_account = self.impersonation_chain
elif len(self.impersonation_chain) == 1:
impersonation_account = self.impersonation_chain[0]
else:
raise AirflowException(
"Chained list of accounts is not supported, please specify only one service account"
)
cmd.extend([
'--impersonate-service-account',
impersonation_account,
])
if self.regional:
cmd.append('--region')
else:
cmd.append('--zone')
cmd.append(self.location)
if self.use_internal_ip:
cmd.append('--internal-ip')
execute_in_subprocess(cmd)
# Tell `KubernetesPodOperator` where the config file is located
self.config_file = os.environ[KUBE_CONFIG_ENV_VAR]
return super().execute(context)
def test_should_print_all_messages1(self):
with self.assertLogs(
"airflow.utils.log.logging_mixin.LoggingMixin") as logs:
process_utils.execute_in_subprocess(
["bash", "-c", "echo CAT; echo KITTY;"])
msgs = [record.getMessage() for record in logs.records]
self.assertEqual([
"Executing cmd: bash -c 'echo CAT; echo KITTY;'", 'Output:', 'CAT',
'KITTY'
], msgs)
def init_module(go_module_name: str, go_module_path: str) -> None:
"""Initialize a Go module. If a ``go.mod`` file already exists, this function
will do nothing.
:param go_module_name: The name of the Go module to initialize.
:param go_module_path: The path to the directory containing the Go module.
:return:
"""
if os.path.isfile(os.path.join(go_module_path, "go.mod")):
return
go_mod_init_cmd = ["go", "mod", "init", go_module_name]
execute_in_subprocess(go_mod_init_cmd, cwd=go_module_path)
def prepare_virtualenv(
venv_directory: str,
python_bin: str,
system_site_packages: bool,
requirements: Optional[List[str]] = None,
requirements_file_path: Optional[str] = None,
) -> str:
"""Creates a virtual environment and installs the additional python packages.
:param venv_directory: The path for directory where the environment will be created.
:type venv_directory: str
:param python_bin: Path to the Python executable.
:type python_bin: str
:param system_site_packages: Whether to include system_site_packages in your virtualenv.
See virtualenv documentation for more information.
:type system_site_packages: bool
:param requirements: List of additional python packages.
:type requirements: List[str]
:param requirements_file_path: Path to the ``requirements.txt`` file.
:type requirements_file_path: str
:return: Path to a binary file with Python in a virtual environment.
:rtype: str
"""
virtualenv_cmd = _generate_virtualenv_cmd(venv_directory, python_bin,
system_site_packages)
execute_in_subprocess(virtualenv_cmd)
if requirements is not None and requirements_file_path is not None:
raise Exception(
"Either requirements OR requirements_file_path has to be passed, but not both"
)
pip_cmd = None
if requirements is not None and len(requirements) != 0:
pip_cmd = _generate_pip_install_cmd_from_list(venv_directory,
requirements)
if requirements_file_path is not None and requirements_file_path:
pip_cmd = _generate_pip_install_cmd_from_file(venv_directory,
requirements_file_path)
if pip_cmd:
execute_in_subprocess(pip_cmd)
return f'{venv_directory}/bin/python'
def test_should_raise_exception(self):
with self.assertRaises(CalledProcessError):
process_utils.execute_in_subprocess(["bash", "-c", "exit 1"])
