Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions src/common/core/constants.py
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
DEFAULT_PROMETHEUS_MULTIPROC_DIR = "/tmp/flagsmith-prometheus"
26 changes: 18 additions & 8 deletions src/common/core/main.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import contextlib
import logging
import os
import shutil
import sys
import typing

Expand All @@ -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__)

Expand Down Expand Up @@ -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"
Expand Down
39 changes: 0 additions & 39 deletions src/common/core/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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()
23 changes: 16 additions & 7 deletions tests/integration/core/test_main.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,4 @@
import os
import subprocess
from pathlib import Path

import django
import pytest
Expand Down Expand Up @@ -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")
Comment on lines +131 to +148
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this use case real? i.e. Is there a reason why we need to keep the temporary files generated by this program?



def test_main__no_django_configured__expected_0(
Expand Down