Merge pull request #202 from MathisVerstrepen/hotfix

Fix: "graphs.node_count does not exist" on new project

Author: MathisVerstrepen

api/app/database/pg/graph_ops/graph_config_crud.py

@@ -41,7 +41,11 @@ async def update_graph_name(
             if not isinstance(db_graph, Graph):
                 raise HTTPException(status_code=404, detail=f"Graph with id {graph_id} not found")
 
-            return db_graph
+        await session.refresh(db_graph, attribute_names=["nodes", "edges"])
+
+        db_graph.node_count = len(db_graph.nodes)
+
+        return db_graph
 
 
 async def toggle_graph_pin(pg_engine: SQLAlchemyAsyncEngine, graph_id: str, pinned: bool) -> Graph:
@@ -72,7 +76,11 @@ async def toggle_graph_pin(pg_engine: SQLAlchemyAsyncEngine, graph_id: str, pinn
             if not isinstance(db_graph, Graph):
                 raise HTTPException(status_code=404, detail=f"Graph with id {graph_id} not found")
 
-            return db_graph
+        await session.refresh(db_graph, attribute_names=["nodes", "edges"])
+
+        db_graph.node_count = len(db_graph.nodes)
+
+        return db_graph
 
 
 class GraphConfigUpdate(BaseModel):
@@ -135,7 +143,11 @@ async def update_graph_config(
             if not isinstance(db_graph, Graph):
                 raise HTTPException(status_code=404, detail=f"Graph with id {graph_id} not found")
 
-            return db_graph
+        await session.refresh(db_graph, attribute_names=["nodes", "edges"])
+
+        db_graph.node_count = len(db_graph.nodes)
+
+        return db_graph
 
 
 async def get_canvas_config(pg_engine: SQLAlchemyAsyncEngine, graph_id: str) -> GraphConfigUpdate:

api/app/database/pg/graph_ops/graph_crud.py

@@ -38,7 +38,7 @@ async def get_all_graphs(engine: SQLAlchemyAsyncEngine, user_id: str) -> list[Gr
             .subquery()
         )
 
-        # Main query to select Graph and the node_count from the subquery
+        # Main query to select the Graph object and the node_count as separate columns
         stmt = (
             select(Graph, node_count_subquery.c.node_count)
             .outerjoin(node_count_subquery, Graph.id == node_count_subquery.c.graph_id)
@@ -48,6 +48,7 @@ async def get_all_graphs(engine: SQLAlchemyAsyncEngine, user_id: str) -> list[Gr
 
         result = await session.exec(stmt)  # type: ignore
 
+        # Manually process the results and assign the count to the property
         graphs = []
         for graph, count in result.all():
             graph.node_count = count or 0
@@ -84,6 +85,8 @@ async def get_graph_by_id(engine: SQLAlchemyAsyncEngine, graph_id: str) -> Compl
         if not db_graph:
             raise HTTPException(status_code=404, detail=f"Graph with id {graph_id} not found")
 
+        db_graph.node_count = len(db_graph.nodes)
+
         complete_graph_response = CompleteGraph(
             graph=db_graph,
             nodes=db_graph.nodes,
@@ -169,6 +172,11 @@ async def persist_temporary_graph(
             raise HTTPException(
                 status_code=500, detail="Unexpected error: Retrieved object is not of type Graph."
             )
+
+        await session.refresh(db_graph, attribute_names=["nodes", "edges"])
+
+        db_graph.node_count = len(db_graph.nodes)
+
         return db_graph
 
 

api/app/database/pg/graph_ops/graph_node_crud.py

@@ -107,6 +107,9 @@ async def update_graph_with_nodes_and_edges(
                 raise e
 
         await session.refresh(db_graph, attribute_names=["nodes", "edges"])
+
+        db_graph.node_count = len(db_graph.nodes)
+
         return db_graph
 
 

api/app/database/pg/models.py

@@ -6,6 +6,7 @@
 from typing import Any, Optional
 
 from models.auth import UserPass
+from pydantic import PrivateAttr, computed_field
 from sqlalchemy import Column, ForeignKeyConstraint, Index, PrimaryKeyConstraint, func, select
 from sqlalchemy.dialects.postgresql import DOUBLE_PRECISION, JSONB, TEXT, TIMESTAMP
 from sqlalchemy.dialects.postgresql import UUID as PG_UUID
@@ -87,9 +88,16 @@ class Graph(SQLModel, table=True):
     nodes: list["Node"] = Relationship(back_populates="graph")
     edges: list["Edge"] = Relationship(back_populates="graph")
 
-    node_count: Optional[int] = (
-        None  # This will be set dynamically in the query, not stored in the database
-    )
+    _node_count: Optional[int] = PrivateAttr(default=None)
+
+    @computed_field  # type: ignore[misc]
+    @property
+    def node_count(self) -> Optional[int]:
+        return self._node_count
+
+    @node_count.setter
+    def node_count(self, value: Optional[int]) -> None:
+        self._node_count = value
 
     __table_args__ = (Index("idx_graphs_user_updated_at", "user_id", "updated_at"),)
 

api/app/database/redis/redis_ops.py

@@ -76,7 +76,6 @@ async def annotation_exists(self, remote_hash: str) -> bool:
         Returns:
             bool: True if the annotation exists, False otherwise.
         """
-        print(remote_hash)
         if not remote_hash:
             return False
         key = f"annotation:{remote_hash}"

api/app/services/stream.py

@@ -200,9 +200,6 @@ async def propagate_stream_to_websocket(
         )
 
         if not is_title_generation:
-            from rich import print
-
-            print(messages)
             openRouterReq = OpenRouterReqChat(
                 api_key=open_router_api_key,
                 model=request_data.model,