diff --git a/src/common/core/main.py b/src/common/core/main.py index 07976148..d1b7bcdb 100644 --- a/src/common/core/main.py +++ b/src/common/core/main.py @@ -2,14 +2,15 @@ import logging import os import sys -import tempfile import typing from django.core.management import ( execute_from_command_line as django_execute_from_command_line, ) +from environs import Env from common.core.cli import healthcheck +from common.core.utils import TemporaryDirectory logger = logging.getLogger(__name__) @@ -30,6 +31,7 @@ def ensure_cli_env() -> typing.Generator[None, None, None]: main() ``` """ + env = Env() ctx = contextlib.ExitStack() # TODO @khvn26 Move logging setup to here @@ -43,9 +45,11 @@ def ensure_cli_env() -> typing.Generator[None, None, None]: os.environ.setdefault("DJANGO_SETTINGS_MODULE", "settings.dev") # Set up Prometheus' multiprocess mode - if "PROMETHEUS_MULTIPROC_DIR" not in os.environ: - prometheus_multiproc_dir_name = tempfile.mkdtemp() - + if not env.str("PROMETHEUS_MULTIPROC_DIR", ""): + delete = not env.bool("PROMETHEUS_MULTIPROC_DIR_KEEP", False) + prometheus_multiproc_dir_name = ctx.enter_context( + TemporaryDirectory(delete=delete) + ) logger.info( "Created %s for Prometheus multi-process mode", prometheus_multiproc_dir_name, diff --git a/src/common/core/utils.py b/src/common/core/utils.py index 301edcd6..861f2b94 100644 --- a/src/common/core/utils.py +++ b/src/common/core/utils.py @@ -2,9 +2,18 @@ import logging import pathlib import random +import sys +import tempfile from functools import lru_cache from itertools import cycle -from typing import Iterator, Literal, NotRequired, TypedDict, TypeVar, get_args +from typing import ( + Iterator, + Literal, + NotRequired, + TypedDict, + TypeVar, + get_args, +) from django.conf import settings from django.contrib.auth import get_user_model @@ -186,3 +195,40 @@ def using_database_replica( return manager return manager.db_manager(chosen_replica) + + +if sys.version_info >= (3, 12): + # Already has the desired behavior; re-export for uniform imports. + TemporaryDirectory = tempfile.TemporaryDirectory +else: + import contextlib + from typing import ContextManager, Generator + + def TemporaryDirectory( + suffix: str | None = None, + prefix: str | None = None, + dir: str | None = None, + *, + delete: bool = True, + ) -> ContextManager[str]: + """ + Create a temporary directory with optional cleanup control. + + This wrapper exists because Python 3.12 changed TemporaryDirectory's behavior + by adding a 'delete' parameter, which doesn't exist in Python 3.11. This + function provides a consistent API across both versions. + + When delete=True, uses the stdlib's TemporaryDirectory (auto-cleanup). + When delete=False, creates a directory with mkdtemp that persists after + the context manager exits, matching Python 3.12's delete=False behavior. + + See https://docs.python.org/3.12/library/tempfile.html#tempfile.TemporaryDirectory for usage details. + """ + if delete: + return tempfile.TemporaryDirectory(suffix, prefix, dir) + + @contextlib.contextmanager + def _tmpdir() -> Generator[str, None, None]: + yield tempfile.mkdtemp(suffix, prefix, dir) + + return _tmpdir() diff --git a/tests/integration/core/test_main.py b/tests/integration/core/test_main.py index 7287b4c3..b302be2a 100644 --- a/tests/integration/core/test_main.py +++ b/tests/integration/core/test_main.py @@ -1,6 +1,10 @@ +import os +from pathlib import Path + import django import pytest from django.core.management import ManagementUtility +from pyfakefs.fake_filesystem import FakeFilesystem from pytest_httpserver import HTTPServer from common.core.main import main @@ -104,3 +108,28 @@ def test_main__healthcheck_http__server_invalid_response__runs_expected( # When & Then with pytest.raises(Exception): main(argv) + + +def test_main__prometheus_multiproc_remove_dir_on_exit_default__expected() -> None: + # Given + os.environ.pop("PROMETHEUS_MULTIPROC_DIR_KEEP", None) + + # When + main(["flagsmith"]) + + # Then + assert not Path(os.environ["PROMETHEUS_MULTIPROC_DIR"]).exists() + + +def test_main__prometheus_multiproc_remove_dir_on_exit_true__expected( + fs: FakeFilesystem, +) -> None: + # Given + os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None) + os.environ["PROMETHEUS_MULTIPROC_DIR_KEEP"] = "true" + + # When + main(["flagsmith"]) + + # Then + assert Path(os.environ["PROMETHEUS_MULTIPROC_DIR"]).exists()