Coverage for src / quber / playground / storage.py: 25%

1"""Where a document's files live. 

2 

3Every file the playground keeps for a document is named by the document's 

4storage key and sits in the local data directory: the served PDF, the parse 

5and figure artifacts the ingest read, and the upload workdir. On the 

6developer host that directory is the only copy. 

7 

8Hosted, the local directory is scratch that a task replacement erases, and 

9the durable copy is an ``s3://`` prefix from settings holding one 

10``<doc_key>/`` prefix per document with the same file names inside it. Two 

11movements keep the two in step: ``publish`` copies files up after every 

12stage of an upload, and ``local_pdf`` fetches the served PDF back down when 

13a task no longer has it. Nothing here ever deletes from the prefix; removing 

14a document from the library removes its database rows and local files only, 

15so a later upload of the same bytes finds every artifact still there. 

16 

17Run as a module, it publishes the whole current library once, which is how 

18the host's documents were moved into the bucket: 

19 

20 uv run python -m quber.playground.storage 

21""" 

22 

23from __future__ import annotations 

24 

25import json 

26import sys 

27from datetime import datetime, timezone 

28from pathlib import Path 

29from typing import Iterable, Optional 

30 

31from cloudpathlib import S3Path 

32from loguru import logger 

33 

34from quber.settings import get_settings 

35 

36DATA_DIR: Path = get_settings().playground.data_dir 

37UPLOADS: Path = DATA_DIR / "uploads" 

38 

39 

40def prefix(doc_key: str) -> Optional[S3Path]: 

41 """The document's prefix in the bucket, or None on the developer host.""" 

42 root = get_settings().playground.artifacts_uri 

43 if not root: 

44 return None 

45 return S3Path(root.rstrip("/") + "/" + doc_key + "/") 

46 

47 

48def publish(doc_key: str, paths: Iterable[Path]) -> int: 

49 """Copy the given local files under the document's prefix. Returns how 

50 many were copied; zero, without touching anything, on the developer host 

51 or when none of the paths exist.""" 

52 dest = prefix(doc_key) 

53 if dest is None: 

54 return 0 

55 copied = 0 

56 for path in paths: 

57 if not path.is_file(): 

58 continue 

59 (dest / path.name).upload_from(path, force_overwrite_to_cloud=True) 

60 copied += 1 

61 if copied: 

62 logger.info("published {} file(s) for {} to {}", copied, doc_key, dest) 

63 return copied 

64 

65 

66def publish_workdir(doc_key: str, workdir: Path) -> int: 

67 """Copy every file the upload workdir holds for the document.""" 

68 if not workdir.is_dir(): 

69 return 0 

70 return publish(doc_key, sorted(p for p in workdir.iterdir() if p.name.startswith(doc_key))) 

71 

72 

73def write_source_manifest( 

74 doc_key: str, *, filename: str, content_hash: str, source_uri: Optional[str] = None 

75) -> Path: 

76 """Record where the document came from beside its PDF: the name it was 

77 uploaded under, its full hash, the S3 object it was taken from if any, and 

78 when. The key carries none of this, so the prefix says it instead.""" 

79 manifest = { 

80 "doc_key": doc_key, 

81 "filename": filename, 

82 "sha256": content_hash, 

83 "source_uri": source_uri, 

84 "uploaded_at": datetime.now(timezone.utc).isoformat(timespec="seconds"), 

85 } 

86 path = UPLOADS / f"{doc_key}.source.json" 

87 path.parent.mkdir(parents=True, exist_ok=True) 

88 path.write_text(json.dumps(manifest, indent=2)) 

89 return path 

90 

91 

92def local_pdf(doc_key: str) -> Optional[Path]: 

93 """The served PDF as a local path, fetched from the bucket when this host 

94 no longer has it. None when it exists nowhere.""" 

95 for candidate in (DATA_DIR / f"{doc_key}.pdf", UPLOADS / f"{doc_key}.pdf"): 

96 if candidate.is_file(): 

97 return candidate 

98 remote = prefix(doc_key) 

99 if remote is None: 

100 return None 

101 source = remote / f"{doc_key}.pdf" 

102 if not source.exists(): 

103 return None 

104 target = DATA_DIR / f"{doc_key}.pdf" 

105 target.parent.mkdir(parents=True, exist_ok=True) 

106 logger.info("fetching {} from {}", target.name, source) 

107 source.download_to(target) 

108 return target 

109 

110 

111def publish_library() -> None: 

112 """Publish every document the database knows about from the local data 

113 directory: the served PDF and the artifacts beside it, plus the upload 

114 workdir when this host still has it.""" 

115 from quber.playground import db 

116 

117 if prefix("probe") is None: 

118 sys.exit("QUBER_PLAYGROUND_ARTIFACTS_URI is not set; nothing to publish to") 

119 with db.connect() as conn: 

120 rows = conn.execute( 

121 "SELECT doc_key, filename, content_hash FROM ade_playground.documents ORDER BY id" 

122 ).fetchall() 

123 total = 0 

124 for doc_key, filename, content_hash in rows: 

125 files = sorted(DATA_DIR.glob(f"{doc_key}.*")) 

126 if not any(f.name == f"{doc_key}.source.json" for f in files): 

127 files.append(write_source_manifest(doc_key, filename=filename, content_hash=content_hash)) 

128 copied = publish(doc_key, files) + publish_workdir(doc_key, UPLOADS / f"{doc_key}-artifacts") 

129 total += copied 

130 logger.info("{}: {} file(s)", doc_key, copied) 

131 logger.success("published {} file(s) for {} document(s)", total, len(rows)) 

132 

133 

134if __name__ == "__main__": 

135 publish_library()