graflo.architecture.backend.writer

Streaming writer for GraFlo file backend directories.

Attributes

Classes

GraFloBackendWriter

Write schema and chunked graph data to a GraFlo backend directory.

Several writers, in one process or several, may write one directory at once: :meth:flush_index adds this writer's chunks to the index on disk under :func:directory_lock, rather than replacing it.

Source code in graflo/architecture/backend/writer.py

class GraFloBackendWriter:
    """Write schema and chunked graph data to a GraFlo backend directory.

    Several writers, in one process or several, may write one directory at
    once: :meth:`flush_index` adds this writer's chunks to the index on disk
    under :func:`directory_lock`, rather than replacing it.
    """

    def __init__(
        self,
        output_dir: Path,
        *,
        chunk_size: int = 50_000,
        resume: bool = False,
    ) -> None:
        self._layout = GraFloLayout(output_dir)
        self._chunk_size = chunk_size
        self._resume = resume
        self._schema: Schema | None = None
        self._index: GraFloIndex | None = None
        self._vertex_writers: dict[str, _CollectionWriter] = {}
        self._edge_writers: dict[tuple[str, str, str | None], _CollectionWriter] = {}
        if resume:
            self._index = _read_index(self._layout.index_path)

    def __enter__(self) -> Self:
        self._layout.ensure_dirs()
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> None:
        if exc is None:
            self.flush_index()

    @property
    def layout(self) -> GraFloLayout:
        return self._layout

    def reset_data(self) -> None:
        """Remove data chunks and index while keeping schema if present."""
        with directory_lock(self._layout.root):
            for path in (self._layout.vertices_dir, self._layout.edges_dir):
                if path.exists():
                    shutil.rmtree(path)
            if self._layout.index_path.exists():
                self._layout.index_path.unlink()
        self._index = None
        self._vertex_writers = {}
        self._edge_writers = {}
        self._layout.ensure_dirs()

    def write_schema(self, schema: Schema) -> None:
        self._schema = schema
        self._layout.ensure_dirs()
        import yaml

        with open(self._layout.schema_path, "w", encoding="utf-8") as fout:
            yaml.safe_dump(
                schema.model_dump(mode="json", by_alias=True, exclude_none=True),
                fout,
                default_flow_style=False,
                sort_keys=False,
            )

    def write_vertex_batch(self, vertex_type: str, docs: list[dict[str, Any]]) -> None:
        if not docs:
            return
        writer = self._vertex_writers.setdefault(
            vertex_type,
            self._make_vertex_writer(vertex_type),
        )
        writer.push_many(docs)

    def write_edge_batch(
        self,
        edge_key: tuple[str, str, str | None],
        docs: list[list[Any]],
    ) -> None:
        if not docs:
            return
        writer = self._edge_writers.setdefault(
            edge_key,
            self._make_edge_writer(edge_key),
        )
        writer.push_many(docs)

    def flush_index(self) -> GraFloIndex:
        if self._schema is None and self._layout.schema_path.exists():
            self._schema = Schema.from_yaml(str(self._layout.schema_path))
        if self._schema is None:
            raise ValueError("Cannot flush GraFlo backend index without schema")

        for writer in self._vertex_writers.values():
            writer.flush()
        for writer in self._edge_writers.values():
            writer.flush()

        with directory_lock(self._layout.root):
            on_disk = _read_index(self._layout.index_path)
            vertices = dict(on_disk.vertices) if on_disk is not None else {}
            for name, writer in self._vertex_writers.items():
                vertices[name] = _merge_entry(vertices.get(name), *writer.unpublished())
            edges = dict(on_disk.edges) if on_disk is not None else {}
            for edge_key, writer in self._edge_writers.items():
                name = GraFloLayout.edge_key_to_index_name(edge_key)
                edges[name] = _merge_entry(edges.get(name), *writer.unpublished())
            index = GraFloIndex(
                graflo_version=_graflo_package_version(),
                schema_hash=backend_schema_hash(self._schema),
                vertices=vertices,
                edges=edges,
            )
            staged = self._layout.index_path.with_name(
                f"{self._layout.index_path.name}.{os.getpid()}.tmp"
            )
            with open(staged, "w", encoding="utf-8") as fout:
                fout.write(
                    json.dumps(
                        index.model_dump(mode="json", by_alias=True, exclude_none=True),
                        indent=2,
                        sort_keys=True,
                    )
                    + "\n"
                )
            os.replace(staged, self._layout.index_path)
        for writer in (*self._vertex_writers.values(), *self._edge_writers.values()):
            writer.published()
        self._index = index
        return index

    def _make_vertex_writer(self, vertex_type: str) -> _CollectionWriter:
        existing = self._index.vertices.get(vertex_type) if self._index else None
        return _CollectionWriter(
            self._layout,
            chunk_size=self._chunk_size,
            vertex_type=vertex_type,
            existing=existing,
        )

    def _make_edge_writer(
        self, edge_key: tuple[str, str, str | None]
    ) -> _CollectionWriter:
        index_name = GraFloLayout.edge_key_to_index_name(edge_key)
        existing = self._index.edges.get(index_name) if self._index else None
        return _CollectionWriter(
            self._layout,
            chunk_size=self._chunk_size,
            edge_key=edge_key,
            existing=existing,
        )

Attributes

layout property

Methods:

__enter__()

Source code in graflo/architecture/backend/writer.py

def __enter__(self) -> Self:
    self._layout.ensure_dirs()
    return self

__exit__(exc_type, exc, tb)

Source code in graflo/architecture/backend/writer.py

def __exit__(
    self,
    exc_type: type[BaseException] | None,
    exc: BaseException | None,
    tb: TracebackType | None,
) -> None:
    if exc is None:
        self.flush_index()

__init__(output_dir, *, chunk_size=50000, resume=False)

Source code in graflo/architecture/backend/writer.py

def __init__(
    self,
    output_dir: Path,
    *,
    chunk_size: int = 50_000,
    resume: bool = False,
) -> None:
    self._layout = GraFloLayout(output_dir)
    self._chunk_size = chunk_size
    self._resume = resume
    self._schema: Schema | None = None
    self._index: GraFloIndex | None = None
    self._vertex_writers: dict[str, _CollectionWriter] = {}
    self._edge_writers: dict[tuple[str, str, str | None], _CollectionWriter] = {}
    if resume:
        self._index = _read_index(self._layout.index_path)

flush_index()

Source code in graflo/architecture/backend/writer.py

def flush_index(self) -> GraFloIndex:
    if self._schema is None and self._layout.schema_path.exists():
        self._schema = Schema.from_yaml(str(self._layout.schema_path))
    if self._schema is None:
        raise ValueError("Cannot flush GraFlo backend index without schema")

    for writer in self._vertex_writers.values():
        writer.flush()
    for writer in self._edge_writers.values():
        writer.flush()

    with directory_lock(self._layout.root):
        on_disk = _read_index(self._layout.index_path)
        vertices = dict(on_disk.vertices) if on_disk is not None else {}
        for name, writer in self._vertex_writers.items():
            vertices[name] = _merge_entry(vertices.get(name), *writer.unpublished())
        edges = dict(on_disk.edges) if on_disk is not None else {}
        for edge_key, writer in self._edge_writers.items():
            name = GraFloLayout.edge_key_to_index_name(edge_key)
            edges[name] = _merge_entry(edges.get(name), *writer.unpublished())
        index = GraFloIndex(
            graflo_version=_graflo_package_version(),
            schema_hash=backend_schema_hash(self._schema),
            vertices=vertices,
            edges=edges,
        )
        staged = self._layout.index_path.with_name(
            f"{self._layout.index_path.name}.{os.getpid()}.tmp"
        )
        with open(staged, "w", encoding="utf-8") as fout:
            fout.write(
                json.dumps(
                    index.model_dump(mode="json", by_alias=True, exclude_none=True),
                    indent=2,
                    sort_keys=True,
                )
                + "\n"
            )
        os.replace(staged, self._layout.index_path)
    for writer in (*self._vertex_writers.values(), *self._edge_writers.values()):
        writer.published()
    self._index = index
    return index

reset_data()

Remove data chunks and index while keeping schema if present.

Source code in graflo/architecture/backend/writer.py

def reset_data(self) -> None:
    """Remove data chunks and index while keeping schema if present."""
    with directory_lock(self._layout.root):
        for path in (self._layout.vertices_dir, self._layout.edges_dir):
            if path.exists():
                shutil.rmtree(path)
        if self._layout.index_path.exists():
            self._layout.index_path.unlink()
    self._index = None
    self._vertex_writers = {}
    self._edge_writers = {}
    self._layout.ensure_dirs()

write_edge_batch(edge_key, docs)

Source code in graflo/architecture/backend/writer.py

def write_edge_batch(
    self,
    edge_key: tuple[str, str, str | None],
    docs: list[list[Any]],
) -> None:
    if not docs:
        return
    writer = self._edge_writers.setdefault(
        edge_key,
        self._make_edge_writer(edge_key),
    )
    writer.push_many(docs)

write_schema(schema)

Source code in graflo/architecture/backend/writer.py

def write_schema(self, schema: Schema) -> None:
    self._schema = schema
    self._layout.ensure_dirs()
    import yaml

    with open(self._layout.schema_path, "w", encoding="utf-8") as fout:
        yaml.safe_dump(
            schema.model_dump(mode="json", by_alias=True, exclude_none=True),
            fout,
            default_flow_style=False,
            sort_keys=False,
        )

write_vertex_batch(vertex_type, docs)

Source code in graflo/architecture/backend/writer.py

def write_vertex_batch(self, vertex_type: str, docs: list[dict[str, Any]]) -> None:
    if not docs:
        return
    writer = self._vertex_writers.setdefault(
        vertex_type,
        self._make_vertex_writer(vertex_type),
    )
    writer.push_many(docs)

Functions:

directory_lock(root)

Hold an exclusive lock on the backend directory root, across processes.

Writers take it only to publish their chunks into INDEX.json; chunk files themselves are claimed by exclusive creation and need no lock.

Source code in graflo/architecture/backend/writer.py

@contextmanager
def directory_lock(root: Path) -> Generator[None, None, None]:
    """Hold an exclusive lock on the backend directory *root*, across processes.

    Writers take it only to publish their chunks into ``INDEX.json``; chunk
    files themselves are claimed by exclusive creation and need no lock.
    """
    root.mkdir(parents=True, exist_ok=True)
    with open(root / LOCK_FILENAME, "a+b") as handle:
        if sys.platform == "win32":
            import msvcrt

            handle.seek(0)
            msvcrt.locking(handle.fileno(), msvcrt.LK_LOCK, 1)
            try:
                yield
            finally:
                handle.seek(0)
                msvcrt.locking(handle.fileno(), msvcrt.LK_UNLCK, 1)
        else:
            import fcntl

            fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
            try:
                yield
            finally:
                fcntl.flock(handle.fileno(), fcntl.LOCK_UN)