Fix model lineage and create Python model UI form

CHANGELOG.md

@@ -1,5 +1,16 @@
 # Change Log
 
+## 2.2.2
+
+### Model Lineage
+
+- **Cleaner lineage for Python models.** Python models now show their real upstream sources in the lineage graph, without the extra placeholder node that used to appear.
+
+### Create Python Model
+
+- **Group options match your project.** The Create Python Model form now lists the groups from the project you've selected, and refreshes when you switch projects.
+- **Fix: creating a Python model works reliably.** Creating a Python model and loading its DAG and group options now work as expected.
+
 ## 2.2.1
 
 ### Security

package.json

@@ -10,7 +10,7 @@
     "type": "git",
     "url": "https://github.com/Workday/dj.git"
   },
-  "version": "2.2.1",
+  "version": "2.2.2",
   "workspaces": [
     "web"
   ],

src/services/api.ts

@@ -91,6 +91,10 @@ export class Api {
       case 'framework-check-model-exists':
       case 'framework-preferences':
       case 'framework-model-cte-analysis':
+      case 'framework-python-model-create':
+      case 'framework-dag-create':
+      case 'framework-get-available-dags':
+      case 'framework-get-python-model-groups':
         return (await this.framework.handleApi(
           payload as any,
         )) as ApiResponse<T>;

src/services/framework/utils/python-model-utils.ts

@@ -475,7 +475,7 @@ function makeMarkdownCell(
  */
 function isNotebookOnlyCell(cell: PythonModelCell): boolean {
   const meta = cell.metadata;
-  if (meta && meta.dj_notebook_only === true) {
+  if (meta?.dj_notebook_only === true) {
     return true;
   }
   return cellText(cell).startsWith('# Python Model:');

src/services/modelLineage.ts

@@ -468,22 +468,9 @@ export class ModelLineage {
               .filter(Boolean)
           : [];
 
-        const pythonModelTable = props['python_model_table'] ?? pythonModelName;
-        const pythonModelId = `python.${pythonModelTable}`;
-        const pythonModelNode: LineageNode = {
-          id: pythonModelId,
-          name: pythonModelTable,
-          type: 'python',
-          description:
-            props['python_model_description'] ??
-            `Python model: ${pythonModelName}`,
-          tags: ['python'],
-          path: '',
-          schema,
-          database: catalog,
-          hasOwnUpstream: upstreamSources.length > 0,
-          hasOwnDownstream: true,
-        };
+        if (upstreamSources.length === 0) {
+          return null;
+        }
 
         const upstreamSourceNodes: LineageNode[] = [];
         for (const sourceId of upstreamSources) {
@@ -506,8 +493,11 @@ export class ModelLineage {
           }
         }
 
+        if (upstreamSourceNodes.length === 0) {
+          return null;
+        }
+
         return {
-          pythonModelNode,
           sourceNodeId: sourceNode.id,
           upstreamSourceNodes,
         };
@@ -522,20 +512,15 @@ export class ModelLineage {
     const results = await Promise.all(queries);
 
     for (const result of results) {
-      if (result) {
-        pythonModelNodes.push(result.pythonModelNode);
+      if (!result) {
+        continue;
+      }
+      for (const upstreamNode of result.upstreamSourceNodes) {
+        pythonModelNodes.push(upstreamNode);
         pythonModelEdges.push({
-          pythonModelNodeId: result.pythonModelNode.id,
+          pythonModelNodeId: upstreamNode.id,
           sourceNodeId: result.sourceNodeId,
         });
-
-        for (const upstreamNode of result.upstreamSourceNodes) {
-          pythonModelNodes.push(upstreamNode);
-          pythonModelEdges.push({
-            pythonModelNodeId: upstreamNode.id,
-            sourceNodeId: result.pythonModelNode.id,
-          });
-        }
       }
     }
 

web/src/pages/ModelLineage/LineageGraph.tsx

@@ -13,6 +13,7 @@ import {
 } from '@xyflow/react';
 import { useCallback, useEffect, useRef } from 'react';
 
+import type { PythonModelEdge } from '../../stores/dataExplorerStore';
 import { useDataExplorerStore } from '../../stores/dataExplorerStore';
 import LightdashNode from './LightdashNode';
 import ModelNode from './ModelNode';
@@ -42,6 +43,7 @@ interface LineageGraphProps {
   currentNode: LineageNode;
   upstreamNodes: LineageNode[];
   downstreamNodes: LineageNode[];
+  pythonModelEdges?: PythonModelEdge[];
   lightdashDownstream?: LightdashLineageNode[];
   projectName: string;
   selectedNodeForQuery: string | null;
@@ -213,6 +215,7 @@ export default function LineageGraph({
   currentNode,
   upstreamNodes,
   downstreamNodes,
+  pythonModelEdges,
   lightdashDownstream,
   projectName,
   selectedNodeForQuery,
@@ -287,6 +290,18 @@ export default function LineageGraph({
       ...additionalNodes,
     ];
 
+    // Redirect map for Python-model upstream sources: child node id -> the
+    const upstreamEdgeTargets = new Map<string, string[]>();
+    (pythonModelEdges ?? []).forEach(({ pythonModelNodeId, sourceNodeId }) => {
+      const existing = upstreamEdgeTargets.get(pythonModelNodeId);
+      if (existing) {
+        existing.push(sourceNodeId);
+      } else {
+        upstreamEdgeTargets.set(pythonModelNodeId, [sourceNodeId]);
+      }
+    });
+    const knownNodeIds = new Set<string>(allLineageNodes.map((n) => n.id));
+
     // Check model outdated status for all models
     const checkStatuses = async () => {
       const statusPromises = allLineageNodes
@@ -396,21 +411,29 @@ export default function LineageGraph({
       };
       newNodes.push(flowNode);
 
-      const edge: Edge = {
-        id: `${node.id}-${currentNode.id}`,
-        source: node.id,
-        target: currentNode.id,
-        sourceHandle: 'output',
-        targetHandle: 'input',
-        style: edgeStyle,
-        markerEnd: {
-          type: MarkerType.ArrowClosed,
-          width: 20,
-          height: 20,
-          color: EDGE_COLOR,
-        },
-      };
-      newEdges.push(edge);
+      const redirectTargets = (upstreamEdgeTargets.get(node.id) ?? []).filter(
+        (targetId) => knownNodeIds.has(targetId),
+      );
+      const edgeTargets =
+        redirectTargets.length > 0 ? redirectTargets : [currentNode.id];
+
+      edgeTargets.forEach((targetId) => {
+        const edge: Edge = {
+          id: `${node.id}-${targetId}`,
+          source: node.id,
+          target: targetId,
+          sourceHandle: 'output',
+          targetHandle: 'input',
+          style: edgeStyle,
+          markerEnd: {
+            type: MarkerType.ArrowClosed,
+            width: 20,
+            height: 20,
+            color: EDGE_COLOR,
+          },
+        };
+        newEdges.push(edge);
+      });
     });
 
     // Create downstream nodes and edges
@@ -473,6 +496,13 @@ export default function LineageGraph({
       // Check if already added
       if (newNodes.some((n) => n.id === node.id)) return;
 
+      const hasVisibleUpstreamEdge = additionalEdges.some(
+        (e) => e.target === node.id,
+      );
+      const hasVisibleDownstreamEdge = additionalEdges.some(
+        (e) => e.source === node.id,
+      );
+
       const flowNode: Node<ModelNodeData> = {
         id: node.id,
         type: 'lineageNode',
@@ -492,8 +522,10 @@ export default function LineageGraph({
           // Use backend values to determine if node has its own upstream/downstream
           hasUpstream: node.hasOwnUpstream === true,
           hasDownstream: node.hasOwnDownstream === true,
-          isUpstreamExpanded: isNodeUpstreamExpanded(node.id),
-          isDownstreamExpanded: isNodeDownstreamExpanded(node.id),
+          isUpstreamExpanded:
+            isNodeUpstreamExpanded(node.id) || hasVisibleUpstreamEdge,
+          isDownstreamExpanded:
+            isNodeDownstreamExpanded(node.id) || hasVisibleDownstreamEdge,
           onRun: onRunQuery,
           onCompile: handleCompile,
           onCompileAndRun: handleCompileAndRun,
@@ -584,6 +616,7 @@ export default function LineageGraph({
     currentNode,
     upstreamNodes,
     downstreamNodes,
+    pythonModelEdges,
     lightdashDownstream,
     projectName,
     onRunQuery,

web/src/pages/ModelLineage/ModelLineage.tsx

@@ -844,6 +844,7 @@ export default function ModelLineage({ onShowAdhocQuery }: ModelLineageProps) {
                 currentNode={lineageData.current}
                 upstreamNodes={lineageData.upstream}
                 downstreamNodes={lineageData.downstream}
+                pythonModelEdges={lineageData.pythonModelEdges}
                 lightdashDownstream={lineageData.lightdashDownstream}
                 projectName={activeModel.projectName}
                 selectedNodeForQuery={selectedNodeForQuery}