fix: make backend tests pass on Windows (#282)
* fix: make backend tests pass on Windows * fix: address Windows locking + lazy graphiti imports * fix: address CodeRabbit review comments * fix: improve file_lock typing * Update apps/backend/runners/github/file_lock.py Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> * fix: satisfy ruff + asyncio loop usage * style: ruff format file_lock * refactor: safer temp file close + async JSON read * style: ruff format file_lock --------- Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Co-authored-by: Andy <119136210+AndyMik90@users.noreply.github.com>
This commit is contained in:
committed by
GitHub
parent
e9782db044
commit
4dcc5afae8
+3
-3
@@ -94,6 +94,9 @@ def create_ollama_embedder(config: "GraphitiConfig") -> Any:
|
||||
ProviderNotInstalled: If graphiti-core is not installed
|
||||
ProviderError: If model is not specified
|
||||
"""
|
||||
if not config.ollama_embedding_model:
|
||||
raise ProviderError("Ollama embedder requires OLLAMA_EMBEDDING_MODEL")
|
||||
|
||||
try:
|
||||
from graphiti_core.embedder.openai import OpenAIEmbedder, OpenAIEmbedderConfig
|
||||
except ImportError as e:
|
||||
@@ -103,9 +106,6 @@ def create_ollama_embedder(config: "GraphitiConfig") -> Any:
|
||||
f"Error: {e}"
|
||||
)
|
||||
|
||||
if not config.ollama_embedding_model:
|
||||
raise ProviderError("Ollama embedder requires OLLAMA_EMBEDDING_MODEL")
|
||||
|
||||
# Get embedding dimension (auto-detect for known models, or use configured value)
|
||||
embedding_dim = get_embedding_dim_for_model(
|
||||
config.ollama_embedding_model,
|
||||
|
||||
@@ -27,6 +27,9 @@ def create_openai_llm_client(config: "GraphitiConfig") -> Any:
|
||||
ProviderNotInstalled: If graphiti-core is not installed
|
||||
ProviderError: If API key is missing
|
||||
"""
|
||||
if not config.openai_api_key:
|
||||
raise ProviderError("OpenAI provider requires OPENAI_API_KEY")
|
||||
|
||||
try:
|
||||
from graphiti_core.llm_client.config import LLMConfig
|
||||
from graphiti_core.llm_client.openai_client import OpenAIClient
|
||||
@@ -37,9 +40,6 @@ def create_openai_llm_client(config: "GraphitiConfig") -> Any:
|
||||
f"Error: {e}"
|
||||
)
|
||||
|
||||
if not config.openai_api_key:
|
||||
raise ProviderError("OpenAI provider requires OPENAI_API_KEY")
|
||||
|
||||
llm_config = LLMConfig(
|
||||
api_key=config.openai_api_key,
|
||||
model=config.openai_model,
|
||||
|
||||
@@ -146,12 +146,12 @@ class GraphitiClient:
|
||||
# The original graphiti-core KuzuDriver has build_indices_and_constraints()
|
||||
# as a no-op, which causes FTS search failures
|
||||
from integrations.graphiti.queries_pkg.kuzu_driver_patched import (
|
||||
PatchedKuzuDriver as KuzuDriver,
|
||||
create_patched_kuzu_driver,
|
||||
)
|
||||
|
||||
db_path = self.config.get_db_path()
|
||||
try:
|
||||
self._driver = KuzuDriver(db=str(db_path))
|
||||
self._driver = create_patched_kuzu_driver(db=str(db_path))
|
||||
except (OSError, PermissionError) as e:
|
||||
logger.warning(
|
||||
f"Failed to initialize LadybugDB driver at {db_path}: {e}"
|
||||
|
||||
@@ -9,6 +9,7 @@ This patched driver fixes both issues for LadybugDB compatibility.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
# Import kuzu (might be real_ladybug via monkeypatch)
|
||||
@@ -17,157 +18,159 @@ try:
|
||||
except ImportError:
|
||||
import real_ladybug as kuzu # type: ignore
|
||||
|
||||
from graphiti_core.driver.driver import GraphProvider
|
||||
from graphiti_core.driver.kuzu_driver import KuzuDriver as OriginalKuzuDriver
|
||||
from graphiti_core.graph_queries import get_fulltext_indices
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PatchedKuzuDriver(OriginalKuzuDriver):
|
||||
"""
|
||||
KuzuDriver with proper FTS index creation and parameter handling.
|
||||
def create_patched_kuzu_driver(db: str = ":memory:", max_concurrent_queries: int = 1):
|
||||
from graphiti_core.driver.driver import GraphProvider
|
||||
from graphiti_core.driver.kuzu_driver import KuzuDriver as OriginalKuzuDriver
|
||||
from graphiti_core.graph_queries import get_fulltext_indices
|
||||
|
||||
Fixes two bugs in graphiti-core:
|
||||
1. FTS indexes are never created (build_indices_and_constraints is a no-op)
|
||||
2. None parameters are filtered out, causing "Parameter not found" errors
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
db: str = ":memory:",
|
||||
max_concurrent_queries: int = 1,
|
||||
):
|
||||
# Store database path before calling parent (which creates the Database)
|
||||
self._database = db # Required by Graphiti for group_id checks
|
||||
super().__init__(db, max_concurrent_queries)
|
||||
|
||||
async def execute_query(
|
||||
self, cypher_query_: str, **kwargs: Any
|
||||
) -> tuple[list[dict[str, Any]] | list[list[dict[str, Any]]], None, None]:
|
||||
class PatchedKuzuDriver(OriginalKuzuDriver):
|
||||
"""
|
||||
Execute a Cypher query with proper None parameter handling.
|
||||
KuzuDriver with proper FTS index creation and parameter handling.
|
||||
|
||||
The original driver filters out None values, but LadybugDB requires
|
||||
all referenced parameters to exist. This override keeps None values
|
||||
in the parameters dict.
|
||||
Fixes two bugs in graphiti-core:
|
||||
1. FTS indexes are never created (build_indices_and_constraints is a no-op)
|
||||
2. None parameters are filtered out, causing "Parameter not found" errors
|
||||
"""
|
||||
# Don't filter out None values - LadybugDB needs them
|
||||
params = {k: v for k, v in kwargs.items()}
|
||||
# Still remove these unsupported parameters
|
||||
params.pop("database_", None)
|
||||
params.pop("routing_", None)
|
||||
|
||||
try:
|
||||
results = await self.client.execute(cypher_query_, parameters=params)
|
||||
except Exception as e:
|
||||
# Truncate long values for logging
|
||||
log_params = {
|
||||
k: (v[:5] if isinstance(v, list) else v) for k, v in params.items()
|
||||
}
|
||||
logger.error(
|
||||
f"Error executing Kuzu query: {e}\n{cypher_query_}\n{log_params}"
|
||||
)
|
||||
raise
|
||||
def __init__(
|
||||
self,
|
||||
db: str = ":memory:",
|
||||
max_concurrent_queries: int = 1,
|
||||
):
|
||||
# Store database path before calling parent (which creates the Database)
|
||||
self._database = db # Required by Graphiti for group_id checks
|
||||
super().__init__(db, max_concurrent_queries)
|
||||
|
||||
if not results:
|
||||
return [], None, None
|
||||
async def execute_query(
|
||||
self, cypher_query_: str, **kwargs: Any
|
||||
) -> tuple[list[dict[str, Any]] | list[list[dict[str, Any]]], None, None]:
|
||||
"""
|
||||
Execute a Cypher query with proper None parameter handling.
|
||||
|
||||
if isinstance(results, list):
|
||||
dict_results = [list(result.rows_as_dict()) for result in results]
|
||||
else:
|
||||
dict_results = list(results.rows_as_dict())
|
||||
return dict_results, None, None # type: ignore
|
||||
The original driver filters out None values, but LadybugDB requires
|
||||
all referenced parameters to exist. This override keeps None values
|
||||
in the parameters dict.
|
||||
"""
|
||||
# Don't filter out None values - LadybugDB needs them
|
||||
params = {k: v for k, v in kwargs.items()}
|
||||
# Still remove these unsupported parameters
|
||||
params.pop("database_", None)
|
||||
params.pop("routing_", None)
|
||||
|
||||
async def build_indices_and_constraints(self, delete_existing: bool = False):
|
||||
"""
|
||||
Build FTS indexes required for Graphiti's hybrid search.
|
||||
try:
|
||||
results = await self.client.execute(cypher_query_, parameters=params)
|
||||
except Exception as e:
|
||||
# Truncate long values for logging
|
||||
log_params = {
|
||||
k: (v[:5] if isinstance(v, list) else v) for k, v in params.items()
|
||||
}
|
||||
logger.error(
|
||||
f"Error executing Kuzu query: {e}\n{cypher_query_}\n{log_params}"
|
||||
)
|
||||
raise
|
||||
|
||||
The original KuzuDriver has this as a no-op, but we need to actually
|
||||
create the FTS indexes for search to work.
|
||||
if not results:
|
||||
return [], None, None
|
||||
|
||||
Args:
|
||||
delete_existing: If True, drop and recreate indexes (default: False)
|
||||
"""
|
||||
logger.info("Building FTS indexes for Kuzu/LadybugDB...")
|
||||
if isinstance(results, list):
|
||||
dict_results = [list(result.rows_as_dict()) for result in results]
|
||||
else:
|
||||
dict_results = list(results.rows_as_dict())
|
||||
return dict_results, None, None # type: ignore
|
||||
|
||||
# Get the FTS index creation queries from Graphiti
|
||||
fts_queries = get_fulltext_indices(GraphProvider.KUZU)
|
||||
async def build_indices_and_constraints(self, delete_existing: bool = False):
|
||||
"""
|
||||
Build FTS indexes required for Graphiti's hybrid search.
|
||||
|
||||
# Create a sync connection for index creation
|
||||
conn = kuzu.Connection(self.db)
|
||||
The original KuzuDriver has this as a no-op, but we need to actually
|
||||
create the FTS indexes for search to work.
|
||||
|
||||
try:
|
||||
for query in fts_queries:
|
||||
try:
|
||||
# Check if we need to drop existing index first
|
||||
if delete_existing:
|
||||
# Extract index name from query
|
||||
# Format: CALL CREATE_FTS_INDEX('TableName', 'index_name', [...])
|
||||
parts = query.split("'")
|
||||
if len(parts) >= 4:
|
||||
table_name = parts[1]
|
||||
index_name = parts[3]
|
||||
drop_query = (
|
||||
f"CALL DROP_FTS_INDEX('{table_name}', '{index_name}')"
|
||||
Args:
|
||||
delete_existing: If True, drop and recreate indexes (default: False)
|
||||
"""
|
||||
logger.info("Building FTS indexes for Kuzu/LadybugDB...")
|
||||
|
||||
# Get the FTS index creation queries from Graphiti
|
||||
fts_queries = get_fulltext_indices(GraphProvider.KUZU)
|
||||
|
||||
# Create a sync connection for index creation
|
||||
conn = kuzu.Connection(self.db)
|
||||
|
||||
try:
|
||||
for query in fts_queries:
|
||||
try:
|
||||
# Check if we need to drop existing index first
|
||||
if delete_existing:
|
||||
# Extract index name from query
|
||||
# Format: CALL CREATE_FTS_INDEX('TableName', 'index_name', [...])
|
||||
match = re.search(
|
||||
r"CREATE_FTS_INDEX\('([^']+)',\s*'([^']+)'", query
|
||||
)
|
||||
try:
|
||||
conn.execute(drop_query)
|
||||
logger.debug(
|
||||
f"Dropped existing FTS index: {index_name}"
|
||||
)
|
||||
except Exception:
|
||||
# Index might not exist, that's fine
|
||||
pass
|
||||
if match:
|
||||
table_name, index_name = match.groups()
|
||||
drop_query = f"CALL DROP_FTS_INDEX('{table_name}', '{index_name}')"
|
||||
try:
|
||||
conn.execute(drop_query)
|
||||
logger.debug(
|
||||
f"Dropped existing FTS index: {index_name}"
|
||||
)
|
||||
except Exception:
|
||||
# Index might not exist, that's fine
|
||||
pass
|
||||
|
||||
# Create the FTS index
|
||||
conn.execute(query)
|
||||
logger.debug(f"Created FTS index: {query[:80]}...")
|
||||
# Create the FTS index
|
||||
conn.execute(query)
|
||||
logger.debug(f"Created FTS index: {query[:80]}...")
|
||||
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
# Handle "index already exists" gracefully
|
||||
if "already exists" in error_msg or "duplicate" in error_msg:
|
||||
logger.debug(
|
||||
f"FTS index already exists (skipping): {query[:60]}..."
|
||||
)
|
||||
else:
|
||||
# Log but don't fail - some indexes might fail in certain Kuzu versions
|
||||
logger.warning(f"Failed to create FTS index: {e}")
|
||||
logger.debug(f"Query was: {query}")
|
||||
|
||||
logger.info("FTS indexes created successfully")
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def setup_schema(self):
|
||||
"""
|
||||
Set up the database schema and install/load the FTS extension.
|
||||
|
||||
Extends the parent setup_schema() to properly set up FTS support.
|
||||
"""
|
||||
conn = kuzu.Connection(self.db)
|
||||
|
||||
try:
|
||||
# First, install the FTS extension (required before loading)
|
||||
try:
|
||||
conn.execute("INSTALL fts")
|
||||
logger.debug("Installed FTS extension")
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
# Handle "index already exists" gracefully
|
||||
if "already exists" in error_msg or "duplicate" in error_msg:
|
||||
logger.debug(
|
||||
f"FTS index already exists (skipping): {query[:60]}..."
|
||||
)
|
||||
else:
|
||||
# Log but don't fail - some indexes might fail in certain Kuzu versions
|
||||
logger.warning(f"Failed to create FTS index: {e}")
|
||||
logger.debug(f"Query was: {query}")
|
||||
if "already" not in error_msg:
|
||||
logger.debug(f"FTS extension install note: {e}")
|
||||
|
||||
logger.info("FTS indexes created successfully")
|
||||
finally:
|
||||
conn.close()
|
||||
# Then load the FTS extension
|
||||
try:
|
||||
conn.execute("LOAD EXTENSION fts")
|
||||
logger.debug("Loaded FTS extension")
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
if "already loaded" not in error_msg:
|
||||
logger.debug(f"FTS extension load note: {e}")
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def setup_schema(self):
|
||||
"""
|
||||
Set up the database schema and install/load the FTS extension.
|
||||
# Run the parent schema setup (creates tables)
|
||||
super().setup_schema()
|
||||
|
||||
Extends the parent setup_schema() to properly set up FTS support.
|
||||
"""
|
||||
conn = kuzu.Connection(self.db)
|
||||
|
||||
try:
|
||||
# First, install the FTS extension (required before loading)
|
||||
try:
|
||||
conn.execute("INSTALL fts")
|
||||
logger.debug("Installed FTS extension")
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
if "already" not in error_msg:
|
||||
logger.debug(f"FTS extension install note: {e}")
|
||||
|
||||
# Then load the FTS extension
|
||||
try:
|
||||
conn.execute("LOAD EXTENSION fts")
|
||||
logger.debug("Loaded FTS extension")
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
if "already loaded" not in error_msg:
|
||||
logger.debug(f"FTS extension load note: {e}")
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
# Run the parent schema setup (creates tables)
|
||||
super().setup_schema()
|
||||
return PatchedKuzuDriver(db=db, max_concurrent_queries=max_concurrent_queries)
|
||||
|
||||
@@ -1,9 +1,10 @@
|
||||
"""
|
||||
File Locking for Concurrent Operations
|
||||
======================================
|
||||
=====================================
|
||||
|
||||
Thread-safe and process-safe file locking utilities for GitHub automation.
|
||||
Uses fcntl.flock() on Unix systems for proper cross-process locking.
|
||||
Uses fcntl.flock() on Unix systems and msvcrt.locking() on Windows for proper
|
||||
cross-process locking.
|
||||
|
||||
Example Usage:
|
||||
# Simple file locking
|
||||
@@ -14,20 +15,79 @@ Example Usage:
|
||||
# Atomic write with locking
|
||||
async with locked_write("path/to/file.json", timeout=5.0) as f:
|
||||
json.dump(data, f)
|
||||
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import fcntl
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
import warnings
|
||||
from collections.abc import Callable
|
||||
from contextlib import asynccontextmanager, contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
_IS_WINDOWS = os.name == "nt"
|
||||
_WINDOWS_LOCK_SIZE = 1024 * 1024
|
||||
|
||||
try:
|
||||
import fcntl # type: ignore
|
||||
except ImportError: # pragma: no cover
|
||||
fcntl = None
|
||||
|
||||
try:
|
||||
import msvcrt # type: ignore
|
||||
except ImportError: # pragma: no cover
|
||||
msvcrt = None
|
||||
|
||||
|
||||
def _try_lock(fd: int, exclusive: bool) -> None:
|
||||
if _IS_WINDOWS:
|
||||
if msvcrt is None:
|
||||
raise FileLockError("msvcrt is required for file locking on Windows")
|
||||
if not exclusive:
|
||||
warnings.warn(
|
||||
"Shared file locks are not supported on Windows; using exclusive lock",
|
||||
RuntimeWarning,
|
||||
stacklevel=3,
|
||||
)
|
||||
msvcrt.locking(fd, msvcrt.LK_NBLCK, _WINDOWS_LOCK_SIZE)
|
||||
return
|
||||
|
||||
if fcntl is None:
|
||||
raise FileLockError(
|
||||
"fcntl is required for file locking on non-Windows platforms"
|
||||
)
|
||||
|
||||
lock_mode = fcntl.LOCK_EX if exclusive else fcntl.LOCK_SH
|
||||
fcntl.flock(fd, lock_mode | fcntl.LOCK_NB)
|
||||
|
||||
|
||||
def _unlock(fd: int) -> None:
|
||||
if _IS_WINDOWS:
|
||||
if msvcrt is None:
|
||||
warnings.warn(
|
||||
"msvcrt unavailable; cannot unlock file descriptor",
|
||||
RuntimeWarning,
|
||||
stacklevel=3,
|
||||
)
|
||||
return
|
||||
msvcrt.locking(fd, msvcrt.LK_UNLCK, _WINDOWS_LOCK_SIZE)
|
||||
return
|
||||
|
||||
if fcntl is None:
|
||||
warnings.warn(
|
||||
"fcntl unavailable; cannot unlock file descriptor",
|
||||
RuntimeWarning,
|
||||
stacklevel=3,
|
||||
)
|
||||
return
|
||||
fcntl.flock(fd, fcntl.LOCK_UN)
|
||||
|
||||
|
||||
class FileLockError(Exception):
|
||||
"""Raised when file locking operations fail."""
|
||||
@@ -43,7 +103,8 @@ class FileLockTimeout(FileLockError):
|
||||
|
||||
class FileLock:
|
||||
"""
|
||||
Cross-process file lock using fcntl.flock().
|
||||
Cross-process file lock using platform-specific locking (fcntl.flock on Unix,
|
||||
msvcrt.locking on Windows).
|
||||
|
||||
Supports both sync and async context managers for flexible usage.
|
||||
|
||||
@@ -89,22 +150,22 @@ class FileLock:
|
||||
self._fd = os.open(str(self._lock_file), os.O_CREAT | os.O_RDWR)
|
||||
|
||||
# Try to acquire lock with timeout
|
||||
lock_mode = fcntl.LOCK_EX if self.exclusive else fcntl.LOCK_SH
|
||||
start_time = time.time()
|
||||
|
||||
while True:
|
||||
try:
|
||||
# Non-blocking lock attempt
|
||||
fcntl.flock(self._fd, lock_mode | fcntl.LOCK_NB)
|
||||
_try_lock(self._fd, self.exclusive)
|
||||
return # Lock acquired
|
||||
except BlockingIOError:
|
||||
except (BlockingIOError, OSError):
|
||||
# Lock held by another process
|
||||
elapsed = time.time() - start_time
|
||||
if elapsed >= self.timeout:
|
||||
os.close(self._fd)
|
||||
self._fd = None
|
||||
raise FileLockTimeout(
|
||||
f"Failed to acquire lock on {self.filepath} within {self.timeout}s"
|
||||
f"Failed to acquire lock on {self.filepath} within "
|
||||
f"{self.timeout}s"
|
||||
)
|
||||
|
||||
# Wait a bit before retrying
|
||||
@@ -114,7 +175,7 @@ class FileLock:
|
||||
"""Release the file lock."""
|
||||
if self._fd is not None:
|
||||
try:
|
||||
fcntl.flock(self._fd, fcntl.LOCK_UN)
|
||||
_unlock(self._fd)
|
||||
os.close(self._fd)
|
||||
except Exception:
|
||||
pass # Best effort cleanup
|
||||
@@ -141,12 +202,12 @@ class FileLock:
|
||||
async def __aenter__(self):
|
||||
"""Async context manager entry."""
|
||||
# Run blocking lock acquisition in thread pool
|
||||
await asyncio.get_event_loop().run_in_executor(None, self._acquire_lock)
|
||||
await asyncio.get_running_loop().run_in_executor(None, self._acquire_lock)
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc_val, exc_tb):
|
||||
"""Async context manager exit."""
|
||||
await asyncio.get_event_loop().run_in_executor(None, self._release_lock)
|
||||
await asyncio.get_running_loop().run_in_executor(None, self._release_lock)
|
||||
return False
|
||||
|
||||
|
||||
@@ -192,7 +253,9 @@ def atomic_write(filepath: str | Path, mode: str = "w"):
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def locked_write(filepath: str | Path, timeout: float = 5.0, mode: str = "w"):
|
||||
async def locked_write(
|
||||
filepath: str | Path, timeout: float = 5.0, mode: str = "w"
|
||||
) -> Any:
|
||||
"""
|
||||
Async context manager combining file locking and atomic writes.
|
||||
|
||||
@@ -219,7 +282,7 @@ async def locked_write(filepath: str | Path, timeout: float = 5.0, mode: str = "
|
||||
|
||||
try:
|
||||
# Atomic write in thread pool (since it uses sync file I/O)
|
||||
fd, tmp_path = await asyncio.get_event_loop().run_in_executor(
|
||||
fd, tmp_path = await asyncio.get_running_loop().run_in_executor(
|
||||
None,
|
||||
lambda: tempfile.mkstemp(
|
||||
dir=filepath.parent, prefix=f".{filepath.name}.tmp.", suffix=""
|
||||
@@ -229,20 +292,20 @@ async def locked_write(filepath: str | Path, timeout: float = 5.0, mode: str = "
|
||||
try:
|
||||
# Open temp file and yield to caller
|
||||
f = os.fdopen(fd, mode)
|
||||
yield f
|
||||
|
||||
# Ensure file is closed before rename
|
||||
f.close()
|
||||
try:
|
||||
yield f
|
||||
finally:
|
||||
f.close()
|
||||
|
||||
# Atomic replace
|
||||
await asyncio.get_event_loop().run_in_executor(
|
||||
await asyncio.get_running_loop().run_in_executor(
|
||||
None, os.replace, tmp_path, filepath
|
||||
)
|
||||
|
||||
except Exception:
|
||||
# Clean up temp file on error
|
||||
try:
|
||||
await asyncio.get_event_loop().run_in_executor(
|
||||
await asyncio.get_running_loop().run_in_executor(
|
||||
None, os.unlink, tmp_path
|
||||
)
|
||||
except Exception:
|
||||
@@ -255,7 +318,7 @@ async def locked_write(filepath: str | Path, timeout: float = 5.0, mode: str = "
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def locked_read(filepath: str | Path, timeout: float = 5.0):
|
||||
async def locked_read(filepath: str | Path, timeout: float = 5.0) -> Any:
|
||||
"""
|
||||
Async context manager for locked file reading.
|
||||
|
||||
@@ -338,7 +401,10 @@ async def locked_json_read(filepath: str | Path, timeout: float = 5.0) -> Any:
|
||||
|
||||
|
||||
async def locked_json_update(
|
||||
filepath: str | Path, updater: callable, timeout: float = 5.0, indent: int = 2
|
||||
filepath: str | Path,
|
||||
updater: Callable[[Any], Any],
|
||||
timeout: float = 5.0,
|
||||
indent: int = 2,
|
||||
) -> Any:
|
||||
"""
|
||||
Helper for atomic read-modify-write of JSON files.
|
||||
@@ -373,17 +439,19 @@ async def locked_json_update(
|
||||
|
||||
try:
|
||||
# Read current data
|
||||
if filepath.exists():
|
||||
with open(filepath) as f:
|
||||
data = json.load(f)
|
||||
else:
|
||||
data = None
|
||||
def _read_json():
|
||||
if filepath.exists():
|
||||
with open(filepath) as f:
|
||||
return json.load(f)
|
||||
return None
|
||||
|
||||
data = await asyncio.get_running_loop().run_in_executor(None, _read_json)
|
||||
|
||||
# Apply update function
|
||||
updated_data = updater(data)
|
||||
|
||||
# Write atomically
|
||||
fd, tmp_path = await asyncio.get_event_loop().run_in_executor(
|
||||
fd, tmp_path = await asyncio.get_running_loop().run_in_executor(
|
||||
None,
|
||||
lambda: tempfile.mkstemp(
|
||||
dir=filepath.parent, prefix=f".{filepath.name}.tmp.", suffix=""
|
||||
@@ -394,13 +462,13 @@ async def locked_json_update(
|
||||
with os.fdopen(fd, "w") as f:
|
||||
json.dump(updated_data, f, indent=indent)
|
||||
|
||||
await asyncio.get_event_loop().run_in_executor(
|
||||
await asyncio.get_running_loop().run_in_executor(
|
||||
None, os.replace, tmp_path, filepath
|
||||
)
|
||||
|
||||
except Exception:
|
||||
try:
|
||||
await asyncio.get_event_loop().run_in_executor(
|
||||
await asyncio.get_running_loop().run_in_executor(
|
||||
None, os.unlink, tmp_path
|
||||
)
|
||||
except Exception:
|
||||
|
||||
Reference in New Issue
Block a user