Move sdk/services from the SDK to VS Code (#852)

## Changes - Move sdk/services from the SDK to VS Code - Get rid of most `@databircks/databricks-sdk/dist/...` imports. Only `context` is still there, which can be cleaned up after the next SDK release. ## Testing - [x] make sure integ tests are run - [x] make sure unit tests are run
databricks · Sep 1, 2023 · f544d6b · f544d6b
1 parent 3b0cadc
commit f544d6b
Show file tree

Hide file tree

Showing 76 changed files with 2,803 additions and 230 deletions.
diff --git a/packages/databricks-vscode-types/index.ts b/packages/databricks-vscode-types/index.ts
@@ -1,4 +1,5 @@
-import {ApiClient, WorkspaceClient, Cluster} from "@databricks/databricks-sdk";
+import {ApiClient, WorkspaceClient} from "@databricks/databricks-sdk";
+import {Cluster} from "databricks/src/sdk-extensions";
 import {Event} from "vscode";
 export type ConnectionState = "CONNECTED" | "CONNECTING" | "DISCONNECTED";
 

diff --git a/packages/databricks-vscode-types/package.json b/packages/databricks-vscode-types/package.json
@@ -29,6 +29,7 @@
         "typescript": "~5.1.6"
     },
     "dependencies": {
-        "@databricks/databricks-sdk": "../../vendor/databricks-sdk.tgz"
+        "@databricks/databricks-sdk": "../../vendor/databricks-sdk.tgz",
+        "databricks": "*"
     }
 }
diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json
@@ -677,7 +677,8 @@
         "test:unit": "yarn run build && node ./out/test/runTest.js",
         "test:integ:prepare": "yarn run package",
         "test:integ:run": "wdio run src/test/e2e/wdio.conf.ts",
-        "test:integ": "yarn run test:integ:prepare && yarn run test:integ:run",
+        "test:integ:sdk": "ts-mocha --type-check 'src/**/*.integ.ts'",
+        "test:integ": "yarn run test:integ:prepare && yarn run test:integ:run && yarn run test:integ:sdk",
         "test:cov": "nyc yarn run test:unit",
         "test": "yarn run test:lint && yarn run test:unit",
         "clean": "rm -rf node_modules out .vscode-test",
@@ -727,6 +728,7 @@
         "nyc": "^15.1.0",
         "prettier": "^3.0.0",
         "tmp-promise": "^3.0.3",
+        "ts-mocha": "^10.0.0",
         "ts-mockito": "^2.6.1",
         "ts-node": "^10.9.1",
         "typescript": "~5.1.6",

diff --git a/packages/databricks-vscode/src/cli/CliWrapper.test.ts b/packages/databricks-vscode/src/cli/CliWrapper.test.ts
@@ -13,7 +13,7 @@ import {writeFile} from "node:fs/promises";
 import {CliWrapper} from "./CliWrapper";
 import path from "node:path";
 import {Context} from "@databricks/databricks-sdk/dist/context";
-import {NamedLogger} from "@databricks/databricks-sdk/dist/logging";
+import {logging} from "@databricks/databricks-sdk";
 
 const execFile = promisify(execFileCb);
 
@@ -152,15 +152,18 @@ host = example.com
             const profiles = await cli.listProfiles(
                 path,
                 new Context({
-                    logger: NamedLogger.getOrCreate("cli-wrapper-test", {
-                        factory: () => {
-                            return {
-                                log: (level, msg, meta) => {
-                                    logs.push({level, msg, meta});
-                                },
-                            };
-                        },
-                    }),
+                    logger: logging.NamedLogger.getOrCreate(
+                        "cli-wrapper-test",
+                        {
+                            factory: () => {
+                                return {
+                                    log: (level, msg, meta) => {
+                                        logs.push({level, msg, meta});
+                                    },
+                                };
+                            },
+                        }
+                    ),
                 })
             );
             assert.equal(profiles.length, 2);

diff --git a/packages/databricks-vscode/src/cli/CliWrapper.ts b/packages/databricks-vscode/src/cli/CliWrapper.ts
@@ -3,11 +3,12 @@ import {ExtensionContext, window, commands} from "vscode";
 import {SyncDestinationMapper} from "../sync/SyncDestination";
 import {workspaceConfigs} from "../vscode-objs/WorkspaceConfigs";
 import {promisify} from "node:util";
-import {withLogContext} from "@databricks/databricks-sdk/dist/logging";
+import {logging} from "@databricks/databricks-sdk";
 import {Loggers} from "../logger";
 import {Context, context} from "@databricks/databricks-sdk/dist/context";
 import {Cloud} from "../utils/constants";
 
+const withLogContext = logging.withLogContext;
 const execFile = promisify(execFileCb);
 
 export interface Command {

diff --git a/packages/databricks-vscode/src/cli/DatabricksCliSyncParser.ts b/packages/databricks-vscode/src/cli/DatabricksCliSyncParser.ts
@@ -1,13 +1,13 @@
-import {LEVELS, NamedLogger} from "@databricks/databricks-sdk/dist/logging";
+import {logging} from "@databricks/databricks-sdk";
 import {EventEmitter} from "vscode";
 import {Loggers} from "../logger";
 import {SyncState} from "../sync";
 
-const databricksLogLevelToSdk = new Map<string, LEVELS>([
-    ["DEBUG", LEVELS.debug],
-    ["INFO", LEVELS.info],
-    ["WARN", LEVELS.warn],
-    ["ERROR", LEVELS.error],
+const databricksLogLevelToSdk = new Map<string, logging.LEVELS>([
+    ["DEBUG", logging.LEVELS.debug],
+    ["INFO", logging.LEVELS.info],
+    ["WARN", logging.LEVELS.warn],
+    ["ERROR", logging.LEVELS.error],
 ]);
 
 type EventBase = {
@@ -96,7 +96,7 @@ export class DatabricksCliSyncParser {
 
     public processStderr(data: string) {
         const logLines = data.split("\n");
-        let currentLogLevel: LEVELS = LEVELS.info;
+        let currentLogLevel: logging.LEVELS = logging.LEVELS.info;
         for (let i = 0; i < logLines.length; i++) {
             const line = logLines[i].trim();
             if (line.length === 0) {
@@ -111,9 +111,13 @@ export class DatabricksCliSyncParser {
                     databricksLogLevelToSdk.get(typeMatch[1]) ??
                     currentLogLevel;
             }
-            NamedLogger.getOrCreate(Loggers.CLI).log(currentLogLevel, line, {
-                outfile: "stderr",
-            });
+            logging.NamedLogger.getOrCreate(Loggers.CLI).log(
+                currentLogLevel,
+                line,
+                {
+                    outfile: "stderr",
+                }
+            );
             this.writeEmitter.fire(line.trim() + "\r\n");
             if (this.matchForErrors(line)) {
                 return;
@@ -150,14 +154,14 @@ export class DatabricksCliSyncParser {
             if (line.length === 0) {
                 continue;
             }
-            NamedLogger.getOrCreate(Loggers.CLI).info(line, {
+            logging.NamedLogger.getOrCreate(Loggers.CLI).info(line, {
                 outfile: "stdout",
             });
 
             try {
                 this.processLine(line);
             } catch (error: any) {
-                NamedLogger.getOrCreate(Loggers.Extension).error(
+                logging.NamedLogger.getOrCreate(Loggers.Extension).error(
                     "Error parsing JSON line from databricks sync stdout: " +
                         error
                 );

diff --git a/packages/databricks-vscode/src/cli/SyncTasks.ts b/packages/databricks-vscode/src/cli/SyncTasks.ts
@@ -14,13 +14,15 @@ import {CliWrapper, Command, SyncType} from "./CliWrapper";
 import {ChildProcess, spawn, SpawnOptions} from "node:child_process";
 import {SyncState} from "../sync/CodeSynchronizer";
 import {DatabricksCliSyncParser} from "./DatabricksCliSyncParser";
-import {withLogContext} from "@databricks/databricks-sdk/dist/logging";
+import {logging} from "@databricks/databricks-sdk";
 import {Loggers} from "../logger";
 import {Context, context} from "@databricks/databricks-sdk/dist/context";
 import {PackageMetaData} from "../utils/packageJsonUtils";
 import {RWLock} from "../locking";
 import {EnvVarGenerators} from "../utils";
 
+const {withLogContext} = logging;
+
 export const TASK_SYNC_TYPE = {
     syncFull: "sync-full",
     sync: "sync",

diff --git a/packages/databricks-vscode/src/cluster/ClusterListDataProvider.test.ts b/packages/databricks-vscode/src/cluster/ClusterListDataProvider.test.ts
@@ -5,7 +5,8 @@ import {mock, when, instance} from "ts-mockito";
 import {ClusterModel} from "./ClusterModel";
 import {Disposable} from "vscode";
 import {ClusterListDataProvider} from "./ClusterListDataProvider";
-import {ApiClient, Cluster, compute} from "@databricks/databricks-sdk";
+import {ApiClient, compute} from "@databricks/databricks-sdk";
+import {Cluster} from "../sdk-extensions/Cluster";
 import {resolveProviderResult} from "../test/utils";
 
 const mockListClustersResponse: compute.ListClustersResponse = {

diff --git a/packages/databricks-vscode/src/cluster/ClusterListDataProvider.ts b/packages/databricks-vscode/src/cluster/ClusterListDataProvider.ts
@@ -9,7 +9,7 @@ import {
     TreeItem,
     TreeItemCollapsibleState,
 } from "vscode";
-import {Cluster} from "@databricks/databricks-sdk";
+import {Cluster} from "../sdk-extensions";
 import {ClusterModel} from "./ClusterModel";
 
 /**

diff --git a/packages/databricks-vscode/src/cluster/ClusterLoader.ts b/packages/databricks-vscode/src/cluster/ClusterLoader.ts
@@ -1,6 +1,6 @@
 /* eslint-disable @typescript-eslint/naming-convention */
-import {Cluster, Time, TimeUnits} from "@databricks/databricks-sdk";
-import {NamedLogger} from "@databricks/databricks-sdk/dist/logging";
+import {Time, TimeUnits, logging} from "@databricks/databricks-sdk";
+import {Cluster} from "../sdk-extensions";
 import {Disposable, Event, EventEmitter} from "vscode";
 import {ConnectionManager} from "../configuration/ConnectionManager";
 import {Loggers} from "../logger";
@@ -142,7 +142,9 @@ export class ClusterLoader implements Disposable {
                             resolve();
                         })
                         .catch((e) => {
-                            NamedLogger.getOrCreate(Loggers.Extension).error(
+                            logging.NamedLogger.getOrCreate(
+                                Loggers.Extension
+                            ).error(
                                 `Error fetching permission for cluster ${c.name}`,
                                 e
                             );
@@ -175,7 +177,7 @@ export class ClusterLoader implements Disposable {
             try {
                 await this._load();
             } catch (e) {
-                NamedLogger.getOrCreate(Loggers.Extension).error(
+                logging.NamedLogger.getOrCreate(Loggers.Extension).error(
                     "Error loading clusters",
                     e
                 );

diff --git a/packages/databricks-vscode/src/cluster/ClusterManager.test.ts b/packages/databricks-vscode/src/cluster/ClusterManager.test.ts
@@ -3,11 +3,11 @@ import {ClusterManager} from "./ClusterManager";
 import {
     ApiClient,
     compute,
-    Cluster,
     Time,
     TimeUnits,
     retries,
 } from "@databricks/databricks-sdk";
+import {Cluster} from "../sdk-extensions";
 import {ClusterFixture} from "@databricks/databricks-sdk/dist/test/fixtures";
 import {
     anything,

diff --git a/packages/databricks-vscode/src/cluster/ClusterManager.ts b/packages/databricks-vscode/src/cluster/ClusterManager.ts
@@ -1,4 +1,5 @@
-import {compute, Cluster, Time, TimeUnits} from "@databricks/databricks-sdk";
+import {compute, Time, TimeUnits} from "@databricks/databricks-sdk";
+import {Cluster} from "../sdk-extensions";
 import {CancellationTokenSource, Disposable} from "vscode";
 
 export class ClusterManager implements Disposable {

diff --git a/packages/databricks-vscode/src/cluster/ClusterModel.test.ts b/packages/databricks-vscode/src/cluster/ClusterModel.test.ts
@@ -2,12 +2,8 @@
 
 import assert from "assert";
 import {mock, when, anything, anyString, instance, spy} from "ts-mockito";
-import {
-    ApiClient,
-    Cluster,
-    WorkspaceClient,
-    compute,
-} from "@databricks/databricks-sdk";
+import {ApiClient, WorkspaceClient, compute} from "@databricks/databricks-sdk";
+import {Cluster} from "../sdk-extensions";
 import {ClusterModel} from "./ClusterModel";
 import {ConnectionManager} from "../configuration/ConnectionManager";
 import {Disposable} from "vscode";

diff --git a/packages/databricks-vscode/src/cluster/ClusterModel.ts b/packages/databricks-vscode/src/cluster/ClusterModel.ts
@@ -1,6 +1,7 @@
 /* eslint-disable @typescript-eslint/naming-convention */
 
-import {Cluster, compute} from "@databricks/databricks-sdk";
+import {compute} from "@databricks/databricks-sdk";
+import {Cluster} from "../sdk-extensions";
 import {Disposable, Event, EventEmitter} from "vscode";
 import {ConnectionManager} from "../configuration/ConnectionManager";
 import {ClusterLoader} from "./ClusterLoader";

diff --git a/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.test.ts b/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.test.ts
@@ -4,7 +4,8 @@ import assert from "assert";
 import {mock, when, instance, anything} from "ts-mockito";
 import {Disposable} from "vscode";
 import {ConfigurationDataProvider} from "./ConfigurationDataProvider";
-import {ApiClient, Cluster} from "@databricks/databricks-sdk";
+import {ApiClient} from "@databricks/databricks-sdk";
+import {Cluster} from "../sdk-extensions";
 import {ConnectionManager} from "./ConnectionManager";
 import {resolveProviderResult} from "../test/utils";
 import {SyncDestinationMapper} from "../sync/SyncDestination";

diff --git a/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts b/packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts
@@ -1,5 +1,6 @@
 /* eslint-disable @typescript-eslint/naming-convention */
-import {ApiClient, Cluster} from "@databricks/databricks-sdk";
+import {ApiClient} from "@databricks/databricks-sdk";
+import {Cluster} from "../sdk-extensions";
 import assert from "assert";
 import {mock} from "ts-mockito";
 import {formatQuickPickClusterDetails} from "./ConnectionCommands";

diff --git a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts
@@ -1,8 +1,4 @@
-import {
-    Cluster,
-    WorkspaceFsEntity,
-    WorkspaceFsUtils,
-} from "@databricks/databricks-sdk";
+import {Cluster, WorkspaceFsEntity, WorkspaceFsUtils} from "../sdk-extensions";
 import {homedir} from "node:os";
 import {
     Disposable,

diff --git a/packages/databricks-vscode/src/configuration/ConnectionManager.ts b/packages/databricks-vscode/src/configuration/ConnectionManager.ts
@@ -1,10 +1,5 @@
-import {
-    WorkspaceClient,
-    Cluster,
-    WorkspaceFsEntity,
-    WorkspaceFsUtils,
-    ApiClient,
-} from "@databricks/databricks-sdk";
+import {WorkspaceClient, ApiClient, logging} from "@databricks/databricks-sdk";
+import {Cluster, WorkspaceFsEntity, WorkspaceFsUtils} from "../sdk-extensions";
 import {
     env,
     EventEmitter,
@@ -26,12 +21,13 @@ import {
 import {configureWorkspaceWizard} from "./configureWorkspaceWizard";
 import {ClusterManager} from "../cluster/ClusterManager";
 import {DatabricksWorkspace} from "./DatabricksWorkspace";
-import {NamedLogger} from "@databricks/databricks-sdk/dist/logging";
 import {Loggers} from "../logger";
 import {CustomWhenContext} from "../vscode-objs/CustomWhenContext";
 import {workspaceConfigs} from "../vscode-objs/WorkspaceConfigs";
 import {WorkspaceStateManager} from "../vscode-objs/WorkspaceState";
 
+// eslint-disable-next-line @typescript-eslint/naming-convention
+const {NamedLogger} = logging;
 export type ConnectionState = "CONNECTED" | "CONNECTING" | "DISCONNECTED";
 
 /**

diff --git a/packages/databricks-vscode/src/configuration/DatabricksWorkspace.ts b/packages/databricks-vscode/src/configuration/DatabricksWorkspace.ts
@@ -1,12 +1,6 @@
-import {
-    WorkspaceClient,
-    Cluster,
-    iam,
-    WorkspaceConf,
-    WorkspaceConfProps,
-} from "@databricks/databricks-sdk";
+import {WorkspaceClient, iam, logging} from "@databricks/databricks-sdk";
+import {Cluster, WorkspaceConf, WorkspaceConfProps} from "../sdk-extensions";
 import {Context, context} from "@databricks/databricks-sdk/dist/context";
-import {withLogContext} from "@databricks/databricks-sdk/dist/logging";
 import {Uri} from "vscode";
 import {Loggers} from "../logger";
 import {workspaceConfigs} from "../vscode-objs/WorkspaceConfigs";
@@ -93,7 +87,7 @@ export class DatabricksWorkspace {
         }
     }
 
-    @withLogContext(Loggers.Extension, "DatabricksWorkspace.load")
+    @logging.withLogContext(Loggers.Extension, "DatabricksWorkspace.load")
     static async load(
         client: WorkspaceClient,
         authProvider: AuthProvider,

diff --git a/packages/databricks-vscode/src/configuration/auth/AzureCliCheck.ts b/packages/databricks-vscode/src/configuration/auth/AzureCliCheck.ts
@@ -2,13 +2,16 @@ import {
     ExecUtils,
     ProductVersion,
     WorkspaceClient,
+    logging,
 } from "@databricks/databricks-sdk";
-import {NamedLogger} from "@databricks/databricks-sdk/dist/logging";
 import {commands, Disposable, Uri, window} from "vscode";
 import {Loggers} from "../../logger";
 import {AzureCliAuthProvider} from "./AuthProvider";
 import {orchestrate, OrchestrationLoopError, Step} from "./orchestrate";
 
+// eslint-disable-next-line @typescript-eslint/naming-convention
+const {NamedLogger} = logging;
+
 // eslint-disable-next-line @typescript-eslint/no-var-requires
 const extensionVersion = require("../../../package.json")
     .version as ProductVersion;
@@ -25,7 +28,7 @@ type AzureStepName =
 export class AzureCliCheck implements Disposable {
     private disposables: Disposable[] = [];
     private isCodeSpaces: boolean;
-    private logger: NamedLogger;
+    private logger: logging.NamedLogger;
 
     tenantId: string | undefined;
     azureLoginAppId: string | undefined;

diff --git a/packages/databricks-vscode/src/configuration/auth/MetadataService.test.ts b/packages/databricks-vscode/src/configuration/auth/MetadataService.test.ts
@@ -1,7 +1,5 @@
 /* eslint-disable @typescript-eslint/naming-convention */
 import {MetadataService} from "./MetadataService";
-import {fetch} from "@databricks/databricks-sdk";
-
 import * as assert from "assert";
 import {
     ApiClient,
@@ -11,8 +9,10 @@ import {
     MetadataServiceVersion,
     MetadataServiceVersionHeader,
     RequestVisitor,
+    logging,
+    fetch,
 } from "@databricks/databricks-sdk";
-import {NamedLogger} from "@databricks/databricks-sdk/dist/logging";
+const {NamedLogger} = logging;
 
 describe(__filename, function () {
     this.timeout(10_000);