fix: Unstable PROMETHEUS_MULTIPROC_DIR (#116)

* fix: Unstable `PROMETHEUS_MULTIPROC_DIR`

* avoid circular imports, fix tests

* remove `PROMETHEUS_MULTIPROC_DIR_KEEP`

* Remove unnecessary `Env()` instantiation

---------

Co-authored-by: Matthew Elwell <[email protected]>

Author: khvn26

src/common/core/constants.py

@@ -0,0 +1 @@
+DEFAULT_PROMETHEUS_MULTIPROC_DIR = "/tmp/flagsmith-prometheus"

src/common/core/main.py

@@ -1,16 +1,16 @@
 import contextlib
 import logging
 import os
+import shutil
 import sys
 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
+from common.core.constants import DEFAULT_PROMETHEUS_MULTIPROC_DIR
 
 logger = logging.getLogger(__name__)
 
@@ -31,7 +31,6 @@ def ensure_cli_env() -> typing.Generator[None, None, None]:
         main()
     ```
     """
-    env = Env()
     ctx = contextlib.ExitStack()
 
     # TODO @khvn26 Move logging setup to here
@@ -45,16 +44,16 @@ 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)
-        )
-        logger.info(
-            "Created %s for Prometheus multi-process mode",
-            prometheus_multiproc_dir_name,
-        )
-        os.environ["PROMETHEUS_MULTIPROC_DIR"] = prometheus_multiproc_dir_name
+    prometheus_multiproc_dir_name = os.environ.setdefault(
+        "PROMETHEUS_MULTIPROC_DIR",
+        DEFAULT_PROMETHEUS_MULTIPROC_DIR,
+    )
+    shutil.rmtree(prometheus_multiproc_dir_name, ignore_errors=True)
+    os.makedirs(prometheus_multiproc_dir_name, exist_ok=True)
+    logger.info(
+        "Re-created %s for Prometheus multi-process mode",
+        prometheus_multiproc_dir_name,
+    )
 
     if "docgen" in sys.argv:
         os.environ["DOCGEN_MODE"] = "true"

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()

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,23 @@ 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)
 
-    # When
-    main(["flagsmith"])
-
-    # Then
-    assert not Path(os.environ["PROMETHEUS_MULTIPROC_DIR"]).exists()
-
-
-def test_main__prometheus_multiproc_remove_dir_on_exit_true__expected(
-    monkeypatch: pytest.MonkeyPatch,
-    fs: FakeFilesystem,
-) -> None:
-    # Given
-    monkeypatch.delenv("PROMETHEUS_MULTIPROC_DIR")
-    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 not fs.exists("/tmp/flagsmith-prometheus/some_metric_file.db")
 
 
 def test_main__no_django_configured__expected_0(