lorebase-store 0.2.3__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (163) hide show
  1. lorebase_store/__init__.py +1 -0
  2. lorebase_store/_migrations/env.py +106 -0
  3. lorebase_store/_migrations/script.py.mako +27 -0
  4. lorebase_store/_migrations/versions/0001_initial_schema.py +1071 -0
  5. lorebase_store/_migrations/versions/0002_merge_attempt_details.py +74 -0
  6. lorebase_store/_migrations/versions/0003_webhooks.py +127 -0
  7. lorebase_store/_migrations/versions/0004_defer_revision_references.py +38 -0
  8. lorebase_store/_migrations/versions/0005_security_definer_reads.py +396 -0
  9. lorebase_store/_migrations/versions/0006_api_key_usage.py +48 -0
  10. lorebase_store/_migrations/versions/0007_document_extraction.py +154 -0
  11. lorebase_store/_migrations/versions/0008_definer_argument_checks.py +178 -0
  12. lorebase_store/_migrations/versions/0009_api_key_networks.py +33 -0
  13. lorebase_store/_migrations/versions/0010_overlay_branches.py +582 -0
  14. lorebase_store/_migrations/versions/0011_branch_scoped_search_index.py +229 -0
  15. lorebase_store/_migrations/versions/0012_overlay_neighborhood_lookups.py +185 -0
  16. lorebase_store/_migrations/versions/0013_web_sessions.py +72 -0
  17. lorebase_store/_migrations/versions/0014_embedding_jobs.py +70 -0
  18. lorebase_store/_migrations/versions/0015_append_only_audit_log.py +83 -0
  19. lorebase_store/_migrations/versions/0016_sealed_webhook_keys.py +55 -0
  20. lorebase_store/_migrations/versions/0017_background_branch_copies.py +77 -0
  21. lorebase_store/_migrations/versions/0018_search_lexicon.py +477 -0
  22. lorebase_store/_migrations/versions/0019_vault_bindings.py +56 -0
  23. lorebase_store/_migrations/versions/0020_outbox_consumers.py +150 -0
  24. lorebase_store/_migrations/versions/0021_schema_suggestions.py +114 -0
  25. lorebase_store/_migrations/versions/0022_read_audit_and_shared_rate_limits.py +133 -0
  26. lorebase_store/_migrations/versions/0023_agent_trace.py +96 -0
  27. lorebase_store/_migrations/versions/0024_embedding_template_keys.py +205 -0
  28. lorebase_store/_migrations/versions/0025_snapshot_refs.py +35 -0
  29. lorebase_store/_migrations/versions/0026_plain_reverse_link_index.py +112 -0
  30. lorebase_store/_migrations/versions/0027_vault_rotation_and_tombstones.py +145 -0
  31. lorebase_store/_migrations/versions/0028_operation_deferred_events.py +71 -0
  32. lorebase_store/_migrations/versions/0029_agent_trace_sessions.py +108 -0
  33. lorebase_store/_migrations/versions/0030_path_fanout.py +31 -0
  34. lorebase_store/_migrations/versions/0031_branch_builds.py +91 -0
  35. lorebase_store/_migrations/versions/0032_lineage_revisions.py +35 -0
  36. lorebase_store/_migrations/versions/0033_vault_moved_tokens.py +30 -0
  37. lorebase_store/_migrations/versions/0034_inference_endpoint_modes.py +73 -0
  38. lorebase_store/_migrations/versions/0035_search_truncation.py +104 -0
  39. lorebase_store/_migrations/versions/0036_lineage_maps.py +325 -0
  40. lorebase_store/_migrations/versions/0037_bm25_search.py +346 -0
  41. lorebase_store/auth/__init__.py +52 -0
  42. lorebase_store/auth/api_keys.py +161 -0
  43. lorebase_store/auth/audit.py +382 -0
  44. lorebase_store/auth/checkpoints.py +185 -0
  45. lorebase_store/auth/grants.py +161 -0
  46. lorebase_store/auth/principals.py +158 -0
  47. lorebase_store/auth/scope.py +70 -0
  48. lorebase_store/auth/sessions.py +153 -0
  49. lorebase_store/auth/unit.py +53 -0
  50. lorebase_store/auth/usage.py +70 -0
  51. lorebase_store/blobs.py +73 -0
  52. lorebase_store/bm25.py +346 -0
  53. lorebase_store/bulk.py +147 -0
  54. lorebase_store/bundle.py +1830 -0
  55. lorebase_store/changefeed.py +261 -0
  56. lorebase_store/commit/__init__.py +42 -0
  57. lorebase_store/commit/changes.py +67 -0
  58. lorebase_store/commit/errors.py +60 -0
  59. lorebase_store/commit/migrations.py +335 -0
  60. lorebase_store/commit/pipeline.py +754 -0
  61. lorebase_store/commit/schema.py +123 -0
  62. lorebase_store/db.py +58 -0
  63. lorebase_store/definer.py +50 -0
  64. lorebase_store/digests.py +99 -0
  65. lorebase_store/embeddings.py +444 -0
  66. lorebase_store/endpoint_modes.py +59 -0
  67. lorebase_store/extraction/__init__.py +61 -0
  68. lorebase_store/extraction/items.py +187 -0
  69. lorebase_store/extraction/materialize.py +364 -0
  70. lorebase_store/extraction/proposal.py +65 -0
  71. lorebase_store/extraction/review.py +168 -0
  72. lorebase_store/gc.py +970 -0
  73. lorebase_store/indexes.py +1033 -0
  74. lorebase_store/jobs.py +651 -0
  75. lorebase_store/keyring.py +72 -0
  76. lorebase_store/layout.py +335 -0
  77. lorebase_store/materialize/__init__.py +103 -0
  78. lorebase_store/materialize/apply.py +436 -0
  79. lorebase_store/materialize/background.py +1184 -0
  80. lorebase_store/materialize/branch.py +268 -0
  81. lorebase_store/materialize/build.py +536 -0
  82. lorebase_store/materialize/derived.py +540 -0
  83. lorebase_store/materialize/lineage_backfill.py +502 -0
  84. lorebase_store/materialize/links.py +266 -0
  85. lorebase_store/materialize/overlay.py +292 -0
  86. lorebase_store/materialize/paths.py +315 -0
  87. lorebase_store/materialize/rebuild.py +326 -0
  88. lorebase_store/materialize/rows.py +153 -0
  89. lorebase_store/materialize/sensitivity.py +66 -0
  90. lorebase_store/materialize/snapshot.py +205 -0
  91. lorebase_store/materialize/stale.py +444 -0
  92. lorebase_store/materialize/verify.py +321 -0
  93. lorebase_store/merging/__init__.py +39 -0
  94. lorebase_store/merging/attempts.py +334 -0
  95. lorebase_store/merging/lineage.py +123 -0
  96. lorebase_store/merging/service.py +652 -0
  97. lorebase_store/migrations.py +133 -0
  98. lorebase_store/models/__init__.py +136 -0
  99. lorebase_store/models/base.py +64 -0
  100. lorebase_store/models/enums.py +132 -0
  101. lorebase_store/models/governance.py +268 -0
  102. lorebase_store/models/ingestion.py +288 -0
  103. lorebase_store/models/materialized.py +277 -0
  104. lorebase_store/models/proposals.py +195 -0
  105. lorebase_store/models/queues.py +87 -0
  106. lorebase_store/models/tracing.py +116 -0
  107. lorebase_store/models/versioning.py +222 -0
  108. lorebase_store/models/webhooks.py +96 -0
  109. lorebase_store/nodestore.py +384 -0
  110. lorebase_store/objectstore.py +162 -0
  111. lorebase_store/outbox_events.py +166 -0
  112. lorebase_store/profiling/__init__.py +58 -0
  113. lorebase_store/profiling/delimited.py +109 -0
  114. lorebase_store/profiling/postgres.py +211 -0
  115. lorebase_store/profiling/record_files.py +127 -0
  116. lorebase_store/profiling/records.py +92 -0
  117. lorebase_store/profiling/sql_databases.py +298 -0
  118. lorebase_store/ratelimit.py +129 -0
  119. lorebase_store/reads.py +1491 -0
  120. lorebase_store/replicas.py +434 -0
  121. lorebase_store/repositories/__init__.py +152 -0
  122. lorebase_store/repositories/changesets.py +59 -0
  123. lorebase_store/repositories/deferred_events.py +142 -0
  124. lorebase_store/repositories/documents.py +305 -0
  125. lorebase_store/repositories/domains.py +97 -0
  126. lorebase_store/repositories/events.py +205 -0
  127. lorebase_store/repositories/index_view.py +489 -0
  128. lorebase_store/repositories/invocations.py +68 -0
  129. lorebase_store/repositories/materializations.py +276 -0
  130. lorebase_store/repositories/materialized.py +1129 -0
  131. lorebase_store/repositories/merges.py +417 -0
  132. lorebase_store/repositories/proposal_items.py +199 -0
  133. lorebase_store/repositories/proposals.py +184 -0
  134. lorebase_store/repositories/provenance.py +77 -0
  135. lorebase_store/repositories/refs.py +452 -0
  136. lorebase_store/repositories/revisions.py +273 -0
  137. lorebase_store/repositories/scope.py +33 -0
  138. lorebase_store/repositories/source_runs.py +123 -0
  139. lorebase_store/repositories/suggestions.py +638 -0
  140. lorebase_store/runtime.py +87 -0
  141. lorebase_store/schema_cache.py +50 -0
  142. lorebase_store/settings.py +879 -0
  143. lorebase_store/storage.py +1478 -0
  144. lorebase_store/telemetry.py +142 -0
  145. lorebase_store/testing.py +108 -0
  146. lorebase_store/tracing.py +505 -0
  147. lorebase_store/traverse.py +578 -0
  148. lorebase_store/vault/__init__.py +50 -0
  149. lorebase_store/vault/keys.py +216 -0
  150. lorebase_store/vault/quotes.py +133 -0
  151. lorebase_store/vault/rotation.py +209 -0
  152. lorebase_store/vault/settings.py +152 -0
  153. lorebase_store/vault/store.py +987 -0
  154. lorebase_store/vault/transfer.py +558 -0
  155. lorebase_store/webhooks/__init__.py +24 -0
  156. lorebase_store/webhooks/delivery.py +637 -0
  157. lorebase_store/webhooks/keys.py +126 -0
  158. lorebase_store/webhooks/repository.py +142 -0
  159. lorebase_store/webhooks/settings.py +138 -0
  160. lorebase_store/webhooks/signing.py +50 -0
  161. lorebase_store-0.2.3.dist-info/METADATA +13 -0
  162. lorebase_store-0.2.3.dist-info/RECORD +163 -0
  163. lorebase_store-0.2.3.dist-info/WHEEL +4 -0
@@ -0,0 +1,1830 @@
1
+ """Portable Domain bundles: export a Domain to a directory or tarball, import it elsewhere.
2
+
3
+ A bundle holds the versioned content of one Domain: every Revision reachable from a ref with its
4
+ parents and changeset, the tree nodes and objects those Revisions reach, the refs and their
5
+ log, the provenance entries, and a readable `schemas/<branch>.yaml` per branch head. It does
6
+ not hold the audit log, principals, grants, proposals, documents, vault or embeddings (see
7
+ docs/operations.md). Materialized heads are not stored: import rebuilds them from the trees.
8
+
9
+ Layout (the `.jsonl` names carry `.gz` when `store.bundle.compression` is gzip)::
10
+
11
+ manifest.json format version, Lorebase version, counts, heads, sha256 per file
12
+ domain.json the Domain row
13
+ refs.json refs (kind, head, version, protection, materialization, materialized)
14
+ ref_log.jsonl head movements of every ref
15
+ revisions.jsonl id, root, generation, parents, author, message, origin, stats, created_at
16
+ changesets.jsonl element changes of each Revision
17
+ provenance.jsonl provenance entries
18
+ nodes.jsonl tree nodes: hash and base64 body
19
+ objects.jsonl objects: hash, type and canonical body
20
+ vault.jsonl only with vault values: token, binding, value or erased (see below)
21
+ vault_lineage.jsonl only with vault values: key changes of vault Records
22
+ vault_tombstones.jsonl only with vault values: fingerprints of erased values
23
+ schemas/<ref>.yaml schema at the head of each branch (informational)
24
+
25
+ Revisions are ordered by generation, so a parent always precedes its children. Everything is
26
+ written and read in batches: memory is bounded by `store.bundle.batch_rows` and by one tree
27
+ level while marking. Export reads one repeatable-read snapshot. Import verifies the sha256 of
28
+ every file, recomputes every Revision id, node hash and object hash, checks that every
29
+ reachable node and object is present, and runs in one transaction: a failing import leaves
30
+ nothing behind.
31
+
32
+ Revision ids include the Domain id, so a bundle can only be imported under the id it was
33
+ exported with.
34
+
35
+ Values of vault Attributes travel as the references Record objects hold: by default the vault
36
+ is not exported, so after an import they read as unknown (masked), and import refuses a Record
37
+ whose vault Attribute holds anything but a well-formed reference (design 13.5).
38
+
39
+ An export with `include_vault_values` also writes the vault sections
40
+ (`lorebase_store.vault.transfer`): for every token the exported Record objects reference, the
41
+ Kind, Attribute and Record key its vault row binds it to and the opened value, or `erased`; the
42
+ Domain's key changes (`vault_lineage`); and its tombstones with the token key they were made
43
+ under. The files are listed in the manifest with their sha256 like the others, and objects keep
44
+ their hashes. Such an export needs the vault keys and holds personal data in plain text (files
45
+ and tarball are created readable by their owner only). It appends `bundle.export_plain_started`
46
+ to the audit log before reading any value and `bundle.export_plain` (Domain, heads and counts,
47
+ never values) when done; when either cannot be written the output is removed. Its snapshot is
48
+ taken while the Domain's vault lock is held shared, so no erasure is in flight when it is taken;
49
+ the lock is let go once the snapshot exists, before any file is written, and an erasure starting
50
+ meanwhile waits at most `vault.erase_lock_timeout_ms`. Entries carry the row's `moved_to`.
51
+ Tombstones recorded before they kept their `fingerprint_key` are exported as made under the
52
+ current token key with the first scheme, and the manifest says how many (`vault_notes`). The
53
+ audit actor is what the caller gives: the command runs with the owner's database credentials and
54
+ vault keys, and is as trustworthy as whoever holds them.
55
+
56
+ Import needs the vault keys and the token key every tombstone was made under (by id). It stores
57
+ the key changes (each between two Records of its Kind the bundle holds) and tombstones, then
58
+ each entry again, sealed under the target's keys, after checking that every Record object
59
+ holding its token holds it in that Attribute under the bound key or one linked to it by the key
60
+ changes, that its `moved_to` follows a key change from its key, and that the value is valid for
61
+ a definition of the Attribute in the bundle's schemas. A value a tombstone holds is handled as
62
+ `vault.erased_values` says; erased entries stay erased. A `bundle.import_vault` audit entry
63
+ records the counts (values, erased, tombstones, key changes, groups of more than
64
+ `MAX_LINKED_KEYS` linked keys) in the same transaction. A target Domain that already has vault
65
+ rows, tombstones or key changes is not empty.
66
+ """
67
+
68
+ from __future__ import annotations
69
+
70
+ import argparse
71
+ import asyncio
72
+ import base64
73
+ import gzip
74
+ import hashlib
75
+ import importlib.metadata
76
+ import itertools
77
+ import json
78
+ import logging
79
+ import os
80
+ import shutil
81
+ import stat
82
+ import sys
83
+ import tarfile
84
+ import tempfile
85
+ import time
86
+ from collections.abc import Collection, Iterable, Iterator, Mapping, Sequence
87
+ from contextlib import AsyncExitStack, contextmanager
88
+ from dataclasses import asdict, dataclass, field
89
+ from datetime import UTC, datetime
90
+ from pathlib import Path, PurePosixPath
91
+ from typing import IO, Any, BinaryIO, Protocol
92
+ from urllib.parse import quote
93
+
94
+ from sqlalchemy import ARRAY, LargeBinary, bindparam, text
95
+ from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
96
+
97
+ from lorebase_core.canonical import ObjectType, from_hex, node_hash
98
+ from lorebase_core.canonical import hex as to_hex
99
+ from lorebase_core.model import Attribute, Kind, dump_schema, vault_token
100
+ from lorebase_core.revision import (
101
+ Author,
102
+ Origin,
103
+ Revision,
104
+ RevisionStats,
105
+ format_timestamp,
106
+ next_generation,
107
+ revision_id,
108
+ )
109
+ from lorebase_core.tree import decode_node
110
+ from lorebase_store.auth import AuditRepository, NewAuditEntry
111
+ from lorebase_store.bulk import OnConflict, bulk_write
112
+ from lorebase_store.gc import Marker
113
+ from lorebase_store.materialize import rebuild, verify
114
+ from lorebase_store.materialize.rebuild import load_schema
115
+ from lorebase_store.models.enums import Materialization, RefKind
116
+ from lorebase_store.nodestore import NodeRepository
117
+ from lorebase_store.objectstore import OBJECT_TYPES_BY_CODE, PgObjectStore, encode_object
118
+ from lorebase_store.repositories.changesets import CHANGESETS
119
+ from lorebase_store.repositories.domains import DomainNotFoundError, DomainRepository, DomainRow
120
+ from lorebase_store.repositories.provenance import PROVENANCE
121
+ from lorebase_store.repositories.refs import REF_LOG, REFS, RefRepository
122
+ from lorebase_store.repositories.revisions import REVISION_PARENTS, REVISIONS, RevisionRepository
123
+ from lorebase_store.repositories.scope import scope_to_domain
124
+ from lorebase_store.runtime import StoreRuntime, default_runtime
125
+ from lorebase_store.settings import BundleSettings
126
+ from lorebase_store.vault import VaultKeys
127
+ from lorebase_store.vault.store import MAX_LINKED_KEYS
128
+ from lorebase_store.vault.transfer import (
129
+ LEGACY_SCHEME,
130
+ Holder,
131
+ KeyLinks,
132
+ LineageEntry,
133
+ TombstoneEntry,
134
+ VaultEntry,
135
+ VaultTransferError,
136
+ holders_problem,
137
+ lineage_of,
138
+ parse_entry,
139
+ parse_lineage,
140
+ parse_tombstone,
141
+ read_entries,
142
+ resolved_legacy,
143
+ restore_entries,
144
+ restore_lineage,
145
+ restore_tombstones,
146
+ tombstones_of,
147
+ unknown_entry,
148
+ value_problem,
149
+ )
150
+
151
+ log = logging.getLogger(__name__)
152
+
153
+ FORMAT = "lorebase-domain-bundle"
154
+ FORMAT_VERSION = 1
155
+ MANIFEST = "manifest.json"
156
+ SCHEMAS_DIR = "schemas"
157
+ TAR_SUFFIX = ".tar"
158
+ VAULT = "vault"
159
+ VAULT_LINEAGE = "vault_lineage"
160
+ VAULT_TOMBSTONES = "vault_tombstones"
161
+ _VAULT_SECTIONS = (VAULT, VAULT_LINEAGE, VAULT_TOMBSTONES)
162
+ EXPORT_PLAIN_STARTED_ACTION = "bundle.export_plain_started"
163
+ """Audit action written before an export holding vault values reads any value."""
164
+ EXPORT_PLAIN_ACTION = "bundle.export_plain"
165
+ """Audit action of a finished export holding vault values."""
166
+ IMPORT_VAULT_ACTION = "bundle.import_vault"
167
+ """Audit action of an import of vault values (counts, never values)."""
168
+ _VAULT_LOCK_SHARED = text(
169
+ "SELECT pg_advisory_xact_lock_shared(hashtextextended('lorebase.vault/' || :domain, 0))"
170
+ )
171
+ BUNDLE_ACTOR = "system:bundle"
172
+ """Default audit actor of the bundle command."""
173
+ PLAIN_TEXT_WARNING = (
174
+ "warning: {path} holds personal data in plain text (vault values); keep it as such and "
175
+ "delete it once imported\n"
176
+ )
177
+
178
+ _JSONL = ("revisions", "changesets", "provenance", "ref_log", "nodes", "objects")
179
+ _PLAIN = ("domain.json", "refs.json")
180
+ _HASH_CHUNK = 1 << 20
181
+ _NOFOLLOW = getattr(os, "O_NOFOLLOW", 0)
182
+ _DIRECTORY = getattr(os, "O_DIRECTORY", 0)
183
+ _NONBLOCK = getattr(os, "O_NONBLOCK", 0)
184
+
185
+
186
+ class BundleError(Exception):
187
+ """The bundle cannot be written, is malformed, or fails verification."""
188
+
189
+
190
+ class DomainNotEmptyError(BundleError):
191
+ """The target Domain already has Revisions or refs."""
192
+
193
+
194
+ class VaultNotConfiguredError(BundleError):
195
+ """Vault values are asked for or carried, but the vault keys are not configured."""
196
+
197
+
198
+ @dataclass(slots=True)
199
+ class BundleReport:
200
+ """What an export or import moved."""
201
+
202
+ domain_id: str
203
+ path: str
204
+ revisions: int = 0
205
+ changesets: int = 0
206
+ provenance: int = 0
207
+ nodes: int = 0
208
+ objects: int = 0
209
+ refs: int = 0
210
+ ref_log: int = 0
211
+ schemas: int = 0
212
+ vault_values: int = 0
213
+ vault_erased: int = 0
214
+ vault_lineage: int = 0
215
+ vault_legacy_tombstones: int = 0
216
+ vault_oversized_links: int = 0
217
+ vault_tombstones: int = 0
218
+ rebuilt_branches: list[str] = field(default_factory=list)
219
+ seconds: float = 0.0
220
+
221
+ def as_dict(self) -> dict[str, Any]:
222
+ return asdict(self)
223
+
224
+
225
+ # --------------------------------------------------------------------------------------
226
+ # Files
227
+
228
+
229
+ def _jsonl_name(base: str, compression: str) -> str:
230
+ return f"{base}.jsonl.gz" if compression == "gzip" else f"{base}.jsonl"
231
+
232
+
233
+ def _dumps(value: Any) -> bytes:
234
+ return json.dumps(value, ensure_ascii=False, separators=(",", ":"), sort_keys=True).encode()
235
+
236
+
237
+ def _create(path: Path) -> BinaryIO:
238
+ """A new file only its owner may read or write."""
239
+ return os.fdopen(os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL | _NOFOLLOW, 0o600), "wb")
240
+
241
+
242
+ def _sha256(path: Path) -> tuple[str, int]:
243
+ digest = hashlib.sha256()
244
+ size = 0
245
+ with path.open("rb") as handle:
246
+ while chunk := handle.read(_HASH_CHUNK):
247
+ digest.update(chunk)
248
+ size += len(chunk)
249
+ return digest.hexdigest(), size
250
+
251
+
252
+ class _Output:
253
+ """The files of a bundle being written into a directory."""
254
+
255
+ def __init__(self, root: Path, compression: str, level: int) -> None:
256
+ self.root = root
257
+ self.compression = compression
258
+ self.level = level
259
+ self.files: dict[str, dict[str, Any]] = {}
260
+
261
+ @contextmanager
262
+ def jsonl(self, base: str, *, private: bool = False) -> Iterator[_LineWriter]:
263
+ """Lines of `base`; `private` creates the file readable by its owner only."""
264
+ name = _jsonl_name(base, self.compression)
265
+ path = self.root / name
266
+ raw = _create(path) if private else path.open("wb")
267
+ stream: _Sink = raw
268
+ if self.compression == "gzip":
269
+ stream = gzip.GzipFile(
270
+ filename="", mode="wb", fileobj=raw, compresslevel=self.level, mtime=0
271
+ )
272
+ writer = _LineWriter(stream)
273
+ try:
274
+ yield writer
275
+ finally:
276
+ stream.close()
277
+ raw.close()
278
+ self._register(name, path, writer.records)
279
+
280
+ def plain(self, name: str, content: bytes, records: int | None = None) -> None:
281
+ path = self.root / name
282
+ path.parent.mkdir(parents=True, exist_ok=True)
283
+ path.write_bytes(content)
284
+ self._register(name, path, records)
285
+
286
+ def _register(self, name: str, path: Path, records: int | None) -> None:
287
+ digest, size = _sha256(path)
288
+ entry: dict[str, Any] = {"sha256": digest, "bytes": size}
289
+ if records is not None:
290
+ entry["records"] = records
291
+ self.files[name] = entry
292
+
293
+
294
+ class _Sink(Protocol):
295
+ def write(self, data: bytes, /) -> Any: ...
296
+
297
+ def close(self) -> None: ...
298
+
299
+
300
+ class _LineWriter:
301
+ def __init__(self, stream: _Sink) -> None:
302
+ self._stream = stream
303
+ self.records = 0
304
+
305
+ def write(self, value: Any) -> None:
306
+ self._stream.write(_dumps(value) + b"\n")
307
+ self.records += 1
308
+
309
+
310
+ def _read_jsonl(root: Path, name: str, limits: BundleSettings) -> Iterator[dict[str, Any]]:
311
+ """The records of a `.jsonl` file, read line by line within the store.bundle limits."""
312
+ with _open_file(root, name) as raw:
313
+ handle: gzip.GzipFile | BinaryIO = (
314
+ gzip.GzipFile(fileobj=raw, mode="rb") if name.endswith(".gz") else raw
315
+ )
316
+ size = 0
317
+ number = 0
318
+ try:
319
+ while line := handle.readline(limits.max_line_bytes + 1):
320
+ number += 1
321
+ size += len(line)
322
+ _check_size(name, size, 0, limits)
323
+ if len(line) > limits.max_line_bytes and not line.endswith(b"\n"):
324
+ raise BundleError(
325
+ f"{name} line {number} is longer than store.bundle.max_line_bytes "
326
+ f"({limits.max_line_bytes})"
327
+ )
328
+ try:
329
+ value = json.loads(line)
330
+ except ValueError as error:
331
+ raise BundleError(f"{name} line {number}: not JSON") from error
332
+ if not isinstance(value, dict):
333
+ raise BundleError(f"{name} line {number}: not an object")
334
+ yield value
335
+ except (OSError, EOFError) as error:
336
+ raise BundleError(f"{name} cannot be read: {error}") from error
337
+ finally:
338
+ handle.close()
339
+
340
+
341
+ def _batches[T](items: Iterable[T], size: int) -> Iterator[tuple[T, ...]]:
342
+ iterator = iter(items)
343
+ while batch := tuple(itertools.islice(iterator, size)):
344
+ yield batch
345
+
346
+
347
+ # --------------------------------------------------------------------------------------
348
+ # Export
349
+
350
+
351
+ def _bytes(name: str) -> Any:
352
+ return bindparam(name, type_=ARRAY(LargeBinary()))
353
+
354
+
355
+ _REVISIONS = text(
356
+ """
357
+ SELECT r.id, r.root, r.generation, r.author, r.on_behalf_of, r.message, r.origin, r.stats,
358
+ r.created_at
359
+ FROM revisions r JOIN lb_gc_revs m ON m.id = r.id
360
+ WHERE r.domain_id = :d AND (r.generation, r.id) > (:g, :i)
361
+ ORDER BY r.generation, r.id LIMIT :n
362
+ """
363
+ )
364
+ _PARENTS = text(
365
+ "SELECT revision_id, parent_id FROM revision_parents "
366
+ "WHERE domain_id = :d AND revision_id = ANY(CAST(:ids AS bytea[])) "
367
+ "ORDER BY revision_id, position"
368
+ ).bindparams(_bytes("ids"))
369
+ _CHANGESETS = text(
370
+ """
371
+ SELECT c.revision_id, c.partition, c.element_key, c.old_hash, c.new_hash, c.provenance
372
+ FROM changesets c JOIN lb_gc_revs m ON m.id = c.revision_id
373
+ WHERE c.domain_id = :d AND (c.revision_id, c.partition, c.element_key) > (:r, :p, :k)
374
+ ORDER BY c.revision_id, c.partition, c.element_key LIMIT :n
375
+ """
376
+ )
377
+ _PROVENANCE = text(
378
+ "SELECT id, origin_type, origin_ref, actor, at FROM provenance "
379
+ "WHERE domain_id = :d AND id > :last ORDER BY id LIMIT :n"
380
+ )
381
+ _REF_LOG = text(
382
+ "SELECT ref_name, seq, old_head, new_head, reason, actor, at FROM ref_log "
383
+ "WHERE domain_id = :d AND (ref_name, seq) > (:r, :s) ORDER BY ref_name, seq LIMIT :n"
384
+ )
385
+ _NODES = text(
386
+ "SELECT n.hash, n.body FROM tree_nodes n JOIN lb_gc_nodes m ON m.hash = n.hash "
387
+ "WHERE n.domain_id = :d AND n.hash > :last ORDER BY n.hash LIMIT :n"
388
+ )
389
+ _OBJECTS = text(
390
+ "SELECT o.hash, o.object_type, o.body FROM objects o JOIN lb_gc_objects m ON m.hash = o.hash "
391
+ "WHERE o.domain_id = :d AND o.hash > :last ORDER BY o.hash LIMIT :n"
392
+ )
393
+ _REFS = text(
394
+ """
395
+ SELECT r.name, r.kind::text AS kind, r.head, r.protected, r.materialization::text AS
396
+ materialization, r.version, r.created_at, r.created_by,
397
+ EXISTS (SELECT 1 FROM mat_schema s WHERE s.domain_id = r.domain_id
398
+ AND s.branch = r.name) AS materialized
399
+ FROM refs r WHERE r.domain_id = :d AND r.deleted_at IS NULL AND r.kind <> 'snapshot'
400
+ ORDER BY r.name
401
+ """
402
+ )
403
+
404
+
405
+ def _hex(value: bytes | None) -> str | None:
406
+ return None if value is None else to_hex(bytes(value))
407
+
408
+
409
+ def _timestamp(value: datetime) -> str:
410
+ return format_timestamp(value)
411
+
412
+
413
+ async def export_domain(
414
+ engine: AsyncEngine,
415
+ domain_id: str,
416
+ destination: str | Path,
417
+ *,
418
+ runtime: StoreRuntime | None = None,
419
+ include_vault_values: bool = False,
420
+ actor: str = BUNDLE_ACTOR,
421
+ ) -> BundleReport:
422
+ """Write a bundle of `domain_id` to `destination`: a new directory, or a `.tar` file.
423
+
424
+ Reads one consistent snapshot; commits made meanwhile are not part of the bundle.
425
+ `include_vault_values` also writes the opened vault values in plain text and appends an
426
+ audit entry for `actor` (module doc); it raises `VaultNotConfiguredError` without the
427
+ vault keys. When such an export fails, nothing of it is left.
428
+ """
429
+ runtime = runtime or default_runtime()
430
+ keys = runtime.vault if include_vault_values else None
431
+ if include_vault_values and keys is None:
432
+ raise VaultNotConfiguredError(
433
+ "exporting vault values needs vault.token_key and vault.encryption_key"
434
+ )
435
+ target = Path(destination)
436
+ started = time.monotonic()
437
+ as_tar = target.suffix == TAR_SUFFIX
438
+ _check_destination(target, as_tar)
439
+ if keys is not None:
440
+ try:
441
+ await _audit(engine, actor, EXPORT_PLAIN_STARTED_ACTION, domain_id, {})
442
+ except BaseException:
443
+ if not as_tar:
444
+ shutil.rmtree(target, ignore_errors=True)
445
+ raise
446
+ if as_tar:
447
+ with tempfile.TemporaryDirectory(dir=target.parent, prefix=".bundle-") as work:
448
+ report, heads = await _export(engine, domain_id, Path(work), runtime, str(target), keys)
449
+ if keys is not None:
450
+ await _audit_plain_export(engine, report, heads, actor)
451
+ _pack(Path(work), target, private=keys is not None)
452
+ else:
453
+ try:
454
+ report, heads = await _export(engine, domain_id, target, runtime, str(target), keys)
455
+ if keys is not None:
456
+ await _audit_plain_export(engine, report, heads, actor)
457
+ except BaseException:
458
+ if keys is not None:
459
+ shutil.rmtree(target, ignore_errors=True)
460
+ raise
461
+ report.seconds = time.monotonic() - started
462
+ return report
463
+
464
+
465
+ async def _audit_plain_export(
466
+ engine: AsyncEngine, report: BundleReport, heads: Mapping[str, str], actor: str
467
+ ) -> None:
468
+ """The audit entry of a finished export holding vault values: counts, never values."""
469
+ details = {
470
+ "heads": dict(heads),
471
+ "values": report.vault_values,
472
+ "erased": report.vault_erased,
473
+ "lineage": report.vault_lineage,
474
+ "tombstones": report.vault_tombstones,
475
+ }
476
+ await _audit(engine, actor, EXPORT_PLAIN_ACTION, report.domain_id, details)
477
+
478
+
479
+ async def _audit(
480
+ engine: AsyncEngine, actor: str, action: str, domain_id: str, details: Mapping[str, Any]
481
+ ) -> None:
482
+ sessions = async_sessionmaker(engine, expire_on_commit=False)
483
+ async with sessions() as session, session.begin():
484
+ await AuditRepository(session).append(
485
+ NewAuditEntry(
486
+ actor=actor, action=action, domain_id=domain_id, target=domain_id, details=details
487
+ )
488
+ )
489
+
490
+
491
+ def _check_destination(target: Path, as_tar: bool) -> None:
492
+ if as_tar:
493
+ if target.exists():
494
+ raise BundleError(f"{target} exists")
495
+ target.parent.mkdir(parents=True, exist_ok=True)
496
+ else:
497
+ if target.exists() and (not target.is_dir() or any(target.iterdir())):
498
+ raise BundleError(f"{target} exists and is not an empty directory")
499
+ target.mkdir(parents=True, exist_ok=True)
500
+
501
+
502
+ def _pack(work: Path, target: Path, *, private: bool = False) -> None:
503
+ """Pack `work` into `target`; `private` makes it readable by its owner only."""
504
+ partial = target.with_name(target.name + ".partial")
505
+ try:
506
+ with (
507
+ _create(partial) if private else partial.open("wb") as sink,
508
+ tarfile.open(fileobj=sink, mode="w") as archive,
509
+ ):
510
+ for path in sorted(work.rglob("*")):
511
+ if path.is_file():
512
+ archive.add(path, arcname=path.relative_to(work).as_posix())
513
+ except BaseException:
514
+ partial.unlink(missing_ok=True)
515
+ raise
516
+ partial.rename(target)
517
+
518
+
519
+ async def _export(
520
+ engine: AsyncEngine,
521
+ domain_id: str,
522
+ root: Path,
523
+ runtime: StoreRuntime,
524
+ label: str,
525
+ keys: VaultKeys | None,
526
+ ) -> tuple[BundleReport, dict[str, str]]:
527
+ settings = runtime.settings
528
+ batch = settings.bundle.batch_rows
529
+ out = _Output(root, settings.bundle.compression, settings.bundle.compression_level)
530
+ report = BundleReport(domain_id, label)
531
+ sessions = async_sessionmaker(engine, expire_on_commit=False)
532
+ async with AsyncExitStack() as stack:
533
+ guard = None
534
+ if keys is not None:
535
+ # The vault lock is held shared from another connection while the snapshot below is
536
+ # taken: an erasure in flight ends first, and one starting meanwhile waits (at most
537
+ # vault.erase_lock_timeout_ms). The snapshot fixes what is read; the lock is let go
538
+ # as soon as it exists, before any file is written.
539
+ guard = await stack.enter_async_context(engine.connect())
540
+ await guard.execute(_VAULT_LOCK_SHARED, {"domain": domain_id})
541
+ session = await stack.enter_async_context(sessions())
542
+ await stack.enter_async_context(session.begin())
543
+ await session.execute(text("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ"))
544
+ await scope_to_domain(session, domain_id) # the first query: it takes the snapshot
545
+ if guard is not None:
546
+ await guard.close()
547
+ domain = await DomainRepository(session).get(domain_id)
548
+ if domain is None:
549
+ raise DomainNotFoundError(domain_id)
550
+ refs = [dict(row) for row in (await session.execute(_REFS, {"d": domain_id})).mappings()]
551
+ marker = Marker(session, domain_id, batch=settings.gc.mark_batch)
552
+ await marker.prepare()
553
+ await marker.mark(attempts=False)
554
+ marked = await marker.stats()
555
+ if marked.missing_nodes:
556
+ raise BundleError(f"{marked.missing_nodes} reachable tree nodes are missing")
557
+
558
+ out.plain("domain.json", _dumps(_domain_json(domain)))
559
+ out.plain("refs.json", _dumps([_ref_json(ref) for ref in refs]), len(refs))
560
+ report.refs = len(refs)
561
+ report.revisions = await _export_revisions(session, out, domain_id, batch)
562
+ report.changesets = await _export_changesets(session, out, domain_id, batch)
563
+ report.provenance = await _export_provenance(session, out, domain_id, batch)
564
+ report.ref_log = await _export_ref_log(session, out, domain_id, batch)
565
+ report.nodes = await _export_nodes(session, out, domain_id, batch)
566
+ holders: dict[str, set[Holder]] | None = None if keys is None else {}
567
+ report.objects = await _export_objects(session, out, domain_id, batch, holders)
568
+ for name, got, want in (
569
+ ("revisions", report.revisions, marked.revisions),
570
+ ("nodes", report.nodes, marked.nodes),
571
+ ("objects", report.objects, marked.objects),
572
+ ):
573
+ if got != want:
574
+ raise BundleError(f"expected {want} {name} but the store holds {got}")
575
+ report.schemas = await _export_schemas(session, out, domain_id, refs, runtime)
576
+ if keys is not None and holders is not None:
577
+ await _export_vault(session, out, domain_id, batch, holders, keys, report)
578
+ await marker.drop()
579
+
580
+ heads = {ref["name"]: to_hex(bytes(ref["head"])) for ref in refs}
581
+ counts = {
582
+ "revisions": report.revisions,
583
+ "changesets": report.changesets,
584
+ "provenance": report.provenance,
585
+ "ref_log": report.ref_log,
586
+ "nodes": report.nodes,
587
+ "objects": report.objects,
588
+ "refs": report.refs,
589
+ }
590
+ if keys is not None:
591
+ counts[VAULT] = report.vault_values + report.vault_erased
592
+ counts[VAULT_LINEAGE] = report.vault_lineage
593
+ counts[VAULT_TOMBSTONES] = report.vault_tombstones
594
+ manifest = {
595
+ "format": FORMAT,
596
+ "format_version": FORMAT_VERSION,
597
+ "lorebase_version": _lorebase_version(),
598
+ "domain_id": domain_id,
599
+ "created_at": format_timestamp(datetime.now(UTC)),
600
+ "compression": out.compression,
601
+ "heads": heads,
602
+ "counts": counts,
603
+ "files": out.files,
604
+ }
605
+ if report.vault_legacy_tombstones:
606
+ manifest["vault_notes"] = {
607
+ "legacy_tombstones": report.vault_legacy_tombstones,
608
+ "legacy_fingerprint_key": f"{keys.token_key_id}/{LEGACY_SCHEME}" if keys else None,
609
+ "note": "tombstones recorded without a fingerprint_key, exported as made under the "
610
+ "current token key with the first fingerprint scheme",
611
+ }
612
+ (root / MANIFEST).write_bytes(json.dumps(manifest, indent=2, sort_keys=True).encode() + b"\n")
613
+ return report, heads
614
+
615
+
616
+ def _lorebase_version() -> str:
617
+ try:
618
+ return importlib.metadata.version("lorebase-store")
619
+ except importlib.metadata.PackageNotFoundError:
620
+ return "unknown"
621
+
622
+
623
+ def _domain_json(domain: DomainRow) -> dict[str, Any]:
624
+ return {
625
+ "id": domain.id,
626
+ "display_name": domain.display_name,
627
+ "description": domain.description,
628
+ "default_branch": domain.default_branch,
629
+ "leaf_size": domain.leaf_size,
630
+ "settings": domain.settings,
631
+ "created_at": _timestamp(domain.created_at),
632
+ "created_by": domain.created_by,
633
+ }
634
+
635
+
636
+ def _ref_json(ref: Mapping[str, Any]) -> dict[str, Any]:
637
+ return {
638
+ "name": ref["name"],
639
+ "kind": ref["kind"],
640
+ "head": to_hex(bytes(ref["head"])),
641
+ "protected": ref["protected"],
642
+ "materialization": ref["materialization"],
643
+ "version": ref["version"],
644
+ "created_at": _timestamp(ref["created_at"]),
645
+ "created_by": ref["created_by"],
646
+ "materialized": ref["materialized"],
647
+ }
648
+
649
+
650
+ async def _export_revisions(session: AsyncSession, out: _Output, domain: str, batch: int) -> int:
651
+ after: tuple[int, bytes] = (0, b"")
652
+ with out.jsonl("revisions") as lines:
653
+ while True:
654
+ rows = (
655
+ await session.execute(
656
+ _REVISIONS, {"d": domain, "g": after[0], "i": after[1], "n": batch}
657
+ )
658
+ ).all()
659
+ if not rows:
660
+ return lines.records
661
+ ids = [bytes(row.id) for row in rows]
662
+ parents: dict[bytes, list[str]] = {}
663
+ for revision, parent in await session.execute(_PARENTS, {"d": domain, "ids": ids}):
664
+ parents.setdefault(bytes(revision), []).append(to_hex(bytes(parent)))
665
+ for row in rows:
666
+ lines.write(
667
+ {
668
+ "id": to_hex(bytes(row.id)),
669
+ "root": to_hex(bytes(row.root)),
670
+ "generation": row.generation,
671
+ "parents": parents.get(bytes(row.id), []),
672
+ "author": row.author,
673
+ "on_behalf_of": row.on_behalf_of,
674
+ "message": row.message,
675
+ "origin": row.origin,
676
+ "stats": row.stats,
677
+ "created_at": _timestamp(row.created_at),
678
+ }
679
+ )
680
+ after = (int(rows[-1].generation), bytes(rows[-1].id))
681
+
682
+
683
+ async def _export_changesets(session: AsyncSession, out: _Output, domain: str, batch: int) -> int:
684
+ after: tuple[bytes, str, str] = (b"", "", "")
685
+ with out.jsonl("changesets") as lines:
686
+ while True:
687
+ rows = (
688
+ await session.execute(
689
+ _CHANGESETS,
690
+ {"d": domain, "r": after[0], "p": after[1], "k": after[2], "n": batch},
691
+ )
692
+ ).all()
693
+ if not rows:
694
+ return lines.records
695
+ for row in rows:
696
+ lines.write(
697
+ {
698
+ "revision": to_hex(bytes(row.revision_id)),
699
+ "partition": row.partition,
700
+ "key": row.element_key,
701
+ "old": _hex(row.old_hash),
702
+ "new": _hex(row.new_hash),
703
+ "provenance": row.provenance,
704
+ }
705
+ )
706
+ last = rows[-1]
707
+ after = (bytes(last.revision_id), last.partition, last.element_key)
708
+
709
+
710
+ async def _export_provenance(session: AsyncSession, out: _Output, domain: str, batch: int) -> int:
711
+ last = ""
712
+ with out.jsonl("provenance") as lines:
713
+ while True:
714
+ rows = (
715
+ await session.execute(_PROVENANCE, {"d": domain, "last": last, "n": batch})
716
+ ).all()
717
+ if not rows:
718
+ return lines.records
719
+ for row in rows:
720
+ lines.write(
721
+ {
722
+ "id": row.id,
723
+ "origin_type": row.origin_type,
724
+ "origin_ref": row.origin_ref,
725
+ "actor": row.actor,
726
+ "at": _timestamp(row.at),
727
+ }
728
+ )
729
+ last = rows[-1].id
730
+
731
+
732
+ async def _export_ref_log(session: AsyncSession, out: _Output, domain: str, batch: int) -> int:
733
+ after: tuple[str, int] = ("", 0)
734
+ with out.jsonl("ref_log") as lines:
735
+ while True:
736
+ rows = (
737
+ await session.execute(
738
+ _REF_LOG, {"d": domain, "r": after[0], "s": after[1], "n": batch}
739
+ )
740
+ ).all()
741
+ if not rows:
742
+ return lines.records
743
+ for row in rows:
744
+ lines.write(
745
+ {
746
+ "ref": row.ref_name,
747
+ "seq": row.seq,
748
+ "old": _hex(row.old_head),
749
+ "new": to_hex(bytes(row.new_head)),
750
+ "reason": row.reason,
751
+ "actor": row.actor,
752
+ "at": _timestamp(row.at),
753
+ }
754
+ )
755
+ after = (rows[-1].ref_name, int(rows[-1].seq))
756
+
757
+
758
+ async def _export_nodes(session: AsyncSession, out: _Output, domain: str, batch: int) -> int:
759
+ last = b""
760
+ with out.jsonl("nodes") as lines:
761
+ while True:
762
+ rows = (await session.execute(_NODES, {"d": domain, "last": last, "n": batch})).all()
763
+ if not rows:
764
+ return lines.records
765
+ for row in rows:
766
+ lines.write(
767
+ {
768
+ "hash": to_hex(bytes(row.hash)),
769
+ "body": base64.b64encode(bytes(row.body)).decode("ascii"),
770
+ }
771
+ )
772
+ last = bytes(rows[-1].hash)
773
+
774
+
775
+ async def _export_objects(
776
+ session: AsyncSession,
777
+ out: _Output,
778
+ domain: str,
779
+ batch: int,
780
+ holders: dict[str, set[Holder]] | None = None,
781
+ ) -> int:
782
+ """Write the marked objects; with `holders`, collect the vault tokens of Record objects."""
783
+ last = b""
784
+ with out.jsonl("objects") as lines:
785
+ while True:
786
+ rows = (await session.execute(_OBJECTS, {"d": domain, "last": last, "n": batch})).all()
787
+ if not rows:
788
+ return lines.records
789
+ for row in rows:
790
+ object_type = OBJECT_TYPES_BY_CODE[row.object_type]
791
+ line = {
792
+ "hash": to_hex(bytes(row.hash)),
793
+ "type": object_type.value,
794
+ "body": row.body,
795
+ }
796
+ if holders is not None and object_type is ObjectType.RECORD:
797
+ _collect_holders(line, holders)
798
+ lines.write(line)
799
+ last = bytes(rows[-1].hash)
800
+
801
+
802
+ def _collect_holders(line: Mapping[str, Any], holders: dict[str, set[Holder]]) -> None:
803
+ """Add the Record of a Record object as a holder of each vault token it holds."""
804
+ body = line.get("body") or {}
805
+ values = body.get("values") or {}
806
+ if not isinstance(values, dict):
807
+ return
808
+ for attribute_id, value in values.items():
809
+ token = vault_token(value)
810
+ if token is not None:
811
+ holders.setdefault(token, set()).add(
812
+ Holder(str(body.get("kind")), str(attribute_id), str(body.get("key")))
813
+ )
814
+
815
+
816
+ async def _export_vault(
817
+ session: AsyncSession,
818
+ out: _Output,
819
+ domain: str,
820
+ batch: int,
821
+ holders: Mapping[str, set[Holder]],
822
+ keys: VaultKeys,
823
+ report: BundleReport,
824
+ ) -> None:
825
+ """Write the key changes, the tombstones and the opened value (or `erased`) of every token
826
+ the exported objects hold, bound as its vault row is. Every Record holding a token must
827
+ hold it in the bound Attribute under the bound key or one linked to it (a key change on an
828
+ erased reference keeps the token)."""
829
+ lineage = await lineage_of(session, domain)
830
+ links = KeyLinks()
831
+ with out.jsonl(VAULT_LINEAGE, private=True) as lines:
832
+ for item in lineage:
833
+ links.link(item)
834
+ lines.write(item.as_json())
835
+ report.vault_lineage = len(lineage)
836
+ with out.jsonl(VAULT_TOMBSTONES, private=True) as lines:
837
+ async for found in tombstones_of(session, domain, batch):
838
+ for tombstone in found:
839
+ if tombstone.fingerprint_key is None:
840
+ report.vault_legacy_tombstones += 1
841
+ lines.write(resolved_legacy(tombstone, keys.token_key_id).as_json())
842
+ report.vault_tombstones += len(found)
843
+ tokens = sorted(holders)
844
+ with out.jsonl(VAULT, private=True) as lines:
845
+ for start in range(0, len(tokens), batch):
846
+ chunk = tokens[start : start + batch]
847
+ try:
848
+ stored = await read_entries(session, domain, chunk, keys)
849
+ except VaultTransferError as error:
850
+ raise BundleError(str(error)) from error
851
+ for token in chunk:
852
+ entry = stored.get(token) or unknown_entry(token, holders[token])
853
+ problem = holders_problem(
854
+ token,
855
+ entry.kind_id,
856
+ entry.attribute_id,
857
+ entry.record_key,
858
+ holders[token],
859
+ links,
860
+ )
861
+ if problem is not None:
862
+ raise BundleError(problem)
863
+ if entry.erased:
864
+ report.vault_erased += 1
865
+ else:
866
+ report.vault_values += 1
867
+ lines.write(entry.as_json())
868
+
869
+
870
+ async def _export_schemas(
871
+ session: AsyncSession,
872
+ out: _Output,
873
+ domain_id: str,
874
+ refs: Sequence[Mapping[str, Any]],
875
+ runtime: StoreRuntime,
876
+ ) -> int:
877
+ settings = runtime.settings
878
+ nodes = NodeRepository(
879
+ session, domain_id, writes=settings.writes, batch=settings.nodes.prefetch_batch
880
+ )
881
+ revisions = RevisionRepository(session, domain_id)
882
+ count = 0
883
+ for ref in refs:
884
+ if ref["kind"] != RefKind.BRANCH.value:
885
+ continue
886
+ head = await revisions.require(bytes(ref["head"]))
887
+ schema, _ = await load_schema(
888
+ session, domain_id, await nodes.load_tree(head.root), runtime=runtime
889
+ )
890
+ name = f"{SCHEMAS_DIR}/{quote(ref['name'], safe='')}.yaml"
891
+ out.plain(name, dump_schema(schema).encode())
892
+ count += 1
893
+ return count
894
+
895
+
896
+ # --------------------------------------------------------------------------------------
897
+ # Import
898
+
899
+
900
+ def _check_size(name: str, size: int, total: int, limits: BundleSettings) -> None:
901
+ if size > limits.max_member_bytes:
902
+ raise BundleError(
903
+ f"{name} is larger than store.bundle.max_member_bytes ({limits.max_member_bytes})"
904
+ )
905
+ if total > limits.max_total_bytes:
906
+ raise BundleError(
907
+ f"the bundle exceeds store.bundle.max_total_bytes ({limits.max_total_bytes}) at {name}"
908
+ )
909
+
910
+
911
+ def _check_entries(name: str, entries: int, limits: BundleSettings) -> None:
912
+ if entries > limits.max_entries:
913
+ raise BundleError(
914
+ f"the bundle has more entries than store.bundle.max_entries ({limits.max_entries}) "
915
+ f"at {name}"
916
+ )
917
+
918
+
919
+ class _Meter:
920
+ """Bytes read from the files of one bundle, checked against the size limits as they come."""
921
+
922
+ def __init__(self, limits: BundleSettings) -> None:
923
+ self.limits = limits
924
+ self.total = 0
925
+
926
+ def chunks(self, name: str, source: IO[bytes]) -> Iterator[bytes]:
927
+ size = 0
928
+ while chunk := source.read(_HASH_CHUNK):
929
+ size += len(chunk)
930
+ self.total += len(chunk)
931
+ _check_size(name, size, self.total, self.limits)
932
+ yield chunk
933
+
934
+
935
+ def _open_file(root: Path, name: str) -> BinaryIO:
936
+ """Open a regular file of a bundle directory without following a link at any level."""
937
+ parts = PurePosixPath(name).parts
938
+ try:
939
+ folder = os.open(root, os.O_RDONLY | _DIRECTORY)
940
+ try:
941
+ for part in parts[:-1]:
942
+ inner = os.open(
943
+ part, os.O_RDONLY | _DIRECTORY | _NOFOLLOW | _NONBLOCK, dir_fd=folder
944
+ )
945
+ os.close(folder)
946
+ folder = inner
947
+ descriptor = os.open(parts[-1], os.O_RDONLY | _NOFOLLOW | _NONBLOCK, dir_fd=folder)
948
+ finally:
949
+ os.close(folder)
950
+ except OSError as error:
951
+ raise BundleError(
952
+ f"{name} cannot be opened (links are not followed): {error.strerror}"
953
+ ) from error
954
+ if not stat.S_ISREG(os.fstat(descriptor).st_mode):
955
+ os.close(descriptor)
956
+ raise BundleError(f"{name} is not a regular file")
957
+ return os.fdopen(descriptor, "rb")
958
+
959
+
960
+ def _read_document(root: Path, name: str, limits: BundleSettings) -> Any:
961
+ """A whole `.json` file of a bundle, at most `store.bundle.max_line_bytes` long."""
962
+ with _open_file(root, name) as handle:
963
+ data = handle.read(limits.max_line_bytes + 1)
964
+ if len(data) > limits.max_line_bytes:
965
+ raise BundleError(
966
+ f"{name} is larger than store.bundle.max_line_bytes ({limits.max_line_bytes})"
967
+ )
968
+ try:
969
+ return json.loads(data)
970
+ except ValueError as error:
971
+ raise BundleError(f"{name} is not valid JSON") from error
972
+
973
+
974
+ def _member_parts(member: tarfile.TarInfo, archive: str) -> tuple[str, ...]:
975
+ """The path of a tarball member, refused unless it is a file or directory of the layout."""
976
+ label = f"{member.name!r} in {archive}"
977
+ if member.issym():
978
+ raise BundleError(f"{label} is a symbolic link")
979
+ if member.islnk():
980
+ raise BundleError(f"{label} is a hard link")
981
+ if not (member.isreg() or member.isdir()):
982
+ raise BundleError(f"{label} is not a regular file or directory")
983
+ path = PurePosixPath(member.name)
984
+ if path.is_absolute():
985
+ raise BundleError(f"{label} has an absolute path")
986
+ if ".." in path.parts:
987
+ raise BundleError(f"{label} leaves the bundle directory")
988
+ parts = path.parts
989
+ expected = (
990
+ parts == (SCHEMAS_DIR,)
991
+ if member.isdir()
992
+ else len(parts) == 1 or (len(parts) == 2 and parts[0] == SCHEMAS_DIR)
993
+ )
994
+ if not expected:
995
+ raise BundleError(f"unexpected entry {label}")
996
+ return parts
997
+
998
+
999
+ def _extract(tarball: Path, into: Path, limits: BundleSettings) -> None:
1000
+ """Unpack a bundle tarball: only files and directories with the bundle's names, within
1001
+ the store.bundle entry and size limits, counting the bytes actually read."""
1002
+ meter = _Meter(limits)
1003
+ entries = 0
1004
+ declared = 0
1005
+ try:
1006
+ with tarfile.open(tarball, "r:") as archive:
1007
+ for member in archive:
1008
+ entries += 1
1009
+ label = f"{member.name!r} in {tarball.name}"
1010
+ _check_entries(label, entries, limits)
1011
+ parts = _member_parts(member, tarball.name)
1012
+ target = into.joinpath(*parts)
1013
+ if member.isdir():
1014
+ target.mkdir(exist_ok=True)
1015
+ continue
1016
+ declared += member.size
1017
+ _check_size(label, member.size, declared, limits)
1018
+ source = archive.extractfile(member)
1019
+ if source is None:
1020
+ raise BundleError(f"{label} cannot be read")
1021
+ target.parent.mkdir(exist_ok=True)
1022
+ try:
1023
+ descriptor = os.open(
1024
+ target, os.O_WRONLY | os.O_CREAT | os.O_EXCL | _NOFOLLOW, 0o600
1025
+ )
1026
+ except FileExistsError as error:
1027
+ raise BundleError(f"{label} appears twice") from error
1028
+ with source, os.fdopen(descriptor, "wb") as sink:
1029
+ for chunk in meter.chunks(label, source):
1030
+ sink.write(chunk)
1031
+ except tarfile.TarError as error:
1032
+ raise BundleError(f"{tarball.name} is not a readable tar file: {error}") from error
1033
+ except OSError as error:
1034
+ raise BundleError(f"{tarball.name} cannot be unpacked: {error.strerror}") from error
1035
+
1036
+
1037
+ def _scan(root: Path, limits: BundleSettings) -> set[str]:
1038
+ """The files of a bundle directory. Links and special files are refused, never followed."""
1039
+ files: set[str] = set()
1040
+ entries = 0
1041
+ total = 0
1042
+ pending = [PurePosixPath()]
1043
+ while pending:
1044
+ folder = pending.pop()
1045
+ with os.scandir(root.joinpath(*folder.parts)) as listing:
1046
+ for entry in listing:
1047
+ relative = folder / entry.name
1048
+ name = relative.as_posix()
1049
+ entries += 1
1050
+ _check_entries(name, entries, limits)
1051
+ if entry.is_symlink():
1052
+ raise BundleError(f"{name} is a symbolic link")
1053
+ if entry.is_dir(follow_symlinks=False):
1054
+ pending.append(relative)
1055
+ continue
1056
+ if not entry.is_file(follow_symlinks=False):
1057
+ raise BundleError(f"{name} is not a regular file or directory")
1058
+ size = entry.stat(follow_symlinks=False).st_size
1059
+ total += size
1060
+ _check_size(name, size, total, limits)
1061
+ files.add(name)
1062
+ return files
1063
+
1064
+
1065
+ @contextmanager
1066
+ def _opened(source: Path, limits: BundleSettings) -> Iterator[Path]:
1067
+ if source.is_dir():
1068
+ yield source
1069
+ elif source.is_file() and source.suffix == TAR_SUFFIX:
1070
+ with tempfile.TemporaryDirectory(prefix="lorebase-bundle-") as work:
1071
+ _extract(source, Path(work), limits)
1072
+ yield Path(work)
1073
+ else:
1074
+ raise BundleError(f"{source} is neither a bundle directory nor a {TAR_SUFFIX} file")
1075
+
1076
+
1077
+ def read_manifest(root: Path, limits: BundleSettings | None = None) -> dict[str, Any]:
1078
+ """The manifest of an opened bundle directory, checked for format and version."""
1079
+ limits = limits or BundleSettings()
1080
+ if not os.path.lexists(root / MANIFEST):
1081
+ raise BundleError(f"{MANIFEST} not found")
1082
+ manifest = _read_document(root, MANIFEST, limits)
1083
+ if not isinstance(manifest, dict):
1084
+ raise BundleError(f"{MANIFEST} is not a JSON object")
1085
+ if manifest.get("format") != FORMAT:
1086
+ raise BundleError(f"not a Lorebase Domain bundle (format {manifest.get('format')!r})")
1087
+ version = manifest.get("format_version")
1088
+ if version != FORMAT_VERSION:
1089
+ raise BundleError(
1090
+ f"bundle format version {version!r} is not supported (this release reads "
1091
+ f"{FORMAT_VERSION})"
1092
+ )
1093
+ for key in ("domain_id", "compression", "counts", "files", "heads"):
1094
+ if key not in manifest:
1095
+ raise BundleError(f"{MANIFEST} lacks {key!r}")
1096
+ return manifest
1097
+
1098
+
1099
+ def _verify_files(root: Path, manifest: Mapping[str, Any], limits: BundleSettings) -> None:
1100
+ files: Mapping[str, Mapping[str, Any]] = manifest["files"]
1101
+ present = _scan(root, limits) - {MANIFEST}
1102
+ expected = set(files)
1103
+ if present != expected:
1104
+ raise BundleError(
1105
+ f"files differ from the manifest: missing {sorted(expected - present)}, "
1106
+ f"unexpected {sorted(present - expected)}"
1107
+ )
1108
+ meter = _Meter(limits)
1109
+ for name, entry in files.items():
1110
+ digest = hashlib.sha256()
1111
+ size = 0
1112
+ with _open_file(root, name) as handle:
1113
+ for chunk in meter.chunks(name, handle):
1114
+ digest.update(chunk)
1115
+ size += len(chunk)
1116
+ if digest.hexdigest() != entry["sha256"] or size != entry["bytes"]:
1117
+ raise BundleError(f"{name} does not match its sha256 in the manifest")
1118
+
1119
+
1120
+ @dataclass(slots=True)
1121
+ class _Import:
1122
+ session: AsyncSession
1123
+ root: Path
1124
+ manifest: Mapping[str, Any]
1125
+ domain_id: str
1126
+ runtime: StoreRuntime
1127
+ batch: int
1128
+ report: BundleReport
1129
+ actor: str = BUNDLE_ACTOR
1130
+
1131
+ def lines(self, base: str) -> Iterator[dict[str, Any]]:
1132
+ name = _jsonl_name(base, self.manifest["compression"])
1133
+ return _read_jsonl(self.root, name, self.runtime.settings.bundle)
1134
+
1135
+ def document(self, name: str) -> Any:
1136
+ return _read_document(self.root, name, self.runtime.settings.bundle)
1137
+
1138
+ def expect(self, name: str, got: int) -> None:
1139
+ want = self.manifest["counts"][name]
1140
+ if got != want:
1141
+ raise BundleError(f"{name}: manifest says {want}, bundle holds {got}")
1142
+
1143
+
1144
+ async def import_bundle(
1145
+ engine: AsyncEngine,
1146
+ source: str | Path,
1147
+ *,
1148
+ domain_id: str | None = None,
1149
+ runtime: StoreRuntime | None = None,
1150
+ verify_heads: bool | None = None,
1151
+ actor: str = BUNDLE_ACTOR,
1152
+ ) -> BundleReport:
1153
+ """Import a bundle (directory or `.tar`) into `domain_id`, which must not exist yet or be
1154
+ empty (no Revisions, no refs). It defaults to, and must equal, the bundle's Domain id.
1155
+ A bundle carrying vault values appends a `bundle.import_vault` audit entry for `actor` in
1156
+ the same transaction.
1157
+
1158
+ Every file hash, Revision id, node hash and object hash is verified, materialized heads are
1159
+ rebuilt for the refs that had one, and with `verify_heads` (default
1160
+ `store.bundle.verify_after_import`) each rebuilt branch is checked against its tree. All
1161
+ in one transaction. Raises `BundleError` (the transaction is rolled back).
1162
+
1163
+ Reading is bounded by the `store.bundle` limits (entries, file and total sizes, line
1164
+ length); links and special files are refused. The content itself is trusted: authors,
1165
+ provenance, the ref log, changesets and Domain settings are imported as written.
1166
+ """
1167
+ runtime = runtime or default_runtime()
1168
+ limits = runtime.settings.bundle
1169
+ started = time.monotonic()
1170
+ path = Path(source)
1171
+ with _opened(path, limits) as root:
1172
+ manifest = read_manifest(root, limits)
1173
+ target = domain_id or manifest["domain_id"]
1174
+ if target != manifest["domain_id"]:
1175
+ raise BundleError(
1176
+ f"the bundle is of Domain {manifest['domain_id']!r}; Revision ids include the "
1177
+ f"Domain id, so it cannot be imported as {target!r}"
1178
+ )
1179
+ _verify_files(root, manifest, limits)
1180
+ settings = runtime.settings
1181
+ report = BundleReport(target, str(path))
1182
+ sessions = async_sessionmaker(engine, expire_on_commit=False)
1183
+ async with sessions() as session, session.begin():
1184
+ await scope_to_domain(session, target)
1185
+ context = _Import(
1186
+ session, root, manifest, target, runtime, settings.bundle.batch_rows, report, actor
1187
+ )
1188
+ await _prepare_domain(context)
1189
+ scanned = await _import_objects(context)
1190
+ if scanned is not None:
1191
+ await _import_vault(context, scanned)
1192
+ await _import_nodes(context)
1193
+ await _import_revisions(context)
1194
+ await _import_changesets(context)
1195
+ await _import_provenance(context)
1196
+ refs = await _import_refs(context)
1197
+ await _import_ref_log(context)
1198
+ await _check_integrity(context)
1199
+ await _materialize(
1200
+ context,
1201
+ refs,
1202
+ settings.bundle.verify_after_import if verify_heads is None else verify_heads,
1203
+ )
1204
+ report.seconds = time.monotonic() - started
1205
+ return report
1206
+
1207
+
1208
+ async def _prepare_domain(c: _Import) -> None:
1209
+ domains = DomainRepository(c.session)
1210
+ data = c.document("domain.json")
1211
+ if data["id"] != c.domain_id:
1212
+ raise BundleError("domain.json names a different Domain than the manifest")
1213
+ if await domains.get(c.domain_id) is not None:
1214
+ counts = (
1215
+ await c.session.execute(
1216
+ text(
1217
+ "SELECT (SELECT count(*) FROM revisions WHERE domain_id = :d), "
1218
+ "(SELECT count(*) FROM refs WHERE domain_id = :d), "
1219
+ "(SELECT count(*) FROM vault WHERE domain_id = :d), "
1220
+ "(SELECT count(*) FROM vault_tombstones WHERE domain_id = :d), "
1221
+ "(SELECT count(*) FROM vault_lineage WHERE domain_id = :d)"
1222
+ ),
1223
+ {"d": c.domain_id},
1224
+ )
1225
+ ).one()
1226
+ if any(counts):
1227
+ raise DomainNotEmptyError(
1228
+ f"Domain {c.domain_id!r} already has {counts[0]} Revisions, {counts[1]} refs, "
1229
+ f"{counts[2]} vault rows, {counts[3]} vault tombstones and {counts[4]} vault "
1230
+ "key changes"
1231
+ )
1232
+ return
1233
+ await domains.create(
1234
+ DomainRow(
1235
+ id=data["id"],
1236
+ display_name=data["display_name"],
1237
+ description=data["description"],
1238
+ default_branch=data["default_branch"],
1239
+ leaf_size=data["leaf_size"],
1240
+ settings=data["settings"],
1241
+ created_at=datetime.fromisoformat(data["created_at"]),
1242
+ created_by=data["created_by"],
1243
+ )
1244
+ )
1245
+
1246
+
1247
+ def _vault_attributes(c: _Import) -> dict[str, set[str]]:
1248
+ """Kind id -> ids of Attributes that are vault Attributes in any schema of the bundle."""
1249
+ found: dict[str, set[str]] = {}
1250
+ for line in c.lines("objects"):
1251
+ if line.get("type") != ObjectType.SCHEMA_KIND.value:
1252
+ continue
1253
+ body = line.get("body") or {}
1254
+ attributes = body.get("attributes") or {}
1255
+ ids = {k for k, a in attributes.items() if isinstance(a, dict) and a.get("vault") is True}
1256
+ if ids:
1257
+ found.setdefault(str(body.get("id")), set()).update(ids)
1258
+ return found
1259
+
1260
+
1261
+ def _check_vault_values(line: dict[str, Any], vaulted: dict[str, set[str]]) -> None:
1262
+ """A Record's vault Attributes hold references only: plain values would never be
1263
+ erasable."""
1264
+ body = line.get("body") or {}
1265
+ values = body.get("values") or {}
1266
+ for attribute_id in vaulted.get(str(body.get("kind")), ()):
1267
+ value = values.get(attribute_id)
1268
+ if value is not None and vault_token(value) is None:
1269
+ raise BundleError(
1270
+ f"object {line.get('hash')}: vault attribute {attribute_id} holds a value, "
1271
+ "not a vault reference"
1272
+ )
1273
+
1274
+
1275
+ def _carries_vault(c: _Import) -> bool:
1276
+ """Whether the bundle holds vault values: the vault, key change and tombstone sections are
1277
+ all listed, each with its count, or none is."""
1278
+ found = set()
1279
+ for section in _VAULT_SECTIONS:
1280
+ listed = _jsonl_name(section, c.manifest["compression"]) in c.manifest["files"]
1281
+ if listed != (section in c.manifest["counts"]):
1282
+ raise BundleError(f"{MANIFEST} lists the {section} file and its count inconsistently")
1283
+ found.add(listed)
1284
+ if len(found) != 1:
1285
+ raise BundleError(f"{MANIFEST} lists some vault sections but not all")
1286
+ return found.pop()
1287
+
1288
+
1289
+ def _vault_definitions(c: _Import) -> dict[tuple[str, str], list[Attribute]]:
1290
+ """(Kind id, Attribute id) -> the distinct definitions of each vault Attribute in the
1291
+ bundle's schemas."""
1292
+ found: dict[tuple[str, str], dict[str, Attribute]] = {}
1293
+ for line in c.lines("objects"):
1294
+ if line.get("type") != ObjectType.SCHEMA_KIND.value:
1295
+ continue
1296
+ try:
1297
+ kind = Kind.model_validate(line.get("body"))
1298
+ except ValueError as error:
1299
+ raise BundleError(f"object {line.get('hash')}: {error}") from error
1300
+ for attribute in kind.attributes.values():
1301
+ if attribute.in_vault:
1302
+ definitions = found.setdefault((kind.id, attribute.id), {})
1303
+ definitions.setdefault(attribute.model_dump_json(), attribute)
1304
+ return {key: list(definitions.values()) for key, definitions in found.items()}
1305
+
1306
+
1307
+ @dataclass(slots=True)
1308
+ class _VaultScan:
1309
+ """What the Record objects of a bundle carrying vault values show: the Records holding
1310
+ each token, and every (Kind id, key) of a Record in any Revision."""
1311
+
1312
+ holders: dict[str, set[Holder]] = field(default_factory=dict)
1313
+ records: set[tuple[str, str]] = field(default_factory=set)
1314
+
1315
+
1316
+ async def _import_objects(c: _Import) -> _VaultScan | None:
1317
+ """Store the objects; when the bundle carries vault values, what its Record objects show
1318
+ (`_VaultScan`)."""
1319
+ store = PgObjectStore(c.session, c.domain_id, writes=c.runtime.settings.writes, batch=c.batch)
1320
+ vaulted = _vault_attributes(c)
1321
+ scan = _VaultScan() if _carries_vault(c) else None
1322
+ total = 0
1323
+ for batch in _batches(c.lines("objects"), c.batch):
1324
+ encoded = []
1325
+ for line in batch:
1326
+ if line.get("type") == ObjectType.RECORD.value:
1327
+ if vaulted:
1328
+ _check_vault_values(line, vaulted)
1329
+ if scan is not None:
1330
+ _collect_holders(line, scan.holders)
1331
+ body = line.get("body") or {}
1332
+ scan.records.add((str(body.get("kind")), str(body.get("key"))))
1333
+ try:
1334
+ item = encode_object(ObjectType(line["type"]), line["body"])
1335
+ except (ValueError, KeyError, TypeError) as error:
1336
+ raise BundleError(f"object {line.get('hash')}: {error}") from error
1337
+ if to_hex(item.hash) != line["hash"]:
1338
+ raise BundleError(f"object {line['hash']} does not match its content")
1339
+ encoded.append(item)
1340
+ await store.put_many(encoded)
1341
+ total += len(batch)
1342
+ c.report.objects = total
1343
+ c.expect("objects", total)
1344
+ return scan
1345
+
1346
+
1347
+ async def _import_vault(c: _Import, scan: _VaultScan) -> None:
1348
+ """Store the key changes, the tombstones and the vault entries (module doc), and append
1349
+ the `bundle.import_vault` audit entry (counts, never values).
1350
+
1351
+ A key change must link two Records of its Kind that the bundle holds. Each entry must be
1352
+ well formed, appear once, hold a value valid for a definition of its Attribute in the
1353
+ bundle's schemas, and every Record holding its token must hold it in that Attribute under
1354
+ its key or a key linked to it by the key changes; its `moved_to` must follow a key change
1355
+ from its key. Values are checked against the tombstones (the bundle's) as
1356
+ `vault.erased_values` says. Groups of more than `MAX_LINKED_KEYS` linked keys are counted
1357
+ for an administrator to look at, as erasure flags them.
1358
+ """
1359
+ holders = scan.holders
1360
+ keys = c.runtime.vault
1361
+ if keys is None:
1362
+ raise VaultNotConfiguredError(
1363
+ "the bundle holds vault values; importing them needs vault.token_key and "
1364
+ "vault.encryption_key"
1365
+ )
1366
+ at = datetime.now(UTC)
1367
+ links, edges = await _import_lineage(c, scan.records, at)
1368
+ c.report.vault_oversized_links = links.oversized(MAX_LINKED_KEYS)
1369
+ await _import_tombstones(c, keys, at)
1370
+ definitions = _vault_definitions(c)
1371
+ seen: set[str] = set()
1372
+ for batch in _batches(enumerate(c.lines(VAULT), start=1), c.batch):
1373
+ entries: list[VaultEntry] = []
1374
+ for number, line in batch:
1375
+ try:
1376
+ entry = parse_entry(line)
1377
+ except VaultTransferError as error:
1378
+ raise BundleError(f"vault line {number}: {error}") from error
1379
+ problem = holders_problem(
1380
+ entry.token,
1381
+ entry.kind_id,
1382
+ entry.attribute_id,
1383
+ entry.record_key,
1384
+ holders.get(entry.token, ()),
1385
+ links,
1386
+ )
1387
+ if problem is None and entry.moved_to is not None:
1388
+ edge = LineageEntry(entry.kind_id, entry.moved_to, entry.record_key)
1389
+ if edge not in edges:
1390
+ problem = (
1391
+ f"{entry.token} moved to {entry.moved_to!r}, but no key change of the "
1392
+ f"bundle leads there from {entry.record_key!r}"
1393
+ )
1394
+ if problem is None and not entry.erased:
1395
+ problem = value_problem(
1396
+ entry, definitions.get((entry.kind_id, entry.attribute_id), [])
1397
+ )
1398
+ if problem is not None:
1399
+ raise BundleError(f"vault line {number}: {problem}")
1400
+ if entry.token in seen:
1401
+ raise BundleError(f"vault line {number}: {entry.token} appears twice")
1402
+ seen.add(entry.token)
1403
+ entries.append(entry)
1404
+ if entry.erased:
1405
+ c.report.vault_erased += 1
1406
+ else:
1407
+ c.report.vault_values += 1
1408
+ try:
1409
+ tombstoned = await restore_entries(
1410
+ c.session,
1411
+ c.domain_id,
1412
+ entries,
1413
+ keys,
1414
+ erased_values=c.runtime.vault_settings.erased_values,
1415
+ at=at,
1416
+ )
1417
+ except VaultTransferError as error:
1418
+ raise BundleError(str(error)) from error
1419
+ c.report.vault_values -= tombstoned
1420
+ c.report.vault_erased += tombstoned
1421
+ c.expect(VAULT, len(seen))
1422
+ report = c.report
1423
+ await AuditRepository(c.session).append(
1424
+ NewAuditEntry(
1425
+ actor=c.actor,
1426
+ action=IMPORT_VAULT_ACTION,
1427
+ domain_id=c.domain_id,
1428
+ target=c.domain_id,
1429
+ details={
1430
+ "values": report.vault_values,
1431
+ "erased": report.vault_erased,
1432
+ "tombstones": report.vault_tombstones,
1433
+ "lineage_edges": report.vault_lineage,
1434
+ "oversized_lineage_groups": report.vault_oversized_links,
1435
+ },
1436
+ at=at,
1437
+ )
1438
+ )
1439
+
1440
+
1441
+ async def _import_lineage(
1442
+ c: _Import, records: Collection[tuple[str, str]], at: datetime
1443
+ ) -> tuple[KeyLinks, set[LineageEntry]]:
1444
+ """The key changes, each between two Records of its Kind that the bundle holds."""
1445
+ links = KeyLinks()
1446
+ seen: set[LineageEntry] = set()
1447
+ for batch in _batches(enumerate(c.lines(VAULT_LINEAGE), start=1), c.batch):
1448
+ entries: list[LineageEntry] = []
1449
+ for number, line in batch:
1450
+ try:
1451
+ entry = parse_lineage(line)
1452
+ except VaultTransferError as error:
1453
+ raise BundleError(f"vault_lineage line {number}: {error}") from error
1454
+ if entry in seen:
1455
+ raise BundleError(f"vault_lineage line {number}: the key change appears twice")
1456
+ missing = [
1457
+ key
1458
+ for key in (entry.previous_key, entry.record_key)
1459
+ if (entry.kind_id, key) not in records
1460
+ ]
1461
+ if missing:
1462
+ raise BundleError(
1463
+ f"vault_lineage line {number}: no Record {missing[0]!r} of Kind "
1464
+ f"{entry.kind_id} in the bundle"
1465
+ )
1466
+ seen.add(entry)
1467
+ links.link(entry)
1468
+ entries.append(entry)
1469
+ await restore_lineage(c.session, c.domain_id, entries, at=at)
1470
+ c.report.vault_lineage = len(seen)
1471
+ c.expect(VAULT_LINEAGE, len(seen))
1472
+ return links, seen
1473
+
1474
+
1475
+ async def _import_tombstones(c: _Import, keys: VaultKeys, at: datetime) -> None:
1476
+ """Tombstones are taken as made: only when the target has the token key they were made
1477
+ under (by id)."""
1478
+ total = 0
1479
+ for batch in _batches(enumerate(c.lines(VAULT_TOMBSTONES), start=1), c.batch):
1480
+ entries: list[TombstoneEntry] = []
1481
+ for number, line in batch:
1482
+ try:
1483
+ entries.append(parse_tombstone(line, keys.token_key_ids))
1484
+ except VaultTransferError as error:
1485
+ raise BundleError(f"vault_tombstones line {number}: {error}") from error
1486
+ await restore_tombstones(c.session, c.domain_id, entries, at=at)
1487
+ total += len(entries)
1488
+ c.report.vault_tombstones = total
1489
+ c.expect(VAULT_TOMBSTONES, total)
1490
+
1491
+
1492
+ async def _import_nodes(c: _Import) -> None:
1493
+ repository = NodeRepository(
1494
+ c.session,
1495
+ c.domain_id,
1496
+ writes=c.runtime.settings.writes,
1497
+ batch=c.runtime.settings.nodes.prefetch_batch,
1498
+ )
1499
+ total = 0
1500
+ for batch in _batches(c.lines("nodes"), c.batch):
1501
+ nodes: dict[bytes, bytes] = {}
1502
+ for line in batch:
1503
+ body = base64.b64decode(line["body"], validate=True)
1504
+ digest = from_hex(line["hash"])
1505
+ if node_hash(body) != digest:
1506
+ raise BundleError(f"tree node {line['hash']} does not match its content")
1507
+ decode_node(body) # rejects malformed nodes
1508
+ nodes[digest] = body
1509
+ await repository.put_many(nodes)
1510
+ total += len(batch)
1511
+ c.report.nodes = total
1512
+ c.expect("nodes", total)
1513
+
1514
+
1515
+ async def _import_revisions(c: _Import) -> None:
1516
+ revisions = RevisionRepository(c.session, c.domain_id)
1517
+ writes = c.runtime.settings.writes
1518
+ total = 0
1519
+ previous = 0
1520
+ for batch in _batches(c.lines("revisions"), c.batch):
1521
+ local: dict[bytes, int] = {}
1522
+ parsed: list[tuple[dict[str, Any], Revision, bytes]] = []
1523
+ for line in batch:
1524
+ revision = _revision(c.domain_id, line)
1525
+ digest = revision_id(revision)
1526
+ if to_hex(digest) != line["id"]:
1527
+ raise BundleError(f"Revision {line['id']} does not match its content")
1528
+ if line["generation"] < previous:
1529
+ raise BundleError("revisions.jsonl is not ordered by generation")
1530
+ previous = line["generation"]
1531
+ parsed.append((line, revision, digest))
1532
+ outside = {p for _, rev, _ in parsed for p in rev.parents} - {d for _, _, d in parsed}
1533
+ known = await revisions.generations(outside)
1534
+ rows: list[dict[str, Any]] = []
1535
+ parent_rows: list[dict[str, Any]] = []
1536
+ for line, revision, digest in parsed:
1537
+ try:
1538
+ generation = next_generation(local.get(p) or known[p] for p in revision.parents)
1539
+ except KeyError as error:
1540
+ raise BundleError(
1541
+ f"Revision {line['id']} comes before its parent {to_hex(error.args[0])}"
1542
+ ) from error
1543
+ if generation != line["generation"]:
1544
+ raise BundleError(f"Revision {line['id']} has the wrong generation")
1545
+ local[digest] = generation
1546
+ rows.append(
1547
+ {
1548
+ "domain_id": c.domain_id,
1549
+ "id": digest,
1550
+ "root": revision.root,
1551
+ "generation": generation,
1552
+ "author": revision.author.principal,
1553
+ "on_behalf_of": revision.author.on_behalf_of,
1554
+ "message": revision.message,
1555
+ "origin": revision.origin.model_dump(),
1556
+ "stats": revision.stats.model_dump(),
1557
+ "created_at": revision.created_at,
1558
+ }
1559
+ )
1560
+ parent_rows.extend(
1561
+ {
1562
+ "domain_id": c.domain_id,
1563
+ "revision_id": digest,
1564
+ "position": position,
1565
+ "parent_id": parent,
1566
+ }
1567
+ for position, parent in enumerate(revision.parents)
1568
+ )
1569
+ await bulk_write(c.session, REVISIONS, rows, writes=writes, on_conflict=OnConflict.ERROR)
1570
+ await bulk_write(
1571
+ c.session, REVISION_PARENTS, parent_rows, writes=writes, on_conflict=OnConflict.ERROR
1572
+ )
1573
+ total += len(batch)
1574
+ c.report.revisions = total
1575
+ c.expect("revisions", total)
1576
+
1577
+
1578
+ def _revision(domain_id: str, line: Mapping[str, Any]) -> Revision:
1579
+ try:
1580
+ return Revision(
1581
+ parents=tuple(from_hex(p) for p in line["parents"]),
1582
+ root=from_hex(line["root"]),
1583
+ domain=domain_id,
1584
+ author=Author(principal=line["author"], on_behalf_of=line["on_behalf_of"]),
1585
+ message=line["message"],
1586
+ origin=Origin(**line["origin"]),
1587
+ created_at=datetime.fromisoformat(line["created_at"]),
1588
+ stats=RevisionStats(**line["stats"]),
1589
+ )
1590
+ except (ValueError, KeyError, TypeError) as error:
1591
+ raise BundleError(f"Revision {line.get('id')}: {error}") from error
1592
+
1593
+
1594
+ def _from_hex(value: str | None) -> bytes | None:
1595
+ return None if value is None else from_hex(value)
1596
+
1597
+
1598
+ async def _import_changesets(c: _Import) -> None:
1599
+ writes = c.runtime.settings.writes
1600
+ total = 0
1601
+ for batch in _batches(c.lines("changesets"), c.batch):
1602
+ rows = [
1603
+ {
1604
+ "domain_id": c.domain_id,
1605
+ "revision_id": from_hex(line["revision"]),
1606
+ "partition": line["partition"],
1607
+ "element_key": line["key"],
1608
+ "old_hash": _from_hex(line["old"]),
1609
+ "new_hash": _from_hex(line["new"]),
1610
+ "provenance": line["provenance"] or None,
1611
+ }
1612
+ for line in batch
1613
+ ]
1614
+ await bulk_write(c.session, CHANGESETS, rows, writes=writes, on_conflict=OnConflict.ERROR)
1615
+ total += len(rows)
1616
+ c.report.changesets = total
1617
+ c.expect("changesets", total)
1618
+
1619
+
1620
+ async def _import_provenance(c: _Import) -> None:
1621
+ writes = c.runtime.settings.writes
1622
+ total = 0
1623
+ for batch in _batches(c.lines("provenance"), c.batch):
1624
+ rows = [
1625
+ {
1626
+ "domain_id": c.domain_id,
1627
+ "id": line["id"],
1628
+ "origin_type": line["origin_type"],
1629
+ "origin_ref": line["origin_ref"],
1630
+ "actor": line["actor"],
1631
+ "at": datetime.fromisoformat(line["at"]),
1632
+ }
1633
+ for line in batch
1634
+ ]
1635
+ await bulk_write(c.session, PROVENANCE, rows, writes=writes, on_conflict=OnConflict.ERROR)
1636
+ total += len(rows)
1637
+ c.report.provenance = total
1638
+ c.expect("provenance", total)
1639
+
1640
+
1641
+ async def _import_refs(c: _Import) -> list[dict[str, Any]]:
1642
+ refs: list[dict[str, Any]] = c.document("refs.json")
1643
+ heads = c.manifest["heads"]
1644
+ rows = []
1645
+ for ref in refs:
1646
+ if heads.get(ref["name"]) != ref["head"]:
1647
+ raise BundleError(f"ref {ref['name']!r} disagrees with the manifest heads")
1648
+ rows.append(
1649
+ {
1650
+ "domain_id": c.domain_id,
1651
+ "name": ref["name"],
1652
+ "kind": RefKind(ref["kind"]),
1653
+ "head": from_hex(ref["head"]),
1654
+ "protected": ref["protected"],
1655
+ "materialization": _imported_materialization(ref["materialization"]),
1656
+ "version": ref["version"],
1657
+ "created_at": datetime.fromisoformat(ref["created_at"]),
1658
+ "created_by": ref["created_by"],
1659
+ }
1660
+ )
1661
+ await bulk_write(
1662
+ c.session, REFS, rows, writes=c.runtime.settings.writes, on_conflict=OnConflict.ERROR
1663
+ )
1664
+ c.report.refs = len(rows)
1665
+ c.expect("refs", len(rows))
1666
+ return refs
1667
+
1668
+
1669
+ def _imported_materialization(value: str) -> Materialization:
1670
+ """Overlay branches come back as full copies: import rebuilds every head from its tree."""
1671
+ found = Materialization(value)
1672
+ return Materialization.FULL if found is Materialization.OVERLAY else found
1673
+
1674
+
1675
+ async def _import_ref_log(c: _Import) -> None:
1676
+ writes = c.runtime.settings.writes
1677
+ total = 0
1678
+ for batch in _batches(c.lines("ref_log"), c.batch):
1679
+ rows = [
1680
+ {
1681
+ "domain_id": c.domain_id,
1682
+ "ref_name": line["ref"],
1683
+ "seq": line["seq"],
1684
+ "old_head": _from_hex(line["old"]),
1685
+ "new_head": from_hex(line["new"]),
1686
+ "reason": line["reason"],
1687
+ "actor": line["actor"],
1688
+ "at": datetime.fromisoformat(line["at"]),
1689
+ }
1690
+ for line in batch
1691
+ ]
1692
+ await bulk_write(c.session, REF_LOG, rows, writes=writes, on_conflict=OnConflict.ERROR)
1693
+ total += len(rows)
1694
+ c.report.ref_log = total
1695
+ c.expect("ref_log", total)
1696
+
1697
+
1698
+ async def _check_integrity(c: _Import) -> None:
1699
+ """Every Revision is reachable, every reachable node and object is stored."""
1700
+ marker = Marker(c.session, c.domain_id, batch=c.runtime.settings.gc.mark_batch)
1701
+ await marker.prepare()
1702
+ await marker.mark(attempts=False)
1703
+ stats = await marker.stats()
1704
+ missing_objects = (
1705
+ await c.session.execute(
1706
+ text(
1707
+ "SELECT count(*) FROM lb_gc_objects m WHERE NOT EXISTS "
1708
+ "(SELECT 1 FROM objects o WHERE o.domain_id = :d AND o.hash = m.hash)"
1709
+ ),
1710
+ {"d": c.domain_id},
1711
+ )
1712
+ ).scalar_one()
1713
+ await marker.drop()
1714
+ if stats.missing_nodes or missing_objects:
1715
+ raise BundleError(
1716
+ f"the bundle lacks {stats.missing_nodes} tree nodes and {missing_objects} objects "
1717
+ "that its Revisions reach"
1718
+ )
1719
+ if stats.revisions != c.report.revisions:
1720
+ raise BundleError(
1721
+ f"{c.report.revisions - stats.revisions} Revisions are not reachable from any ref"
1722
+ )
1723
+
1724
+
1725
+ async def _materialize(c: _Import, refs: Sequence[Mapping[str, Any]], check: bool) -> None:
1726
+ for ref in refs:
1727
+ if ref["kind"] != RefKind.BRANCH.value or not ref["materialized"]:
1728
+ continue
1729
+ await rebuild(c.session, c.domain_id, ref["name"], runtime=c.runtime)
1730
+ c.report.rebuilt_branches.append(ref["name"])
1731
+ if check:
1732
+ result = await verify(c.session, c.domain_id, ref["name"], runtime=c.runtime)
1733
+ if not result.ok:
1734
+ raise BundleError(
1735
+ f"branch {ref['name']!r} fails verification after rebuild: "
1736
+ f"{result.divergences[:3]}"
1737
+ )
1738
+ heads = RefRepository(c.session, c.domain_id)
1739
+ for ref in refs: # the refs must still be exactly the bundle's
1740
+ stored = await heads.require(ref["name"])
1741
+ if stored.head != from_hex(ref["head"]) or stored.version != ref["version"]:
1742
+ raise BundleError(f"ref {ref['name']!r} changed while importing")
1743
+ schemas = len([name for name in c.manifest["files"] if name.startswith(SCHEMAS_DIR + "/")])
1744
+ c.report.schemas = schemas
1745
+
1746
+
1747
+ # --------------------------------------------------------------------------------------
1748
+ # Command line
1749
+
1750
+
1751
+ def _parser() -> argparse.ArgumentParser:
1752
+ parser = argparse.ArgumentParser(
1753
+ prog="python -m lorebase_store.bundle",
1754
+ description="Export and import portable Domain bundles. Reads the database and "
1755
+ "store.bundle settings from the Lorebase configuration.",
1756
+ )
1757
+ commands = parser.add_subparsers(dest="command", required=True)
1758
+ export = commands.add_parser("export", help="Write a bundle of a Domain.")
1759
+ export.add_argument("--domain", required=True, help="Domain id.")
1760
+ export.add_argument(
1761
+ "--out", required=True, type=Path, help="New directory, or a path ending in .tar."
1762
+ )
1763
+ export.add_argument(
1764
+ "--include-vault-values",
1765
+ action="store_true",
1766
+ help="Also write the vault values, in plain text (needs the vault keys; audited).",
1767
+ )
1768
+ export.add_argument(
1769
+ "--actor",
1770
+ default=BUNDLE_ACTOR,
1771
+ help="Who exports, recorded in the audit entry of --include-vault-values.",
1772
+ )
1773
+ load = commands.add_parser("import", help="Import a bundle into an empty Domain.")
1774
+ load.add_argument("--in", dest="source", required=True, type=Path, help="Directory or .tar.")
1775
+ load.add_argument("--domain", default=None, help="Target Domain id (must equal the bundle's).")
1776
+ load.add_argument(
1777
+ "--no-verify", action="store_true", help="Skip verifying rebuilt branches against trees."
1778
+ )
1779
+ load.add_argument(
1780
+ "--actor",
1781
+ default=BUNDLE_ACTOR,
1782
+ help="Who imports, recorded in the audit entry of a bundle with vault values.",
1783
+ )
1784
+ return parser
1785
+
1786
+
1787
+ async def _main(args: argparse.Namespace) -> int:
1788
+ from lorebase_platform.settings import load_settings
1789
+ from lorebase_store.db import create_engine
1790
+ from lorebase_store.settings import database_settings
1791
+
1792
+ settings = load_settings()
1793
+ runtime = StoreRuntime.from_loaded(settings)
1794
+ engine = create_engine(database_settings(settings), pooled=False, statement_timeout_ms=0)
1795
+ try:
1796
+ if args.command == "export":
1797
+ report = await export_domain(
1798
+ engine,
1799
+ args.domain,
1800
+ args.out,
1801
+ runtime=runtime,
1802
+ include_vault_values=args.include_vault_values,
1803
+ actor=args.actor,
1804
+ )
1805
+ if args.include_vault_values:
1806
+ sys.stderr.write(PLAIN_TEXT_WARNING.format(path=args.out))
1807
+ else:
1808
+ report = await import_bundle(
1809
+ engine,
1810
+ args.source,
1811
+ domain_id=args.domain,
1812
+ runtime=runtime,
1813
+ verify_heads=False if args.no_verify else None,
1814
+ actor=args.actor,
1815
+ )
1816
+ except (BundleError, DomainNotFoundError) as error:
1817
+ sys.stderr.write(f"{error}\n")
1818
+ return 1
1819
+ finally:
1820
+ await engine.dispose()
1821
+ sys.stdout.write(json.dumps(report.as_dict()) + "\n")
1822
+ return 0
1823
+
1824
+
1825
+ def main(argv: Sequence[str] | None = None) -> int:
1826
+ return asyncio.run(_main(_parser().parse_args(argv)))
1827
+
1828
+
1829
+ if __name__ == "__main__":
1830
+ raise SystemExit(main())