feat: add DBT_LOOM_CONFIG_PATH environment variable support (#1761)

Co-authored-by: Claude <[email protected]>

Author: mdesmet

src/dbt_client/dbtCoreIntegration.ts

@@ -280,18 +280,14 @@ export class DBTCoreProjectIntegration
       "DBTCoreProjectIntegration",
       `Registering dbt core project at ${this.projectRoot}`,
     );
-    this.python = this.executionInfrastructure.createPythonBridge(
-      this.projectRoot.fsPath,
-    );
+    this.python = this.createPythonBridge();
     this.executionInfrastructure.createQueue(
       DBTCoreProjectIntegration.QUEUE_ALL,
     );
 
     this.disposables.push(
       this.pythonEnvironment.onPythonEnvironmentChanged(() => {
-        this.python = this.executionInfrastructure.createPythonBridge(
-          this.projectRoot.fsPath,
-        );
+        this.python = this.createPythonBridge();
       }),
       this.rebuildManifestDiagnostics,
       this.pythonBridgeDiagnostics,
@@ -409,6 +405,24 @@ export class DBTCoreProjectIntegration
     );
   }
 
+  private createPythonBridge(): PythonBridge {
+    const env: { [key: string]: string | undefined } = {};
+
+    // Check if dbt_loom.config.yml exists and add it to environment
+    const dbtLoomConfigPath = path.join(
+      this.projectRoot.fsPath,
+      "dbt_loom.config.yml",
+    );
+    if (existsSync(dbtLoomConfigPath)) {
+      env.DBT_LOOM_CONFIG_PATH = dbtLoomConfigPath;
+    }
+
+    return this.executionInfrastructure.createPythonBridge(
+      this.projectRoot.fsPath,
+      env,
+    );
+  }
+
   async executeSQL(
     query: string,
     limit: number,
@@ -417,9 +431,7 @@ export class DBTCoreProjectIntegration
     this.throwBridgeErrorIfAvailable();
     const { limitQuery } = await this.getQuery(query, limit);
 
-    const queryThread = this.executionInfrastructure.createPythonBridge(
-      this.projectRoot.fsPath,
-    );
+    const queryThread = this.createPythonBridge();
     return new QueryExecution(
       async () => {
         queryThread.kill(2);
@@ -1185,9 +1197,7 @@ export class DBTCoreProjectIntegration
     configPath,
   }: HealthcheckArgs): Promise<ProjectHealthcheck> {
     this.throwBridgeErrorIfAvailable();
-    const healthCheckThread = this.executionInfrastructure.createPythonBridge(
-      this.projectRoot.fsPath,
-    );
+    const healthCheckThread = this.createPythonBridge();
     try {
       await this.createPythonDbtProject(healthCheckThread);
       await healthCheckThread.ex`from dbt_healthcheck import *`;

src/dbt_client/dbtIntegration.ts

@@ -413,7 +413,10 @@ export class DBTCommandExecutionInfrastructure {
     private terminal: DBTTerminal,
   ) {}
 
-  createPythonBridge(cwd: string): PythonBridge {
+  createPythonBridge(
+    cwd: string,
+    env?: { [key: string]: string | undefined },
+  ): PythonBridge {
     let pythonPath = this.pythonEnvironment.pythonPath;
     const envVars = this.pythonEnvironment.environmentVariables;
 
@@ -440,6 +443,7 @@ export class DBTCommandExecutionInfrastructure {
       python: pythonPath,
       cwd: cwd,
       env: {
+        ...env,
         ...envVars,
         PYTHONPATH: __dirname,
       },

src/test/suite/dbtCoreIntegration.test.ts

@@ -7,6 +7,7 @@ import { Container } from "inversify";
 import {
   CLIDBTCommandExecutionStrategy,
   PythonDBTCommandExecutionStrategy,
+  DBTCommandExecutionInfrastructure,
 } from "../../dbt_client/dbtIntegration";
 import { TelemetryService } from "../../telemetry";
 import { DBTProjectContainer } from "../../manifest/dbtProjectContainer";
@@ -18,6 +19,9 @@ import {
   CommandProcessExecutionFactory,
   CommandProcessExecution,
 } from "../../commandProcessExecution";
+import * as fs from "fs";
+import * as path from "path";
+import * as os from "os";
 
 // Mock workspace folders
 jest.mock("vscode", () => ({
@@ -256,3 +260,132 @@ describe("DBTCoreDetection Tests", () => {
     });
   });
 });
+
+describe("DBTCoreProjectIntegration - dbt-loom config", () => {
+  let mockExecutionInfrastructure: jest.Mocked<DBTCommandExecutionInfrastructure>;
+  let mockTelemetry: jest.Mocked<TelemetryService>;
+  let mockPythonBridge: jest.Mocked<any>;
+  let tempDir: string;
+
+  beforeEach(() => {
+    // Create a temporary directory for testing
+    tempDir = fs.mkdtempSync(path.join(os.tmpdir(), "dbt-test-"));
+
+    // Create mock dependencies
+    mockPythonBridge = {
+      lock: jest.fn(),
+      ex: jest.fn(),
+      connected: true,
+    };
+
+    mockTelemetry = {
+      sendTelemetryEvent: jest.fn(),
+      sendTelemetryError: jest.fn(),
+      setTelemetryCustomAttribute: jest.fn(),
+    } as any;
+
+    mockExecutionInfrastructure = {
+      createPythonBridge: jest.fn().mockReturnValue(mockPythonBridge),
+      createQueue: jest.fn(),
+    } as any;
+  });
+
+  afterEach(() => {
+    // Clean up temporary directory
+    if (tempDir && fs.existsSync(tempDir)) {
+      fs.rmSync(tempDir, { recursive: true, force: true });
+    }
+    jest.clearAllMocks();
+  });
+
+  it("should set DBT_LOOM_CONFIG_PATH when dbt_loom.config.yml exists", () => {
+    // Arrange
+    const projectRoot = Uri.file(tempDir);
+    const dbtLoomConfigPath = path.join(tempDir, "dbt_loom.config.yml");
+
+    // Create the dbt_loom.config.yml file
+    fs.writeFileSync(dbtLoomConfigPath, "# Test config file");
+
+    const mockDiagnosticCollection =
+      languages.createDiagnosticCollection("dbt-project-config");
+    const mockPythonEnvChangeEmitter = new EventEmitter<Uri | undefined>();
+
+    // Act
+    new DBTCoreProjectIntegration(
+      mockExecutionInfrastructure as any,
+      {
+        onPythonEnvironmentChanged: mockPythonEnvChangeEmitter.event,
+        pythonPath: "/usr/bin/python3",
+        environmentVariables: {},
+      } as any,
+      mockTelemetry,
+      {} as PythonDBTCommandExecutionStrategy,
+      {} as (
+        projectRoot: Uri,
+        dbtPath: string,
+      ) => CLIDBTCommandExecutionStrategy,
+      {} as DBTProjectContainer,
+      {} as AltimateRequest,
+      {
+        debug: jest.fn(),
+        error: jest.fn(),
+        log: jest.fn(),
+      } as any,
+      {} as ValidationProvider,
+      {} as DeferToProdService,
+      projectRoot,
+      mockDiagnosticCollection,
+    );
+
+    // Assert
+    expect(mockExecutionInfrastructure.createPythonBridge).toHaveBeenCalledWith(
+      tempDir,
+      {
+        DBT_LOOM_CONFIG_PATH: dbtLoomConfigPath,
+      },
+    );
+  });
+
+  it("should not set DBT_LOOM_CONFIG_PATH when dbt_loom.config.yml does not exist", () => {
+    // Arrange
+    const projectRoot = Uri.file(tempDir);
+    // Explicitly NOT creating the dbt_loom.config.yml file
+
+    const mockDiagnosticCollection =
+      languages.createDiagnosticCollection("dbt-project-config");
+    const mockPythonEnvChangeEmitter = new EventEmitter<Uri | undefined>();
+
+    // Act
+    new DBTCoreProjectIntegration(
+      mockExecutionInfrastructure as any,
+      {
+        onPythonEnvironmentChanged: mockPythonEnvChangeEmitter.event,
+        pythonPath: "/usr/bin/python3",
+        environmentVariables: {},
+      } as any,
+      mockTelemetry,
+      {} as PythonDBTCommandExecutionStrategy,
+      {} as (
+        projectRoot: Uri,
+        dbtPath: string,
+      ) => CLIDBTCommandExecutionStrategy,
+      {} as DBTProjectContainer,
+      {} as AltimateRequest,
+      {
+        debug: jest.fn(),
+        error: jest.fn(),
+        log: jest.fn(),
+      } as any,
+      {} as ValidationProvider,
+      {} as DeferToProdService,
+      projectRoot,
+      mockDiagnosticCollection,
+    );
+
+    // Assert
+    expect(mockExecutionInfrastructure.createPythonBridge).toHaveBeenCalledWith(
+      tempDir,
+      {}, // Empty env object when file doesn't exist
+    );
+  });
+});