azure-kusto-data
Version:
Azure Data Explorer Query SDK
250 lines • 10.8 kB
JavaScript
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
import { isNodeLike } from "@azure/core-util";
import axios from "axios";
import http from "http";
import https from "https";
import { v4 as uuidv4 } from "uuid";
import ClientRequestProperties from "./clientRequestProperties.js";
import CloudSettings from "./cloudSettings.js";
import ConnectionStringBuilder from "./connectionBuilder.js";
import { ThrottlingError } from "./errors.js";
import { kustoTrustedEndpoints } from "./kustoTrustedEndpoints.js";
import { KustoResponseDataSetV1, KustoResponseDataSetV2 } from "./response.js";
import AadHelper from "./security.js";
import { toMilliseconds } from "./timeUtils.js";
const COMMAND_TIMEOUT_IN_MILLISECS = toMilliseconds(0, 10, 30);
const QUERY_TIMEOUT_IN_MILLISECS = toMilliseconds(0, 4, 30);
const CLIENT_SERVER_DELTA_IN_MILLISECS = toMilliseconds(0, 0, 30);
const MGMT_PREFIX = ".";
var ExecutionType;
(function (ExecutionType) {
ExecutionType["Mgmt"] = "mgmt";
ExecutionType["Query"] = "query";
ExecutionType["Ingest"] = "ingest";
ExecutionType["QueryV1"] = "queryv1";
})(ExecutionType || (ExecutionType = {}));
export class KustoClient {
constructor(kcsb) {
this.cancelToken = axios.CancelToken.source();
this._isClosed = false;
this.connectionString = typeof kcsb === "string" ? new ConnectionStringBuilder(kcsb) : kcsb;
if (!this.connectionString.dataSource) {
throw new Error("Cluster url is required");
}
const url = new URL(this.connectionString.dataSource);
this.cluster = url.toString();
if (this.cluster.endsWith("/")) {
this.cluster = this.cluster.slice(0, -1);
}
this.defaultDatabase = this.connectionString.initialCatalog;
this.endpoints = {
[ExecutionType.Mgmt]: `${this.cluster}/v1/rest/mgmt`,
[ExecutionType.Query]: `${this.cluster}/v2/rest/query`,
[ExecutionType.Ingest]: `${this.cluster}/v1/rest/ingest`,
[ExecutionType.QueryV1]: `${this.cluster}/v1/rest/query`,
};
this.aadHelper = new AadHelper(this.connectionString);
let headers = {
Accept: "application/json",
};
if (isNodeLike) {
headers = Object.assign(Object.assign({}, headers), { "Accept-Encoding": "gzip,deflate", Connection: "Keep-Alive" });
}
const axiosProps = {
headers,
validateStatus: (status) => status === 200,
maxBodyLength: Infinity,
maxContentLength: Infinity,
maxRedirects: 0,
};
// http and https are Node modules and are not found in browsers
if (isNodeLike) {
// keepAlive pools and reuses TCP connections, so it's faster
axiosProps.httpAgent = new http.Agent({ keepAlive: true });
axiosProps.httpsAgent = new https.Agent({ keepAlive: true });
}
axiosProps.cancelToken = this.cancelToken.token;
this.axiosInstance = axios.create(axiosProps);
}
async execute(db, query, properties) {
query = query.trim();
if (query.startsWith(MGMT_PREFIX)) {
return this.executeMgmt(db, query, properties);
}
return this.executeQuery(db, query, properties);
}
async executeQuery(db, query, properties) {
return this._execute(this.endpoints[ExecutionType.Query], ExecutionType.Query, db, { query }, properties);
}
async executeQueryV1(db, query, properties) {
return this._execute(this.endpoints[ExecutionType.QueryV1], ExecutionType.QueryV1, db, { query }, properties);
}
async executeMgmt(db, query, properties) {
return this._execute(this.endpoints[ExecutionType.Mgmt], ExecutionType.Mgmt, db, { query }, properties);
}
async executeStreamingIngest(db, table, stream, streamFormat, mappingName, blob, clientRequestId) {
let endpoint = `${this.endpoints[ExecutionType.Ingest]}/${this.getDb(db)}/${table}?streamFormat=${streamFormat}`;
if (mappingName != null) {
endpoint += `&mappingName=${mappingName}`;
}
if (blob) {
endpoint += `&sourceKind=uri`;
}
let properties = null;
if (clientRequestId) {
properties = new ClientRequestProperties();
properties.clientRequestId = clientRequestId;
}
return this._execute(endpoint, ExecutionType.Ingest, db, blob ? { blob } : { stream }, properties);
}
async _execute(endpoint, executionType, db, entity, properties) {
this.ensureOpen();
kustoTrustedEndpoints.validateTrustedEndpoint(endpoint, (await CloudSettings.getCloudInfoForCluster(this.cluster)).LoginEndpoint);
db = this.getDb(db);
const headers = {};
let payload;
let clientRequestPrefix = "";
const timeout = this._getClientTimeout(executionType, properties);
let payloadContent = "";
if ("query" in entity) {
payload = {
db,
csl: entity.query,
};
if (properties != null) {
payload.properties = properties.toJSON();
}
payloadContent = JSON.stringify(payload);
headers["Content-Type"] = "application/json; charset=utf-8";
clientRequestPrefix = "KNC.execute;";
}
else if ("stream" in entity) {
payloadContent = entity.stream;
clientRequestPrefix = "KNC.executeStreamingIngest;";
if (isNodeLike) {
headers["Content-Encoding"] = "gzip";
headers["Content-Type"] = "application/octet-stream";
}
else {
headers["Content-Type"] = "application/json";
}
}
else if ("blob" in entity) {
payloadContent = {
sourceUri: entity.blob,
};
clientRequestPrefix = "KNC.executeStreamingIngestFromBlob;";
headers["Content-Type"] = "application/json";
}
else {
throw new Error("Invalid parameters - expected query or streaming ingest");
}
let kustoHeaders = this.connectionString.clientDetails().getHeaders();
kustoHeaders["x-ms-client-request-id"] = `${clientRequestPrefix}${uuidv4()}`;
if (properties != null) {
kustoHeaders = Object.assign(Object.assign({}, kustoHeaders), properties.getHeaders());
}
for (const key of Object.keys(kustoHeaders)) {
if (kustoHeaders[key]) {
headers[key] = kustoHeaders[key];
}
}
const authHeader = await this.aadHelper.getAuthHeader();
if (authHeader != null) {
headers.Authorization = authHeader;
}
return this._doRequest(endpoint, executionType, headers, payloadContent, timeout, properties);
}
getDb(db) {
if (db == null) {
if (this.defaultDatabase == null) {
throw new Error("No database provided, and no default database specified in connection string");
}
db = this.defaultDatabase;
}
return db;
}
async _doRequest(endpoint, executionType, headers, payload, timeout, properties) {
var _a, _b, _c;
// replace non-ascii characters with ? in headers
for (const key of Object.keys(headers)) {
headers[key] = headers[key].replace(/[^\x00-\x7F]+/g, "?");
}
const axiosConfig = {
headers,
timeout,
};
let axiosResponse;
try {
axiosResponse = await this.axiosInstance.post(endpoint, payload, axiosConfig);
}
catch (error) {
if (axios.isAxiosError(error)) {
// Since it's impossible to modify the error request object, the only way to censor the Authorization header is to remove it.
error.request = undefined;
if ((_a = error === null || error === void 0 ? void 0 : error.config) === null || _a === void 0 ? void 0 : _a.headers) {
error.config.headers.Authorization = "<REDACTED>";
}
// eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
if ((_c = (_b = error === null || error === void 0 ? void 0 : error.response) === null || _b === void 0 ? void 0 : _b.request) === null || _c === void 0 ? void 0 : _c._header) {
// eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
error.response.request._header = "<REDACTED>";
}
if (error.response && error.response.status === 429) {
throw new ThrottlingError("POST request failed with status 429 (Too Many Requests)", error);
}
}
throw error;
}
return this._parseResponse(axiosResponse.data, executionType, properties, axiosResponse.status);
}
_parseResponse(response, executionType, properties, status) {
const { raw } = properties || {};
if (raw === true || executionType === ExecutionType.Ingest) {
return response;
}
let kustoResponse = null;
try {
if (executionType === ExecutionType.Query) {
kustoResponse = new KustoResponseDataSetV2(response);
}
else {
kustoResponse = new KustoResponseDataSetV1(response);
}
}
catch (ex) {
throw new Error(`Failed to parse response ({${status}}) with the following error [${ex}].`);
}
if (kustoResponse.getErrorsCount().errors > 0) {
throw new Error(`Kusto request had errors. ${kustoResponse.getExceptions()}`);
}
return kustoResponse;
}
_getClientTimeout(executionType, properties) {
if (properties != null) {
const clientTimeout = properties.getClientTimeout();
if (clientTimeout) {
return clientTimeout;
}
const serverTimeout = properties.getTimeout();
if (serverTimeout) {
return serverTimeout + CLIENT_SERVER_DELTA_IN_MILLISECS;
}
}
return executionType === ExecutionType.Query || executionType === ExecutionType.QueryV1 ? QUERY_TIMEOUT_IN_MILLISECS : COMMAND_TIMEOUT_IN_MILLISECS;
}
close() {
if (!this._isClosed) {
this.cancelToken.cancel("Client Closed");
}
this._isClosed = true;
}
ensureOpen() {
if (this._isClosed) {
throw new Error("Client is closed");
}
}
}
export default KustoClient;
//# sourceMappingURL=client.js.map