"""Root-last publication of immutable version 0.1 datasets."""
from __future__ import annotations
import hashlib
from collections.abc import Iterable, Mapping, Sequence
from signal_dataset._internal.constants import ROOT_DOCUMENT
from signal_dataset._internal.json import Json
from signal_dataset._internal.retry import RetryPolicy, retrying
from signal_dataset.config import PublicationOptions, StorageOptions
from signal_dataset.dataset.control_io import create_document
from signal_dataset.dataset.layout import (
DatasetRoot,
Manifest,
ManifestReference,
ShardEntry,
Snapshot,
)
from signal_dataset.dataset.reader import Dataset, open
from signal_dataset.record.model import PublishedShard
from signal_dataset.storage import object_store_for
from signal_dataset.storage.contracts import ObjectStore, ShardStore
from signal_dataset.storage.paths import (
join_uri,
relative_uri,
validate_root_uri,
)
from signal_dataset.storage.reference import ObjectReference
def _fanout(identifier: str) -> str:
return hashlib.sha256(identifier.encode("utf-8")).hexdigest()[:2]
def _object_reference(root: str, uri: str, size: int) -> ObjectReference:
"""Describe a stored object relative to the dataset root.
New references record no object version.
A generation identified one particular write of an object, and only GCS
produced one, so the same dataset described itself differently on each
backend and its control documents differed on every run. Every object a
reference can name is create-only and immutable, so the current version
is the only version: reading a path unpinned returns the same bytes the
pin would have selected. Dropping it makes a published document a pure
function of its inputs, which is what lets an interrupted publication be
re-run. References written before this change still carry the field and
the read path still honours it.
"""
return ObjectReference(relative_uri(root, uri), None, size)
def _winners(shards: Iterable[PublishedShard]) -> list[PublishedShard]:
winners: dict[str, PublishedShard] = {}
for shard in shards:
previous = winners.get(shard.work_id)
if previous is None or shard.attempt > previous.attempt:
winners[shard.work_id] = shard
return [winners[name] for name in sorted(winners)]
[docs]
def publish(
root: str,
shards: Iterable[PublishedShard],
*,
dataset_id: str,
snapshot_id: str,
expected_work_ids: Sequence[str] | None = None,
metadata: Mapping[str, Json] | None = None,
object_store: ObjectStore | None = None,
shard_store: ShardStore | None = None,
storage_options: StorageOptions | None = None,
publication_options: PublicationOptions | None = None,
retry: int | RetryPolicy | None = None,
) -> Dataset:
"""Publish shards as an immutable snapshot, writing root.json last.
Safe to re-run. Every write is create-only and resumes on an identical
object, so an attempt interrupted anywhere completes when repeated -- which
is also what makes `retry` safe: a transient failure re-enters the same
body and finds its own earlier writes rather than colliding with them.
"""
policy = RetryPolicy.coerce(retry)
root = validate_root_uri(root)
objects = (
object_store_for(root, root_is_directory=True) if object_store is None else object_store
)
settings = PublicationOptions() if publication_options is None else publication_options
selected = _winners(shards)
if expected_work_ids is not None:
expected = set(expected_work_ids)
if len(expected) != len(expected_work_ids):
raise ValueError("expected_work_ids contains duplicates")
observed = {item.work_id for item in selected}
missing = expected - observed
if missing:
raise ValueError(f"missing work IDs: {sorted(missing)!r}")
unexpected = observed - expected
if unexpected:
raise ValueError(f"unexpected work IDs: {sorted(unexpected)!r}")
def attempt() -> Dataset:
return _publish_once(
root,
objects,
selected,
dataset_id=dataset_id,
snapshot_id=snapshot_id,
metadata=metadata,
shard_store=shard_store,
storage_options=storage_options,
settings=settings,
)
return retrying(policy, attempt, description=f"publishing {root}")
def _publish_once(
root: str,
objects: ObjectStore,
selected: Sequence[PublishedShard],
*,
dataset_id: str,
snapshot_id: str,
metadata: Mapping[str, Json] | None,
shard_store: ShardStore | None,
storage_options: StorageOptions | None,
settings: PublicationOptions,
) -> Dataset:
_verify_shards(
objects,
root,
selected,
shard_store,
storage_options,
settings.verify_record_counts,
)
entries: list[ShardEntry] = []
ordinal = 0
for item in selected:
entries.append(
ShardEntry(
ordinal,
item.record_count,
_object_reference(root, item.data_uri, item.data_bytes),
_object_reference(root, item.metadata_uri, item.metadata_bytes),
)
)
ordinal += item.record_count
references: list[ManifestReference] = []
for page_number, start in enumerate(range(0, len(entries), settings.shards_per_manifest)):
page_entries = tuple(entries[start : start + settings.shards_per_manifest])
page = Manifest(
page_entries[0].first_ordinal,
sum(item.record_count for item in page_entries),
page_entries,
)
page_id = f"{snapshot_id}-{page_number:08d}"
path = f"manifests/{_fanout(page_id)}/{page_id}.json"
size = create_document(objects, join_uri(root, path), page.to_dict())
references.append(
ManifestReference(
page.first_ordinal,
page.record_count,
len(page_entries),
sum(item.data.stored_bytes + item.metadata.stored_bytes for item in page_entries),
ObjectReference(path, None, size),
)
)
stored_bytes = sum(item.data_bytes + item.metadata_bytes for item in selected)
snapshot = Snapshot(
snapshot_id,
ordinal,
len(selected),
stored_bytes,
tuple(references),
)
snapshot_path = f"snapshots/{snapshot_id}/snapshot.json"
snapshot_size = create_document(objects, join_uri(root, snapshot_path), snapshot.to_dict())
dataset_root = DatasetRoot(
dataset_id,
ObjectReference(snapshot_path, None, snapshot_size),
dict(metadata or {}),
)
create_document(objects, join_uri(root, ROOT_DOCUMENT), dataset_root.to_dict())
return open(root, object_store=objects, options=storage_options, shard_store=shard_store)
def _verify_shards(
objects: ObjectStore,
root: str,
shards: Sequence[PublishedShard],
shard_store: ShardStore | None,
storage_options: StorageOptions | None,
verify_record_counts: bool = False,
) -> None:
from signal_dataset.storage import ArrayRecordShardStore
from signal_dataset.storage.staging import default_staging
options = StorageOptions() if storage_options is None else storage_options
indexed = (
ArrayRecordShardStore(
staging=default_staging(
objects, root=root, max_open_shards=options.max_open_shards
)
)
if shard_store is None
else shard_store
)
for shard in shards:
if shard.record_count < 0:
raise ValueError("shard record_count must be nonnegative")
for uri, size, generation in (
(shard.data_uri, shard.data_bytes, shard.data_generation),
(shard.metadata_uri, shard.metadata_bytes, shard.metadata_generation),
):
info = objects.info(uri, generation=generation)
if info.size != size:
raise ValueError(f"published shard size mismatch: {uri}")
if generation is not None and info.generation != generation:
raise ValueError(f"published shard generation mismatch: {uri}")
if not verify_record_counts:
continue
count = indexed.count(
uri,
generation=generation,
options=options.reader_options,
file_reader_buffer_size=options.file_reader_buffer_size,
)
if count != shard.record_count:
raise ValueError(f"published shard record count mismatch: {uri}")