/*--------------------------------------------------------------------------------------------- * Copyright (c) Microsoft Corporation. All rights reserved. * Licensed under the Source EULA. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ import * as azdata from 'azdata'; import * as vscode from 'vscode'; import * as fspath from 'path'; import * as fs from 'fs'; import * as nls from 'vscode-nls'; const localize = nls.loadMessageBundle(); import * as Constants from '../constants'; import { IFileSource, IFile, File, FileType } from './fileSources'; import { CancelableStream } from './cancelableStream'; import { TreeNode } from './treeNodes'; import * as utils from '../utils'; import { IFileNode } from './types'; import { MountStatus } from '../hdfs/mount'; import { SqlClusterSession } from './objectExplorerNodeProvider'; export interface ITreeChangeHandler { notifyNodeChanged(node: TreeNode): void; } export class TreeDataContext { constructor(public extensionContext: vscode.ExtensionContext, public changeHandler: ITreeChangeHandler) { } } export abstract class HdfsFileSourceNode extends TreeNode { constructor(protected context: TreeDataContext, protected _path: string, fileSource: IFileSource | undefined, protected mountStatus?: MountStatus) { super(fileSource); } public get hdfsPath(): string { return this._path; } public get nodePathValue(): string { return this.getDisplayName(); } protected isMounted(): boolean { return this.mountStatus === MountStatus.Mount || this.mountStatus === MountStatus.Mount_Child; } getDisplayName(): string { return fspath.basename(this._path); } public async delete(recursive: boolean = false): Promise { const fileSource = await this.getFileSource(); await fileSource.delete(this.hdfsPath, recursive); // Notify parent should be updated. If at top, will return undefined which will refresh whole tree (this.parent).onChildRemoved(); this.context.changeHandler.notifyNodeChanged(this.parent); } public abstract onChildRemoved(): void; } export class FolderNode extends HdfsFileSourceNode { private children: TreeNode[] = []; protected _nodeType: string; constructor(context: TreeDataContext, path: string, fileSource: IFileSource | undefined, nodeType?: string, mountStatus?: MountStatus) { super(context, path, fileSource, mountStatus); this._nodeType = nodeType ? nodeType : Constants.MssqlClusterItems.Folder; } public onChildRemoved(): void { this.children = undefined; } async getChildren(refreshChildren: boolean): Promise { if (refreshChildren || !this.children) { try { const fileSource = await this.getFileSource(); let files: IFile[] = await fileSource.enumerateFiles(this._path); if (files) { // Note: for now, assuming HDFS-provided sorting is sufficient this.children = files.map((file) => { let node: TreeNode = file.fileType === FileType.File ? new FileNode(this.context, file.path, fileSource, this.getChildMountStatus(file)) : new FolderNode(this.context, file.path, fileSource, Constants.MssqlClusterItems.Folder, this.getChildMountStatus(file)); node.parent = this; return node; }); } } catch (error) { this.children = [ErrorNode.create(localize('errorExpanding', "Error: {0}", utils.getErrorMessage(error)), this, error.statusCode)]; } } return this.children; } private getChildMountStatus(file: IFile): MountStatus { if (file.mountStatus !== undefined && file.mountStatus !== MountStatus.None) { return file.mountStatus; } else if (this.mountStatus !== undefined && this.mountStatus !== MountStatus.None) { // Any child node of a mount (or subtree) must be a mount child return MountStatus.Mount_Child; } return MountStatus.None; } getTreeItem(): vscode.TreeItem | Promise { let item = new vscode.TreeItem(this.getDisplayName(), vscode.TreeItemCollapsibleState.Collapsed); // For now, folder always looks the same. We're using SQL icons to differentiate remote vs local files item.iconPath = { dark: this.context.extensionContext.asAbsolutePath('resources/light/Folder.svg'), light: this.context.extensionContext.asAbsolutePath('resources/light/Folder.svg') }; item.contextValue = this._nodeType; return item; } getNodeInfo(): azdata.NodeInfo { // TODO handle error message case by returning it in the OE API // TODO support better mapping of node type let nodeInfo: azdata.NodeInfo = { label: this.getDisplayName(), isLeaf: false, errorMessage: undefined, metadata: undefined, nodePath: this.generateNodePath(), nodeStatus: undefined, nodeType: this._nodeType, nodeSubType: this.getSubType(), iconType: this.isMounted() ? 'Folder_mounted' : 'Folder' }; return nodeInfo; } private getSubType(): string | undefined { if (this.mountStatus === MountStatus.Mount) { return Constants.MssqlClusterItemsSubType.Mount; } else if (this.mountStatus === MountStatus.Mount_Child) { return Constants.MssqlClusterItemsSubType.MountChild; } return undefined; } public async writeFile(localFile: IFile): Promise { return this.runChildAddAction(() => this.writeFileAsync(localFile)); } private async writeFileAsync(localFile: IFile): Promise { const fileSource = await this.getFileSource(); await fileSource.writeFile(localFile, this._path); let fileNode = new FileNode(this.context, File.createPath(this._path, File.getBasename(localFile)), fileSource); return fileNode; } public async mkdir(name: string): Promise { return this.runChildAddAction(() => this.mkdirAsync(name)); } private async mkdirAsync(name: string): Promise { const fileSource = await this.getFileSource(); await fileSource.mkdir(name, this._path); let subDir = new FolderNode(this.context, File.createPath(this._path, name), fileSource); return subDir; } private async runChildAddAction(action: () => Promise): Promise { let node = await action(); await this.getChildren(true); if (this.children) { // Find the child matching the node. This is necessary // since writing can add duplicates. node = this.children.find(n => n.nodePathValue === node.nodePathValue) as T; this.context.changeHandler.notifyNodeChanged(this); } else { // Failed to retrieve children from server so something went wrong node = undefined; } return node; } } export class ConnectionNode extends FolderNode { constructor(context: TreeDataContext, private displayName: string, private clusterSession: SqlClusterSession) { super(context, '/', undefined, Constants.MssqlClusterItems.Connection); } override getDisplayName(): string { return this.displayName; } public override async delete(): Promise { throw new Error(localize('errDeleteConnectionNode', "Cannot delete a connection. Only subfolders and files can be deleted.")); } override async getTreeItem(): Promise { let item = await super.getTreeItem(); item.contextValue = this._nodeType; return item; } public override async getFileSource(): Promise { // The node is initially created without a filesource and then one is created only once an action is // taken that requires a connection const fileSource = await super.getFileSource(); if (!fileSource) { await this.updateFileSource(await this.clusterSession.getSqlClusterConnection()); } return super.getFileSource(); } override getNodeInfo(): azdata.NodeInfo { // TODO handle error message case by returning it in the OE API // TODO support better mapping of node type let nodeInfo: azdata.NodeInfo = { label: this.getDisplayName(), isLeaf: false, errorMessage: undefined, metadata: undefined, nodePath: this.generateNodePath(), nodeStatus: undefined, nodeType: 'mssqlCluster:hdfs', nodeSubType: undefined, iconType: 'Folder' }; return nodeInfo; } } export class FileNode extends HdfsFileSourceNode implements IFileNode { constructor(context: TreeDataContext, path: string, fileSource: IFileSource, mountStatus?: MountStatus) { super(context, path, fileSource, mountStatus); } public onChildRemoved(): void { // do nothing } getChildren(refreshChildren: boolean): TreeNode[] | Promise { return []; } getTreeItem(): vscode.TreeItem | Promise { let item = new vscode.TreeItem(this.getDisplayName(), vscode.TreeItemCollapsibleState.None); item.iconPath = { dark: this.context.extensionContext.asAbsolutePath('resources/dark/file_inverse.svg'), light: this.context.extensionContext.asAbsolutePath('resources/light/file.svg') }; item.contextValue = Constants.MssqlClusterItems.File; return item; } getNodeInfo(): azdata.NodeInfo { // TODO improve node type handling so it's not tied to SQL Server types let nodeInfo: azdata.NodeInfo = { label: this.getDisplayName(), isLeaf: true, errorMessage: undefined, metadata: undefined, nodePath: this.generateNodePath(), nodeStatus: undefined, nodeType: Constants.MssqlClusterItems.File, nodeSubType: this.getSubType(), iconType: this.isMounted() ? 'FileGroupFile_mounted' : 'FileGroupFile' }; return nodeInfo; } public async getFileContentsAsString(maxBytes?: number): Promise { const fileSource = await this.getFileSource(); let contents: Buffer = await fileSource.readFile(this.hdfsPath, maxBytes); return contents ? contents.toString('utf8') : ''; } public async getFileLinesAsString(maxLines: number): Promise { const fileSource = await this.getFileSource(); let contents: Buffer = await fileSource.readFileLines(this.hdfsPath, maxLines); return contents ? contents.toString('utf8') : ''; } public async writeFileContentsToDisk(localPath: string, cancelToken?: vscode.CancellationTokenSource): Promise { const fileSource = await this.getFileSource(); return new Promise((resolve, reject) => { let readStream: fs.ReadStream = fileSource.createReadStream(this.hdfsPath); readStream.on('error', (err) => { reject(err); }); let error: string | Error = undefined; let writeStream = fs.createWriteStream(localPath, { encoding: 'utf8' }); writeStream.on('error', (err) => { error = err; reject(error); }); writeStream.on('finish', () => { if (!error) { resolve(vscode.Uri.file(localPath)); } }); let cancelable = new CancelableStream(cancelToken); cancelable.on('error', (err) => { reject(err); }); readStream.pipe(cancelable).pipe(writeStream); }); } private getSubType(): string | undefined { let subType = ''; if (this.getDisplayName().toLowerCase().endsWith('.jar') || this.getDisplayName().toLowerCase().endsWith('.py')) { subType += Constants.MssqlClusterItemsSubType.Spark; } else if (this.mountStatus === MountStatus.Mount_Child) { subType += Constants.MssqlClusterItemsSubType.MountChild; } return subType.length > 0 ? subType : undefined; } } class ErrorNode extends TreeNode { static messageNum: number = 0; private _nodePathValue: string; constructor(private message: string) { super(undefined); } public static create(message: string, parent: TreeNode, errorCode?: number): ErrorNode { let node = new ErrorNode(message); node.parent = parent; if (errorCode) { node.errorStatusCode = errorCode; } return node; } private ensureNodePathValue(): void { if (!this._nodePathValue) { this._nodePathValue = `message_${ErrorNode.messageNum++}`; } } public get nodePathValue(): string { this.ensureNodePathValue(); return this._nodePathValue; } public getChildren(refreshChildren: boolean): TreeNode[] | Promise { return []; } public getTreeItem(): vscode.TreeItem | Promise { let item = new vscode.TreeItem(this.message, vscode.TreeItemCollapsibleState.None); item.contextValue = Constants.MssqlClusterItems.Error; return item; } getNodeInfo(): azdata.NodeInfo { let nodeInfo: azdata.NodeInfo = { label: this.message, isLeaf: false, errorMessage: undefined, metadata: undefined, nodePath: this.generateNodePath(), nodeStatus: undefined, nodeType: Constants.MssqlClusterItems.Error, nodeSubType: undefined, iconType: 'MessageType' }; return nodeInfo; } }