/*--------------------------------------------------------------------------------------------- * 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, IHdfsOptions, IFile, File, FileSourceFactory } from './fileSources'; import { CancelableStream } from './cancelableStream'; import { TreeNode } from './treeNodes'; import * as utils from '../utils'; import { IFileNode } from './types'; export interface ITreeChangeHandler { notifyNodeChanged(node: TreeNode): void; } export class TreeDataContext { constructor(public extensionContext: vscode.ExtensionContext, public changeHandler: ITreeChangeHandler) { } } export class HdfsProvider implements vscode.TreeDataProvider, ITreeChangeHandler { static readonly NoConnectionsMessage = 'No connections added'; static readonly ConnectionsLabel = 'Connections'; private connections: ConnectionNode[]; private _onDidChangeTreeData = new vscode.EventEmitter(); private context: TreeDataContext; constructor(extensionContext: vscode.ExtensionContext) { this.connections = []; this.context = new TreeDataContext(extensionContext, this); } public get onDidChangeTreeData(): vscode.Event { return this._onDidChangeTreeData.event; } getTreeItem(element: TreeNode): vscode.TreeItem | Thenable { return element.getTreeItem(); } getChildren(element?: TreeNode): vscode.ProviderResult { if (element) { return element.getChildren(false); } else { return this.connections.length > 0 ? this.connections : [ErrorNode.create(HdfsProvider.NoConnectionsMessage, element)]; } } addConnection(displayName: string, fileSource: IFileSource): void { if (!this.connections.find(c => c.getDisplayName() === displayName)) { this.connections.push(new ConnectionNode(this.context, displayName, fileSource)); this._onDidChangeTreeData.fire(); } } addHdfsConnection(options: IHdfsOptions): void { let displayName = `${options.user}@${options.host}:${options.port}`; let fileSource = FileSourceFactory.instance.createHdfsFileSource(options); this.addConnection(displayName, fileSource); } notifyNodeChanged(node: TreeNode): void { this._onDidChangeTreeData.fire(node); } } export abstract class HdfsFileSourceNode extends TreeNode { constructor(protected context: TreeDataContext, protected _path: string, protected fileSource: IFileSource) { super(); } public get hdfsPath(): string { return this._path; } public get nodePathValue(): string { return this.getDisplayName(); } getDisplayName(): string { return fspath.basename(this._path); } public async delete(recursive: boolean = false): Promise { await this.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, nodeType?: string) { super(context, path, fileSource); this._nodeType = nodeType ? nodeType : Constants.MssqlClusterItems.Folder; } private ensureChildrenExist(): void { if (!this.children) { this.children = []; } } public onChildRemoved(): void { this.children = undefined; } async getChildren(refreshChildren: boolean): Promise { if (refreshChildren || !this.children) { this.ensureChildrenExist(); try { let files: IFile[] = await this.fileSource.enumerateFiles(this._path); if (files) { // Note: for now, assuming HDFS-provided sorting is sufficient this.children = files.map((file) => { let node: TreeNode = file.isDirectory ? new FolderNode(this.context, file.path, this.fileSource) : new FileNode(this.context, file.path, this.fileSource); node.parent = this; return node; }); } } catch (error) { this.children = [ErrorNode.create(localize('errorExpanding', 'Error: {0}', utils.getErrorMessage(error)), this, error.statusCode)]; } } return this.children; } 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: undefined, iconType: 'Folder' }; return nodeInfo; } public async writeFile(localFile: IFile): Promise { return this.runChildAddAction(() => this.writeFileAsync(localFile)); } private async writeFileAsync(localFile: IFile): Promise { await this.fileSource.writeFile(localFile, this._path); let fileNode = new FileNode(this.context, File.createPath(this._path, File.getBasename(localFile)), this.fileSource); return fileNode; } public async mkdir(name: string): Promise { return this.runChildAddAction(() => this.mkdirAsync(name)); } private async mkdirAsync(name: string): Promise { await this.fileSource.mkdir(name, this._path); let subDir = new FolderNode(this.context, File.createPath(this._path, name), this.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, fileSource: IFileSource) { super(context, '/', fileSource, Constants.MssqlClusterItems.Connection); } getDisplayName(): string { return this.displayName; } public async delete(): Promise { throw new Error(localize('errDeleteConnectionNode', 'Cannot delete a connection. Only subfolders and files can be deleted.')); } async getTreeItem(): Promise { let item = await super.getTreeItem(); 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: undefined, iconType: 'HDFSFolder' }; return nodeInfo; } } export class FileNode extends HdfsFileSourceNode implements IFileNode { constructor(context: TreeDataContext, path: string, fileSource: IFileSource) { super(context, path, fileSource); } 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: 'FileGroupFile' }; return nodeInfo; } public async getFileContentsAsString(maxBytes?: number): Promise { let contents: Buffer = await this.fileSource.readFile(this.hdfsPath, maxBytes); return contents ? contents.toString('utf8') : ''; } public async getFileLinesAsString(maxLines: number): Promise { let contents: Buffer = await this.fileSource.readFileLines(this.hdfsPath, maxLines); return contents ? contents.toString('utf8') : ''; } public writeFileContentsToDisk(localPath: string, cancelToken?: vscode.CancellationTokenSource): Promise { return new Promise((resolve, reject) => { let readStream: fs.ReadStream = this.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', (location) => { 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 { if (this.getDisplayName().toLowerCase().endsWith('.jar') || this.getDisplayName().toLowerCase().endsWith('.py')) { return Constants.MssqlClusterItemsSubType.Spark; } return undefined; } } export class ErrorNode extends TreeNode { static messageNum: number = 0; private _nodePathValue: string; constructor(private message: string) { super(); } 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; } }