Support target selector (#1424)

Author: mdesmet

dbt_core_integration.py

@@ -344,7 +344,7 @@ class DbtProject:
 
     def __init__(
         self,
-        target: Optional[str] = None,
+        target_name: Optional[str] = None,
         profiles_dir: Optional[str] = None,
         project_dir: Optional[str] = None,
         threads: Optional[int] = 1,
@@ -357,7 +357,7 @@ def __init__(
     ):
         self.args = ConfigInterface(
             threads=threads,
-            target=target,
+            target=target_name,
             profiles_dir=profiles_dir,
             project_dir=project_dir,
             profile=profile,
@@ -398,7 +398,9 @@ def init_config(self):
             set_invocation_context(os.environ)
             set_from_args(self.args, None)
             # Copy over global_flags
-            self.args.__dict__.update(get_flags().__dict__)
+            for key, value in get_flags().__dict__.items():
+                if key not in self.args.__dict__:
+                    self.args.__dict__[key] = value
         else:
             set_from_args(self.args, self.args)
         self.config = RuntimeConfig.from_args(self.args)
@@ -882,3 +884,15 @@ def validate_sql_dry_run(self, compiled_sql: str):
             return self.adapter.validate_sql(compiled_sql)
         except Exception as e:
             raise Exception(str(e))
+
+    def get_target_names(self):
+        from dbt.config.profile import read_profile
+        profile = read_profile(self.args.profiles_dir)
+        project = profile[self.project_name]
+        if "outputs" in project:
+            outputs = project["outputs"]
+            return outputs.keys()
+        return []
+    
+    def set_selected_target(self, target: str):
+        self.args.target = target

src/commands/index.ts

@@ -15,6 +15,7 @@ import {
   ProgressLocation,
   TextEditorDecorationType,
   DecorationRangeBehavior,
+  QuickPickItem,
 } from "vscode";
 import { SqlPreviewContentProvider } from "../content_provider/sqlPreviewContentProvider";
 import { RunModelType } from "../domain";
@@ -568,6 +569,44 @@ export class VSCodeCommands implements Disposable {
           this.notebookController.createNotebook(args);
         },
       ),
+      commands.registerCommand(
+        "dbtPowerUser.openTargetSelector",
+        async (targets, project: DBTProject, statusBar) => {
+          try {
+            if (!targets) {
+              return;
+            }
+            this.dbtTerminal.debug(
+              "OpenTargetSelector",
+              "Showing following targets",
+              targets,
+            );
+            const target = await window.showQuickPick(targets, {
+              title: "Select your target",
+              canPickMany: false,
+            });
+            if (target) {
+              await project.setSelectedTarget(target);
+              await statusBar.updateStatusBar();
+              this.dbtTerminal.info(
+                "OpenTargetSelector",
+                "Selecting target",
+                true,
+                target,
+              );
+            }
+          } catch (error) {
+            this.dbtTerminal.error(
+              "OpenTargetSelector",
+              "An error occurred while changing target",
+              error,
+            );
+            window.showErrorMessage(
+              "An error occurred while changing target: " + error,
+            );
+          }
+        },
+      ),
       commands.registerCommand(
         "dbtPowerUser.createSqlFile",
         async (args: { code?: string; fileName?: string } | undefined) => {

src/dbt_client/dbtCloudIntegration.ts

@@ -336,6 +336,18 @@ export class DBTCloudProjectIntegration
     this.dbtPath = getDBTPath(this.pythonEnvironment, this.terminal);
   }
 
+  async setSelectedTarget(targetName: string): Promise<void> {
+    throw new Error("Method not implemented.");
+  }
+
+  async getTargetNames(): Promise<Array<string>> {
+    throw new Error("Method not implemented.");
+  }
+
+  getSelectedTarget(): string | undefined {
+    throw new Error("Method not implemented.");
+  }
+
   getTargetPath(): string | undefined {
     return this.targetPath;
   }

src/dbt_client/dbtCoreIntegration.ts

@@ -241,6 +241,7 @@ export class DBTCoreProjectIntegration
   private profilesDir?: string;
   private targetPath?: string;
   private adapterType?: string;
+  private targetName?: string;
   private version?: number[];
   private packagesInstallPath?: string;
   private modelPaths?: string[];
@@ -370,6 +371,7 @@ export class DBTCoreProjectIntegration
   async refreshProjectConfig(): Promise<void> {
     await this.createPythonDbtProject(this.python);
     await this.python.ex`project.init_project()`;
+    this.targetName = await this.findSelectedTarget();
     this.targetPath = await this.findTargetPath();
     this.modelPaths = await this.findModelPaths();
     this.seedPaths = await this.findSeedPaths();
@@ -379,6 +381,25 @@ export class DBTCoreProjectIntegration
     this.adapterType = await this.findAdapterType();
   }
 
+  async findSelectedTarget(): Promise<string> {
+    return await this.python.lock(
+      (python) => python`to_dict(project.config.target_name)`,
+    );
+  }
+
+  async setSelectedTarget(targetName: string): Promise<void> {
+    await this.python.lock(
+      (python) => python`project.set_selected_target(${targetName})`,
+    );
+    await this.refreshProjectConfig();
+  }
+
+  async getTargetNames(): Promise<Array<string>> {
+    return await this.python.lock(
+      (python) => python`to_dict(project.get_target_names())`,
+    );
+  }
+
   async executeSQL(
     query: string,
     limit: number,
@@ -507,6 +528,10 @@ export class DBTCoreProjectIntegration
     }
   }
 
+  getSelectedTarget() {
+    return this.targetName;
+  }
+
   getTargetPath(): string | undefined {
     return this.targetPath;
   }
@@ -798,6 +823,10 @@ export class DBTCoreProjectIntegration
       command.addArgument("--profiles-dir");
       command.addArgument(this.profilesDir);
     }
+    if (this.targetName) {
+      command.addArgument("--target");
+      command.addArgument(this.targetName);
+    }
     command.setExecutionStrategy(this.pythonDBTCommandExecutionStrategy);
     return command;
   }

src/dbt_client/dbtIntegration.ts

@@ -325,6 +325,9 @@ export interface DBTProjectIntegration extends Disposable {
   initializeProject(): Promise<void>;
   // called when project configuration is changed
   refreshProjectConfig(): Promise<void>;
+  // Change target
+  setSelectedTarget(targetName: string): Promise<void>;
+  getTargetNames(): Promise<Array<string>>;
   // retrieve dbt configs
   getTargetPath(): string | undefined;
   getModelPaths(): string[] | undefined;
@@ -333,6 +336,7 @@ export interface DBTProjectIntegration extends Disposable {
   getPackageInstallPath(): string | undefined;
   getAdapterType(): string | undefined;
   getVersion(): number[] | undefined;
+  getSelectedTarget(): string | undefined;
   // parse manifest
   rebuildManifest(): Promise<void>;
   // execute queries