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)