azure-kusto-ingest
Version:
Azure Data Explorer Ingestion SDK
246 lines • 11.4 kB
JavaScript
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
import { KustoDataErrors, TimeUtils } from "azure-kusto-data";
import { ExponentialRetry } from "./retry.js";
import { ContainerClient } from "@azure/storage-blob";
import { TableClient } from "@azure/data-tables";
import { RankedStorageAccountSet } from "./rankedStorageAccountSet.js";
import { QueueClient } from "@azure/storage-queue";
const ATTEMPT_COUNT = 4;
export var ResourceType;
(function (ResourceType) {
ResourceType[ResourceType["Queue"] = 0] = "Queue";
ResourceType[ResourceType["Container"] = 1] = "Container";
ResourceType[ResourceType["Table"] = 2] = "Table";
})(ResourceType || (ResourceType = {}));
export class ResourceURI {
constructor(uri, accountName, resourceType) {
this.uri = uri;
this.accountName = accountName;
this.resourceType = resourceType;
}
}
export class IngestClientResources {
constructor(securedReadyForAggregationQueues = null, failedIngestionsQueues = null, successfulIngestionsQueues = null, containers = null, statusTables = null) {
this.securedReadyForAggregationQueues = securedReadyForAggregationQueues;
this.failedIngestionsQueues = failedIngestionsQueues;
this.successfulIngestionsQueues = successfulIngestionsQueues;
this.containers = containers;
this.statusTables = statusTables;
}
valid() {
const resources = [this.securedReadyForAggregationQueues, this.failedIngestionsQueues, this.failedIngestionsQueues, this.containers, this.statusTables];
return resources.reduce((prev, current) => !!(prev && current), true);
}
}
export class ResourceManager {
constructor(kustoClient, isBrowser = false) {
this.kustoClient = kustoClient;
this.isBrowser = isBrowser;
this.baseSleepTimeSecs = 1;
this.baseJitterSecs = 1;
this.refreshPeriod = TimeUtils.toMilliseconds(1, 0, 0);
this.refreshPeriodOnError = TimeUtils.toMilliseconds(0, 10, 0);
this.ingestClientResources = null;
this.ingestClientResourcesLastUpdate = null;
this.authorizationContext = null;
this.authorizationContextLastUpdate = null;
this.rankedStorageAccountSet = new RankedStorageAccountSet();
}
async refreshIngestClientResources() {
const error = await this.tryRefresh(false);
if (!this.ingestClientResources) {
throw new Error(`Failed to fetch ingestion resources from service. ${error === null || error === void 0 ? void 0 : error.message}.\n ${error === null || error === void 0 ? void 0 : error.stack}`);
}
return this.ingestClientResources;
}
async getIngestClientResourcesFromService() {
const retry = new ExponentialRetry(ATTEMPT_COUNT, this.baseSleepTimeSecs, this.baseJitterSecs);
while (retry.shouldTry()) {
try {
const cmd = `.get ingestion resources ${this.isBrowser ? `with (EnableBlobCors='true', EnableQueueCors='true', EnableTableCors='true')` : ""}`;
const response = await this.kustoClient.execute("NetDefaultDB", cmd);
const table = response.primaryResults[0];
const resoures = new IngestClientResources(this.getResourceByName(table, "SecuredReadyForAggregationQueue", ResourceType.Queue), this.getResourceByName(table, "FailedIngestionsQueue", ResourceType.Queue), this.getResourceByName(table, "SuccessfulIngestionsQueue", ResourceType.Queue), this.getResourceByName(table, "TempStorage", ResourceType.Container), this.getResourceByName(table, "IngestionsStatusTable", ResourceType.Table));
if (!resoures.valid()) {
throw new Error("Unexpected error occured - fetched data returned missing resource");
}
return resoures;
}
catch (error) {
if (!(error instanceof KustoDataErrors.ThrottlingError)) {
throw error;
}
await retry.backoff();
}
}
throw new Error(`Failed to get ingestion resources from server - the request was throttled ${ATTEMPT_COUNT} times.`);
}
getResourceByName(table, resourceName, resourceType) {
const result = [];
for (const row of table.rows()) {
const typedRow = row;
if (typedRow.ResourceTypeName === resourceName) {
let accountName = "";
if (resourceType === ResourceType.Queue) {
accountName = new QueueClient(typedRow.StorageRoot).accountName;
}
else if (resourceType === ResourceType.Container) {
accountName = new ContainerClient(typedRow.StorageRoot).accountName;
}
result.push(new ResourceURI(typedRow.StorageRoot, accountName, resourceType));
}
}
return result;
}
pupulateStorageAccounts() {
if (this.ingestClientResources == null) {
return;
}
// containers
const accounts = new Set();
if (this.ingestClientResources.containers != null) {
for (const container of this.ingestClientResources.containers) {
accounts.add(container.accountName);
}
}
// queues
if (this.ingestClientResources.securedReadyForAggregationQueues != null) {
for (const queue of this.ingestClientResources.securedReadyForAggregationQueues) {
accounts.add(queue.accountName);
}
}
for (const account of accounts) {
this.rankedStorageAccountSet.registerStorageAccount(account);
}
}
groupResourcesByStorageAccount(resources) {
var _a;
const result = new Map();
for (const resource of resources) {
if (!result.has(resource.accountName)) {
result.set(resource.accountName, []);
}
(_a = result.get(resource.accountName)) === null || _a === void 0 ? void 0 : _a.push(resource);
}
return result;
}
getRankedAndShuffledStorageAccounts(resources) {
const resourcesByAccount = this.groupResourcesByStorageAccount(resources);
const rankedStorageAccounts = this.rankedStorageAccountSet.getRankedShuffledAccounts();
const result = new Array();
for (const account of rankedStorageAccounts) {
const accountName = account.getAccountName();
if (resourcesByAccount.has(accountName)) {
result.push(resourcesByAccount.get(accountName));
}
}
return result;
}
getRoundRobinRankedAndShuffledResources(resources) {
const rankedAccounts = this.getRankedAndShuffledStorageAccounts(resources);
const result = new Array();
let index = 0;
while (result.length < resources.length) {
const account = rankedAccounts[index % rankedAccounts.length];
if (account.length > 0) {
result.push(account.shift());
}
index++;
}
return result;
}
async refreshAuthorizationContext() {
const error = await this.tryRefresh(true);
if (this.authorizationContext == null) {
throw new Error(`Failed to fetch Authorization context from service. ${error === null || error === void 0 ? void 0 : error.message}.\n ${error === null || error === void 0 ? void 0 : error.stack}`);
}
return this.authorizationContext;
}
async tryRefresh(isAuthContextFetch) {
var _a;
const resource = isAuthContextFetch ? (_a = this.authorizationContext) === null || _a === void 0 ? void 0 : _a.trim() : this.ingestClientResources;
const lastRefresh = isAuthContextFetch ? this.authorizationContextLastUpdate : this.ingestClientResourcesLastUpdate;
const now = Date.now();
let error = null;
if (!resource || !lastRefresh || lastRefresh + this.refreshPeriod <= now) {
try {
if (isAuthContextFetch) {
this.authorizationContext = await this.getAuthorizationContextFromService();
this.authorizationContextLastUpdate = now;
}
else {
this.ingestClientResources = await this.getIngestClientResourcesFromService();
this.ingestClientResourcesLastUpdate = now;
this.pupulateStorageAccounts();
}
}
catch (e) {
error = e;
}
}
return error;
}
async getAuthorizationContextFromService() {
const retry = new ExponentialRetry(ATTEMPT_COUNT, this.baseSleepTimeSecs, this.baseJitterSecs);
while (retry.shouldTry()) {
try {
const response = await this.kustoClient.execute("NetDefaultDB", ".get kusto identity token");
const next = response.primaryResults[0].rows().next();
if (next.done) {
throw new Error("Failed to get authorization context - got empty results");
}
return next.value.toJSON().AuthorizationContext;
}
catch (error) {
if (!(error instanceof KustoDataErrors.ThrottlingError)) {
throw error;
}
await retry.backoff();
}
}
throw new Error(`Failed to get identity token from server - the request was throttled ${ATTEMPT_COUNT} times.`);
}
async getIngestionQueues() {
const queues = (await this.refreshIngestClientResources()).securedReadyForAggregationQueues;
return queues ? this.getRoundRobinRankedAndShuffledResources(queues) : null;
}
async getFailedIngestionsQueues() {
return (await this.refreshIngestClientResources()).failedIngestionsQueues;
}
async getSuccessfulIngestionsQueues() {
return (await this.refreshIngestClientResources()).successfulIngestionsQueues;
}
async getContainers() {
const containers = (await this.refreshIngestClientResources()).containers;
return containers ? this.getRoundRobinRankedAndShuffledResources(containers) : null;
}
async getAuthorizationContext() {
return this.refreshAuthorizationContext();
}
async getStatusTables() {
return (await this.refreshIngestClientResources()).statusTables;
}
async createStatusTable() {
const statusTables = await this.getStatusTables();
if (!statusTables) {
throw new Error("Failed to get status table");
}
return createStatusTableClient(statusTables[0].uri);
}
close() {
this.kustoClient.close();
}
reportResourceUsageResult(accountName, success) {
this.rankedStorageAccountSet.logResultToAccount(accountName, success);
}
}
export const createStatusTableClient = (uri) => {
const tableUrl = new URL(uri);
const origin = tableUrl.origin;
const sasToken = tableUrl.search;
const tableName = tableUrl.pathname.replace("/", "");
return new TableClient(origin + sasToken, tableName);
};
export default ResourceManager;
//# sourceMappingURL=resourceManager.js.map