dbgraph 0.1.0__py3-none-any.whl

dbgraph/__init__.py

@@ -0,0 +1,21 @@
+from dbgraph.entity.asset import Asset
+from dbgraph.entity.asset_type import AssetType
+from dbgraph.entity.dbgraph import DatabaseGraph
+from dbgraph.entity.link import Link
+from dbgraph.entity.link_type import LinkType
+from dbgraph.builder.sql_graph_builder import SQLGraphBuilder
+from dbgraph.persistent.sql_graph_persistent import SQLGraphPersistent
+from dbgraph.descriptor.langchain_graph_descriptor import LangchainGraphDescriptor
+from dbgraph.render.markdown_renderer import MarkdownRenderer
+
+__all__ = [
+    "Asset",
+    "AssetType",
+    "DatabaseGraph",
+    "Link",
+    "LinkType",
+    "SQLGraphBuilder",
+    "SQLGraphPersistent",
+    "LangchainGraphDescriptor",
+    "MarkdownRenderer"
+]

dbgraph/application.py

@@ -0,0 +1,70 @@
+from dataclasses import dataclass
+from uuid import uuid4
+
+from dbgraph.builder.graph_builder import GraphBuilder
+from dbgraph.builder.sql_graph_builder import SQLGraphBuilder
+from dbgraph.descriptor.graph_descriptor import GraphDescriptor
+from dbgraph.entity.asset import Asset
+from dbgraph.entity.dbgraph import DatabaseGraph
+from dbgraph.entity.link import Link
+from dbgraph.entity.rdbgraph import RDatabaseGraph
+from dbgraph.persistent.graph_persistent import GraphPersistent
+from dbgraph.persistent.sql_graph_persistent import SQLGraphPersistent
+
+
+@dataclass
+class Application:
+    """God class to control the logical flow within dbgraph"""
+
+    source_database_type: str
+    graph_persistent_uri: str
+    graph_persistent_type: str
+    descriptor: GraphDescriptor | None = None
+
+    def __post_init__(self):
+        self.graph_persistent = self.create_graph_persistent()
+
+    def create_graph_builder(self, database_uri: str) -> GraphBuilder:
+        match self.source_database_type:
+            case "sql":
+                builder = SQLGraphBuilder(database_uri)
+            case _:
+                raise NotImplementedError()
+        return builder
+
+    def create_graph_persistent(self) -> GraphPersistent:
+        match self.graph_persistent_type:
+            case "sql":
+                persistent = SQLGraphPersistent(self.graph_persistent_uri)
+            case _:
+                raise NotImplementedError()
+        return persistent
+
+    def build_graph(self, name: str, database_uri: str, fill_semantic: bool) -> str:
+        """Build the graph, return its ID in the database"""
+        builder = self.create_graph_builder(database_uri)
+        graph = builder.build_graph()
+        if fill_semantic and self.descriptor is not None:
+            graph = self.descriptor.rfill_semantic_aspects(
+                RDatabaseGraph.from_graph(graph)
+            )
+        graph_id = str(uuid4())
+        self.graph_persistent.create_graph(graph_id, name)
+        self.graph_persistent.save_graph(graph, name)
+        return graph_id
+
+    def load_graph(self, graph_id: str) -> DatabaseGraph:
+        """Load the graph using its ID"""
+        return self.graph_persistent.load_graph(graph_id)
+
+    def delete_graph(self, graph_id: str):
+        """Delete the graph using its ID"""
+        self.graph_persistent.delete_graph(graph_id)
+
+    def get_asset(self, asset_id: str, graph_id: str) -> Asset:
+        """Find an asset using its ID"""
+        return self.graph_persistent.get_asset(asset_id, graph_id)
+
+    def get_link(self, link_id: str, graph_id: str) -> Link:
+        """Find a link using its ID"""
+        return self.graph_persistent.get_link(link_id, graph_id)

dbgraph/builder/graph_builder.py

@@ -0,0 +1,22 @@
+from abc import ABC, abstractmethod
+
+from dbgraph.entity.asset import Asset
+from dbgraph.entity.dbgraph import DatabaseGraph
+from dbgraph.entity.link import Link
+
+
+class GraphBuilder(ABC):
+    """Interface for relational database graph builder"""
+
+    @abstractmethod
+    def _build_assets(self) -> list[Asset]:
+        pass
+
+    @abstractmethod
+    def _build_links(self) -> list[Link]:
+        pass
+
+    def build_graph(self) -> DatabaseGraph:
+        assets = self._build_assets()
+        links = self._build_links()
+        return DatabaseGraph(assets, links)

dbgraph/builder/sql_graph_builder.py

@@ -0,0 +1,245 @@
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+from uuid import uuid4
+
+from sqlalchemy import (
+    Column,
+    MetaData,
+    Numeric,
+    PrimaryKeyConstraint,
+    String,
+    Table,
+    create_engine,
+    func,
+    inspect,
+    select,
+)
+
+from dbgraph.builder.graph_builder import GraphBuilder
+from dbgraph.entity.aspect import (
+    RCategoricalStatistics,
+    RColumnSchemaAspect,
+    RColumnStatisticsAspect,
+    RForeignKeyAspect,
+    RNumericalStatistics,
+    RTableSchemaAspect,
+    RTableStatisticsAspect,
+)
+from dbgraph.entity.asset import Asset
+from dbgraph.entity.asset_type import AssetType
+from dbgraph.entity.link import Link
+from dbgraph.entity.link_type import LinkType
+
+
+@dataclass
+class SQLGraphBuilder(GraphBuilder):
+    """Build the database graph using SQL Alchemy"""
+
+    db_uri: str
+    max_worker: int = 4
+
+    def __post_init__(self):
+        self.engine = create_engine(self.db_uri)
+        self.columns_assets: dict[str, list[Asset]] = {}
+        self.tables_assets: dict[str, Asset] = {}
+        self._tables_orm: dict[str, Table] = {
+            table.name: table for table in self._orm_tables
+        }
+
+    @property
+    def _orm_tables(self) -> list[Table]:
+        orm_tables = []
+        for table_name in self.table_names:
+            metadata = MetaData()
+            table = Table(table_name, metadata, autoload_with=self.engine)
+            orm_tables.append(table)
+        return orm_tables
+
+    @property
+    def table_names(self) -> list[str]:
+        inspector = inspect(self.engine)
+        return inspector.get_table_names()
+
+    def _get_numerical_stats_aspect(self, col: Column) -> RNumericalStatistics:
+        stmt = select(
+            func.min(col).label("min"),
+            func.max(col).label("max"),
+            func.avg(col).label("mean"),
+        )
+        with self.engine.connect() as conn:
+            result = conn.execute(stmt).one()
+        return RNumericalStatistics(
+            min=float(result.min), max=float(result.max), mean=float(result.mean)
+        )
+
+    def _get_null_info(self, col: Column) -> tuple[int, int]:
+        stmt = select(
+            func.count(col).label("non_null_count"),
+            (func.count() - func.count(col)).label("null_count"),
+        )
+        with self.engine.connect() as conn:
+            result = conn.execute(stmt).one()
+        return int(result.non_null_count), int(result.null_count)
+
+    def _get_cat_stats_aspect(self, col: Column) -> RCategoricalStatistics:
+        stmt = (
+            select(
+                col.label("value"),
+                func.count().label("value_count"),
+            )
+            .where(col.is_not(None))
+            .group_by(col)
+            .order_by(func.count().desc())
+            .limit(10)
+        )
+        with self.engine.connect() as conn:
+            rows = conn.execute(stmt).all()
+        value_counts = {str(row[0]): int(row[1]) for row in rows}
+        return RCategoricalStatistics(value_counts)
+
+    def _get_stats_aspect(self, col: Column) -> RColumnStatisticsAspect:
+        non_null_count, null_count = self._get_null_info(col)
+        if isinstance(col.type, Numeric):
+            num_aspect = self._get_numerical_stats_aspect(col)
+            cat_aspect = None
+        elif isinstance(col.type, String):
+            num_aspect = None
+            cat_aspect = self._get_cat_stats_aspect(col)
+        else:
+            num_aspect = None
+            cat_aspect = None
+        return RColumnStatisticsAspect(
+            name=f"{col.name}_column_stats",
+            numerical_stats=num_aspect,
+            categorical_stats=cat_aspect,
+            non_null_count=non_null_count,
+            null_count=null_count,
+        )
+
+    def _make_column_asset(self, col: Column, pks: PrimaryKeyConstraint) -> Asset:
+        schema_aspect = RColumnSchemaAspect(
+            name=f"{col.name}_column_schema",
+            dtype=str(col.type),
+            is_nullable=col.nullable or False,
+            is_pk=col.name in [c.name for c in pks.columns],
+        )
+        stats_aspect = self._get_stats_aspect(col)
+        return Asset(
+            asset_id=str(uuid4()),
+            name=col.name,
+            type=AssetType.RCOLUMN,
+            aspects={
+                "schema_properties": schema_aspect,
+                "statistical_properties": stats_aspect,
+            },
+        )
+
+    def _get_columns_assets(self, table_name: str) -> list[Asset]:
+        table = self._tables_orm[table_name]
+        assets = []
+        pks = table.primary_key
+        with ThreadPoolExecutor(max_workers=self.max_worker) as executor:
+            assets = list(
+                executor.map(
+                    self._make_column_asset, table.columns, [pks] * len(table.columns)
+                )
+            )
+        return assets
+
+    def _get_table_stats_aspect(self, table_name: str) -> RTableStatisticsAspect:
+        table = self._tables_orm[table_name]
+        nrow = self._get_table_nrow(table_name)
+        ncol = len(table.columns)
+        return RTableStatisticsAspect(
+            name=f"{table.name}_table_stats", num_rows=nrow, num_columns=ncol
+        )
+
+    def _get_table_nrow(self, table_name: str) -> int:
+        table = self._tables_orm[table_name]
+
+        stmt = select(func.count()).select_from(table)
+
+        with self.engine.connect() as conn:
+            nrow = conn.execute(stmt).scalar()
+        if nrow is None:
+            raise RuntimeError("Can't get number of row")
+        return int(nrow)
+
+    def _get_table_schema_aspect(self, table_name: str) -> RTableSchemaAspect:
+        table = self._tables_orm[table_name]
+        pks = [c.name for c in table.primary_key.columns]
+        indices = {str(i.name): [c.name for c in i.columns] for i in table.indexes}
+
+        return RTableSchemaAspect(
+            name=f"{table_name}_table_schema", pks=pks, indices=indices
+        )
+
+    def _make_table_asset(self, table_name: str) -> Asset:
+        stats_aspect = self._get_table_stats_aspect(table_name)
+        schema_aspect = self._get_table_schema_aspect(table_name)
+        table_asset = Asset(
+            asset_id=str(uuid4()),
+            name=table_name,
+            type=AssetType.RTABLE,
+            aspects={
+                "schema_properties": schema_aspect,
+                "statistical_properties": stats_aspect,
+            },
+        )
+        self.columns_assets[table_asset.name] = self._get_columns_assets(table_name)
+        self.tables_assets[table_asset.name] = table_asset
+        return table_asset
+
+    def _build_assets(self) -> list[Asset]:
+        with ThreadPoolExecutor(max_workers=self.max_worker) as executor:
+            assets = list(executor.map(self._make_table_asset, self.table_names))
+        for assets_list in self.columns_assets.values():
+            assets.extend(assets_list)
+        return assets
+
+    def _build_contain_links(self) -> list[Link]:
+        contain_links = []
+        for table_name, columns_assets in self.columns_assets.items():
+            table_asset = self.tables_assets[table_name]
+            for column_asset in columns_assets:
+                link = Link(
+                    link_id=str(uuid4()),
+                    name=f"{table_asset.name}_{column_asset.name}",
+                    source_id=table_asset.asset_id,
+                    destination_id=column_asset.asset_id,
+                    type=LinkType.CONTAIN,
+                )
+                contain_links.append(link)
+        return contain_links
+
+    def _build_fk_links(self) -> list[Link]:
+        links = []
+        for table_name in self.tables_assets:
+            table = self._tables_orm[table_name]
+            for fk in table.foreign_keys:
+                to_table = fk.column.table.name
+                from_table = table.name
+                name = f"{from_table}_{to_table}_fk"
+                link = Link(
+                    link_id=str(uuid4()),
+                    name=name,
+                    source_id=self.tables_assets[from_table].asset_id,
+                    destination_id=self.tables_assets[to_table].asset_id,
+                    type=LinkType.FOREIGN_KEY,
+                    aspects={
+                        "foreign_key_properties": RForeignKeyAspect(
+                            name=name,
+                            from_column=fk.parent.name,
+                            to_column=fk.column.name,
+                            on_delete=fk.ondelete or "",
+                            on_update=fk.onupdate or "",
+                        )
+                    },
+                )
+                links.append(link)
+        return links
+
+    def _build_links(self) -> list[Link]:
+        contain_links = self._build_contain_links()
+        fk_links = self._build_fk_links()
+        return contain_links + fk_links

dbgraph/descriptor/graph_descriptor.py

@@ -0,0 +1,68 @@
+from abc import ABC, abstractmethod
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+
+from dbgraph.entity.aspect import SemanticAspect
+from dbgraph.entity.asset import Asset
+from dbgraph.entity.rdbgraph import RDatabaseGraph
+
+
+@dataclass
+class GraphDescriptor(ABC):
+    """Base class that generate semantic aspects for a database graph or a sub-graph"""
+
+    max_workers: int
+
+    @abstractmethod
+    def get_semantic_aspect(
+        self, asset: Asset, context: RDatabaseGraph
+    ) -> SemanticAspect:
+        """Get the description of this asset.
+
+        Args:
+            asset: Target asset
+            context: an instance of `RDatabaseGraph`, could be the whole database graph or a sub-graph
+
+        Returns:
+            semantic aspect for this asset
+        """
+
+    def rfill_semantic_aspects(self, graph: RDatabaseGraph) -> RDatabaseGraph:
+        tables = graph.get_tables()
+        assets = tables
+        contexts = [graph.select_connected_tables(table.asset_id) for table in tables]
+        for table in tables:
+            context = RDatabaseGraph(assets=[table], links=[])
+            columns = graph.get_columns(table.asset_id)
+            for column in columns:
+                assets.append(column)
+                contexts.append(context)
+        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
+            results = {
+                asset.asset_id: future
+                for asset, future in zip(
+                    assets, executor.map(self.get_semantic_aspect, assets, contexts)
+                )
+            }
+        for key, aspect in results.items():
+            asset = graph.get_asset(key)
+            asset.aspects["semantic_properties"] = aspect
+        return graph
+
+    def rfill_semantic_aspects_seq(self, graph: RDatabaseGraph) -> RDatabaseGraph:
+        tables = graph.get_tables()
+        assets = tables
+        contexts = [graph.select_connected_tables(table.asset_id) for table in tables]
+        for table in tables:
+            context = RDatabaseGraph(assets=[table], links=[])
+            columns = graph.get_columns(table.asset_id)
+            for column in columns:
+                assets.append(column)
+                contexts.append(context)
+        results = {}
+        for asset, context in zip(assets, contexts):
+            results[asset.asset_id] = self.get_semantic_aspect(asset, context)
+        for key, aspect in results.items():
+            asset = graph.get_asset(key)
+            asset.aspects["semantic_properties"] = aspect
+        return graph

dbgraph/descriptor/langchain_graph_descriptor.py

@@ -0,0 +1,74 @@
+from dataclasses import dataclass
+
+from langchain.chat_models import BaseChatModel
+from langchain_core.messages import SystemMessage
+from langchain_core.messages.human import HumanMessage
+
+from dbgraph import AssetType
+from dbgraph.descriptor.graph_descriptor import GraphDescriptor
+from dbgraph.entity.aspect import SemanticAspect
+from dbgraph.entity.asset import Asset
+from dbgraph.entity.rdbgraph import RDatabaseGraph
+from dbgraph.render.markdown_renderer import MarkdownRenderer
+
+
+@dataclass
+class LangchainGraphDescriptor(GraphDescriptor):
+    """Implementation of `GraphDescriptor` using Langchain"""
+
+    model: BaseChatModel
+    table_system_prompt: str
+    table_question_prompt: str
+    column_system_prompt: str
+    column_question_prompt: str
+    markdown_renderer: MarkdownRenderer
+
+    def _generate(self, system_prompt: str, question: str) -> str:
+        response = self.model.invoke(
+            [SystemMessage(content=system_prompt), HumanMessage(content=question)]
+        )
+        if response.content is None:
+            raise ValueError("Couldn't generate description")
+        return str(response.content)
+
+    def _get_semantic_aspect_table(
+        self, asset: Asset, context: RDatabaseGraph
+    ) -> SemanticAspect:
+        self.markdown_renderer.render(context)
+        system_prompt = (
+            self.table_system_prompt + "\n" + self.markdown_renderer.get_content()
+        )
+        question = self.table_question_prompt + " " + asset.name
+        description = self._generate(system_prompt, question)
+        return SemanticAspect(
+            name=f"{asset.name}_semantic_properties",
+            description=description,
+            keywords=[],
+        )
+
+    def _get_semantic_aspect_column(
+        self, asset: Asset, context: RDatabaseGraph
+    ) -> SemanticAspect:
+        self.markdown_renderer.render(context)
+        system_prompt = (
+            self.column_system_prompt + "\n" + self.markdown_renderer.get_content()
+        )
+        question = self.column_question_prompt + " " + asset.name
+        description = self._generate(system_prompt, question)
+        return SemanticAspect(
+            name=f"{asset.name}_semantic_properties",
+            description=description,
+            keywords=[],
+        )
+
+    def get_semantic_aspect(
+        self, asset: Asset, context: RDatabaseGraph
+    ) -> SemanticAspect:
+        if asset.type == AssetType.RCOLUMN:
+            return self._get_semantic_aspect_column(asset, context)
+        elif asset.type == AssetType.RTABLE:
+            return self._get_semantic_aspect_table(asset, context)
+        else:
+            raise NotImplementedError(
+                "Only support for table and column in relational databases"
+            )

dbgraph/descriptor/prompt_templates.py

@@ -0,0 +1,80 @@
+TABLE_SYSTEM_PROMPT = """
+You are an expert in SQL database schema understanding.
+
+Generate a concise semantic description of the database table from the provided schema.
+
+Your description should capture:
+1. WHAT the table represents.
+2. WHAT a single row represents.
+3. The table's primary business or technical purpose.
+4. Important relationships to other entities when they can be inferred from foreign keys.
+5. Whether the table represents an entity, relationship/mapping, transaction, event, log, audit record, or reference data when this is clear from the schema.
+
+Inference rules:
+- Use the table name, column names, data types, primary keys, foreign keys, unique constraints, and relationships.
+- Prefer semantic meaning over simply repeating column names.
+- Do not infer information that is not reasonably supported by the schema.
+- When uncertain, use conservative wording rather than guessing.
+- For junction tables, describe the relationship between the referenced entities.
+- For tables with timestamps/status fields, mention their temporal or state-tracking role only when supported by the schema.
+
+Output rules:
+- Return ONLY the table description.
+- Use plain text, with no Markdown.
+- Use 1 concise sentence, or 2 sentences if necessary.
+- Do not start with "Description:".
+- Do not list columns unless necessary to explain the table.
+- Do not include example values.
+- Do not include SQL, comments, explanations, or reasoning.
+- Do not ask for additional information.
+
+Relevant schema:
+"""
+
+TABLE_QUESTION_PROMPT = """
+Generate a short description for the following table:
+"""
+
+COLUMN_SYSTEM_PROMPT = """
+You are an expert in SQL database schema understanding.
+
+Your task is to generate a concise semantic description of the given column based only on the provided schema.
+
+The description should explain the semantic meaning and role of the column within its table, rather than merely restating its name or data type.
+
+Focus on:
+- WHAT the column represents.
+- WHAT the value means in the context of the table.
+- The role of the column when it can be inferred, such as identifier, attribute, status, flag, timestamp, date, amount, quantity, metric, category, or foreign key.
+- The entity or concept that the column refers to.
+- For foreign keys, describe what entity the column identifies or references.
+- For status/type/category columns, describe what concept or state they represent.
+- For boolean/flag columns, describe the condition or property represented by the flag.
+- For date/time columns, describe the business or technical event represented by the timestamp when it can be inferred.
+- For numeric columns, describe the semantic meaning of the number rather than only saying it is a numeric value.
+- Use surrounding columns, primary keys, foreign keys, constraints, and the table's purpose to disambiguate the column's meaning.
+
+Inference rules:
+- Interpret the column in the context of the table, not in isolation.
+- Prefer semantic meaning over literal column-name expansion.
+- Use relationships and constraints to infer meaning when possible.
+- Do not invent business rules, meanings, units, or semantics that are not reasonably supported by the schema.
+- If the meaning is ambiguous, provide the most conservative description supported by the available schema.
+- Do not assume that a column name has a standard meaning if the schema provides evidence otherwise.
+
+Output rules:
+- Return ONLY the column description.
+- Use plain text with no Markdown.
+- Use 1 concise sentence, preferably under 25 words.
+- Do not start with "Description:".
+- Do not include example values.
+- Do not list the column's data type unless it is essential to its semantic meaning.
+- Do not include explanations, reasoning, comments, or SQL.
+- Do not ask questions or request additional information.
+
+Relevant schema:
+"""
+
+COLUMN_QUESTION_PROMPT = """
+Generate a short description for the following column:
+"""