dbgraph 0.1.0__py3-none-any.whl

dbgraph/io/json_graph_loader.py

@@ -0,0 +1,160 @@
+import json
+from dataclasses import dataclass
+from pathlib import Path
+
+from dbgraph.entity.aspect import (
+    RCategoricalStatistics,
+    RColumnSchemaAspect,
+    RColumnStatisticsAspect,
+    RForeignKeyAspect,
+    RNumericalStatistics,
+    RTableSchemaAspect,
+    RTableStatisticsAspect,
+    SemanticAspect,
+)
+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.io.graph_loader import GraphLoader
+
+
+@dataclass
+class JSONGraphLoader(GraphLoader):
+    json_path: Path
+
+    def load(self) -> DatabaseGraph:
+        with self.json_path.open("r") as f:
+            data = json.load(f)
+            assets_data = data["assets"]
+            links_data = data["links"]
+        assets = []
+        links = []
+        for data in assets_data:
+            if data["type"] == "table":
+                asset = self._parse_table_asset(data)
+            elif data["type"] == "column":
+                asset = self._parse_column_asset(data)
+            else:
+                raise NotImplementedError(f"Invalid asset type '{data['type']}'")
+            assets.append(asset)
+        for data in links_data:
+            if data["type"] == "contain":
+                link = self._parse_contain_link(data)
+            elif data["type"] == "foreign-key":
+                link = self._parse_fk_link(data)
+            else:
+                raise NotImplementedError(f"Invalid link type '{data['type']}'")
+            links.append(link)
+        return DatabaseGraph(assets, links)
+
+    def _parse_fk_link(self, data: dict) -> Link:
+        fk_data = data["aspects"]["foreign_key_properties"]
+        fk_aspect = RForeignKeyAspect(
+            name=fk_data["name"],
+            from_column=fk_data["from_column"],
+            to_column=fk_data["to_column"],
+            on_delete=fk_data["on_delete"],
+            on_update=fk_data["on_update"],
+        )
+        return Link(
+            link_id=data["link_id"],
+            name=data["name"],
+            type=LinkType.FOREIGN_KEY,
+            source_id=data["source_id"],
+            destination_id=data["destination_id"],
+            aspects={"foreign_key_properties": fk_aspect},
+        )
+
+    def _parse_contain_link(self, data: dict) -> Link:
+        return Link(
+            link_id=data["link_id"],
+            name=data["name"],
+            type=LinkType.CONTAIN,
+            source_id=data["source_id"],
+            destination_id=data["destination_id"],
+        )
+
+    def _parse_table_asset(self, data: dict) -> Asset:
+        aspects_data = data["aspects"]
+        aspects = {}
+        if "schema_properties" in aspects_data:
+            properties = aspects_data["schema_properties"]
+            aspects["schema_properties"] = RTableSchemaAspect(
+                name=properties["name"],
+                pks=properties["pks"],
+                indices=properties["indices"],
+            )
+        if "statistical_properties" in aspects_data:
+            properties = aspects_data["statistical_properties"]
+            aspects["statistical_properties"] = RTableStatisticsAspect(
+                name=properties["name"],
+                num_columns=properties["num_columns"],
+                num_rows=properties["num_rows"],
+            )
+        if "semantic_properties" in aspects_data:
+            properties = aspects_data["semantic_properties"]
+            aspects["semantic_properties"] = SemanticAspect(
+                name=properties["name"],
+                description=properties["description"],
+                keywords=properties["keywords"],
+            )
+        return Asset(
+            asset_id=data["asset_id"],
+            name=data["name"],
+            type=AssetType.RTABLE,
+            aspects=aspects,
+        )
+
+    def _parse_column_asset(self, data: dict) -> Asset:
+        aspects_data = data["aspects"]
+        aspects = {}
+        if "schema_properties" in aspects_data:
+            properties = aspects_data["schema_properties"]
+            aspects["schema_properties"] = RColumnSchemaAspect(
+                name=properties["name"],
+                dtype=properties["dtype"],
+                is_pk=properties["is_pk"],
+                is_nullable=properties["is_nullable"],
+            )
+        if "statistical_properties" in aspects_data:
+            properties = aspects_data["statistical_properties"]
+            numerical_stats = properties["numerical_stats"]
+            categorical_stats = properties["categorical_stats"]
+            rnumerical_stats = rcategorical_stats = None
+            if numerical_stats:
+                rnumerical_stats = RNumericalStatistics(
+                    mean=numerical_stats["mean"],
+                    std=numerical_stats["std"],
+                    min=numerical_stats["min"],
+                    q2=numerical_stats["q2"],
+                    median=numerical_stats["median"],
+                    q3=numerical_stats["q3"],
+                    max=numerical_stats["max"],
+                )
+            if categorical_stats:
+                rcategorical_stats = RCategoricalStatistics(
+                    value_counts=categorical_stats["value_counts"]
+                )
+            aspects["statistical_properties"] = RColumnStatisticsAspect(
+                name=properties["name"],
+                non_null_count=properties["non_null_count"],
+                null_count=properties["null_count"],
+                numerical_stats=rnumerical_stats,
+                categorical_stats=rcategorical_stats,
+            )
+        if "semantic_properties" in aspects_data:
+            properties = aspects_data["semantic_properties"]
+            aspects["semantic_properties"] = SemanticAspect(
+                name=properties["name"],
+                description=properties["description"],
+                keywords=properties["keywords"],
+            )
+
+        return Asset(
+            asset_id=data["asset_id"],
+            name=data["name"],
+            type=AssetType.RCOLUMN,
+            aspects=aspects,
+        )

dbgraph/io/json_graph_writer.py

@@ -0,0 +1,17 @@
+import json
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+from dbgraph.entity.dbgraph import DatabaseGraph
+from dbgraph.io.graph_writer import GraphWriter
+
+
+@dataclass
+class JSONGraphWriter(GraphWriter):
+    json_path: Path
+    indent: int | None = None
+
+    def write(self, graph: DatabaseGraph):
+        graph_data = asdict(graph)
+        with self.json_path.open("w") as f:
+            json.dump(graph_data, f, indent=self.indent)

dbgraph/persistent/graph_persistent.py

@@ -0,0 +1,98 @@
+from abc import ABC, abstractmethod
+
+from dbgraph.entity.aspect import Aspect
+from dbgraph.entity.asset import Asset
+from dbgraph.entity.dbgraph import DatabaseGraph
+from dbgraph.entity.link import Link
+
+
+class GraphPersistent(ABC):
+    """Interface for graph persistent component, which save and load the graphs"""
+
+    @abstractmethod
+    def create_graph(self, graph_id: str, graph_name: str):
+        """Insert the graph in the database"""
+
+    @abstractmethod
+    def delete_graph(self, graph_id: str):
+        """Delete a graph and return it"""
+
+    @abstractmethod
+    def insert_assets(self, assets: list[Asset], graph_id: str):
+        """Insert the assets into database"""
+
+    @abstractmethod
+    def insert_asset(self, asset: Asset, graph_id: str):
+        """Insert a single asset"""
+
+    @abstractmethod
+    def insert_links(self, links: list[Link], graph_id: str):
+        """Insert links into database"""
+
+    @abstractmethod
+    def insert_link(self, link: Link, graph_id: str):
+        """Insert a single link into the database"""
+
+    @abstractmethod
+    def remove_assets(self, assets_ids: list[str], graph_id: str):
+        """Remove the assets from the database"""
+
+    @abstractmethod
+    def remove_asset(self, asset_id: str, graph_id: str):
+        """Remove the asset from the database"""
+
+    @abstractmethod
+    def remove_links(self, links_ids: list[str], graph_id: str):
+        """Remove the links"""
+
+    @abstractmethod
+    def remove_link(self, link_id: str, graph_id: str):
+        """Remove the link"""
+
+    @abstractmethod
+    def insert_asset_aspects(self, asset_id: str, aspects: dict[str, Aspect]):
+        """Insert aspects into asset"""
+
+    @abstractmethod
+    def get_asset(self, asset_id: str, graph_id: str) -> Asset:
+        """Find asset by its ID"""
+
+    @abstractmethod
+    def get_link(self, link_id: str, graph_id: str) -> Link:
+        """Find link by its ID"""
+
+    @abstractmethod
+    def insert_link_aspects(self, link_id: str, aspects: dict[str, Aspect]):
+        """Insert aspects into link"""
+
+    @abstractmethod
+    def get_aspects_of_asset(self, asset_id: str) -> dict[str, Aspect]:
+        """Get the aspects of an asset"""
+
+    @abstractmethod
+    def get_aspects_of_link(self, link_id: str) -> dict[str, Aspect]:
+        """Get the aspects of a link"""
+
+    @abstractmethod
+    def get_assets(self, graph_id: str) -> list[Asset]:
+        """Return list of assets in a graph"""
+
+    @abstractmethod
+    def get_links(self, graph_id: str) -> list[Link]:
+        """Return list of links in a graph"""
+
+    @abstractmethod
+    def load_graph(self, graph_id: str) -> DatabaseGraph:
+        """Load the full database graph"""
+
+    @abstractmethod
+    def save_graph(self, graph: DatabaseGraph, graph_id: str):
+        """Save the graph into database"""
+
+    @abstractmethod
+    def set_graph_state(self, graph_id: str, completed: bool):
+        """Mark the graph as completely built and saved in the database or hasn't been finished building yet"""
+
+    @abstractmethod
+    def get_graph_info(self, graph_id: str) -> tuple[str, str, bool]:
+        """Return the graph's info, which is a tuple containing [graph_id, graph_name, completed]"""

dbgraph/persistent/models.py

@@ -0,0 +1,222 @@
+from sqlalchemy import (
+    JSON,
+    Boolean,
+    ForeignKey,
+    String,
+    Text,
+)
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
+
+
+class Base(DeclarativeBase):
+    pass
+
+
+class GraphModel(Base):
+    __tablename__ = "graphs"
+
+    graph_id: Mapped[str] = mapped_column(
+        String(36),
+        primary_key=True,
+    )
+
+    name: Mapped[str] = mapped_column(
+        Text,
+        nullable=False,
+    )
+
+    completed: Mapped[bool] = mapped_column(
+        Boolean,
+        nullable=False,
+        default=False,
+    )
+
+    assets: Mapped[list["AssetModel"]] = relationship(
+        back_populates="graph",
+        cascade="all, delete-orphan",
+        passive_deletes=True,
+    )
+
+    links: Mapped[list["LinkModel"]] = relationship(
+        back_populates="graph",
+        cascade="all, delete-orphan",
+        passive_deletes=True,
+    )
+
+
+class AssetModel(Base):
+    __tablename__ = "assets"
+
+    asset_id: Mapped[str] = mapped_column(
+        String(36),
+        primary_key=True,
+    )
+
+    name: Mapped[str] = mapped_column(
+        Text,
+        nullable=False,
+    )
+
+    asset_type: Mapped[str] = mapped_column(
+        String(30),
+        nullable=False,
+    )
+
+    graph_id: Mapped[str] = mapped_column(
+        ForeignKey(
+            "graphs.graph_id",
+            ondelete="CASCADE",
+        ),
+        nullable=False,
+        index=True,
+    )
+
+    graph: Mapped["GraphModel"] = relationship(
+        back_populates="assets",
+    )
+
+    outgoing_links: Mapped[list["LinkModel"]] = relationship(
+        foreign_keys="LinkModel.source_id",
+        back_populates="source",
+        cascade="all, delete-orphan",
+    )
+
+    incoming_links: Mapped[list["LinkModel"]] = relationship(
+        foreign_keys="LinkModel.destination_id",
+        back_populates="destination",
+        cascade="all, delete-orphan",
+    )
+
+    asset_aspects: Mapped[list["AssetAspectModel"]] = relationship(
+        back_populates="asset",
+        cascade="all, delete-orphan",
+        passive_deletes=True,
+    )
+
+
+class LinkModel(Base):
+    __tablename__ = "links"
+
+    link_id: Mapped[str] = mapped_column(
+        String(36),
+        primary_key=True,
+    )
+
+    name: Mapped[str] = mapped_column(
+        Text,
+        nullable=False,
+    )
+
+    link_type: Mapped[str] = mapped_column(
+        String(30),
+        nullable=False,
+    )
+
+    source_id: Mapped[str] = mapped_column(
+        ForeignKey(
+            "assets.asset_id",
+            ondelete="CASCADE",
+        ),
+        nullable=False,
+        index=True,
+    )
+
+    destination_id: Mapped[str] = mapped_column(
+        ForeignKey(
+            "assets.asset_id",
+            ondelete="CASCADE",
+        ),
+        nullable=False,
+        index=True,
+    )
+
+    graph_id: Mapped[str] = mapped_column(
+        ForeignKey(
+            "graphs.graph_id",
+            ondelete="CASCADE",
+        ),
+        nullable=False,
+        index=True,
+    )
+
+    source: Mapped["AssetModel"] = relationship(
+        foreign_keys=[source_id],
+        back_populates="outgoing_links",
+    )
+
+    destination: Mapped["AssetModel"] = relationship(
+        foreign_keys=[destination_id],
+        back_populates="incoming_links",
+    )
+
+    graph: Mapped["GraphModel"] = relationship(
+        back_populates="links",
+    )
+
+    link_aspects: Mapped[list["LinkAspectModel"]] = relationship(
+        back_populates="link",
+        cascade="all, delete-orphan",
+        passive_deletes=True,
+    )
+
+
+class AssetAspectModel(Base):
+    __tablename__ = "asset_aspects"
+
+    asset_id: Mapped[str] = mapped_column(
+        ForeignKey(
+            "assets.asset_id",
+            ondelete="CASCADE",
+        ),
+        primary_key=True,
+    )
+
+    name: Mapped[str] = mapped_column(
+        String(100),
+        primary_key=True,
+    )
+
+    type: Mapped[str] = mapped_column(
+        String(30),
+        nullable=False,
+    )
+
+    json_aspect: Mapped[dict] = mapped_column(
+        JSON,
+        nullable=False,
+    )
+
+    asset: Mapped["AssetModel"] = relationship(
+        back_populates="asset_aspects",
+    )
+
+
+class LinkAspectModel(Base):
+    __tablename__ = "link_aspects"
+
+    link_id: Mapped[str] = mapped_column(
+        ForeignKey(
+            "links.link_id",
+            ondelete="CASCADE",
+        ),
+        primary_key=True,
+    )
+
+    name: Mapped[str] = mapped_column(
+        String(100),
+        primary_key=True,
+    )
+
+    type: Mapped[str] = mapped_column(
+        String(30),
+        nullable=False,
+    )
+
+    json_aspect: Mapped[dict] = mapped_column(
+        JSON,
+        nullable=False,
+    )
+
+    link: Mapped["LinkModel"] = relationship(
+        back_populates="link_aspects",
+    )