Metadata-Version: 2.4
Name: dbgraph-sdk
Version: 0.1.0
Summary: Build LLM-assisted schema graphs over relational databases: introspect, describe, search and render.
Project-URL: Homepage, https://github.com/hainamnguyen192/dbgraph-sdk
Project-URL: Repository, https://github.com/hainamnguyen192/dbgraph-sdk
Author: minhdenthedev, hainamnguyen192
License: MIT
License-File: LICENSE
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Database
Classifier: Topic :: Software Development :: Libraries
Requires-Python: >=3.11
Requires-Dist: rustworkx>=0.18.0
Requires-Dist: tqdm>=4.68.4
Provides-Extra: all
Requires-Dist: bm25s>=0.3.9; extra == 'all'
Requires-Dist: openai>=2.45.0; extra == 'all'
Requires-Dist: psycopg[binary]>=3.3.4; extra == 'all'
Requires-Dist: pymysql>=1.2.0; extra == 'all'
Requires-Dist: pystemmer>=3.1.0; extra == 'all'
Requires-Dist: trino>=0.338.0; extra == 'all'
Provides-Extra: mysql
Requires-Dist: pymysql>=1.2.0; extra == 'mysql'
Provides-Extra: openai
Requires-Dist: openai>=2.45.0; extra == 'openai'
Provides-Extra: postgres
Requires-Dist: psycopg[binary]>=3.3.4; extra == 'postgres'
Provides-Extra: search
Requires-Dist: bm25s>=0.3.9; extra == 'search'
Requires-Dist: pystemmer>=3.1.0; extra == 'search'
Provides-Extra: trino
Requires-Dist: trino>=0.338.0; extra == 'trino'
Description-Content-Type: text/markdown

# 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):

```bash
pip install dbgraph-sdk
```

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

```bash
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

```python
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:

```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

![Use cases of DBGraph](https://raw.githubusercontent.com/hainamnguyen192/dbgraph-sdk/main/diagrams/usecase.png)

- **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

![Class diagram of DBGraph](https://raw.githubusercontent.com/hainamnguyen192/dbgraph-sdk/main/diagrams/entity.png)

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

```bash
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.
