dbgraph 0.1.0__py3-none-any.whl

dbgraph/entity/aspect.py

@@ -0,0 +1,153 @@
+from dataclasses import dataclass
+from enum import StrEnum
+
+
+@dataclass
+class Aspect:
+    """Represent a set of properties that are related to an aspect
+
+    Attributes:
+        name: The name of this aspect
+    """
+
+    name: str
+
+
+@dataclass
+class SemanticAspect(Aspect):
+    """Represent properties related to semantic meaning of the asset
+
+    Attributes:
+        description: The natural language description for this asset
+        keywords: list of relevant keywords
+    """
+
+    description: str
+    keywords: list[str]
+
+
+@dataclass
+class RTableSchemaAspect(Aspect):
+    """Group of properties relevant to relational databases' table schema
+
+    Attributes:
+        pks: list of column names belong to the primary key
+        indices: a dictionary which map the table's indices with the columns belong to corresponding indices
+    """
+
+    pks: list[str]
+    indices: dict[str, list[str]]
+
+
+@dataclass
+class RTableStatisticsAspect(Aspect):
+    """Group of properties relevant to relational databases' table statistics
+
+    Attributes:
+        num_columns: number of columns
+        num_rows: number of rows
+    """
+
+    num_columns: int
+    num_rows: int
+
+
+@dataclass
+class RColumnSchemaAspect(Aspect):
+    """Group of properties relevant to RDB's column schema
+
+    Attributes:
+        dtype: data type
+        is_pk: whether or not this column belongs to primary key
+        is_nullable: whether or not this column is nullable
+    """
+
+    dtype: str
+    is_pk: bool
+    is_nullable: bool
+
+
+@dataclass
+class RNumericalStatistics:
+    """Statistics for numerical columns
+
+    Attributes:
+        mean: Mean value
+        std: Standard deviation
+        min: Minimum value
+        q2: 2nd quantile
+        median: Median value
+        q3: 3rd quantile
+        max: Maximum value
+
+    """
+
+    mean: float | None = None
+    std: float | None = None
+    min: float | None = None
+    q2: float | None = None
+    median: float | None = None
+    q3: float | None = None
+    max: float | None = None
+
+
+@dataclass
+class RCategoricalStatistics:
+    """Statistics for categorical columns
+
+    Attributes:
+        value_counts: top 10 most frequent values in a column and
+            its corresponding frequency
+    """
+
+    value_counts: dict[str, int]
+
+
+@dataclass
+class RColumnStatisticsAspect(Aspect):
+    """Group of properties relevant to RDB's column statistics
+
+    Attributes:
+        non_null_count: Number of non null records
+        null_count: Number of null records
+        numerical_stats: numerical statistics of this column if its data type is numeric
+        categorical_stats: categorical statistics of this column if its data type is string-like
+    """
+
+    non_null_count: int
+    null_count: int
+    numerical_stats: RNumericalStatistics | None = None
+    categorical_stats: RCategoricalStatistics | None = None
+
+
+class FKBehavior(StrEnum):
+    """Defines the behavior of a foreign key when its referenced row changes.
+
+    Attributes:
+        CASCADE: Automatically propagates the change to referencing rows.
+        SET_NULL: Sets the foreign key value to ``NULL``.
+        RESTRICT: Prevents the operation if referencing rows exist.
+        NO_ACTION: Does not perform any action on referencing rows.
+    """
+
+    CASCADE = "CASCADE"
+    SET_NULL = "SET NULL"
+    RESTRICT = "RESTRICT"
+    NO_ACTION = "NO ACTION"
+
+
+@dataclass
+class RForeignKeyAspect(Aspect):
+    """Group of properties relevant to foreign keys
+
+    Attributes:
+        from_column: the source column
+        to_column: referred column
+        on_delete: foreign key behavior on delete event
+        on_update: foreign key behavior on update event
+    """
+
+    from_column: str
+    to_column: str
+    on_delete: FKBehavior | str
+    on_update: FKBehavior | str

dbgraph/entity/asset.py

@@ -0,0 +1,26 @@
+from dataclasses import dataclass, field
+
+from dbgraph.entity.aspect import (
+    Aspect,
+)
+from dbgraph.entity.asset_type import AssetType
+
+
+@dataclass
+class Asset:
+    """
+    Represent a data asset, such as a table in SQL databases or a document in MongoDB.
+
+    It acts as a vertex in the graph.
+
+    Attributes:
+        asset_id: ID of this asset
+        name: original name of this asset
+        type: type of asset
+        aspects: properties aspects belong to this asset
+    """
+
+    asset_id: str
+    name: str
+    type: AssetType
+    aspects: dict[str, Aspect] = field(default_factory=dict)

dbgraph/entity/asset_type.py

@@ -0,0 +1,13 @@
+from enum import StrEnum
+
+
+class AssetType(StrEnum):
+    """Type of `Asset`.
+
+    Attributes:
+        table: table in SQL databases
+        column: column in SQL databases
+    """
+
+    RTABLE = "table"
+    RCOLUMN = "column"

dbgraph/entity/dbgraph.py

@@ -0,0 +1,152 @@
+from __future__ import annotations
+
+import itertools
+from dataclasses import dataclass
+
+import rustworkx as rx
+
+from dbgraph.entity.asset import Asset
+from dbgraph.entity.link import Link
+
+
+@dataclass
+class DatabaseGraph:
+    """Represents a database as a graph of assets and relationships.
+
+    A `DatabaseGraph` models the structure of a database using
+    assets as nodes and links as relationships between those assets.
+    It provides methods for creating, querying, and traversing the
+    database structure in a graph-oriented way.
+
+    Attributes:
+        assets: The collection of `Asset` objects in the graph.
+        links: The collection of `Link` objects representing
+            relationships between assets.
+    """
+
+    assets: list[Asset]
+    links: list[Link]
+
+    def __post_init__(self):
+        self._asset_id_to_node_idx = {a.asset_id: i for i, a in enumerate(self.assets)}
+        self._nodes_data = {i: asset for i, asset in enumerate(self.assets)}
+        self._edges_data = {i: link for i, link in enumerate(self.links)}
+        links_tuples = [
+            (
+                self._asset_id_to_node_idx[link.source_id],
+                self._asset_id_to_node_idx[link.destination_id],
+                link,
+            )
+            for link in self.links
+        ]
+        self._graph = rx.PyGraph()
+        self._graph.add_nodes_from(self.assets)
+        self._graph.add_edges_from(links_tuples)
+
+    def _node_idx(self, asset_id: str) -> int:
+        try:
+            return self._asset_id_to_node_idx[asset_id]
+        except KeyError:
+            raise KeyError(f"Asset with id={asset_id} not found!")
+
+    def _neighbors(self, asset_id: str) -> tuple[set[int], list[Link]]:
+        """Return list of asset's indices in the graph and those links"""
+        node_idx = self._node_idx(asset_id)
+
+        # Find links
+        links: dict[int, Link] = self._graph.adj(node_idx)
+
+        # Find assets indices
+        asset_indices = {self._node_idx(link.destination_id) for link in links.values()}
+        asset_indices = asset_indices.union(
+            {self._node_idx(link.source_id) for link in links.values()}
+        )
+
+        return asset_indices, list(links.values())
+
+    def get_asset(self, asset_id: str) -> Asset:
+        """Return the asset using asset_id"""
+        return self._graph.get_node_data(self._node_idx(asset_id))
+
+    def get_link(self, src_id: str, dst_id: str) -> Link:
+        """Return the link connect these two Assets"""
+        return self._graph.get_edge_data(self._node_idx(src_id), self._node_idx(dst_id))
+
+    def select_neighbors(self, asset_id: str) -> DatabaseGraph:
+        """Get this node and its neighbors.
+
+        Args:
+            asset_id: ID of the asset to be queried
+
+        Return:
+            subgraph: an instance of `DatabaseGraph` as subgraph
+        """
+
+        asset_indices, links = self._neighbors(asset_id)
+        assets = [self._graph.get_node_data(i) for i in asset_indices]
+        return DatabaseGraph(assets, links)
+
+    def select_shortest_paths(self, src_id: str, dst_id: str) -> list[DatabaseGraph]:
+        """
+        Find shortest paths between two Assets.
+
+        Args:
+            src_id: ID of the source asset
+            dst_id: ID of the destination asset
+
+        Returns:
+            subgraphs: list of `DatabaseGraph` which are the shortest paths
+                between source and destination Assets
+        """
+        src_idx = self._node_idx(src_id)
+        dst_idx = self._node_idx(dst_id)
+
+        paths = rx.all_shortest_paths(self._graph, src_idx, dst_idx)
+
+        # Get a list of lists of assets
+        assets_lists = [[self._graph.get_node_data(i) for i in path] for path in paths]
+
+        # Get a list of lists of links
+        pairs_lists = [itertools.pairwise(path) for path in paths]
+        links_lists = [
+            [self._graph.get_edge_data(pair[0], pair[1]) for pair in pairs]
+            for pairs in pairs_lists
+        ]
+        return [
+            DatabaseGraph(assets, links)
+            for assets, links in zip(assets_lists, links_lists)
+        ]
+
+    def get_neighbors(self, asset_id: str) -> list[Asset]:
+        """Get neighbors assets of an asset
+
+        Args:
+            asset_id: target asset
+
+        Return:
+            list of neighbor `Asset`s
+        """
+        asset_indices, _ = self._neighbors(asset_id)
+        # Remove itself
+        try:
+            asset_indices.remove(self._node_idx(asset_id))
+        except KeyError:
+            pass
+        return [self._graph.get_node_data(i) for i in asset_indices]
+
+    def find_shortest_paths(self, src_id: str, dst_id: str) -> list[list[Asset]]:
+        """Find shortest paths between two Assets
+
+        Args:
+            src_id: ID of the source asset
+            dst_id: ID of the destination asset
+
+        Returns:
+            assets: list of lists of assets. Each list is a shortest path
+                from source Asset to destination Asset
+        """
+        src_idx = self._node_idx(src_id)
+        dst_idx = self._node_idx(dst_id)
+        paths = rx.all_shortest_paths(self._graph, src_idx, dst_idx)
+
+        return [[self._graph.get_node_data(i) for i in path] for path in paths]

dbgraph/entity/link.py

@@ -0,0 +1,37 @@
+from dataclasses import dataclass, field
+from typing import cast
+
+from dbgraph.entity.aspect import Aspect, RForeignKeyAspect
+from dbgraph.entity.link_type import LinkType
+
+
+@dataclass
+class Link:
+    """
+    Represent relationship between two assets.
+
+    Attributes:
+        link_id: ID of this link
+        name: name
+        type: type of link
+        source_id: ID of the source `Asset`
+        destination_id: ID of the destination `Asset`
+        aspects: properties aspects belong to this link
+    """
+
+    link_id: str
+    name: str
+    type: LinkType
+    source_id: str
+    destination_id: str
+    aspects: dict[str, Aspect] = field(default_factory=dict)
+
+    def to_markdown(self, src_name: str, dst_name: str) -> str:
+        match self.type:
+            case LinkType.CONTAIN:
+                return f"- Column `{dst_name}` belongs to table `{src_name}`\n"
+            case LinkType.FOREIGN_KEY:
+                fk_aspect = cast(
+                    RForeignKeyAspect, self.aspects["foreign_key_properties"]
+                )
+                return f"- `{src_name}`.`{fk_aspect.from_column}` -> `{dst_name}`.`{fk_aspect.to_column}`\n"

dbgraph/entity/link_type.py

@@ -0,0 +1,13 @@
+from enum import StrEnum
+
+
+class LinkType(StrEnum):
+    """Type of link.
+
+    Attributes:
+        foreign-key: Foreign key constraint
+        contain: containment relationship (table - columns)
+    """
+
+    FOREIGN_KEY = "foreign-key"
+    CONTAIN = "contain"

dbgraph/entity/rdbgraph.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+from dbgraph import LinkType
+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
+
+
+class RDatabaseGraph(DatabaseGraph):
+    """Graph for Relational Database"""
+
+    @staticmethod
+    def from_graph(graph: DatabaseGraph) -> RDatabaseGraph:
+        return RDatabaseGraph(graph.assets, graph.links)
+
+    def get_columns(self, table_id: str) -> list[Asset]:
+        """Get list of assets which are columns belong to this table
+
+        Args:
+            table_id: ID of the target table
+
+        Returns:
+            assets: columns belong to this table
+        """
+        neighbors = self.get_neighbors(table_id)
+        return [n for n in neighbors if n.type == AssetType.RCOLUMN]
+
+    def get_tables(self) -> list[Asset]:
+        """Get all the tables in the graph"""
+        return [a for a in self.assets if a.type == AssetType.RTABLE]
+
+    def get_foreign_keys(self) -> list[Link]:
+        """Get all foreign keys in the graph"""
+        return [link for link in self.links if link.type == LinkType.FOREIGN_KEY]
+
+    def get_connected_tables(self, table_id: str) -> list[Asset]:
+        """Get all connected table
+
+        Args:
+            table_id: target table
+
+        Returns:
+            assets: list of references table
+        """
+        neighbors = self.get_neighbors(table_id)
+        return [n for n in neighbors if n.type == AssetType.RTABLE]
+
+    def select_connected_tables(self, table_id: str) -> RDatabaseGraph:
+        """Get references tables with their columns
+
+        Args:
+            table_id: target table
+
+        Returns:
+            subgraph: the subgraph of which include connected tables and their columns
+        """
+        refs_tables = self.get_connected_tables(table_id)
+        assets: list[Asset] = [table for table in refs_tables]
+        target_table: Asset = self.get_asset(table_id)
+        target_table_columns = self.get_columns(table_id)
+        assets.append(target_table)
+        assets.extend(target_table_columns)
+        links: list[Link] = [
+            self.get_link(target_table.asset_id, column.asset_id)
+            for column in target_table_columns
+        ]
+        for table in refs_tables:
+            fk_link = self.get_link(table_id, table.asset_id)
+            links.append(fk_link)
+            columns = self.get_columns(table.asset_id)
+            for column in columns:
+                assets.append(column)
+                link = self.get_link(table.asset_id, column.asset_id)
+                links.append(link)
+        return RDatabaseGraph(assets, links)

dbgraph/event/event.py

@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class Event:
+    """Base event used in DBGraph"""
+
+    message: str
+
+
+@dataclass
+class ProgressEvent(Event):
+    """Event with progress"""
+
+    target: float
+    current: float
+
+    def get_percentage(self) -> float:
+        """Return the progress of the task in percentage"""
+        return self.current / self.target

dbgraph/event/event_bus.py

@@ -0,0 +1,61 @@
+import queue
+import threading
+
+from dbgraph.event.event import Event
+from dbgraph.event.event_handler import EventHandler
+
+
+class EventBus:
+    """Store events in a queue and manage a background thread to distribute events"""
+
+    def __init__(self) -> None:
+        self.queue = queue.Queue()
+        self.handlers: dict[type[Event], set[EventHandler]] = {}
+        self._thread: threading.Thread | None = None
+        self._stop_event = threading.Event()
+
+    def subscribe(self, event_type: type[Event], handler: EventHandler):
+        """Subscribe a handler to specific type of event"""
+        if event_type in self.handlers:
+            self.handlers[event_type].add(handler)
+        else:
+            self.handlers[event_type] = {handler}
+
+    def unsubscribe(self, event_type: type[Event], handler: EventHandler):
+        """Unsubscribe a handler from listening to specific event type"""
+        self.handlers[event_type].remove(handler)
+
+    def publish(self, event: Event):
+        """Push an event into the event queue"""
+        self.queue.put(event)
+
+    def _dispatch(self):
+        """Get the latest event and let the appropriate handler process it"""
+        event = self.queue.get()
+        for handler in self.handlers[type(event)]:
+            handler.handle(event)
+
+    def _run(self):
+        """Run the event loop"""
+        while not self._stop_event.is_set():
+            try:
+                self._dispatch()
+            except queue.Empty:
+                continue
+
+    def start(self):
+        """Start the event background thread"""
+        if self._thread is not None and self._thread.is_alive():
+            return
+        self._stop_event.clear()
+        self._thread = threading.Thread(
+            target=self._run, name="dbgraph-event-bus", daemon=True
+        )
+        self._thread.start()
+
+    def stop(self):
+        """Stop the event background thread"""
+        self._stop_event.set()
+        if self._thread is not None:
+            self._thread.join()
+            self._thread = None

dbgraph/event/event_handler.py

@@ -0,0 +1,11 @@
+from abc import ABC, abstractmethod
+
+from dbgraph.event.event import Event
+
+
+class EventHandler(ABC):
+    """Interface for event handler"""
+
+    @abstractmethod
+    def handle(self, event: Event):
+        """Handle the event"""

dbgraph/io/graph_loader.py

@@ -0,0 +1,11 @@
+from abc import ABC, abstractmethod
+
+from dbgraph.entity.dbgraph import DatabaseGraph
+
+
+class GraphLoader(ABC):
+    """Interface for classes that read graphs from disk"""
+
+    @abstractmethod
+    def load(self) -> DatabaseGraph:
+        pass

dbgraph/io/graph_writer.py

@@ -0,0 +1,11 @@
+from abc import ABC, abstractmethod
+
+from dbgraph.entity.dbgraph import DatabaseGraph
+
+
+class GraphWriter(ABC):
+    """Interface for classes that write graph to disk"""
+
+    @abstractmethod
+    def write(self, graph: DatabaseGraph):
+        pass