65 lines
2.4 KiB
JavaScript
65 lines
2.4 KiB
JavaScript
|
|
import { AthenaClient, StartQueryExecutionCommand, GetQueryExecutionCommand, GetQueryResultsCommand } from "@aws-sdk/client-athena";
|
|
const athenaClient = new AthenaClient({ region: process.env.AWS_REGION, profile: 'default' });
|
|
|
|
const executionTimeout = 1000;
|
|
|
|
const productQuery = `SELECT DISTINCT
|
|
line_item_resource_id AS resourceId,
|
|
line_item_product_code AS productCode,
|
|
line_item_usage_account_id AS accountId
|
|
FROM ${process.env.ATHENA_CU_TABLE};`;
|
|
|
|
export const executeQueryAsync = async (sqlQuery) => {
|
|
const startQueryCommand = new StartQueryExecutionCommand({
|
|
QueryString: sqlQuery,
|
|
QueryExecutionContext: { Database: process.env.ATHENA_CU_DATABASE },
|
|
ResultConfiguration: { OutputLocation: process.env.ATHENA_OUTPUT_S3_BUCKET },
|
|
});
|
|
|
|
const startQueryResponse = await athenaClient.send(startQueryCommand);
|
|
const queryExecutionId = startQueryResponse.QueryExecutionId;
|
|
|
|
let queryExecutionStatus;
|
|
do {
|
|
const getQueryExecutionCommand = new GetQueryExecutionCommand({
|
|
QueryExecutionId: queryExecutionId,
|
|
});
|
|
const queryExecutionResponse = await athenaClient.send(getQueryExecutionCommand);
|
|
queryExecutionStatus = queryExecutionResponse.QueryExecution.Status.State;
|
|
|
|
if (queryExecutionStatus === "FAILED") {
|
|
console.error(`Query Failed: ${JSON.stringify(queryExecutionResponse)}`);
|
|
break;
|
|
}
|
|
await new Promise((resolve) => setTimeout(resolve, executionTimeout));
|
|
} while (queryExecutionStatus !== "SUCCEEDED");
|
|
|
|
return queryExecutionId;
|
|
}
|
|
|
|
export const retrieveResultsAsync = async (queryExecutionId) => {
|
|
const getQueryResultsCommand = new GetQueryResultsCommand({
|
|
QueryExecutionId: queryExecutionId,
|
|
});
|
|
|
|
const result = await athenaClient.send(getQueryResultsCommand);
|
|
if (!result || !result.ResultSet || !result.ResultSet.Rows) {
|
|
throw new Error('No results');
|
|
}
|
|
let rows = result.ResultSet.Rows;
|
|
if (0 == rows.length || 1 == rows.length) {
|
|
throw new Error('No result data');
|
|
}
|
|
const columnNames = rows[0].Data.map(item => item.VarCharValue);
|
|
const items = [];
|
|
rows.slice(1).forEach(cells => {
|
|
const item = {};
|
|
for (let i = 0; i < cells.Data.length; i++) {
|
|
item[columnNames[i]] = cells.Data[i].VarCharValue;
|
|
}
|
|
items.push(item);
|
|
});
|
|
return items;
|
|
};
|