Skip to content
Merged
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
44 changes: 31 additions & 13 deletions src/quant_data_kit/normalized_v3.py
Original file line number Diff line number Diff line change
Expand Up @@ -498,17 +498,34 @@ def _partition_identity(
def _streaming_stage(root: Path) -> Iterator[Path]:
staging_root = _mkdir_in_lake(root, root / "normalized" / "staging")
owners_root = _mkdir_in_lake(root, root / "normalized" / ".stage-owners")
for stale in sorted(staging_root.glob("normalized-batch-stream-*")):
checked = _validate_lake_path(root, stale, allow_missing=False)
if not checked.is_dir():
raise ValidationError(f"Normalized staging entry is not a directory: {checked}")
owner = owners_root / f"{checked.name}.lock"
try:
with process_file_lock(owner, timeout_seconds=0.01):
if checked.exists():
shutil.rmtree(checked)
except TimeoutError:
continue
gc_lock = owners_root / ".gc.lock"
with process_file_lock(gc_lock):
for stale in sorted(staging_root.glob("normalized-batch-stream-*")):
if not stale.exists():
continue
checked = _validate_lake_path(root, stale, allow_missing=False)
if not checked.is_dir():
raise ValidationError(f"Normalized staging entry is not a directory: {checked}")
owner = owners_root / f"{checked.name}.lock"
try:
with process_file_lock(owner, timeout_seconds=0.01):
if checked.exists():
shutil.rmtree(checked)
if owner.exists():
owner.unlink()
except TimeoutError:
continue
for owner in sorted(owners_root.glob("normalized-batch-stream-*.lock")):
stage = staging_root / owner.name.removesuffix(".lock")
if stage.exists():
continue
try:
with process_file_lock(owner, timeout_seconds=0.01):
removable = not stage.exists()
if removable and owner.exists():
owner.unlink()
except TimeoutError:
continue
operation = uuid.uuid4().hex
stage = staging_root / f"normalized-batch-stream-{operation}"
owner = owners_root / f"{stage.name}.lock"
Expand All @@ -519,8 +536,9 @@ def _streaming_stage(root: Path) -> Iterator[Path]:
finally:
if stage.exists():
shutil.rmtree(stage)
if owner.exists():
owner.unlink()
with process_file_lock(gc_lock):
if owner.exists():
owner.unlink()


def _sql_text(value: str | Path) -> str:
Expand Down
33 changes: 33 additions & 0 deletions tests/test_normalized_v3_failures.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import json
from collections.abc import Callable
from contextlib import contextmanager
from copy import deepcopy
from dataclasses import replace
from datetime import date, datetime, timedelta, timezone
Expand Down Expand Up @@ -166,6 +167,38 @@ def test_digest_and_staging_corruption_guards(
pass


def test_streaming_stage_serializes_owner_file_cleanup(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
root = lake_module._resolved_lake_root(tmp_path / "lake", create=True)
real_lock = normalized_v3.process_file_lock
real_unlink = Path.unlink
held_locks: set[Path] = set()

@contextmanager
def tracked_lock(path: Path, *, timeout_seconds: float = 60.0):
checked = Path(path)
with real_lock(checked, timeout_seconds=timeout_seconds):
held_locks.add(checked)
try:
yield
finally:
held_locks.remove(checked)

def guarded_unlink(path: Path, *args: object, **kwargs: object) -> None:
if path.parent.name == ".stage-owners" and path.name.startswith("normalized-batch-stream-"):
assert any(item.name == ".gc.lock" for item in held_locks)
real_unlink(path, *args, **kwargs)

monkeypatch.setattr(normalized_v3, "process_file_lock", tracked_lock)
monkeypatch.setattr(Path, "unlink", guarded_unlink)
with normalized_v3._streaming_stage(root) as stage:
assert stage.is_dir()
owners_root = root / "normalized" / ".stage-owners"
assert sorted(item.name for item in owners_root.iterdir()) == [".gc.lock"]


def _rewrite_index_manifest(
manifest_path: Path,
mutate: Callable[[dict[str, Any]], None],
Expand Down
Loading