From ef59e410f8f52e99a2d6be2bcecab9496277375f Mon Sep 17 00:00:00 2001
From: Lars Reimann <mail@larsreimann.com>
Date: Mon, 22 Apr 2024 21:26:26 +0200
Subject: [PATCH] fix: import error when exploring tables in pipelines that
 call segments

---
 packages/safe-ds-lang/src/language/index.ts     |  7 ++++++-
 .../lsp/safe-ds-rename-provider.test.ts         |  9 +++------
 .../src/extension/eda/apis/runnerApi.ts         | 17 +++++++++++------
 3 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/packages/safe-ds-lang/src/language/index.ts b/packages/safe-ds-lang/src/language/index.ts
index 030392ddc..f6afcd620 100644
--- a/packages/safe-ds-lang/src/language/index.ts
+++ b/packages/safe-ds-lang/src/language/index.ts
@@ -20,10 +20,15 @@ export { locationToString, positionToString, rangeToString } from '../helpers/lo
 // Messages
 export * as messages from './runtime/messages.js';
 
-// Constants
+// Commands
 export * as commands from './communication/commands.js';
+
+// RPC
 export * as rpc from './communication/rpc.js';
 
+// Generation
+export { CODEGEN_PREFIX } from './generation/safe-ds-python-generator.js';
+
 // Dependencies
 export const dependencies = {
     'safe-ds-runner': {
diff --git a/packages/safe-ds-lang/tests/language/lsp/safe-ds-rename-provider.test.ts b/packages/safe-ds-lang/tests/language/lsp/safe-ds-rename-provider.test.ts
index 84ace157f..29f9b3ae6 100644
--- a/packages/safe-ds-lang/tests/language/lsp/safe-ds-rename-provider.test.ts
+++ b/packages/safe-ds-lang/tests/language/lsp/safe-ds-rename-provider.test.ts
@@ -9,7 +9,6 @@ import { TextDocument } from 'vscode-languageserver-textdocument';
 const services = (await createSafeDsServices(NodeFileSystem)).SafeDs;
 const documentBuilder = services.shared.workspace.DocumentBuilder;
 const langiumDocuments = services.shared.workspace.LangiumDocuments;
-const langiumDocumentFactory = services.shared.workspace.LangiumDocumentFactory;
 const nameProvider = services.references.NameProvider;
 const renameProvider = services.lsp.RenameProvider!;
 
@@ -258,11 +257,9 @@ describe('SafeDsRenameProvider', async () => {
         }
 
         // Add documents to workspace
-        const documents = descriptions.map((description) => {
-            const document = langiumDocumentFactory.fromString(description.originalContent, URI.parse(description.uri));
-            langiumDocuments.addDocument(document);
-            return document;
-        });
+        const documents = descriptions.map((description) =>
+            langiumDocuments.createDocument(URI.parse(description.uri), description.originalContent),
+        );
         await documentBuilder.build(documents);
 
         // Get the element to rename
diff --git a/packages/safe-ds-vscode/src/extension/eda/apis/runnerApi.ts b/packages/safe-ds-vscode/src/extension/eda/apis/runnerApi.ts
index 466b8f07b..a15ac4dd7 100644
--- a/packages/safe-ds-vscode/src/extension/eda/apis/runnerApi.ts
+++ b/packages/safe-ds-vscode/src/extension/eda/apis/runnerApi.ts
@@ -1,10 +1,9 @@
 import { Base64Image, Column, Profiling, ProfilingDetailStatistical, Table } from '@safe-ds/eda/types/state.js';
-import { ast, messages, SafeDsServices } from '@safe-ds/lang';
-import { AstNode, LangiumDocument } from 'langium';
+import { ast, CODEGEN_PREFIX, messages, SafeDsServices } from '@safe-ds/lang';
+import { LangiumDocument } from 'langium';
 import * as vscode from 'vscode';
 import crypto from 'crypto';
 import { getPipelineDocument } from '../../mainClient.ts';
-import { CODEGEN_PREFIX } from '../../../../../safe-ds-lang/src/language/generation/safe-ds-python-generator.ts';
 import { safeDsLogger } from '../../helpers/logging.js';
 
 export class RunnerApi {
@@ -12,7 +11,7 @@ export class RunnerApi {
     pipelinePath: vscode.Uri;
     pipelineName: string;
     pipelineNode: ast.SdsPipeline;
-    baseDocument: LangiumDocument<AstNode> | undefined;
+    baseDocument: LangiumDocument | undefined;
     placeholderCounter = 0;
 
     constructor(
@@ -52,12 +51,18 @@ export class RunnerApi {
             const afterPipelineEnd = documentText.substring(endOfPipeline - 1);
             const newDocumentText = beforePipelineEnd + addedLines + afterPipelineEnd;
 
-            const newDoc = this.services.shared.workspace.LangiumDocumentFactory.fromString(
-                newDocumentText,
+            this.services.shared.workspace.LangiumDocuments.deleteDocument(this.pipelinePath);
+            const newDoc = this.services.shared.workspace.LangiumDocuments.createDocument(
                 this.pipelinePath,
+                newDocumentText,
             );
+            await this.services.shared.workspace.DocumentBuilder.build([newDoc]);
+
             await this.services.runtime.Runner.executePipeline(pipelineExecutionId, newDoc, this.pipelineName);
 
+            this.services.shared.workspace.LangiumDocuments.deleteDocument(this.pipelinePath);
+            this.services.shared.workspace.LangiumDocuments.addDocument(this.baseDocument);
+
             const runtimeCallback = (message: messages.RuntimeProgressMessage) => {
                 if (message.id !== pipelineExecutionId) {
                     return;