UNPKG

azure-kusto-ingest

Version:
124 lines 6.51 kB
// Copyright (c) Microsoft Corporation. // Licensed under the MIT License. import { Client as KustoClient, KustoConnectionStringBuilder } from "azure-kusto-data"; import ResourceManager from "./resourceManager.js"; import IngestionBlobInfo from "./ingestionBlobInfo.js"; import { ContainerClient } from "@azure/storage-blob"; import { QueueClient } from "@azure/storage-queue"; import { ReportLevel, ReportMethod } from "./ingestionProperties.js"; import { AbstractKustoClient } from "./abstractKustoClient.js"; import { TableReportIngestionResult, IngestionStatusInTableDescription, IngestionStatusResult, OperationStatus, putRecordInTable, } from "./ingestionResult.js"; import { BlobDescriptor, StreamDescriptor } from "./descriptors.js"; export class KustoIngestClientBase extends AbstractKustoClient { constructor(kcsb, defaultProps, autoCorrectEndpoint = true, isBrowser) { super(defaultProps); if (typeof kcsb === "string") { kcsb = new KustoConnectionStringBuilder(kcsb); } if (autoCorrectEndpoint) { kcsb.dataSource = this.getIngestionEndpoint(kcsb.dataSource); } const kustoClient = new KustoClient(kcsb); this.resourceManager = new ResourceManager(kustoClient, isBrowser); this.defaultDatabase = kustoClient.defaultDatabase; const clientDetails = kcsb.clientDetails(); this.applicationForTracing = clientDetails.applicationNameForTracing; this.clientVersionForTracing = clientDetails.versionForTracing; } async ingestFromBlob(blob, ingestionProperties, maxRetries = KustoIngestClientBase.MaxNumberOfRetryAttempts) { this.ensureOpen(); const props = this._getMergedProps(ingestionProperties); const descriptor = blob instanceof BlobDescriptor ? blob : new BlobDescriptor(blob); const authorizationContext = await this.resourceManager.getAuthorizationContext(); const ingestionBlobInfo = new IngestionBlobInfo(descriptor, props, authorizationContext, this.applicationForTracing, this.clientVersionForTracing); const reportToTable = props.reportLevel !== ReportLevel.DoNotReport && props.reportMethod !== ReportMethod.Queue; if (reportToTable) { const statusTableClient = await this.resourceManager.createStatusTable(); const status = this.createStatusObject(props, OperationStatus.Pending, ingestionBlobInfo); await putRecordInTable(statusTableClient, Object.assign(Object.assign({}, status), { partitionKey: ingestionBlobInfo.Id, rowKey: ingestionBlobInfo.Id })); const desc = new IngestionStatusInTableDescription(statusTableClient.url, ingestionBlobInfo.Id, ingestionBlobInfo.Id); ingestionBlobInfo.IngestionStatusInTable = desc; await this.sendQueueMessage(maxRetries, ingestionBlobInfo); return new TableReportIngestionResult(desc, statusTableClient); } await this.sendQueueMessage(maxRetries, ingestionBlobInfo); return new IngestionStatusResult(this.createStatusObject(props, OperationStatus.Queued, ingestionBlobInfo)); } async sendQueueMessage(maxRetries, blobInfo) { const queues = await this.resourceManager.getIngestionQueues(); if (queues == null) { throw new Error("Failed to get queues"); } const ingestionBlobInfoJson = JSON.stringify(blobInfo); const encoded = Buffer.from(ingestionBlobInfoJson).toString("base64"); const retryCount = Math.min(maxRetries, queues.length); for (let i = 0; i < retryCount; i++) { const queueClient = new QueueClient(queues[i].uri); try { const queueResponse = await queueClient.sendMessage(encoded); this.resourceManager.reportResourceUsageResult(queueClient.accountName, true); return queueResponse; } catch (_) { this.resourceManager.reportResourceUsageResult(queueClient.accountName, false); } } throw new Error("Failed to send message to queue."); } createStatusObject(props, status, ingestionBlobInfo) { const time = Date.now().toString(); return { Status: status, Timestamp: time, IngestionSourceId: ingestionBlobInfo.Id, IngestionSourcePath: ingestionBlobInfo.BlobPath.split(/[?;]/)[0], Database: props.database, Table: props.table, UpdatedOn: time, Details: "", }; } async uploadToBlobWithRetry(descriptor, blobName, maxRetries = KustoIngestClientBase.MaxNumberOfRetryAttempts) { const containers = await this.resourceManager.getContainers(); if (containers == null || containers.length === 0) { throw new Error("Failed to get containers"); } const retryCount = Math.min(maxRetries, containers.length); // Go over all containers and try to upload the file to the first one that succeeds for (let i = 0; i < retryCount; i++) { const containerClient = new ContainerClient(containers[i].uri); const blockBlobClient = containerClient.getBlockBlobClient(blobName); try { if (typeof descriptor == "string") { await blockBlobClient.uploadFile(descriptor); } else if (descriptor instanceof StreamDescriptor) { if (descriptor.stream instanceof Buffer) { await blockBlobClient.uploadData(descriptor.stream); } else { await blockBlobClient.uploadStream(descriptor.stream); } } else { await blockBlobClient.uploadData(descriptor); } this.resourceManager.reportResourceUsageResult(containerClient.accountName, true); return blockBlobClient.url; } catch (ex) { this.resourceManager.reportResourceUsageResult(containerClient.accountName, false); } } throw new Error("Failed to upload to blob."); } close() { if (!this._isClosed) { this.resourceManager.close(); } super.close(); } } KustoIngestClientBase.MaxNumberOfRetryAttempts = 3; export default KustoIngestClientBase; //# sourceMappingURL=ingestClientBase.js.map