From 983e1c146a4dd9c3842a403c938ab88a9cbe926d Mon Sep 17 00:00:00 2001
From: Milho001 <l.milhomemfrancochristino@uu.nl>
Date: Wed, 15 Nov 2023 11:01:19 +0000
Subject: [PATCH] feat(qb): connector drop suggests nodes

---
 .../lib/data-access/store/sessionSlice.ts     |  19 +-
 .../lib/mock-data/schema/moviesSchemaRaw.ts   |   6 +
 .../mock-data/schema/northwindSchemaRaw.ts    |   4 +
 .../lib/mock-data/schema/simpleAirportRaw.ts  |   1 +
 libs/shared/lib/mock-data/schema/simpleRaw.ts |  10 +-
 .../lib/mock-data/schema/twitterSchemaRaw.ts  |  21 ++
 .../querybuilder/model/graphology/model.ts    |   1 +
 .../querybuilder/model/graphology/utils.ts    |   3 +
 .../querybuilder/model/reactflow/handles.tsx  |  67 +-----
 .../querybuilder/model/reactflow/model.tsx    |   1 -
 .../lib/querybuilder/panel/querybuilder.tsx   |  95 ++++-----
 .../queryBuilderLogicPillsPanel.tsx           |  70 +++++-
 .../queryBuilderRelatedNodesPanel.tsx         | 199 ++++++++++++++++++
 .../stories/querybuilder-simple.stories.tsx   |   1 +
 .../querybuilder/panel/utils/connectorDrop.ts |  17 ++
 libs/shared/lib/schema/model/FromBackend.ts   |   6 +
 .../lib/schema/panel/schema.stories.tsx       |   1 +
 .../schema/pills/nodes/entity/entity-node.tsx |   8 +-
 .../nodes/relation/relation-node.stories.tsx  |   1 +
 .../pills/nodes/relation/relation-node.tsx    |  24 ++-
 .../lib/vis/paohvis/paohvis.stories.tsx       |   1 +
 .../semanticsubstrates.stories.tsx            |   1 +
 22 files changed, 404 insertions(+), 153 deletions(-)
 create mode 100644 libs/shared/lib/querybuilder/panel/querysidepanel/queryBuilderRelatedNodesPanel.tsx
 create mode 100644 libs/shared/lib/querybuilder/panel/utils/connectorDrop.ts

diff --git a/libs/shared/lib/data-access/store/sessionSlice.ts b/libs/shared/lib/data-access/store/sessionSlice.ts
index 02c7201e9..d77dff617 100644
--- a/libs/shared/lib/data-access/store/sessionSlice.ts
+++ b/libs/shared/lib/data-access/store/sessionSlice.ts
@@ -12,14 +12,12 @@ export type ErrorMessage = {
 export type SessionCacheI = {
   currentDatabase?: string;
   databases: string[];
-  error?: ErrorMessage;
 };
 
 // Define the initial state using that type
 export const initialState: SessionCacheI = {
   currentDatabase: undefined,
   databases: [],
-  error: undefined,
 };
 
 export const sessionSlice = createSlice({
@@ -38,25 +36,10 @@ export const sessionSlice = createSlice({
         else state.currentDatabase = undefined;
       }
     },
-    displayError(state, action: PayloadAction<ErrorMessage>): void {
-      // use a switch, in case certain errors will have side effects in the future
-      switch (action.payload.errorStatus) {
-        case 'Bad request':
-        case 'Bad credentials':
-        case 'Translation error':
-        case 'Database error':
-        case 'ML bad request':
-          state.error = action.payload;
-          break;
-      }
-    },
-    closeError(state): void {
-      state.error = undefined;
-    },
   },
 });
 
-export const { updateCurrentDatabase, updateDatabaseList, displayError, closeError } = sessionSlice.actions;
+export const { updateCurrentDatabase, updateDatabaseList } = sessionSlice.actions;
 
 // Other code such as selectors can use the imported `RootState` type
 export const sessionCacheState = (state: RootState) => state.sessionCache;
diff --git a/libs/shared/lib/mock-data/schema/moviesSchemaRaw.ts b/libs/shared/lib/mock-data/schema/moviesSchemaRaw.ts
index 8a69595d8..d94c4a40c 100644
--- a/libs/shared/lib/mock-data/schema/moviesSchemaRaw.ts
+++ b/libs/shared/lib/mock-data/schema/moviesSchemaRaw.ts
@@ -41,6 +41,7 @@ export const movieSchemaRaw: SchemaFromBackend = {
   edges: [
     {
       name: 'ACTED_IN',
+      label: 'ACTED_IN',
       collection: 'ACTED_IN',
       from: 'Person',
       to: 'Movie',
@@ -53,6 +54,7 @@ export const movieSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'REVIEWED',
+      label: 'REVIEWED',
       collection: 'REVIEWED',
       from: 'Person',
       to: 'Movie',
@@ -69,6 +71,7 @@ export const movieSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'PRODUCED',
+      label: 'PRODUCED',
       collection: 'PRODUCED',
       from: 'Person',
       to: 'Movie',
@@ -76,6 +79,7 @@ export const movieSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'WROTE',
+      label: 'WROTE',
       collection: 'WROTE',
       from: 'Person',
       to: 'Movie',
@@ -83,6 +87,7 @@ export const movieSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'FOLLOWS',
+      label: 'FOLLOWS',
       collection: 'FOLLOWS',
       from: 'Person',
       to: 'Person',
@@ -90,6 +95,7 @@ export const movieSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'DIRECTED',
+      label: 'DIRECTED',
       collection: 'DIRECTED',
       from: 'Person',
       to: 'Movie',
diff --git a/libs/shared/lib/mock-data/schema/northwindSchemaRaw.ts b/libs/shared/lib/mock-data/schema/northwindSchemaRaw.ts
index 7b665842a..3c4c78eec 100644
--- a/libs/shared/lib/mock-data/schema/northwindSchemaRaw.ts
+++ b/libs/shared/lib/mock-data/schema/northwindSchemaRaw.ts
@@ -236,6 +236,7 @@ export const northwindSchemaRaw: SchemaFromBackend = {
   edges: [
     {
       name: 'ORDERS',
+      label: 'ORDERS',
       collection: 'ORDERS',
       from: 'Order',
       to: 'Product',
@@ -264,6 +265,7 @@ export const northwindSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'PART_OF',
+      label: 'PART_OF',
       collection: 'PART_OF',
       from: 'Product',
       to: 'Category',
@@ -271,6 +273,7 @@ export const northwindSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'SUPPLIES',
+      label: 'SUPPLIES',
       collection: 'SUPPLIES',
       from: 'Supplier',
       to: 'Product',
@@ -278,6 +281,7 @@ export const northwindSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'PURCHASED',
+      label: 'PURCHASED',
       collection: 'PURCHASED',
       from: 'Customer',
       to: 'Order',
diff --git a/libs/shared/lib/mock-data/schema/simpleAirportRaw.ts b/libs/shared/lib/mock-data/schema/simpleAirportRaw.ts
index 7a877d85f..5a2fdbe30 100644
--- a/libs/shared/lib/mock-data/schema/simpleAirportRaw.ts
+++ b/libs/shared/lib/mock-data/schema/simpleAirportRaw.ts
@@ -19,6 +19,7 @@ export const simpleSchemaAirportRaw: SchemaFromBackend = {
   edges: [
     {
       name: 'flights',
+      label: 'flights',
       from: 'airports',
       to: 'airports',
       collection: 'flights',
diff --git a/libs/shared/lib/mock-data/schema/simpleRaw.ts b/libs/shared/lib/mock-data/schema/simpleRaw.ts
index 9ccfa362f..41a7fbf84 100644
--- a/libs/shared/lib/mock-data/schema/simpleRaw.ts
+++ b/libs/shared/lib/mock-data/schema/simpleRaw.ts
@@ -35,6 +35,7 @@ export const simpleSchemaRaw: SchemaFromBackend = {
   edges: [
     {
       name: 'Airport2:Airport',
+      label: 'Airport2:Airport',
       from: 'Airport2',
       to: 'Airport',
       collection: 'flights',
@@ -45,6 +46,7 @@ export const simpleSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'Airport:Staff',
+      label: 'Airport:Staff',
       from: 'Airport',
       to: 'Staff',
       collection: 'flights',
@@ -52,6 +54,7 @@ export const simpleSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'Plane:Airport',
+      label: 'Plane:Airport',
       from: 'Plane',
       to: 'Airport',
       collection: 'flights',
@@ -59,6 +62,7 @@ export const simpleSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'Airport:Thijs',
+      label: 'Airport:Thijs',
       from: 'Airport',
       to: 'Thijs',
       collection: 'flights',
@@ -66,6 +70,7 @@ export const simpleSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'Thijs:Airport',
+      label: 'Thijs:Airport',
       from: 'Thijs',
       to: 'Airport',
       collection: 'flights',
@@ -73,6 +78,7 @@ export const simpleSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'Staff:Plane',
+      label: 'Staff:Plane',
       from: 'Staff',
       to: 'Plane',
       collection: 'flights',
@@ -80,6 +86,7 @@ export const simpleSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'Staff:Airport2',
+      label: 'Staff:Airport2',
       from: 'Staff',
       to: 'Airport2',
       collection: 'flights',
@@ -87,14 +94,15 @@ export const simpleSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'Airport2:Plane',
+      label: 'Airport2:Plane',
       from: 'Airport2',
       to: 'Plane',
       collection: 'flights',
       attributes: [{ name: 'hallo', type: 'string' }],
     },
-
     {
       name: 'Airport:Airport',
+      label: 'Airport:Airport',
       from: 'Airport',
       to: 'Airport',
       collection: 'flights',
diff --git a/libs/shared/lib/mock-data/schema/twitterSchemaRaw.ts b/libs/shared/lib/mock-data/schema/twitterSchemaRaw.ts
index 7637bdb26..07eed928b 100644
--- a/libs/shared/lib/mock-data/schema/twitterSchemaRaw.ts
+++ b/libs/shared/lib/mock-data/schema/twitterSchemaRaw.ts
@@ -157,6 +157,7 @@ export const twitterSchemaRaw: SchemaFromBackend = {
   edges: [
     {
       name: 'USING',
+      label: 'USING',
       collection: 'USING',
       from: 'Tweet',
       to: 'Source',
@@ -164,6 +165,7 @@ export const twitterSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'SIMILAR_TO',
+      label: 'SIMILAR_TO',
       collection: 'SIMILAR_TO',
       from: 'User',
       to: 'User',
@@ -176,6 +178,7 @@ export const twitterSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'SIMILAR_TO',
+      label: 'SIMILAR_TO',
       collection: 'SIMILAR_TO',
       from: 'User',
       to: 'Me',
@@ -188,6 +191,7 @@ export const twitterSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'AMPLIFIES',
+      label: 'AMPLIFIES',
       collection: 'AMPLIFIES',
       from: 'Me',
       to: 'User',
@@ -195,6 +199,7 @@ export const twitterSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'AMPLIFIES',
+      label: 'AMPLIFIES',
       collection: 'AMPLIFIES',
       from: 'User',
       to: 'User',
@@ -202,6 +207,7 @@ export const twitterSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'RT_MENTIONS',
+      label: 'RT_MENTIONS',
       collection: 'RT_MENTIONS',
       from: 'Me',
       to: 'User',
@@ -209,6 +215,7 @@ export const twitterSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'RT_MENTIONS',
+      label: 'RT_MENTIONS',
       collection: 'RT_MENTIONS',
       from: 'User',
       to: 'User',
@@ -216,6 +223,7 @@ export const twitterSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'FOLLOWS',
+      label: 'FOLLOWS',
       collection: 'FOLLOWS',
       from: 'User',
       to: 'Me',
@@ -223,6 +231,7 @@ export const twitterSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'FOLLOWS',
+      label: 'FOLLOWS',
       collection: 'FOLLOWS',
       from: 'Me',
       to: 'User',
@@ -230,6 +239,7 @@ export const twitterSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'FOLLOWS',
+      label: 'FOLLOWS',
       collection: 'FOLLOWS',
       from: 'User',
       to: 'User',
@@ -237,6 +247,7 @@ export const twitterSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'FOLLOWS',
+      label: 'FOLLOWS',
       collection: 'FOLLOWS',
       from: 'Me',
       to: 'Me',
@@ -244,6 +255,7 @@ export const twitterSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'INTERACTS_WITH',
+      label: 'INTERACTS_WITH',
       collection: 'INTERACTS_WITH',
       from: 'User',
       to: 'User',
@@ -251,6 +263,7 @@ export const twitterSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'INTERACTS_WITH',
+      label: 'INTERACTS_WITH',
       collection: 'INTERACTS_WITH',
       from: 'Me',
       to: 'User',
@@ -258,6 +271,7 @@ export const twitterSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'RETWEETS',
+      label: 'RETWEETS',
       collection: 'RETWEETS',
       from: 'Tweet',
       to: 'Tweet',
@@ -265,6 +279,7 @@ export const twitterSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'REPLY_TO',
+      label: 'REPLY_TO',
       collection: 'REPLY_TO',
       from: 'Tweet',
       to: 'Tweet',
@@ -272,6 +287,7 @@ export const twitterSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'CONTAINS',
+      label: 'CONTAINS',
       collection: 'CONTAINS',
       from: 'Tweet',
       to: 'Link',
@@ -279,6 +295,7 @@ export const twitterSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'MENTIONS',
+      label: 'MENTIONS',
       collection: 'MENTIONS',
       from: 'Tweet',
       to: 'User',
@@ -286,6 +303,7 @@ export const twitterSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'MENTIONS',
+      label: 'MENTIONS',
       collection: 'MENTIONS',
       from: 'Tweet',
       to: 'Me',
@@ -293,6 +311,7 @@ export const twitterSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'TAGS',
+      label: 'TAGS',
       collection: 'TAGS',
       from: 'Tweet',
       to: 'Hashtag',
@@ -300,6 +319,7 @@ export const twitterSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'POSTS',
+      label: 'POSTS',
       collection: 'POSTS',
       from: 'User',
       to: 'Tweet',
@@ -307,6 +327,7 @@ export const twitterSchemaRaw: SchemaFromBackend = {
     },
     {
       name: 'POSTS',
+      label: 'POSTS',
       collection: 'POSTS',
       from: 'Me',
       to: 'Tweet',
diff --git a/libs/shared/lib/querybuilder/model/graphology/model.ts b/libs/shared/lib/querybuilder/model/graphology/model.ts
index 6c14d6d14..1f851ab78 100644
--- a/libs/shared/lib/querybuilder/model/graphology/model.ts
+++ b/libs/shared/lib/querybuilder/model/graphology/model.ts
@@ -19,6 +19,7 @@ export interface NodeAttribute {
 
 export type NodeDefaults = {
   id?: string;
+  schemaKey?: string;
   type: QueryElementTypes;
   width?: number;
   height?: number;
diff --git a/libs/shared/lib/querybuilder/model/graphology/utils.ts b/libs/shared/lib/querybuilder/model/graphology/utils.ts
index e67436f70..e46dd9ac7 100644
--- a/libs/shared/lib/querybuilder/model/graphology/utils.ts
+++ b/libs/shared/lib/querybuilder/model/graphology/utils.ts
@@ -50,6 +50,7 @@ export class QueryMultiGraphology extends Graph<QueryGraphNodes, QueryGraphEdges
     attributes.y = y;
     attributes.width = width;
     attributes.height = height;
+    attributes.schemaKey = attributes.schemaKey;
 
     if (!attributes.id) attributes.id = 'id_' + (Date.now() + Math.floor(Math.random() * 1000)).toString();
 
@@ -101,6 +102,8 @@ export class QueryMultiGraphology extends Graph<QueryGraphNodes, QueryGraphEdges
       throw Error('using wrong function! use addLogicPill2Graphology instead');
     }
 
+    console.log(attributes);
+
     // Add a node to the graphology object
     this.addNode(attributes.id, { ...attributes });
 
diff --git a/libs/shared/lib/querybuilder/model/reactflow/handles.tsx b/libs/shared/lib/querybuilder/model/reactflow/handles.tsx
index 3bdb84114..676ef3207 100644
--- a/libs/shared/lib/querybuilder/model/reactflow/handles.tsx
+++ b/libs/shared/lib/querybuilder/model/reactflow/handles.tsx
@@ -16,71 +16,26 @@ export enum Handles {
   RelationLeft = 'relationLeftHandle', //target
   RelationRight = 'relationRightHandle', //source
   RelationAttribute = 'relationAttributeHandle', //source
-  ToAttribute = 'attributesHandle', //target
   EntityLeft = 'entityLeftHandle', //source
   EntityRight = 'entityRightHandle', //target
   EntityAttribute = 'entityAttributeHandle', //source
-  OnAttribute = 'onAttributeHandle', //source
-  ReceiveFunction = 'receiveFunctionHandle', //target
-  FunctionBase = 'functionHandle_', // + name from FunctionTypes args //source
-  FromAttribute = 'fromAttributeHandle', //source
   LogicLeft = 'leftLogicHandle',
   LogicRight = 'rightLogicHandle',
 }
 
-/** returns a boolean that check whether the handle is a function handle */
-export function isFunctionHandle(handle: string): boolean {
-  return handle.startsWith(Handles.FunctionBase);
+export function isRelationHandle(handle: Handles): boolean {
+  return handle.startsWith(Handles.RelationLeft) || handle.startsWith(Handles.RelationRight);
 }
 
-// /**
-//  *  returns the functionargumenttype
-//  *  Currently only working for groupby but made that in the future other functions can use this as well.
-//  */
-// export function functionHandleToType(handle: string): FunctionArgTypes {
-//   if (isFunctionHandle(handle))
-//     return handle.slice(Handles.FunctionBase.length) as FunctionArgTypes;
-//   else
-//     throw new Error('Incorrectly trying to assert handle to function handle');
-// }
-// /** Creates a handle from a functiontype */
-// export function typeToFunctionHandle(type: FunctionArgTypes): string {
-//   return Handles.FunctionBase + type;
-// }
-
-/**
- * Return a list of handles to which a connection can be made by dragging a node nearby
- */
-export function nodeToHandlesThatCanReceiveDragconnect(node: SchemaReactflowNode): string[] {
-  switch (node.type) {
-    case QueryElementTypes.Entity:
-      return [Handles.ToAttribute];
-    case QueryElementTypes.Relation:
-      return [Handles.RelationLeft, Handles.RelationRight, Handles.ToAttribute];
-    // case QueryElementTypes.Function:
-    //   return [Handles.ToAttribute];
-    // case QueryElementTypes.Attribute:
-    //   return [];
-    default:
-      throw new Error('Unsupported node');
-  }
+export function isEntityHandle(handle: Handles): boolean {
+  return handle.startsWith(Handles.EntityLeft) || handle.startsWith(Handles.EntityRight);
 }
 
-/**
- * Return a list of handles from which a connection can be made while dragging the node they are on
- * @deprecated
- */
-export function nodeToHandlesThatCanSendDragconnect(node: SchemaReactflowNode): string[] {
-  switch (node.type) {
-    // case QueryElementTypes.Entity:
-    //   return [Handles.ToRelation];
-    // case QueryElementTypes.Relation:
-    //   return [];
-    // case QueryElementTypes.Function:
-    //   return [];
-    // case QueryElementTypes.Attribute:
-    //   return [Handles.OnAttribute];
-    default:
-      throw new Error('Unsupported node');
-  }
+export function isLogicHandle(handle: Handles): boolean {
+  return (
+    handle.startsWith(Handles.LogicLeft) ||
+    handle.startsWith(Handles.LogicRight) ||
+    handle.startsWith(Handles.RelationAttribute) ||
+    handle.startsWith(Handles.EntityAttribute)
+  );
 }
diff --git a/libs/shared/lib/querybuilder/model/reactflow/model.tsx b/libs/shared/lib/querybuilder/model/reactflow/model.tsx
index ed03cfb64..a35239236 100644
--- a/libs/shared/lib/querybuilder/model/reactflow/model.tsx
+++ b/libs/shared/lib/querybuilder/model/reactflow/model.tsx
@@ -5,7 +5,6 @@
  */
 import { Edge as ReactEdge, NodeProps } from 'reactflow';
 import { EntityData, EntityNodeAttributes, LogicData, LogicNodeAttributes, RelationData, RelationNodeAttributes } from '../graphology';
-import { Schema } from 'inspector';
 
 /** Enums for the possible types of query elements */
 export enum QueryElementTypes {
diff --git a/libs/shared/lib/querybuilder/panel/querybuilder.tsx b/libs/shared/lib/querybuilder/panel/querybuilder.tsx
index 12f975c72..c0ca3d3aa 100644
--- a/libs/shared/lib/querybuilder/panel/querybuilder.tsx
+++ b/libs/shared/lib/querybuilder/panel/querybuilder.tsx
@@ -36,6 +36,7 @@ import {
   QueryElementTypes,
   QueryGraphNodes,
   createReactFlowElements,
+  isLogicHandle,
   toHandleData,
 } from '../model';
 import { InputNodeType } from '../model/logic/general';
@@ -51,12 +52,15 @@ import { QueryMLDialog } from './querysidepanel/queryMLDialog';
 import { QuerySettingsDialog } from './querysidepanel/querySettingsDialog';
 import { toSchemaGraphology } from '../../data-access/store/schemaSlice';
 import { LayoutFactory } from '../../graph-layout';
+import { ConnectingNodeDataI } from './utils/connectorDrop';
+import { QueryBuilderRelatedNodesPanel } from './querysidepanel/queryBuilderRelatedNodesPanel';
+import { addError } from '../../data-access/store/configSlice';
 
 export type QueryBuilderProps = {
   onRunQuery?: () => void;
 };
 
-type SettingsPanel = 'settings' | 'ml' | 'logic' | undefined;
+type SettingsPanel = 'settings' | 'ml' | 'logic' | 'relatedNodes' | undefined;
 
 /**
  * This is the main querybuilder component. It is responsible for holding all pills and fire off the visual part of the querybuilder panel logic
@@ -83,12 +87,7 @@ export const QueryBuilderInner = (props: QueryBuilderProps) => {
   const config = useConfig();
   const dispatch = useDispatch();
   const isDraggingPill = useRef(false);
-  const connectingNodeId = useRef<{
-    params: OnConnectStartParams;
-    position: XYPosition;
-    node: QueryGraphNodes;
-    attribute: NodeAttribute;
-  } | null>(null);
+  const connectingNodeId = useRef<ConnectingNodeDataI | null>(null);
   const reactFlow = useReactFlow();
   const isEdgeUpdating = useRef(false);
   const isOnConnect = useRef(false);
@@ -200,9 +199,11 @@ export const QueryBuilderInner = (props: QueryBuilderProps) => {
             x: position.x,
             y: position.y,
             name: dragData.name,
+            schemaKey: dragData.name,
           },
           schema.getNodeAttribute(dragData.name, 'attributes')
         );
+
         dispatch(setQuerybuilderGraphology(graphologyGraph));
         break;
       // Creates a relation element and will also create the 2 related entities together with the connections
@@ -214,6 +215,7 @@ export const QueryBuilderInner = (props: QueryBuilderProps) => {
             y: position.y,
             depth: { min: queryBuilderSettings.depth.min, max: queryBuilderSettings.depth.max },
             name: dragData.collection,
+            schemaKey: dragData.label,
             collection: dragData.collection,
           },
           schema.getEdgeAttribute(dragData.label, 'attributes')
@@ -315,57 +317,21 @@ export const QueryBuilderInner = (props: QueryBuilderProps) => {
         const position = reactFlow.project({ x: clientX, y: clientY });
         if (connectingNodeId?.current) connectingNodeId.current.position = position;
 
-        setToggleSettings('logic');
+        if (!connectingNodeId?.current?.params?.handleId) {
+          dispatch(addError('Connection has no source or target handle id'));
+          return;
+        } else {
+          const data = toHandleData(connectingNodeId.current.params.handleId);
+          if (isLogicHandle(data.handleType)) setToggleSettings('logic');
+          else setToggleSettings('relatedNodes');
+        }
+
+        // setToggleSettings('logic');
       }
     },
     [reactFlow.project]
   );
 
-  const onNewNodeFromPopup = (value: AllLogicDescriptions) => {
-    const logic = AllLogicMap[value.key];
-    const firstLeftLogicInput = logic.inputs?.[0];
-    if (!firstLeftLogicInput) return;
-
-    if (connectingNodeId.current === null || connectingNodeId.current?.params?.handleId == null) {
-      const bounds = reactFlowWrapper?.current?.getBoundingClientRect() || { x: 0, y: 0, width: 0, height: 0 };
-
-      const logicNode = graphologyGraph.addLogicPill2Graphology({
-        name: value.name,
-        type: QueryElementTypes.Logic,
-        x: bounds.width / 2,
-        y: bounds.height / 2,
-        logic: logic,
-      });
-    } else {
-      const params = connectingNodeId.current.params;
-      const position = connectingNodeId.current.position;
-
-      const logicNode = graphologyGraph.addLogicPill2Graphology({
-        name: value.name,
-        type: QueryElementTypes.Logic,
-        x: position.x,
-        y: position.y,
-        logic: logic,
-      });
-
-      if (!logicNode?.id) throw new Error('Logic node has no id');
-      if (!logicNode?.name) throw new Error('Logic node has no name');
-      if (!params.handleId) throw new Error('Connection has no source or target');
-
-      const sourceHandleData = toHandleData(params.handleId);
-      graphologyGraph.addEdge2Graphology(
-        graphologyGraph.getNodeAttributes(params.nodeId),
-        graphologyGraph.getNodeAttributes(logicNode.id),
-        { type: 'connection' },
-        { sourceHandleName: sourceHandleData.attributeName, targetHandleName: firstLeftLogicInput.name }
-      );
-    }
-
-    dispatch(setQuerybuilderGraphology(graphologyGraph));
-    setToggleSettings(undefined);
-    connectingNodeId.current = null;
-  };
-
   const onEdgeUpdateStart = useCallback(() => {
     isEdgeUpdating.current = true;
   }, []);
@@ -438,13 +404,34 @@ export const QueryBuilderInner = (props: QueryBuilderProps) => {
         }}
       >
         <QueryBuilderLogicPillsPanel
-          onClick={onNewNodeFromPopup}
+          onClick={(v) => {
+            connectingNodeId.current = null;
+            setToggleSettings(undefined);
+          }}
+          reactFlowWrapper={reactFlowWrapper.current}
           title="Logic Pills usable by the node"
           className="min-h-[75vh] max-h-[75vh]"
           onDrag={() => {}}
           connection={connectingNodeId?.current}
         />
       </Dialog>
+      <Dialog
+        open={toggleSettings === 'relatedNodes'}
+        onClose={() => {
+          setToggleSettings(undefined);
+        }}
+      >
+        <QueryBuilderRelatedNodesPanel
+          onFinished={() => {
+            connectingNodeId.current = null;
+            setToggleSettings(undefined);
+          }}
+          reactFlowWrapper={reactFlowWrapper.current}
+          title="Related nodes available to add to the query"
+          className="min-h-[75vh] max-h-[75vh]"
+          connection={connectingNodeId?.current}
+        />
+      </Dialog>
       <svg height={0}>
         <defs>
           <marker id="arrowIn" markerWidth="9" markerHeight="10" refX={0} refY="5" orient="auto">
diff --git a/libs/shared/lib/querybuilder/panel/querysidepanel/queryBuilderLogicPillsPanel.tsx b/libs/shared/lib/querybuilder/panel/querysidepanel/queryBuilderLogicPillsPanel.tsx
index 9f8b9b1c8..5d60445a5 100644
--- a/libs/shared/lib/querybuilder/panel/querysidepanel/queryBuilderLogicPillsPanel.tsx
+++ b/libs/shared/lib/querybuilder/panel/querysidepanel/queryBuilderLogicPillsPanel.tsx
@@ -3,22 +3,28 @@ import FunctionsIcon from '@mui/icons-material/Functions';
 import GridOnIcon from '@mui/icons-material/GridOn';
 import NumbersIcon from '@mui/icons-material/Numbers';
 import AbcIcon from '@mui/icons-material/Abc';
-import { useState } from 'react';
-import { AllLogicDescriptions, AllLogicMap, NodeAttribute, QueryGraphNodes, toHandleData } from '../../model';
+import { useMemo, useState } from 'react';
+import { AllLogicDescriptions, AllLogicMap, NodeAttribute, QueryElementTypes, QueryGraphNodes, toHandleData } from '../../model';
 import { OnConnectStartParams, XYPosition } from 'reactflow';
+import { ConnectingNodeDataI } from '../utils/connectorDrop';
+import { useQuerybuilderGraph } from '@graphpolaris/shared/lib/data-access';
+import { toQuerybuilderGraphology, setQuerybuilderGraphology } from '@graphpolaris/shared/lib/data-access/store/querybuilderSlice';
+import { useDispatch } from 'react-redux';
 
 export const QueryBuilderLogicPillsPanel = (props: {
+  reactFlowWrapper: HTMLDivElement | null;
   className?: string;
   title?: string;
   onClick: (item: AllLogicDescriptions) => void;
   onDrag?: (item: AllLogicDescriptions) => void;
-  connection?: {
-    params: OnConnectStartParams;
-    position: XYPosition;
-    node: QueryGraphNodes;
-    attribute: NodeAttribute;
-  } | null;
+  connection?: ConnectingNodeDataI | null;
 }) => {
+  const graph = useQuerybuilderGraph();
+  const dispatch = useDispatch();
+  const graphologyGraph = toQuerybuilderGraphology(graph);
+  const [selectedOp, setSelectedOp] = useState(-1);
+  const [selectedType, setSelectedType] = useState(-1);
+
   let filterType = (props.connection?.params?.handleId ? toHandleData(props.connection.params.handleId).attributeType : null) as string;
   if (!filterType) return <></>;
   else if (filterType === 'string') filterType = 'string';
@@ -53,8 +59,6 @@ export const QueryBuilderLogicPillsPanel = (props: {
       icon: <AbcIcon fontSize="small" />,
     },
   ].filter((item) => !filterType || item.title === filterType);
-  const [selectedOp, setSelectedOp] = useState(-1);
-  const [selectedType, setSelectedType] = useState(-1);
 
   const onDragStart = (event: React.DragEvent, value: AllLogicDescriptions) => {
     console.log('drag start');
@@ -64,6 +68,50 @@ export const QueryBuilderLogicPillsPanel = (props: {
     if (props.onDrag) props.onDrag(value);
   };
 
+  const onNewNodeFromPopup = (value: AllLogicDescriptions) => {
+    const logic = AllLogicMap[value.key];
+    const firstLeftLogicInput = logic.inputs?.[0];
+    if (!firstLeftLogicInput) return;
+
+    if (props.connection === null || props.connection?.params?.handleId == null) {
+      const bounds = props.reactFlowWrapper?.getBoundingClientRect() || { x: 0, y: 0, width: 0, height: 0 };
+
+      const logicNode = graphologyGraph.addLogicPill2Graphology({
+        name: value.name,
+        type: QueryElementTypes.Logic,
+        x: bounds.width / 2,
+        y: bounds.height / 2,
+        logic: logic,
+      });
+    } else {
+      const params = props.connection.params;
+      const position = props.connection.position;
+
+      const logicNode = graphologyGraph.addLogicPill2Graphology({
+        name: value.name,
+        type: QueryElementTypes.Logic,
+        x: position.x,
+        y: position.y,
+        logic: logic,
+      });
+
+      if (!logicNode?.id) throw new Error('Logic node has no id');
+      if (!logicNode?.name) throw new Error('Logic node has no name');
+      if (!params.handleId) throw new Error('Connection has no source or target');
+
+      const sourceHandleData = toHandleData(params.handleId);
+      graphologyGraph.addEdge2Graphology(
+        graphologyGraph.getNodeAttributes(params.nodeId),
+        graphologyGraph.getNodeAttributes(logicNode.id),
+        { type: 'connection' },
+        { sourceHandleName: sourceHandleData.attributeName, targetHandleName: firstLeftLogicInput.name }
+      );
+    }
+
+    dispatch(setQuerybuilderGraphology(graphologyGraph));
+    props.onClick(value);
+  };
+
   return (
     <div className={props.className + ' card'}>
       {props.title && <h1 className="card-title mb-7">{props.title}</h1>}
@@ -110,7 +158,7 @@ export const QueryBuilderLogicPillsPanel = (props: {
                   onDragStart={(e) => onDragStart(e, item)}
                   draggable={true}
                   onClick={() => {
-                    props.onClick(item);
+                    onNewNodeFromPopup(item);
                   }}
                 >
                   {item.icon && (
diff --git a/libs/shared/lib/querybuilder/panel/querysidepanel/queryBuilderRelatedNodesPanel.tsx b/libs/shared/lib/querybuilder/panel/querysidepanel/queryBuilderRelatedNodesPanel.tsx
new file mode 100644
index 000000000..004bf405f
--- /dev/null
+++ b/libs/shared/lib/querybuilder/panel/querysidepanel/queryBuilderRelatedNodesPanel.tsx
@@ -0,0 +1,199 @@
+import FilterAltIcon from '@mui/icons-material/FilterAlt';
+import FunctionsIcon from '@mui/icons-material/Functions';
+import GridOnIcon from '@mui/icons-material/GridOn';
+import NumbersIcon from '@mui/icons-material/Numbers';
+import AbcIcon from '@mui/icons-material/Abc';
+import { useMemo, useState } from 'react';
+import { AllLogicDescriptions, AllLogicMap, Handles, NodeAttribute, QueryElementTypes, QueryGraphNodes, toHandleData } from '../../model';
+import { OnConnectStartParams, XYPosition } from 'reactflow';
+import { ConnectingNodeDataI } from '../utils/connectorDrop';
+import { useQuerybuilderGraph, useQuerybuilderSettings, useSchemaGraph } from '@graphpolaris/shared/lib/data-access';
+import { toQuerybuilderGraphology, setQuerybuilderGraphology } from '@graphpolaris/shared/lib/data-access/store/querybuilderSlice';
+import { useDispatch } from 'react-redux';
+import { toSchemaGraphology } from '@graphpolaris/shared/lib/data-access/store/schemaSlice';
+import { SchemaAttribute, SchemaEdge, SchemaNode } from '@graphpolaris/shared/lib/schema';
+import { schemaExpandRelation } from '@graphpolaris/shared/lib/schema/schema-utils';
+
+export const QueryBuilderRelatedNodesPanel = (props: {
+  reactFlowWrapper: HTMLDivElement | null;
+  className?: string;
+  title?: string;
+  onFinished: () => void;
+  connection?: ConnectingNodeDataI | null;
+}) => {
+  const schema = useSchemaGraph();
+  const graph = useQuerybuilderGraph();
+  const dispatch = useDispatch();
+  const graphologyGraph = toQuerybuilderGraphology(graph);
+  const schemaGraph = toSchemaGraphology(schema);
+  const queryBuilderSettings = useQuerybuilderSettings();
+  const [selectedType, setSelectedType] = useState<'entity' | 'relation' | 'all'>('all');
+
+  const handleData = props.connection?.params?.handleId ? toHandleData(props.connection.params.handleId) : null;
+
+  const relatedEntities = useMemo<SchemaNode[]>(() => {
+    if (!props.connection) return [];
+    const type = props.connection.node.type;
+    if (type !== QueryElementTypes.Entity && type !== QueryElementTypes.Relation)
+      throw new Error('Connection is not an entity or relation');
+
+    if (type === QueryElementTypes.Entity) {
+      // find entities that are connected to the current entity
+      const nodes =
+        handleData?.handleType === Handles.EntityRight
+          ? schemaGraph.outboundNeighbors(props.connection.node.schemaKey)
+          : schemaGraph.inboundNeighbors(props.connection.node.schemaKey);
+      return nodes.map((node) => schemaGraph.getNodeAttributes(node) as SchemaNode);
+    } else if (type === QueryElementTypes.Relation) {
+      // find relations of the current entity
+      const attributes = schemaGraph.getEdgeAttributes(props.connection.node.schemaKey);
+      return handleData?.handleType === Handles.RelationRight
+        ? [schemaGraph.getNodeAttributes(attributes.to)]
+        : [schemaGraph.getNodeAttributes(attributes.from)];
+    }
+
+    return [];
+  }, [schema, graph, props]);
+
+  const relatedRelations = useMemo<SchemaEdge[]>(() => {
+    if (!props.connection) return [];
+    const type = props.connection.node.type;
+    if (type !== QueryElementTypes.Entity && type !== QueryElementTypes.Relation)
+      throw new Error('Connection is not an entity or relation');
+
+    if (type === QueryElementTypes.Entity) {
+      // find entities on the edge of the current relation
+      const edges =
+        handleData?.handleType === Handles.EntityRight
+          ? schemaGraph.outboundEdges(props.connection.node.schemaKey)
+          : schemaGraph.inboundEdges(props.connection.node.schemaKey);
+      return edges.map((edge) => ({ ...schemaGraph.getEdgeAttributes(edge), label: edge } as SchemaEdge));
+    } else if (type === QueryElementTypes.Relation) {
+      // find relations that are connected to the proper neighboring entity of current relation
+      const attributes = schemaGraph.getEdgeAttributes(props.connection.node.schemaKey);
+      const edges =
+        handleData?.handleType === Handles.RelationRight
+          ? schemaGraph.outboundEdges(attributes.to)
+          : schemaGraph.inboundEdges(attributes.from);
+      return edges.map((edge) => ({ ...schemaGraph.getEdgeAttributes(edge), label: edge } as SchemaEdge));
+    }
+
+    return [];
+  }, [schema, graph, props]);
+
+  const newEntity = (entity: SchemaNode) => {
+    if (!props.connection) {
+      props.onFinished();
+      return;
+    }
+
+    const params = props.connection.params;
+    const position = props.connection.position;
+
+    const newNode = graphologyGraph.addPill2Graphology(
+      {
+        type: QueryElementTypes.Entity,
+        x: position.x,
+        y: position.y,
+        name: entity.name,
+        schemaKey: entity.name,
+      },
+      schemaGraph.getNodeAttribute(entity.name, 'attributes')
+    );
+
+    if (!newNode?.id) throw new Error('Logic node has no id');
+    if (!newNode?.name) throw new Error('Logic node has no name');
+    if (!params.handleId) throw new Error('Connection has no source or target');
+
+    if (handleData?.handleType === Handles.EntityRight || handleData?.handleType === Handles.RelationRight)
+      graphologyGraph.addEdge2Graphology(graphologyGraph.getNodeAttributes(params.nodeId), newNode);
+    else graphologyGraph.addEdge2Graphology(newNode, graphologyGraph.getNodeAttributes(params.nodeId));
+
+    dispatch(setQuerybuilderGraphology(graphologyGraph));
+    props.onFinished();
+  };
+  const newRelation = (relation: SchemaEdge) => {
+    if (!props.connection) {
+      props.onFinished();
+      return;
+    }
+
+    const params = props.connection.params;
+    const position = props.connection.position;
+
+    const newNode = graphologyGraph.addPill2Graphology(
+      {
+        type: QueryElementTypes.Relation,
+        x: position.x,
+        y: position.y,
+        schemaKey: relation.label,
+        depth: { min: queryBuilderSettings.depth.min, max: queryBuilderSettings.depth.max },
+        name: relation.collection,
+        collection: relation.collection,
+      },
+      schemaGraph.getEdgeAttribute(relation.label, 'attributes')
+    );
+
+    if (!newNode?.id) throw new Error('Logic node has no id');
+    if (!newNode?.name) throw new Error('Logic node has no name');
+    if (!params.handleId) throw new Error('Connection has no source or target');
+
+    if (handleData?.handleType === Handles.EntityRight || handleData?.handleType === Handles.RelationRight)
+      graphologyGraph.addEdge2Graphology(graphologyGraph.getNodeAttributes(params.nodeId), newNode);
+    else graphologyGraph.addEdge2Graphology(newNode, graphologyGraph.getNodeAttributes(params.nodeId));
+    dispatch(setQuerybuilderGraphology(graphologyGraph));
+    props.onFinished();
+  };
+
+  return (
+    <div className={props.className + ' card'}>
+      {props.title && <h1 className="card-title mb-7">{props.title}</h1>}
+      <div className="overflow-x-hidden h-[75rem] w-full gap-2 pt-3">
+        <div className="w-full">
+          <button
+            className={'btn w-[calc(50%-0.25rem)] mr-2 normal-case text-lg mb-3 ' + (selectedType !== 'relation' ? 'btn-active' : '')}
+            onClick={(e) => {
+              e.preventDefault();
+              selectedType === 'entity' ? setSelectedType('all') : setSelectedType('entity');
+            }}
+          >
+            Related Entities
+          </button>
+          <button
+            className={'btn w-[calc(50%-0.25rem)] text-lg normal-case mb-3 ' + (selectedType !== 'entity' ? 'btn-active' : '')}
+            onClick={(e) => {
+              e.preventDefault();
+              selectedType === 'relation' ? setSelectedType('all') : setSelectedType('relation');
+            }}
+          >
+            Related Relations
+          </button>
+        </div>
+        <div className="w-full flex flex-row gap-2">
+          {selectedType !== 'relation' && (
+            <ul className="menu p-0 [&_li>*]:rounded-none w-full pb-10">
+              {relatedEntities.map((entity) => {
+                return (
+                  <button className="btn btn-sm border-entity-400 normal-case" key={entity.name} onClick={() => newEntity(entity)}>
+                    {entity.name}
+                  </button>
+                );
+              })}
+            </ul>
+          )}
+          {selectedType !== 'entity' && (
+            <ul className="menu p-0 [&_li>*]:rounded-none w-full pb-10">
+              {relatedRelations.map((relation) => {
+                return (
+                  <button className="btn btn-sm border-relation-400 normal-case" key={relation.name} onClick={() => newRelation(relation)}>
+                    {relation.name}
+                  </button>
+                );
+              })}
+            </ul>
+          )}
+        </div>
+      </div>
+    </div>
+  );
+};
diff --git a/libs/shared/lib/querybuilder/panel/stories/querybuilder-simple.stories.tsx b/libs/shared/lib/querybuilder/panel/stories/querybuilder-simple.stories.tsx
index f55aec918..8f40b7593 100644
--- a/libs/shared/lib/querybuilder/panel/stories/querybuilder-simple.stories.tsx
+++ b/libs/shared/lib/querybuilder/panel/stories/querybuilder-simple.stories.tsx
@@ -46,6 +46,7 @@ export const Simple = {
       edges: [
         {
           name: 'entity:entity',
+          label: 'entity:entity',
           from: 'entity',
           to: 'entity',
           collection: 'entity2entity',
diff --git a/libs/shared/lib/querybuilder/panel/utils/connectorDrop.ts b/libs/shared/lib/querybuilder/panel/utils/connectorDrop.ts
new file mode 100644
index 000000000..7a4820693
--- /dev/null
+++ b/libs/shared/lib/querybuilder/panel/utils/connectorDrop.ts
@@ -0,0 +1,17 @@
+import { OnConnectStartParams, XYPosition } from 'reactflow';
+import {
+  AllLogicDescriptions,
+  AllLogicMap,
+  NodeAttribute,
+  QueryElementTypes,
+  QueryGraphNodes,
+  QueryMultiGraphology,
+  toHandleData,
+} from '../../model';
+
+export type ConnectingNodeDataI = {
+  params: OnConnectStartParams;
+  position: XYPosition;
+  node: QueryGraphNodes;
+  attribute: NodeAttribute;
+};
diff --git a/libs/shared/lib/schema/model/FromBackend.ts b/libs/shared/lib/schema/model/FromBackend.ts
index 058ce78f1..c080591de 100644
--- a/libs/shared/lib/schema/model/FromBackend.ts
+++ b/libs/shared/lib/schema/model/FromBackend.ts
@@ -1,4 +1,7 @@
 /*************** schema format from the backend *************** */
+
+import { QueryElementTypes } from '../../querybuilder';
+
 /** Schema type, consist of nodes and edges */
 export type SchemaFromBackend = {
   edges: SchemaEdge[];
@@ -17,6 +20,7 @@ export type SchemaAttribute = {
 export type SchemaNode = {
   name: string;
   attributes: SchemaAttribute[];
+  type?: string;
 };
 
 /** Edge type, consist of a name, start point, end point and a list of attributes */
@@ -26,4 +30,6 @@ export type SchemaEdge = {
   from: string;
   collection: string;
   attributes: SchemaAttribute[];
+  label: string;
+  type?: string;
 };
diff --git a/libs/shared/lib/schema/panel/schema.stories.tsx b/libs/shared/lib/schema/panel/schema.stories.tsx
index 8a91d9337..61457bbf4 100644
--- a/libs/shared/lib/schema/panel/schema.stories.tsx
+++ b/libs/shared/lib/schema/panel/schema.stories.tsx
@@ -75,6 +75,7 @@ export const TestSimple = {
       edges: [
         {
           name: 'Thijs:Airport',
+          label: 'Thijs:Airport',
           from: 'Thijs',
           to: 'Airport',
           collection: 'flights',
diff --git a/libs/shared/lib/schema/pills/nodes/entity/entity-node.tsx b/libs/shared/lib/schema/pills/nodes/entity/entity-node.tsx
index af97fd7cd..f5ce86800 100644
--- a/libs/shared/lib/schema/pills/nodes/entity/entity-node.tsx
+++ b/libs/shared/lib/schema/pills/nodes/entity/entity-node.tsx
@@ -15,6 +15,7 @@ import { SchemaReactflowNodeWithFunctions } from '../../../model/reactflow';
 import { QueryElementTypes } from '@graphpolaris/shared/lib/querybuilder';
 import { SchemaEntityPopup } from './SchemaEntityPopup';
 import { Popup } from '@graphpolaris/shared/lib/components/Popup';
+import { SchemaNode } from '../../../model';
 
 /**
  * EntityNode is the node that represents the database entities.
@@ -30,7 +31,12 @@ export const EntityNode = React.memo(({ id, selected, data }: NodeProps<SchemaRe
    * @param event React Mouse drag event
    */
   const onDragStart = (event: React.DragEvent<HTMLDivElement>) => {
-    event.dataTransfer.setData('application/reactflow', JSON.stringify({ type: QueryElementTypes.Entity, name: id }));
+    const eventData: SchemaNode = {
+      type: QueryElementTypes.Entity,
+      name: id,
+      attributes: data.attributes.map((attribute) => attribute.attributes).flat(), // TODO: this seems wrong
+    };
+    event.dataTransfer.setData('application/reactflow', JSON.stringify(eventData));
     event.dataTransfer.effectAllowed = 'move';
   };
 
diff --git a/libs/shared/lib/schema/pills/nodes/relation/relation-node.stories.tsx b/libs/shared/lib/schema/pills/nodes/relation/relation-node.stories.tsx
index fc6012764..bb8bc3f1b 100644
--- a/libs/shared/lib/schema/pills/nodes/relation/relation-node.stories.tsx
+++ b/libs/shared/lib/schema/pills/nodes/relation/relation-node.stories.tsx
@@ -57,6 +57,7 @@ export const Default = {
       edges: [
         {
           name: 'Thijs:Thijs',
+          label: 'Thijs:Thijs',
           from: 'Thijs',
           to: 'Thijs',
           collection: 'flights',
diff --git a/libs/shared/lib/schema/pills/nodes/relation/relation-node.tsx b/libs/shared/lib/schema/pills/nodes/relation/relation-node.tsx
index 6f42caf44..e36cd36ff 100644
--- a/libs/shared/lib/schema/pills/nodes/relation/relation-node.tsx
+++ b/libs/shared/lib/schema/pills/nodes/relation/relation-node.tsx
@@ -15,6 +15,7 @@ import { SchemaReactflowRelation, SchemaReactflowRelationWithFunctions } from '.
 import { QueryElementTypes } from '@graphpolaris/shared/lib/querybuilder';
 import { Popup } from '@graphpolaris/shared/lib/components/Popup';
 import { SchemaRelationshipPopup } from './SchemaRelationshipPopup';
+import { SchemaEdge } from '../../../model';
 
 /**
  * Relation node component that renders a relation node for the schema.
@@ -29,17 +30,18 @@ export const RelationNode = React.memo(({ id, selected, data }: NodeProps<Schema
    * @param event React Mouse drag event.
    */
   const onDragStart = (event: React.DragEvent<HTMLDivElement>) => {
-    event.dataTransfer.setData(
-      'application/reactflow',
-      JSON.stringify({
-        type: QueryElementTypes.Relation,
-        name: id, //TODO id?
-        from: data.from,
-        to: data.to,
-        collection: data.collection,
-        label: data.label,
-      })
-    );
+    console.log(data);
+
+    const eventData: SchemaEdge = {
+      type: QueryElementTypes.Relation,
+      name: id, //TODO id?
+      from: data.from,
+      to: data.to,
+      collection: data.collection,
+      label: data.label,
+      attributes: data.attributes.map((attribute) => attribute.attributes).flat(), // TODO this seems wrong
+    };
+    event.dataTransfer.setData('application/reactflow', JSON.stringify(eventData));
     event.dataTransfer.effectAllowed = 'move';
   };
 
diff --git a/libs/shared/lib/vis/paohvis/paohvis.stories.tsx b/libs/shared/lib/vis/paohvis/paohvis.stories.tsx
index 1feba2010..36815e586 100644
--- a/libs/shared/lib/vis/paohvis/paohvis.stories.tsx
+++ b/libs/shared/lib/vis/paohvis/paohvis.stories.tsx
@@ -55,6 +55,7 @@ export const TestWithData = {
       edges: [
         {
           name: '12',
+          label: '12',
           from: '1',
           to: '1',
           collection: '1c',
diff --git a/libs/shared/lib/vis/semanticsubstrates/semanticsubstrates.stories.tsx b/libs/shared/lib/vis/semanticsubstrates/semanticsubstrates.stories.tsx
index f1f381671..a984c344d 100644
--- a/libs/shared/lib/vis/semanticsubstrates/semanticsubstrates.stories.tsx
+++ b/libs/shared/lib/vis/semanticsubstrates/semanticsubstrates.stories.tsx
@@ -41,6 +41,7 @@ export const TestWithData = {
       edges: [
         {
           name: '12',
+          label: '12',
           from: '1',
           to: '1',
           collection: '1c',
-- 
GitLab