diff --git a/src/common/core/constants.py b/src/common/core/constants.py new file mode 100644 index 0000000..d4d74ea --- /dev/null +++ b/src/common/core/constants.py @@ -0,0 +1 @@ +DEFAULT_PROMETHEUS_MULTIPROC_DIR = "/tmp/flagsmith-prometheus" diff --git a/src/common/core/main.py b/src/common/core/main.py index d1b7bcd..2346d70 100644 --- a/src/common/core/main.py +++ b/src/common/core/main.py @@ -1,6 +1,7 @@ import contextlib import logging import os +import shutil import sys import typing @@ -10,7 +11,7 @@ from environs import Env from common.core.cli import healthcheck -from common.core.utils import TemporaryDirectory +from common.core.constants import DEFAULT_PROMETHEUS_MULTIPROC_DIR logger = logging.getLogger(__name__) @@ -45,16 +46,25 @@ def ensure_cli_env() -> typing.Generator[None, None, None]: os.environ.setdefault("DJANGO_SETTINGS_MODULE", "settings.dev") # Set up Prometheus' multiprocess mode - 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) - ) + prometheus_multiproc_dir_name = os.environ.setdefault( + "PROMETHEUS_MULTIPROC_DIR", + DEFAULT_PROMETHEUS_MULTIPROC_DIR, + ) + prometheus_multiproc_dir_keep = env.bool( + "PROMETHEUS_MULTIPROC_DIR_KEEP", + default=False, + ) + if not prometheus_multiproc_dir_keep: + shutil.rmtree(prometheus_multiproc_dir_name, ignore_errors=True) logger.info( - "Created %s for Prometheus multi-process mode", + "Removed %s to ensure a clean state for Prometheus multi-process mode", prometheus_multiproc_dir_name, ) - os.environ["PROMETHEUS_MULTIPROC_DIR"] = prometheus_multiproc_dir_name + os.makedirs(prometheus_multiproc_dir_name, exist_ok=True) + logger.info( + "Created %s for Prometheus multi-process mode", + prometheus_multiproc_dir_name, + ) if "docgen" in sys.argv: os.environ["DOCGEN_MODE"] = "true" diff --git a/src/common/core/utils.py b/src/common/core/utils.py index bb5217e..cf747ab 100644 --- a/src/common/core/utils.py +++ b/src/common/core/utils.py @@ -2,8 +2,6 @@ import logging import pathlib import random -import sys -import tempfile from functools import lru_cache from itertools import cycle from typing import ( @@ -200,40 +198,3 @@ 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 6c44edd..df07b40 100644 --- a/tests/integration/core/test_main.py +++ b/tests/integration/core/test_main.py @@ -1,6 +1,4 @@ -import os import subprocess -from pathlib import Path import django import pytest @@ -111,32 +109,43 @@ def test_main__healthcheck_http__server_invalid_response__runs_expected( main(argv) -def test_main__prometheus_multiproc_remove_dir_on_exit_default__expected( +def test_main__prometheus_multiproc_remove_dir_on_start_default__expected( monkeypatch: pytest.MonkeyPatch, + fs: FakeFilesystem, ) -> None: # Given monkeypatch.delenv("PROMETHEUS_MULTIPROC_DIR_KEEP", raising=False) + fs.create_file( + "/tmp/flagsmith-prometheus/some_metric_file.db", + create_missing_dirs=True, + ) + # When main(["flagsmith"]) # Then - assert not Path(os.environ["PROMETHEUS_MULTIPROC_DIR"]).exists() + assert not fs.exists("/tmp/flagsmith-prometheus/some_metric_file.db") -def test_main__prometheus_multiproc_remove_dir_on_exit_true__expected( +def test_main__prometheus_multiproc_remove_dir_on_start_true__expected( monkeypatch: pytest.MonkeyPatch, fs: FakeFilesystem, ) -> None: # Given - monkeypatch.delenv("PROMETHEUS_MULTIPROC_DIR") + monkeypatch.delenv("PROMETHEUS_MULTIPROC_DIR", raising=False) monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR_KEEP", "true") + fs.create_file( + "/tmp/flagsmith-prometheus/some_metric_file.db", + create_missing_dirs=True, + ) + # When main(["flagsmith"]) # Then - assert Path(os.environ["PROMETHEUS_MULTIPROC_DIR"]).exists() + assert fs.exists("/tmp/flagsmith-prometheus/some_metric_file.db") def test_main__no_django_configured__expected_0(