Skip to content

Commit

Permalink
Chef - Use temporary folder for stateful_shell artifacts (#19362)
Browse files Browse the repository at this point in the history
  • Loading branch information
cpagravel authored Jun 10, 2022
1 parent 96f83a0 commit 8b16650
Showing 1 changed file with 67 additions and 73 deletions.
140 changes: 67 additions & 73 deletions examples/chef/stateful_shell.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
import shlex
import subprocess
import sys
import tempfile
import time
from typing import Dict, Optional

Expand Down Expand Up @@ -49,11 +50,6 @@ def __init__(self) -> None:
self.env: Dict[str, str] = os.environ.copy()
self.cwd: str = self.env["PWD"]

# This file holds the env after running a command. This is a better approach
# than writing to stdout because commands could redirect the stdout.
self._envfile_path: str = os.path.join(_HERE, _ENV_FILENAME)
self._cmd_output_path: str = os.path.join(_HERE, _OUTPUT_FILENAME)

def print_env(self) -> None:
"""Print environment variables in commandline friendly format for export.
Expand Down Expand Up @@ -87,72 +83,70 @@ def run_cmd(
Returns:
Output of command if return_cmd_output set to True.
"""
env_dict = {}
# Set OLDPWD at beginning because opening the shell clears this. This handles 'cd -'.
# env -0 prints the env variables separated by null characters for easy parsing.

if return_cmd_output:
# Piping won't work here because piping will affect how environment variables
# are propagated. This solution uses tee without piping to preserve env variables.
redirect = f" > >(tee \"{self._cmd_output_path}\") 2>&1 " # include stderr

# Delete the file before running the command so we can later check if the file
# exists as a signal that tee has finished writing to the file.
if os.path.isfile(self._cmd_output_path):
os.remove(self._cmd_output_path)
else:
redirect = ""

# TODO: Use env -0 when `macos-latest` refers to macos-12 in github actions.
# env -0 is ideal because it will support cases where an env variable that has newline
# characters. The flag "-0" is requires MacOS 12 which is still in beta in Github Actions.
# The less ideal `env` command is used by itself, with the caveat that newline chars
# are unsupported in env variables.
save_env_cmd = f"env > {self._envfile_path}"

command_with_state = (
f"OLDPWD={self.env.get('OLDPWD', '')}; {cmd} {redirect}; RETCODE=$?; "
f"{save_env_cmd}; exit $RETCODE")
with subprocess.Popen(
[command_with_state],
env=self.env, cwd=self.cwd,
shell=True, executable=self.shell_app
) as proc:
returncode = proc.wait()

# Load env state from envfile.
with open(self._envfile_path, encoding="latin1") as f:
# TODO: Split on null char after updating to env -0 - requires MacOS 12.
env_entries = f.read().split("\n")
for entry in env_entries:
parts = entry.split("=")
# Handle case where an env variable contains text with '='.
env_dict[parts[0]] = "=".join(parts[1:])
self.env = env_dict
self.cwd = self.env["PWD"]

if raise_on_returncode and returncode != 0:
raise RuntimeError(
"Error. Nonzero return code."
f"\nReturncode: {returncode}"
f"\nCmd: {cmd}")

if return_cmd_output:
# Poll for file due to give 'tee' time to close.
# This is necessary because 'tee' waits for all subshells to finish before writing.
start_time = time.time()
while time.time() - start_time < _TEE_WAIT_TIMEOUT:
try:
with open(self._cmd_output_path, encoding="latin1") as f:
output = f.read()
break
except FileNotFoundError:
pass
time.sleep(0.1)
with tempfile.TemporaryDirectory(dir=os.path.dirname(_HERE)) as temp_dir:
envfile_path: str = os.path.join(temp_dir, _ENV_FILENAME)
cmd_output_path: str = os.path.join(temp_dir, _OUTPUT_FILENAME)

env_dict = {}
# Set OLDPWD at beginning because opening the shell clears this. This handles 'cd -'.
# env -0 prints the env variables separated by null characters for easy parsing.

if return_cmd_output:
# Piping won't work here because piping will affect how environment variables
# are propagated. This solution uses tee without piping to preserve env variables.
redirect = f" > >(tee \"{cmd_output_path}\") 2>&1 " # include stderr
else:
raise TimeoutError(
f"Error. Output file: {self._cmd_output_path} not created within "
f"the alloted time of: {_TEE_WAIT_TIMEOUT}s"
)

return output
redirect = ""

# TODO: Use env -0 when `macos-latest` refers to macos-12 in github actions.
# env -0 is ideal because it will support cases where an env variable that has newline
# characters. The flag "-0" is requires MacOS 12 which is still in beta in Github Actions.
# The less ideal `env` command is used by itself, with the caveat that newline chars
# are unsupported in env variables.
save_env_cmd = f"env > {envfile_path}"

command_with_state = (
f"OLDPWD={self.env.get('OLDPWD', '')}; {cmd} {redirect}; RETCODE=$?; "
f"{save_env_cmd}; exit $RETCODE")
with subprocess.Popen(
[command_with_state],
env=self.env, cwd=self.cwd,
shell=True, executable=self.shell_app
) as proc:
returncode = proc.wait()

# Load env state from envfile.
with open(envfile_path, encoding="latin1") as f:
# TODO: Split on null char after updating to env -0 - requires MacOS 12.
env_entries = f.read().split("\n")
for entry in env_entries:
parts = entry.split("=")
# Handle case where an env variable contains text with '='.
env_dict[parts[0]] = "=".join(parts[1:])
self.env = env_dict
self.cwd = self.env["PWD"]

if raise_on_returncode and returncode != 0:
raise RuntimeError(
"Error. Nonzero return code."
f"\nReturncode: {returncode}"
f"\nCmd: {cmd}")

if return_cmd_output:
# Poll for file due to give 'tee' time to close.
# This is necessary because 'tee' waits for all subshells to finish before writing.
start_time = time.time()
while time.time() - start_time < _TEE_WAIT_TIMEOUT:
if os.path.isfile(cmd_output_path):
with open(cmd_output_path, encoding="latin1") as f:
output = f.read()
if output: # Ensure that file has been written to.
break
time.sleep(0.1)
else:
raise TimeoutError(
f"Error. Output file: {cmd_output_path} not created within "
f"the alloted time of: {_TEE_WAIT_TIMEOUT}s"
)

return output

0 comments on commit 8b16650

Please sign in to comment.