diff --git a/src/common/core/constants.py b/src/common/core/constants.py index d4d74ea..b55a6ca 100644 --- a/src/common/core/constants.py +++ b/src/common/core/constants.py @@ -1 +1 @@ -DEFAULT_PROMETHEUS_MULTIPROC_DIR = "/tmp/flagsmith-prometheus" +DEFAULT_PROMETHEUS_MULTIPROC_DIR_NAME = "flagsmith-prometheus" diff --git a/src/common/core/main.py b/src/common/core/main.py index 47cd10e..20149e3 100644 --- a/src/common/core/main.py +++ b/src/common/core/main.py @@ -1,16 +1,17 @@ import contextlib import logging import os -import shutil import sys import typing +from tempfile import gettempdir from django.core.management import ( execute_from_command_line as django_execute_from_command_line, ) from common.core.cli import healthcheck -from common.core.constants import DEFAULT_PROMETHEUS_MULTIPROC_DIR +from common.core.constants import DEFAULT_PROMETHEUS_MULTIPROC_DIR_NAME +from common.core.utils import clear_directory, make_writable_directory logger = logging.getLogger(__name__) @@ -35,6 +36,15 @@ def ensure_cli_env() -> typing.Generator[None, None, None]: # TODO @khvn26 Move logging setup to here + # Prometheus multiproc support + prom_dir = os.environ.setdefault( + "PROMETHEUS_MULTIPROC_DIR", + os.path.join(gettempdir(), DEFAULT_PROMETHEUS_MULTIPROC_DIR_NAME), + ) + if os.path.exists(prom_dir): + clear_directory(prom_dir) + make_writable_directory(prom_dir) + # Currently we don't install Flagsmith modules as a package, so we need to add # $CWD to the Python path to be able to import them sys.path.append(os.getcwd()) @@ -43,18 +53,6 @@ def ensure_cli_env() -> typing.Generator[None, None, None]: # without resorting to it being set outside of the application os.environ.setdefault("DJANGO_SETTINGS_MODULE", "settings.dev") - # Set up Prometheus' multiprocess mode - 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" diff --git a/src/common/core/utils.py b/src/common/core/utils.py index cf747ab..c73f4d3 100644 --- a/src/common/core/utils.py +++ b/src/common/core/utils.py @@ -1,7 +1,9 @@ import json import logging +import os import pathlib import random +import shutil from functools import lru_cache from itertools import cycle from typing import ( @@ -198,3 +200,30 @@ def using_database_replica( return manager return manager.db_manager(chosen_replica) + + +def clear_directory(directory_path: str) -> None: + """ + Safely clear a directory including all subdirectories and files. + """ + for p in pathlib.Path(directory_path).rglob("*"): + try: + # Ensure that the cleanup doesn't silently fail on + # files and subdirs created by other users. + p.chmod(0o777) + except (PermissionError, FileNotFoundError): # pragma: no cover + pass + + shutil.rmtree(directory_path, ignore_errors=True) + + +def make_writable_directory(directory_path: str) -> None: + os.makedirs(directory_path, exist_ok=True) + + try: + # While `mkdir` sets mode=0o777 by default, this can be affected by umask + # resulting in lesser permissions for other users. This step ensures the + # directory is writable for all users. + os.chmod(directory_path, 0o777) + except PermissionError: # pragma: no cover + pass