From fdbc529307e3c64845786dfa6b3d586f5db8ef20 Mon Sep 17 00:00:00 2001 From: John Sirois Date: Wed, 12 Oct 2022 18:32:46 -0700 Subject: [PATCH] Support injecting args and env vars in a PEX. Pex now supports setting defaults for environment variables and injecting default arguments for the entry point chosen at PEX build time. This eliminates the need for projects like https://github.com/kwlzn/pyuwsgi_pex and makes it easy to package a PEXed framework app where the main entry point is fixed in the framework and the hooks are configured to point to user supplied entry points. This removes the last known legitimate API use of Pex and paves the way for Pex 3 to be cut in good conscience with a CLI-only API that does not hinder any known use cases. Fixes #987 --- docs/recipes.rst | 34 ++ pex/bin/pex.py | 37 ++- pex/pex.py | 25 +- pex/pex_info.py | 20 ++ pex/venv/pex.py | 9 + tests/integration/test_inject_env_and_args.py | 301 ++++++++++++++++++ 6 files changed, 418 insertions(+), 8 deletions(-) create mode 100644 tests/integration/test_inject_env_and_args.py diff --git a/docs/recipes.rst b/docs/recipes.rst index 850d88c42..f87b96750 100644 --- a/docs/recipes.rst +++ b/docs/recipes.rst @@ -3,6 +3,40 @@ PEX Recipes and Notes ===================== +Uvicorn and other customizable application servers +-------------------------------------------------- + +Often you want to run a third-party application server and have it use your code. You can always do +this by writing a shim bit of python code that starts the application server configured to use your +code. It may be simpler though to use ``--inject-env`` and ``--inject-args`` to seal this +configuration into a PEX file without needing to write a shim. + +For example, to package up a uvicorn-powered server of your app coroutine in ``example.py`` that ran +on port 8888 by default you could: + +.. code-block:: bash + + $ pex "uvicorn[standard]" -c uvicorn --inject-args 'example:app --port 8888' -oexample-app.pex + $ ./example-app.pex + INFO: Started server process [2014] + INFO: Waiting for application startup. + INFO: ASGI 'lifespan' protocol appears unsupported. + INFO: Application startup complete. + INFO: Uvicorn running on http://127.0.0.1:8888 (Press CTRL+C to quit) + ^CINFO: Shutting down + INFO: Finished server process [2014] + +You could then over-ride the port with: + +.. code-block:: bash + + $ ./example-app.pex --port 0 + INFO: Started server process [2248] + INFO: Waiting for application startup. + INFO: ASGI 'lifespan' protocol appears unsupported. + INFO: Application startup complete. + INFO: Uvicorn running on http://127.0.0.1:45751 (Press CTRL+C to quit) + Long running PEX applications and daemons ----------------------------------------- diff --git a/pex/bin/pex.py b/pex/bin/pex.py index 5803b98a4..1beaba8df 100755 --- a/pex/bin/pex.py +++ b/pex/bin/pex.py @@ -11,8 +11,9 @@ import itertools import json import os +import shlex import sys -from argparse import Action, ArgumentDefaultsHelpFormatter, ArgumentParser +from argparse import Action, ArgumentDefaultsHelpFormatter, ArgumentError, ArgumentParser from textwrap import TextWrapper from pex import pex_warnings @@ -45,7 +46,6 @@ from pex.resolve.resolvers import Unsatisfiable from pex.resolver import resolve from pex.result import catch, try_ -from pex.sh_boot import create_sh_boot_script from pex.targets import Targets from pex.tracer import TRACER from pex.typing import TYPE_CHECKING, cast @@ -405,6 +405,37 @@ def configure_clp_pex_entry_points(parser): "`from a.b.c import m` during validation.", ) + class InjectEnvAction(Action): + def __call__(self, parser, namespace, value, option_str=None): + components = value.split("=", 1) + if len(components) != 2: + raise ArgumentError( + self, + "Environment variable values must be of the form `name=value`. " + "Given: {value}".format(value=value), + ) + self.default.append(tuple(components)) + + group.add_argument( + "--inject-env", + dest="inject_env", + default=[], + action=InjectEnvAction, + help="Environment variables to freeze in to the application environment.", + ) + + class InjectArgAction(Action): + def __call__(self, parser, namespace, value, option_str=None): + self.default.extend(shlex.split(value)) + + group.add_argument( + "--inject-args", + dest="inject_args", + default=[], + action=InjectArgAction, + help="Command line arguments to the application to freeze in.", + ) + class Seed(Enum["Seed.Value"]): class Value(Enum.Value): @@ -603,6 +634,8 @@ def build_pex( pex_builder.add_source(src_file_path, dst_path) pex_info = pex_builder.info + pex_info.inject_env = dict(options.inject_env) + pex_info.inject_args = options.inject_args pex_info.venv = bool(options.venv) pex_info.venv_bin_path = options.venv or BinPath.FALSE pex_info.venv_copies = options.venv_copies diff --git a/pex/pex.py b/pex/pex.py index 3e75eea45..c019bf541 100644 --- a/pex/pex.py +++ b/pex/pex.py @@ -560,6 +560,17 @@ def _execute(self): TRACER.log("PEX_INTERPRETER specified, dropping into interpreter") return self.execute_interpreter() + if not any( + ( + self._pex_info_overrides.script, + self._pex_info_overrides.entry_point, + self._pex_info.script, + self._pex_info.entry_point, + ) + ): + TRACER.log("No entry point specified, dropping into interpreter") + return self.execute_interpreter() + if self._pex_info_overrides.script and self._pex_info_overrides.entry_point: return "Cannot specify both script and entry_point for a PEX!" @@ -568,19 +579,21 @@ def _execute(self): if self._pex_info_overrides.script: return self.execute_script(self._pex_info_overrides.script) - elif self._pex_info_overrides.entry_point: + if self._pex_info_overrides.entry_point: return self.execute_entry( EntryPoint.parse("run = {}".format(self._pex_info_overrides.entry_point)) ) - elif self._pex_info.script: + + for name, value in self._pex_info.inject_env.items(): + os.environ.setdefault(name, value) + sys.argv[1:1] = list(self._pex_info.inject_args) + + if self._pex_info.script: return self.execute_script(self._pex_info.script) - elif self._pex_info.entry_point: + else: return self.execute_entry( EntryPoint.parse("run = {}".format(self._pex_info.entry_point)) ) - else: - TRACER.log("No entry point specified, dropping into interpreter") - return self.execute_interpreter() @classmethod def demote_bootstrap(cls): diff --git a/pex/pex_info.py b/pex/pex_info.py index e3e92fd34..dfb462b44 100644 --- a/pex/pex_info.py +++ b/pex/pex_info.py @@ -165,6 +165,26 @@ def build_properties(self, value): self._pex_info["build_properties"] = self.make_build_properties() self._pex_info["build_properties"].update(value) + @property + def inject_env(self): + # type: () -> Dict[str, str] + return dict(self._pex_info.get("inject_env", {})) + + @inject_env.setter + def inject_env(self, value): + # type: (Mapping[str, str]) -> None + self._pex_info["inject_env"] = dict(value) + + @property + def inject_args(self): + # type: () -> Tuple[str, ...] + return tuple(self._pex_info.get("inject_args", ())) + + @inject_args.setter + def inject_args(self, value): + # type: (Iterable[str]) -> None + self._pex_info["inject_args"] = tuple(value) + @property def venv(self): # type: () -> bool diff --git a/pex/venv/pex.py b/pex/venv/pex.py index 95fa1e21f..100918401 100644 --- a/pex/venv/pex.py +++ b/pex/venv/pex.py @@ -470,6 +470,8 @@ def sys_executable_paths(): ) ) sys.exit(1) + is_exec_override = len(pex_overrides) == 1 + if {strip_pex_env!r}: for key in list(os.environ): if key.startswith("PEX_"): @@ -559,6 +561,11 @@ def sys_executable_paths(): {exec_ast} sys.exit(0) + if not is_exec_override: + for name, value in {inject_env!r}: + os.environ.setdefault(name, value) + sys.argv[1:1] = {inject_args!r} + module_name, _, function = entry_point.partition(":") if not function: import runpy @@ -578,6 +585,8 @@ def sys_executable_paths(): shebang_python=venv_python, bin_path=bin_path, strip_pex_env=pex_info.strip_pex_env, + inject_env=tuple(pex_info.inject_env.items()), + inject_args=list(pex_info.inject_args), entry_point=pex_info.entry_point, script=pex_info.script, exec_ast=( diff --git a/tests/integration/test_inject_env_and_args.py b/tests/integration/test_inject_env_and_args.py new file mode 100644 index 000000000..28b13b839 --- /dev/null +++ b/tests/integration/test_inject_env_and_args.py @@ -0,0 +1,301 @@ +# Copyright 2022 Pants project contributors (see CONTRIBUTORS.md). +# Licensed under the Apache License, Version 2.0 (see LICENSE). + +import hashlib +import json +import os.path +import re +import signal +import socket +import subprocess +from contextlib import closing +from textwrap import dedent + +import pytest + +from pex.common import safe_open +from pex.fetcher import URLFetcher +from pex.testing import IS_PYPY, PY_VER, make_env, run_pex_command +from pex.typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing import Any, Iterable, List, Optional + + +def test_inject_env_invalid(): + # type: () -> None + result = run_pex_command(args=["--inject-env", "FOO"]) + result.assert_failure() + assert "--inject-env" in result.error + assert ( + "Environment variable values must be of the form `name=value`. Given: FOO" in result.error + ) + + +parametrize_execution_mode_args = pytest.mark.parametrize( + "execution_mode_args", + [ + pytest.param([], id="ZIPAPP"), + pytest.param(["--venv"], id="VENV"), + ], +) + + +@parametrize_execution_mode_args +def test_inject_env( + tmpdir, # type: Any + execution_mode_args, # type: List[str] +): + # type: (...) -> None + + print_FOO_env_code = "import os; print(os.environ.get('FOO', ''))" + + pex = os.path.join(str(tmpdir), "pex") + with open(os.path.join(str(tmpdir), "exe.py"), "w") as fp: + fp.write(print_FOO_env_code) + run_pex_command( + args=["--inject-env", "FOO=bar", "--exe", fp.name, "-o", pex] + execution_mode_args + ).assert_success() + + def assert_FOO( + expected_env_value, # type: str + runtime_env_value=None, # type: Optional[str] + ): + assert ( + expected_env_value + == subprocess.check_output(args=[pex], env=make_env(FOO=runtime_env_value)) + .decode("utf-8") + .strip() + ) + + assert_FOO(expected_env_value="bar") + assert_FOO(expected_env_value="baz", runtime_env_value="baz") + assert_FOO(expected_env_value="", runtime_env_value="") + + # Switching away from the built-in entrypoint should disable the injected env. + assert ( + "" + == subprocess.check_output( + args=[pex, "-c", print_FOO_env_code], env=make_env(PEX_INTERPRETER=1) + ) + .decode("utf-8") + .strip() + ) + + +DUMP_ARGS_CODE = "import json, sys; print(json.dumps(sys.argv[1:]))" + + +def create_inject_args_pex( + tmpdir, # type: Any + execution_mode_args, # type: Iterable[str] + *inject_args # type: str +): + # type: (...) -> str + pex = os.path.join( + str(tmpdir), + "pex-{}".format(hashlib.sha256(json.dumps(inject_args).encode("utf-8")).hexdigest()), + ) + with open(os.path.join(str(tmpdir), "exe.py"), "w") as fp: + fp.write(DUMP_ARGS_CODE) + argv = ["--exe", fp.name, "-o", pex] + for inject in inject_args: + argv.append("--inject-args") + argv.append(inject) + argv.extend(execution_mode_args) + run_pex_command(args=argv).assert_success() + return pex + + +@parametrize_execution_mode_args +def test_inject_args( + tmpdir, # type: Any + execution_mode_args, # type: List[str] +): + # type: (...) -> None + + pex_individual = create_inject_args_pex(tmpdir, execution_mode_args, "foo", "bar") + pex_shlex = create_inject_args_pex(tmpdir, execution_mode_args, "foo bar") + pex_combined = create_inject_args_pex(tmpdir, execution_mode_args, "foo bar", "baz") + + def assert_argv( + pex, # type: str + expected_argv, # type: List[str] + runtime_args=(), # type: Iterable[str] + **env + ): + assert expected_argv == json.loads( + subprocess.check_output(args=[pex] + list(runtime_args), env=make_env(**env)) + ) + + assert_argv(pex_individual, expected_argv=["foo", "bar"]) + assert_argv(pex_individual, expected_argv=["foo", "bar", "baz"], runtime_args=["baz"]) + assert_argv(pex_shlex, expected_argv=["foo", "bar"]) + assert_argv(pex_shlex, expected_argv=["foo", "bar", "baz"], runtime_args=["baz"]) + assert_argv(pex_combined, expected_argv=["foo", "bar", "baz"]) + assert_argv(pex_combined, expected_argv=["foo", "bar", "baz", "baz"], runtime_args=["baz"]) + + # Switching away from the built-in entrypoint should disable injected args. + assert_argv( + pex_individual, expected_argv=[], runtime_args=["-c", DUMP_ARGS_CODE], PEX_INTERPRETER=1 + ) + assert_argv(pex_shlex, expected_argv=[], runtime_args=["-c", DUMP_ARGS_CODE], PEX_INTERPRETER=1) + assert_argv( + pex_combined, expected_argv=[], runtime_args=["-c", DUMP_ARGS_CODE], PEX_INTERPRETER=1 + ) + + +@pytest.mark.skipif(PY_VER < (3, 7), reason="Uvicorn only support Python 3.7+.") +@parametrize_execution_mode_args +def test_complex( + tmpdir, # type: Any + execution_mode_args, # type: List[str] +): + # type: (...) -> None + + pex = os.path.join(str(tmpdir), "pex") + src = os.path.join(str(tmpdir), "src") + with safe_open(os.path.join(src, "example.py"), "w") as fp: + fp.write( + dedent( + """\ + import json + import os + import sys + + + async def app(scope, receive, send): + assert scope['type'] == 'http' + + await send({ + 'type': 'http.response.start', + 'status': 200, + 'headers': [ + [b'content-type', b'text/plain'], + ], + }) + await send({ + 'type': 'http.response.body', + 'body': os.environb.get(b'MESSAGE') or b'', + }) + + if __name__ == "__main__": + json.dump( + {"args": sys.argv[1:], "MESSAGE": os.environ.get("MESSAGE")}, sys.stdout + ) + """ + ) + ) + run_pex_command( + args=[ + "-D", + src, + "uvicorn[standard]==0.18.3", + "-c", + "uvicorn", + "--inject-args", + "example:app", + "--inject-env", + "MESSAGE=Hello, world!", + "-o", + pex, + ] + + execution_mode_args + ).assert_success() + + def assert_message( + expected, # type: bytes + **env # type: str + ): + # type: (...) -> None + with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock: + sock.bind(("127.0.0.1", 0)) + stderr_read_fd, stderr_write_fd = os.pipe() + # Python 2.7 doesn't support pass_fds, but we don't test against Python2.7. + process = subprocess.Popen( # type: ignore[call-arg] + args=[pex, "--fd", str(sock.fileno())], + stderr=stderr_write_fd, + pass_fds=[sock.fileno()], + env=make_env(**env), + ) + with os.fdopen(stderr_read_fd, "r") as stderr_fp: + for line in stderr_fp: + if "Uvicorn running" in line: + break + + host, port = sock.getsockname() + with URLFetcher().get_body_stream( + "http://{host}:{port}".format(host=host, port=port) + ) as fp: + assert expected == fp.read() + process.send_signal(signal.SIGINT) + process.kill() + os.close(stderr_write_fd) + + assert_message(b"Hello, world!") + assert_message(b"42", MESSAGE="42") + + # Switching away from the built-in entrypoint should disable injected args and env. + assert {"args": ["foo", "bar"], "MESSAGE": None} == json.loads( + subprocess.check_output(args=[pex, "foo", "bar"], env=make_env(PEX_MODULE="example")) + ) + + +@pytest.mark.skipif( + IS_PYPY or PY_VER > (3, 10) or PY_VER < (3, 6), + reason="The pyuwsgi distribution only has wheels for Linux and Mac for Python 3.6 through 3.10", +) +@parametrize_execution_mode_args +def test_pyuwsgi( + tmpdir, # type: Any + execution_mode_args, # type: List[str] +): + # type: (...) -> None + + pex = os.path.join(str(tmpdir), "pex") + src = os.path.join(str(tmpdir), "src") + with safe_open(os.path.join(src, "myflaskapp.py"), "w") as fp: + fp.write( + dedent( + """\ + from flask import Flask + + app = Flask(__name__) + + @app.route('/') + def index(): + return "I am app 1" + """ + ) + ) + run_pex_command( + args=[ + "-D", + src, + "pyuwsgi", + "flask", + "-c", + "uwsgi", + "--inject-args", + "--master --module myflaskapp:app", + "-o", + pex, + ] + + execution_mode_args + ).assert_success() + + stderr_read_fd, stderr_write_fd = os.pipe() + process = subprocess.Popen(args=[pex, "--http-socket", "127.0.0.1:0"], stderr=stderr_write_fd) + port = None # type: Optional[str] + with os.fdopen(stderr_read_fd, "r") as stderr_fp: + for line in stderr_fp: + match = re.search(r"bound to TCP address 127.0.0.1:(?P\d+)", line) + if match: + port = match.group("port") + break + assert port is not None, "Could not determine uwsgi server port." + with URLFetcher().get_body_stream("http://127.0.0.1:{port}".format(port=port)) as fp: + assert b"I am app 1" == fp.read() + process.send_signal(signal.SIGINT) + process.kill() + os.close(stderr_write_fd)