Copyright (c) 2026 Microsoft Corporation. - Graphrag

python

# Copyright (c) 2026 Microsoft Corporation.
# Licensed under the MIT License.

Custom storage example

Here we create a custom storage implementation by extending the base Storage class and registering it with the GraphRAG storage system. Once registered, the custom storage can be instantiated through the factory pattern using either StorageConfig or directly via storage_factory, enabling extensible storage solutions for different backends.

python

import re
from collections.abc import Iterator
from typing import Any

from graphrag_storage import Storage, StorageConfig, create_storage, register_storage


class MyStorage(Storage):
    """Custom storage implementation example."""

    def __init__(
        self,
        some_setting: str,
        optional_setting: str = "default setting",
        **kwargs: Any,
    ):
        # Validate settings and initialize
        self._storage: dict[str, Any] = {}
        self._some_setting = some_setting
        self._optional_setting = optional_setting

    async def get(
        self, key: str, as_bytes: bool | None = None, encoding: str | None = None
    ) -> Any:
        """Retrieve a value from the cache by key."""
        return self._storage.get(key)

    async def set(self, key: str, value: Any, encoding: str | None = None) -> None:
        """Store a value in the cache with the specified key."""
        self._storage[key] = value

    async def has(self, key: str) -> bool:
        """Check if a key exists in the cache."""
        return key in self._storage

    async def delete(self, key: str) -> None:
        """Remove a key and its value from the cache."""
        self._storage.pop(key, None)

    async def clear(self) -> None:
        """Clear all items from the cache."""
        self._storage.clear()

    def find(
        self,
        file_pattern: re.Pattern[str] | None = None,
        prefix: str | None = None,
        suffix: str | None = None,
    ) -> Iterator[str]:
        """Find keys matching the given pattern, prefix, or suffix."""
        keys = list(self._storage.keys())
        if file_pattern:
            keys = [k for k in keys if file_pattern.match(k)]
        if prefix:
            keys = [k for k in keys if k.startswith(prefix)]
        if suffix:
            keys = [k for k in keys if k.endswith(suffix)]
        return iter(keys)

    def keys(self) -> list[str]:
        """Return all keys in the storage."""
        return list(self._storage.keys())

    def child(self, name: str | None = None) -> "Storage":
        """Create a child storage with the given path prefix."""
        return self

    async def get_creation_date(self, key: str) -> str:
        """Get the creation date of a key."""
        return ""


register_storage("MyStorage", MyStorage)


async def run():
    """Run the custom storage example."""
    storage = create_storage(
        StorageConfig(
            type="MyStorage",
            some_setting="My Setting",  # type: ignore
        )
    )

    print("Saving and retrieving a value from storage...")
    print("Setting key 'my_key' to 'value'")
    await storage.set("my_key", "value")
    print("Getting key 'my_key':")
    print(await storage.get("my_key"))


if __name__ == "__main__":
    await run()