microsoft--graphrag/packages/graphrag-storage/example_notebooks/custom_storage_example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1addaf67",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright (c) 2026 Microsoft Corporation.\n",
    "# Licensed under the MIT License."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f42d0ae8",
   "metadata": {},
   "source": [
    "## Custom storage example\n",
    "\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "55c64a48",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving and retrieving a value from storage...\n",
      "Setting key 'my_key' to 'value'\n",
      "Getting key 'my_key':\n",
      "value\n"
     ]
    }
   ],
   "source": [
    "import re\n",
    "from collections.abc import Iterator\n",
    "from typing import Any\n",
    "\n",
    "from graphrag_storage import Storage, StorageConfig, create_storage, register_storage\n",
    "\n",
    "\n",
    "class MyStorage(Storage):\n",
    "    \"\"\"Custom storage implementation example.\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        some_setting: str,\n",
    "        optional_setting: str = \"default setting\",\n",
    "        **kwargs: Any,\n",
    "    ):\n",
    "        # Validate settings and initialize\n",
    "        self._storage: dict[str, Any] = {}\n",
    "        self._some_setting = some_setting\n",
    "        self._optional_setting = optional_setting\n",
    "\n",
    "    async def get(\n",
    "        self, key: str, as_bytes: bool | None = None, encoding: str | None = None\n",
    "    ) -> Any:\n",
    "        \"\"\"Retrieve a value from the cache by key.\"\"\"\n",
    "        return self._storage.get(key)\n",
    "\n",
    "    async def set(self, key: str, value: Any, encoding: str | None = None) -> None:\n",
    "        \"\"\"Store a value in the cache with the specified key.\"\"\"\n",
    "        self._storage[key] = value\n",
    "\n",
    "    async def has(self, key: str) -> bool:\n",
    "        \"\"\"Check if a key exists in the cache.\"\"\"\n",
    "        return key in self._storage\n",
    "\n",
    "    async def delete(self, key: str) -> None:\n",
    "        \"\"\"Remove a key and its value from the cache.\"\"\"\n",
    "        self._storage.pop(key, None)\n",
    "\n",
    "    async def clear(self) -> None:\n",
    "        \"\"\"Clear all items from the cache.\"\"\"\n",
    "        self._storage.clear()\n",
    "\n",
    "    def find(\n",
    "        self,\n",
    "        file_pattern: re.Pattern[str] | None = None,\n",
    "        prefix: str | None = None,\n",
    "        suffix: str | None = None,\n",
    "    ) -> Iterator[str]:\n",
    "        \"\"\"Find keys matching the given pattern, prefix, or suffix.\"\"\"\n",
    "        keys = list(self._storage.keys())\n",
    "        if file_pattern:\n",
    "            keys = [k for k in keys if file_pattern.match(k)]\n",
    "        if prefix:\n",
    "            keys = [k for k in keys if k.startswith(prefix)]\n",
    "        if suffix:\n",
    "            keys = [k for k in keys if k.endswith(suffix)]\n",
    "        return iter(keys)\n",
    "\n",
    "    def keys(self) -> list[str]:\n",
    "        \"\"\"Return all keys in the storage.\"\"\"\n",
    "        return list(self._storage.keys())\n",
    "\n",
    "    def child(self, name: str | None = None) -> \"Storage\":\n",
    "        \"\"\"Create a child storage with the given path prefix.\"\"\"\n",
    "        return self\n",
    "\n",
    "    async def get_creation_date(self, key: str) -> str:\n",
    "        \"\"\"Get the creation date of a key.\"\"\"\n",
    "        return \"\"\n",
    "\n",
    "\n",
    "register_storage(\"MyStorage\", MyStorage)\n",
    "\n",
    "\n",
    "async def run():\n",
    "    \"\"\"Run the custom storage example.\"\"\"\n",
    "    storage = create_storage(\n",
    "        StorageConfig(\n",
    "            type=\"MyStorage\",\n",
    "            some_setting=\"My Setting\",  # type: ignore\n",
    "        )\n",
    "    )\n",
    "\n",
    "    print(\"Saving and retrieving a value from storage...\")\n",
    "    print(\"Setting key 'my_key' to 'value'\")\n",
    "    await storage.set(\"my_key\", \"value\")\n",
    "    print(\"Getting key 'my_key':\")\n",
    "    print(await storage.get(\"my_key\"))\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    await run()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}