/*--------------------------------------------------------------------------------------------- * Copyright (c) Microsoft Corporation. All rights reserved. * Licensed under the Source EULA. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ import { SqlOpsDataClient, SqlOpsFeature } from 'dataprotocol-client'; import { ClientCapabilities, StaticFeature, RPCMessageType, ServerCapabilities } from 'vscode-languageclient'; import { Disposable } from 'vscode'; import { Telemetry } from './telemetry'; import * as contracts from './contracts'; import * as azdata from 'azdata'; import * as Utils from './utils'; import * as UUID from 'vscode-languageclient/lib/utils/uuid'; export class TelemetryFeature implements StaticFeature { constructor(private _client: SqlOpsDataClient) { } fillClientCapabilities(capabilities: ClientCapabilities): void { Utils.ensure(capabilities, 'telemetry')!.telemetry = true; } initialize(): void { this._client.onNotification(contracts.TelemetryNotification.type, e => { Telemetry.sendTelemetryEvent(e.params.eventName, e.params.properties, e.params.measures); }); } } export class AgentServicesFeature extends SqlOpsFeature { private static readonly messagesTypes: RPCMessageType[] = [ contracts.AgentJobsRequest.type, contracts.AgentJobHistoryRequest.type, contracts.AgentJobActionRequest.type ]; private onUpdatedHandler: () => any; constructor(client: SqlOpsDataClient) { super(client, AgentServicesFeature.messagesTypes); } public fillClientCapabilities(capabilities: ClientCapabilities): void { // this isn't explicitly necessary // ensure(ensure(capabilities, 'connection')!, 'agentServices')!.dynamicRegistration = true; } public initialize(capabilities: ServerCapabilities): void { this.register(this.messages, { id: UUID.generateUuid(), registerOptions: undefined }); } protected registerProvider(options: undefined): Disposable { const client = this._client; let self = this; // On updated registration let registerOnUpdated = (handler: () => any): void => { self.onUpdatedHandler = handler; }; let fireOnUpdated = (): void => { if (self.onUpdatedHandler) { self.onUpdatedHandler(); } }; // Job management methods let getJobs = (ownerUri: string): Thenable => { let params: contracts.AgentJobsParams = { ownerUri: ownerUri, jobId: null }; return client.sendRequest(contracts.AgentJobsRequest.type, params).then( r => r, e => { client.logFailedRequest(contracts.AgentJobsRequest.type, e); return Promise.resolve(undefined); } ); }; let getJobHistory = (ownerUri: string, jobID: string, jobName: string): Thenable => { let params: contracts.AgentJobHistoryParams = { ownerUri: ownerUri, jobId: jobID, jobName: jobName }; return client.sendRequest(contracts.AgentJobHistoryRequest.type, params).then( r => r, e => { client.logFailedRequest(contracts.AgentJobHistoryRequest.type, e); return Promise.resolve(undefined); } ); }; let jobAction = (ownerUri: string, jobName: string, action: string): Thenable => { let params: contracts.AgentJobActionParams = { ownerUri: ownerUri, jobName: jobName, action: action }; return client.sendRequest(contracts.AgentJobActionRequest.type, params).then( r => r, e => { client.logFailedRequest(contracts.AgentJobActionRequest.type, e); return Promise.resolve(undefined); } ); }; let createJob = (ownerUri: string, jobInfo: azdata.AgentJobInfo): Thenable => { let params: contracts.CreateAgentJobParams = { ownerUri: ownerUri, job: jobInfo }; let requestType = contracts.CreateAgentJobRequest.type; return client.sendRequest(requestType, params).then( r => { fireOnUpdated(); return r; }, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; let updateJob = (ownerUri: string, originalJobName: string, jobInfo: azdata.AgentJobInfo): Thenable => { let params: contracts.UpdateAgentJobParams = { ownerUri: ownerUri, originalJobName: originalJobName, job: jobInfo }; let requestType = contracts.UpdateAgentJobRequest.type; return client.sendRequest(requestType, params).then( r => { fireOnUpdated(); return r; }, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; let deleteJob = (ownerUri: string, jobInfo: azdata.AgentJobInfo): Thenable => { let params: contracts.DeleteAgentJobParams = { ownerUri: ownerUri, job: jobInfo }; let requestType = contracts.DeleteAgentJobRequest.type; return client.sendRequest(requestType, params).then( r => { fireOnUpdated(); return r; }, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; let getJobDefaults = (ownerUri: string): Thenable => { let params: contracts.AgentJobDefaultsParams = { ownerUri: ownerUri }; let requestType = contracts.AgentJobDefaultsRequest.type; return client.sendRequest(requestType, params).then( r => r, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; // Job Step management methods let createJobStep = (ownerUri: string, stepInfo: azdata.AgentJobStepInfo): Thenable => { let params: contracts.CreateAgentJobStepParams = { ownerUri: ownerUri, step: stepInfo }; let requestType = contracts.CreateAgentJobStepRequest.type; return client.sendRequest(requestType, params).then( r => { fireOnUpdated(); return r; }, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; let updateJobStep = (ownerUri: string, originalJobStepName: string, stepInfo: azdata.AgentJobStepInfo): Thenable => { let params: contracts.UpdateAgentJobStepParams = { ownerUri: ownerUri, originalJobStepName: originalJobStepName, step: stepInfo }; let requestType = contracts.UpdateAgentJobStepRequest.type; return client.sendRequest(requestType, params).then( r => { fireOnUpdated(); return r; }, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; let deleteJobStep = (ownerUri: string, stepInfo: azdata.AgentJobStepInfo): Thenable => { let params: contracts.DeleteAgentJobStepParams = { ownerUri: ownerUri, step: stepInfo }; let requestType = contracts.DeleteAgentJobStepRequest.type; return client.sendRequest(requestType, params).then( r => { fireOnUpdated(); return r; }, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; // Notebook Management methods const getNotebooks = (ownerUri: string): Thenable => { let params: contracts.AgentNotebookParams = { ownerUri: ownerUri }; return client.sendRequest(contracts.AgentNotebooksRequest.type, params).then( r => r, e => { client.logFailedRequest(contracts.AgentNotebooksRequest.type, e); return Promise.resolve(undefined); } ); }; const getNotebookHistory = (ownerUri: string, jobID: string, jobName: string, targetDatabase: string): Thenable => { let params: contracts.AgentNotebookHistoryParams = { ownerUri: ownerUri, jobId: jobID, jobName: jobName, targetDatabase: targetDatabase }; return client.sendRequest(contracts.AgentNotebookHistoryRequest .type, params).then( r => r, e => { client.logFailedRequest(contracts.AgentNotebookHistoryRequest.type, e); return Promise.resolve(undefined); } ); }; const getMaterializedNotebook = (ownerUri: string, targetDatabase: string, notebookMaterializedId: number): Thenable => { let params: contracts.AgentNotebookMaterializedParams = { ownerUri: ownerUri, targetDatabase: targetDatabase, notebookMaterializedId: notebookMaterializedId }; return client.sendRequest(contracts.AgentNotebookMaterializedRequest .type, params).then( r => r, e => { client.logFailedRequest(contracts.AgentNotebookMaterializedRequest.type, e); return Promise.resolve(undefined); } ); }; const getTemplateNotebook = (ownerUri: string, targetDatabase: string, jobId: string): Thenable => { let params: contracts.AgentNotebookTemplateParams = { ownerUri: ownerUri, targetDatabase: targetDatabase, jobId: jobId }; return client.sendRequest(contracts.AgentNotebookTemplateRequest .type, params).then( r => r, e => { client.logFailedRequest(contracts.AgentNotebookTemplateRequest.type, e); return Promise.resolve(undefined); } ); }; const createNotebook = (ownerUri: string, notebookInfo: azdata.AgentNotebookInfo, templateFilePath: string): Thenable => { let params: contracts.CreateAgentNotebookParams = { ownerUri: ownerUri, notebook: notebookInfo, templateFilePath: templateFilePath }; let requestType = contracts.CreateAgentNotebookRequest.type; return client.sendRequest(requestType, params).then( r => { fireOnUpdated(); return r; }, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; const updateNotebook = (ownerUri: string, originalNotebookName: string, notebookInfo: azdata.AgentNotebookInfo, templateFilePath: string): Thenable => { let params: contracts.UpdateAgentNotebookParams = { ownerUri: ownerUri, originalNotebookName: originalNotebookName, notebook: notebookInfo, templateFilePath: templateFilePath }; let requestType = contracts.UpdateAgentNotebookRequest.type; return client.sendRequest(requestType, params).then( r => { fireOnUpdated(); return r; }, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; const deleteNotebook = (ownerUri: string, notebookInfo: azdata.AgentNotebookInfo): Thenable => { let params: contracts.DeleteAgentNotebookParams = { ownerUri: ownerUri, notebook: notebookInfo }; let requestType = contracts.DeleteAgentNotebookRequest.type; return client.sendRequest(requestType, params).then( r => { fireOnUpdated(); return r; }, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; const deleteMaterializedNotebook = (ownerUri: string, agentNotebookHistory: azdata.AgentNotebookHistoryInfo, targetDatabase: string): Thenable => { let params: contracts.DeleteAgentMaterializedNotebookParams = { ownerUri: ownerUri, targetDatabase: targetDatabase, agentNotebookHistory: agentNotebookHistory }; return client.sendRequest(contracts.DeleteMaterializedNotebookRequest .type, params).then( r => r, e => { client.logFailedRequest(contracts.DeleteMaterializedNotebookRequest.type, e); return Promise.resolve(undefined); } ); }; const updateNotebookMaterializedName = (ownerUri: string, agentNotebookHistory: azdata.AgentNotebookHistoryInfo, targetDatabase: string, name: string): Thenable => { let params: contracts.UpdateAgentNotebookRunNameParams = { ownerUri: ownerUri, targetDatabase: targetDatabase, agentNotebookHistory: agentNotebookHistory, materializedNotebookName: name }; return client.sendRequest(contracts.UpdateAgentNotebookRunNameRequest .type, params).then( r => r, e => { client.logFailedRequest(contracts.UpdateAgentNotebookRunNameRequest.type, e); return Promise.resolve(undefined); } ); }; const updateNotebookMaterializedPin = (ownerUri: string, agentNotebookHistory: azdata.AgentNotebookHistoryInfo, targetDatabase: string, pin: boolean): Thenable => { let params: contracts.UpdateAgentNotebookRunPinParams = { ownerUri: ownerUri, targetDatabase: targetDatabase, agentNotebookHistory: agentNotebookHistory, materializedNotebookPin: pin }; return client.sendRequest(contracts.UpdateAgentNotebookRunPinRequest .type, params).then( r => r, e => { client.logFailedRequest(contracts.UpdateAgentNotebookRunPinRequest.type, e); return Promise.resolve(undefined); } ); }; // Alert management methods let getAlerts = (ownerUri: string): Thenable => { let params: contracts.AgentAlertsParams = { ownerUri: ownerUri }; let requestType = contracts.AgentAlertsRequest.type; return client.sendRequest(requestType, params).then( r => r, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; let createAlert = (ownerUri: string, alertInfo: azdata.AgentAlertInfo): Thenable => { let params: contracts.CreateAgentAlertParams = { ownerUri: ownerUri, alert: alertInfo }; let requestType = contracts.CreateAgentAlertRequest.type; return client.sendRequest(requestType, params).then( r => { fireOnUpdated(); return r; }, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; let updateAlert = (ownerUri: string, originalAlertName: string, alertInfo: azdata.AgentAlertInfo): Thenable => { let params: contracts.UpdateAgentAlertParams = { ownerUri: ownerUri, originalAlertName: originalAlertName, alert: alertInfo }; let requestType = contracts.UpdateAgentAlertRequest.type; return client.sendRequest(requestType, params).then( r => { fireOnUpdated(); return r; }, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; let deleteAlert = (ownerUri: string, alertInfo: azdata.AgentAlertInfo): Thenable => { let params: contracts.DeleteAgentAlertParams = { ownerUri: ownerUri, alert: alertInfo }; let requestType = contracts.DeleteAgentAlertRequest.type; return client.sendRequest(requestType, params).then( r => { fireOnUpdated(); return r; }, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; // Operator management methods let getOperators = (ownerUri: string): Thenable => { let params: contracts.AgentOperatorsParams = { ownerUri: ownerUri }; let requestType = contracts.AgentOperatorsRequest.type; return client.sendRequest(requestType, params).then( r => r, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; let createOperator = (ownerUri: string, operatorInfo: azdata.AgentOperatorInfo): Thenable => { let params: contracts.CreateAgentOperatorParams = { ownerUri: ownerUri, operator: operatorInfo }; let requestType = contracts.CreateAgentOperatorRequest.type; return client.sendRequest(requestType, params).then( r => { fireOnUpdated(); return r; }, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; let updateOperator = (ownerUri: string, originalOperatorName: string, operatorInfo: azdata.AgentOperatorInfo): Thenable => { let params: contracts.UpdateAgentOperatorParams = { ownerUri: ownerUri, originalOperatorName: originalOperatorName, operator: operatorInfo }; let requestType = contracts.UpdateAgentOperatorRequest.type; return client.sendRequest(requestType, params).then( r => { fireOnUpdated(); return r; }, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; let deleteOperator = (ownerUri: string, operatorInfo: azdata.AgentOperatorInfo): Thenable => { let params: contracts.DeleteAgentOperatorParams = { ownerUri: ownerUri, operator: operatorInfo }; let requestType = contracts.DeleteAgentOperatorRequest.type; return client.sendRequest(requestType, params).then( r => { fireOnUpdated(); return r; }, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; // Proxy management methods let getProxies = (ownerUri: string): Thenable => { let params: contracts.AgentProxiesParams = { ownerUri: ownerUri }; let requestType = contracts.AgentProxiesRequest.type; return client.sendRequest(requestType, params).then( r => r, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; let createProxy = (ownerUri: string, proxyInfo: azdata.AgentProxyInfo): Thenable => { let params: contracts.CreateAgentProxyParams = { ownerUri: ownerUri, proxy: proxyInfo }; let requestType = contracts.CreateAgentProxyRequest.type; return client.sendRequest(requestType, params).then( r => { fireOnUpdated(); return r; }, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; let updateProxy = (ownerUri: string, originalProxyName: string, proxyInfo: azdata.AgentProxyInfo): Thenable => { let params: contracts.UpdateAgentProxyParams = { ownerUri: ownerUri, originalProxyName: originalProxyName, proxy: proxyInfo }; let requestType = contracts.UpdateAgentProxyRequest.type; return client.sendRequest(requestType, params).then( r => { fireOnUpdated(); return r; }, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; let deleteProxy = (ownerUri: string, proxyInfo: azdata.AgentProxyInfo): Thenable => { let params: contracts.DeleteAgentProxyParams = { ownerUri: ownerUri, proxy: proxyInfo }; let requestType = contracts.DeleteAgentProxyRequest.type; return client.sendRequest(requestType, params).then( r => { fireOnUpdated(); return r; }, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; // Agent Credential Method let getCredentials = (ownerUri: string): Thenable => { let params: contracts.GetCredentialsParams = { ownerUri: ownerUri }; let requestType = contracts.AgentCredentialsRequest.type; return client.sendRequest(requestType, params).then( r => r, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; // Job Schedule management methods let getJobSchedules = (ownerUri: string): Thenable => { let params: contracts.AgentJobScheduleParams = { ownerUri: ownerUri }; let requestType = contracts.AgentJobSchedulesRequest.type; return client.sendRequest(requestType, params).then( r => r, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; let createJobSchedule = (ownerUri: string, scheduleInfo: azdata.AgentJobScheduleInfo): Thenable => { let params: contracts.CreateAgentJobScheduleParams = { ownerUri: ownerUri, schedule: scheduleInfo }; let requestType = contracts.CreateAgentJobScheduleRequest.type; return client.sendRequest(requestType, params).then( r => { fireOnUpdated(); return r; }, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; let updateJobSchedule = (ownerUri: string, originalScheduleName: string, scheduleInfo: azdata.AgentJobScheduleInfo): Thenable => { let params: contracts.UpdateAgentJobScheduleParams = { ownerUri: ownerUri, originalScheduleName: originalScheduleName, schedule: scheduleInfo }; let requestType = contracts.UpdateAgentJobScheduleRequest.type; return client.sendRequest(requestType, params).then( r => { fireOnUpdated(); return r; }, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; let deleteJobSchedule = (ownerUri: string, scheduleInfo: azdata.AgentJobScheduleInfo): Thenable => { let params: contracts.DeleteAgentJobScheduleParams = { ownerUri: ownerUri, schedule: scheduleInfo }; let requestType = contracts.DeleteAgentJobScheduleRequest.type; return client.sendRequest(requestType, params).then( r => { fireOnUpdated(); return r; }, e => { client.logFailedRequest(requestType, e); return Promise.resolve(undefined); } ); }; return azdata.dataprotocol.registerAgentServicesProvider({ providerId: client.providerId, getJobs, getJobHistory, jobAction, createJob, updateJob, deleteJob, getJobDefaults, createJobStep, updateJobStep, deleteJobStep, getNotebooks, getNotebookHistory, getMaterializedNotebook, getTemplateNotebook, createNotebook, updateNotebook, deleteMaterializedNotebook, updateNotebookMaterializedName, updateNotebookMaterializedPin, deleteNotebook, getAlerts, createAlert, updateAlert, deleteAlert, getOperators, createOperator, updateOperator, deleteOperator, getProxies, createProxy, updateProxy, deleteProxy, getCredentials, getJobSchedules, createJobSchedule, updateJobSchedule, deleteJobSchedule, registerOnUpdated }); } } export class SerializationFeature extends SqlOpsFeature { private static readonly messageTypes: RPCMessageType[] = [ contracts.SerializeDataStartRequest.type, contracts.SerializeDataContinueRequest.type, ]; constructor(client: SqlOpsDataClient) { super(client, SerializationFeature.messageTypes); } public fillClientCapabilities(capabilities: ClientCapabilities): void { } public initialize(capabilities: ServerCapabilities): void { this.register(this.messages, { id: UUID.generateUuid(), registerOptions: undefined }); } protected registerProvider(options: undefined): Disposable { const client = this._client; let startSerialization = (requestParams: azdata.SerializeDataStartRequestParams): Thenable => { return client.sendRequest(contracts.SerializeDataStartRequest.type, requestParams).then( r => { return r; }, e => { client.logFailedRequest(contracts.SerializeDataStartRequest.type, e); return Promise.resolve({ succeeded: false, messages: Utils.getErrorMessage(e) }); } ); }; let continueSerialization = (requestParams: azdata.SerializeDataContinueRequestParams): Thenable => { return client.sendRequest(contracts.SerializeDataContinueRequest.type, requestParams).then( r => { return r; }, e => { client.logFailedRequest(contracts.SerializeDataContinueRequest.type, e); return Promise.resolve({ succeeded: false, messages: Utils.getErrorMessage(e) }); } ); }; return azdata.dataprotocol.registerSerializationProvider({ providerId: client.providerId, startSerialization, continueSerialization }); } }