dbgraph-sdk

DBGraph helps you explore and find relevant data assets in a large, complex relational database. It introspects a schema, builds a navigable graph of tables/columns, optionally enriches it with LLM-generated descriptions, and lets you search/render/traverse that graph.

This package is a standalone SDK extracted from the internal dbgraph project, packaged for reuse across teams/services (originally built by minhdenthedev, packaged as an SDK by hainamnguyen192).

Install

Core install (graph building + traversal only, no DB drivers or LLM client):

pip install dbgraph-sdk

Pick the extras you actually need — each one only pulls in the dependency for that piece:

pip install "dbgraph-sdk[postgres]"        # PostgresDataGateway
pip install "dbgraph-sdk[mysql]"           # MySQLDataGateway
pip install "dbgraph-sdk[trino]"           # TrinoDataGateway
pip install "dbgraph-sdk[openai]"          # OAICompatibleLLM
pip install "dbgraph-sdk[search]"          # BM25SearchEngine
pip install "dbgraph-sdk[all]"             # everything above

SQLite is supported out of the box (Python's built-in sqlite3), no extra needed.

Quick start

from pathlib import Path
from typing import cast

from dbgraph import (
    BM25SearchEngine,
    JSONGraphLoader,
    JSONGraphWriter,
    RGraphBuilder,
    SemanticAspect,
    SqliteDataGateway,
)

# RGraphBuilder depends only on the RDataGateway interface, never on a
# specific database driver — swap in PostgresDataGateway, MySQLDataGateway
# or TrinoDataGateway to point it at a different database.
graph_builder = RGraphBuilder(SqliteDataGateway(Path("data/northwind.db")))

# build the schema graph (introspects tables/columns + profiles them)
graph = graph_builder.build_graph()

# optional: generate semantic descriptions for assets via an LLM
# from dbgraph import GraphDescriptorV1, OAICompatibleLLM
#
# graph_descriptor = GraphDescriptorV1(
#     llm=OAICompatibleLLM(model=..., base_url=..., api_key=...),
#     system_prompt=..., formating_prompt=..., target_prompt=...,
# )
# graph = graph_descriptor.rfill_semantic_aspects(graph)

# save the graph
JSONGraphWriter(json_path=Path("data/northwind-graph.json"), indent=2).write(graph)

# load it back later
graph = JSONGraphLoader(json_path=Path("data/northwind-graph.json")).load()

# index + search it with BM25 (requires the `search` extra)
search_engine = BM25SearchEngine(Path("data/northwind-index"))
semantic_aspects = {
    a.asset_id: cast(SemanticAspect, a.aspects["semantic_properties"])
    for a in graph.assets
    if "semantic_properties" in a.aspects
}
search_engine.index(semantic_aspects)

asset_ids = search_engine.search("Give me the total count of orders in each category")

For visualization purposes, here is a graph saved as JSON:

{
  "assets": [
    {
      "asset_id": "8ab5a624-0596-497e-a0ee-3996d95dbe63",
      "name": "Categories",
      "type": "table",
      "aspects": {
        "schema_properties": { "name": "Categories_table_schema", "pks": ["CategoryID"], "indices": {} },
        "statistical_properties": { "name": "Categories_table_stats", "num_columns": 4, "num_rows": 8 },
        "semantic_properties": {
          "name": "Categories_semantic",
          "description": "Stores product category definitions and metadata, serving as a lookup table for classifying products in the inventory system.",
          "keywords": ["categories", "product classification", "category definitions", "inventory groups", "product types"]
        }
      }
    }
  ],
  "links": [
    {
      "link_id": "db6bea93-a02c-4426-a2db-449e4a7bba8f",
      "name": "Categories_CategoryID",
      "type": "contain",
      "source_id": "8ab5a624-0596-497e-a0ee-3996d95dbe63",
      "destination_id": "04c20046-2808-4021-bbf1-99876e0eea6e",
      "aspects": {}
    }
  ]
}

Use cases

• Manipulating database schema — build the schema graph, store it, and use it to traverse the database, find join paths, get referenced tables, ...
• Profiling database — the Aspect concept represents different kinds of properties attached to a data asset (schema, statistics, semantics, ...).
• Render graph — output a schema graph as Markdown/text to feed as LLM context.
• LLM assistance — use an LLM to generate data assets' descriptions/keywords, and as input for downstream SQL generation.
• Search for data assets — search assets by description via BM25 indexing/retrieval.

Architecture

DBGraph is designed to be easy to extend:

1. Core classes (entities) hold the shared business logic of database graphs (traversal, neighborhoods, ...) and core operations (building graphs, profiling databases, ...). The prefix R... stands for "Relational" (the only paradigm currently supported); D..., V..., G... are reserved for Document/Vector/Graph paradigms.
2. Interfaces (extensions) mark the parts of the system meant to be pluggable:
   • RGraphBuilder works against any RDBMS via the RDataGateway abstraction — implementations ship for SQLite, PostgreSQL, MySQL and Trino.
   • LLM abstracts the model provider — OAICompatibleLLM is the bundled implementation ([openai] extra); bring your own by implementing LLM.generate/agenerate.
   • GraphWriter/GraphLoader abstract graph persistence — JSONGraphWriter/JSONGraphLoader are the bundled implementation.
   • SearchEngine abstracts indexing/retrieval — BM25SearchEngine is the bundled implementation ([search] extra).

Development

uv sync --group dev --all-extras
uv run pytest
uv run pylint dbgraph
uv run mypy dbgraph
uv run flake8 dbgraph

Tests that talk to Postgres/MySQL/Trino/OpenAI need real credentials (see tests/) and are skipped/fail without them; the SQLite, entity and loader/writer tests run standalone.