From 9b8faec48519e7b68134d44c8a5acfdf7270172c Mon Sep 17 00:00:00 2001
From: Fabian Jakobs <fabian.jakobs@databricks.com>
Date: Tue, 25 Jul 2023 10:14:25 +0200
Subject: [PATCH] Add DABs code completion for includes

---
 packages/databricks-vscode/package.json       |  3 +-
 .../src/bundle/GenerateBundle.ts              | 67 +++++++++++++++----
 packages/databricks-vscode/src/extension.ts   |  6 +-
 yarn.lock                                     |  8 +++
 4 files changed, 69 insertions(+), 15 deletions(-)

diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json
index 6ce55baa0..fd708127f 100644
--- a/packages/databricks-vscode/package.json
+++ b/packages/databricks-vscode/package.json
@@ -688,7 +688,8 @@
         "ansi-to-html": "^0.7.2",
         "bcryptjs": "^2.4.3",
         "triple-beam": "^1.4.1",
-        "winston": "^3.10.0"
+        "winston": "^3.10.0",
+        "yaml": "^2.3.1"
     },
     "devDependencies": {
         "@istanbuljs/nyc-config-typescript": "^1.0.2",
diff --git a/packages/databricks-vscode/src/bundle/GenerateBundle.ts b/packages/databricks-vscode/src/bundle/GenerateBundle.ts
index 55a948ef0..38d54820c 100644
--- a/packages/databricks-vscode/src/bundle/GenerateBundle.ts
+++ b/packages/databricks-vscode/src/bundle/GenerateBundle.ts
@@ -1,8 +1,11 @@
 import {CliWrapper} from "../cli/CliWrapper";
 import {extensions, Uri} from "vscode";
-import path from "node:path";
+import {workspace, RelativePattern, GlobPattern, Disposable} from "vscode";
+import YAML from "yaml";
 
-export async function generateBundleSchema(cli: CliWrapper) {
+export async function generateBundleSchema(
+    cli: CliWrapper
+): Promise<Disposable> {
     // get freshly generated bundle schema
     const bundleSchema = await cli.getBundleSchema();
 
@@ -12,6 +15,51 @@ export async function generateBundleSchema(cli: CliWrapper) {
     // URI for bundle root config json schema
     const rootConfigSchemaUri = `${dabsUriScheme}:///root.json`;
 
+    const folder = workspace.workspaceFolders?.[0];
+    const configFilePattern = new RelativePattern(
+        folder!,
+        "{databricks,bundle}.{yml,yaml}"
+    );
+
+    // file watcher on all YAML files
+    const watcher = workspace.createFileSystemWatcher("**/*.{yml,yaml}");
+    watcher.onDidChange(async () => {
+        await updateFileGlobs();
+    });
+
+    let configFiles = new Set<string>();
+    await updateFileGlobs();
+
+    async function updateFileGlobs() {
+        const fileGlobs: GlobPattern[] = [configFilePattern];
+
+        // find all YAML files that are included in the root config file
+        for (const configFile of await workspace.findFiles(configFilePattern)) {
+            try {
+                const fileContents = await workspace.fs.readFile(configFile);
+                const config = YAML.parse(fileContents.toString());
+                if (config.include) {
+                    fileGlobs.push(
+                        ...config.include.map(
+                            (g: string) => new RelativePattern(folder!, g)
+                        )
+                    );
+                }
+            } catch (e) {
+                // ignore errors
+            }
+        }
+
+        // expand globs to find all config files
+        const newConfigFiles = new Set<string>();
+        for (const glob of fileGlobs) {
+            for (const file of await workspace.findFiles(glob)) {
+                newConfigFiles.add(file.path);
+            }
+        }
+        configFiles = newConfigFiles;
+    }
+
     const extensionYaml = extensions.getExtension("redhat.vscode-yaml");
     if (extensionYaml) {
         const redHatYamlSchemaApi = await extensionYaml.activate();
@@ -21,16 +69,9 @@ export async function generateBundleSchema(cli: CliWrapper) {
         redHatYamlSchemaApi.registerContributor(
             "dabs",
             (resource: string) => {
-                const validFileNames: string[] = [
-                    "databricks.yml",
-                    "databricks.yaml",
-                    "bundle.yml",
-                    "bundle.yaml",
-                ];
-                for (const name of validFileNames) {
-                    if (path.basename(resource) === name) {
-                        return rootConfigSchemaUri;
-                    }
+                const resourceUri = Uri.parse(resource);
+                if (configFiles.has(resourceUri.path)) {
+                    return rootConfigSchemaUri;
                 }
                 return undefined;
             },
@@ -43,4 +84,6 @@ export async function generateBundleSchema(cli: CliWrapper) {
             }
         );
     }
+
+    return watcher;
 }
diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts
index daeef7435..6135d2483 100644
--- a/packages/databricks-vscode/src/extension.ts
+++ b/packages/databricks-vscode/src/extension.ts
@@ -537,12 +537,14 @@ export async function activate(
 
     // generate a json schema for bundle root and load a custom provider into
     // redhat.vscode-yaml extension to validate bundle config files with this schema
-    generateBundleSchema(cli).catch((e) => {
+    try {
+        context.subscriptions.push(await generateBundleSchema(cli));
+    } catch (e) {
         NamedLogger.getOrCreate("Extension").error(
             "Failed to load bundle schema: ",
             e
         );
-    });
+    }
 
     connectionManager.login(false).catch((e) => {
         NamedLogger.getOrCreate(Loggers.Extension).error("Login error", e);
diff --git a/yarn.lock b/yarn.lock
index b3bf484b7..2bfd9c4b7 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -3715,6 +3715,7 @@ __metadata:
     wdio-video-reporter: ^4.0.3
     wdio-vscode-service: ^5.2.0
     winston: ^3.10.0
+    yaml: ^2.3.1
     yargs: ^17.7.2
   languageName: unknown
   linkType: soft
@@ -11198,6 +11199,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"yaml@npm:^2.3.1":
+  version: 2.3.1
+  resolution: "yaml@npm:2.3.1"
+  checksum: 2c7bc9a7cd4c9f40d3b0b0a98e370781b68b8b7c4515720869aced2b00d92f5da1762b4ffa947f9e795d6cd6b19f410bd4d15fdd38aca7bd96df59bd9486fb54
+  languageName: node
+  linkType: hard
+
 "yargs-parser@npm:20.2.4":
   version: 20.2.4
   resolution: "yargs-parser@npm:20.2.4"