Skip to content

Commit

Permalink
Refactored temp dir handling
Browse files Browse the repository at this point in the history
Made unique temp dir per process so they don't clobber each other
  • Loading branch information
mdellabitta committed Nov 13, 2024
1 parent 324de39 commit 1156519
Showing 1 changed file with 22 additions and 15 deletions.
37 changes: 22 additions & 15 deletions common.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@
import logging
import os
import re
import shutil
import sys
import tempfile
from datetime import datetime
Expand Down Expand Up @@ -52,15 +51,15 @@
PROVIDER_FIELD_NAME,
RIGHTS_CATEGORY_FIELD_NAME,
S3_RETRIES,
TMP_DIR_BASE,
UNLIMITED_RE_USE,
UPLOAD_FIELD_NAME,
WIKIDATA_FIELD_NAME,
S3_BUCKET,
EDM_IS_SHOWN_AT,
)

__http_session = None
__http_session: requests.Session | None = None
__temp_dir: tempfile.TemporaryDirectory | None = None


def load_ids(ids_file: IO) -> list[str]:
Expand Down Expand Up @@ -262,16 +261,32 @@ def s3_file_exists(path: str, s3: S3ServiceResource):


def setup_temp_dir() -> None:
if not os.path.isdir(TMP_DIR_BASE):
os.mkdir(TMP_DIR_BASE)
global __temp_dir
if __temp_dir is not None:
__temp_dir = tempfile.TemporaryDirectory(
"tmp", "wiki", dir=".", ignore_cleanup_errors=True, delete=False
)


def cleanup_temp_dir() -> None:
shutil.rmtree(TMP_DIR_BASE)
global __temp_dir
if __temp_dir is not None:
__temp_dir.cleanup()


def get_temp_file():
return tempfile.NamedTemporaryFile(delete=False, dir=TMP_DIR_BASE)
global __temp_dir
if __temp_dir is None:
raise Exception("Temp dir not initialized.")
return tempfile.NamedTemporaryFile(delete=False, dir=__temp_dir.name)


def clean_up_tmp_file(temp_file) -> None:
try:
if temp_file:
os.unlink(temp_file.name)
except Exception as e:
logging.warning("Temp file unlink failed.", exc_info=e)


def get_s3() -> S3ServiceResource:
Expand Down Expand Up @@ -317,14 +332,6 @@ def setup_logging(partner: str, event_type: str, level: int = logging.INFO) -> N
logging.getLogger(d).setLevel(logging.ERROR)


def clean_up_tmp_file(temp_file) -> None:
try:
if temp_file:
os.unlink(temp_file.name)
except Exception as e:
logging.warning("Temp file unlink failed.", exc_info=e)


Result = Enum("Result", ["DOWNLOADED", "FAILED", "SKIPPED", "UPLOADED", "BYTES"])


Expand Down

0 comments on commit 1156519

Please sign in to comment.