diff --git a/CHANGELOG.md b/CHANGELOG.md index 37370f95623..62b3fd66496 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -71,6 +71,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). is now executed automatically by the DB migration Lambda handler post-migration. - **CUMULUS-4986** - Added `storage_type` variable to `tf-modules/cumulus-rds-tf` module with default value `aurora`. +- **CUMULUS-4954** + - Added `AthenaQueryClient` to `packages/aws-client` that can interact with and run queries in Athena. ### Changed diff --git a/packages/aws-client/package.json b/packages/aws-client/package.json index 07765dd5e95..e4dff406b2c 100644 --- a/packages/aws-client/package.json +++ b/packages/aws-client/package.json @@ -48,6 +48,7 @@ "license": "Apache-2.0", "dependencies": { "@aws-sdk/client-api-gateway": "^3.993.0", + "@aws-sdk/client-athena": "^3.993.0", "@aws-sdk/client-cloudformation": "^3.993.0", "@aws-sdk/client-cloudwatch-events": "^3.993.0", "@aws-sdk/client-dynamodb": "^3.993.0", diff --git a/packages/aws-client/src/AthenaQueryClient.ts b/packages/aws-client/src/AthenaQueryClient.ts new file mode 100644 index 00000000000..44aa19b4c50 --- /dev/null +++ b/packages/aws-client/src/AthenaQueryClient.ts @@ -0,0 +1,237 @@ +/** + * module AthenaQueryClient + */ + +import { + AthenaClient, + AthenaClientConfig, + StartQueryExecutionCommand, + GetQueryExecutionCommand, + QueryExecutionState, + GetQueryExecutionCommandOutput, + GetQueryResultsCommand, + ResultSet, +} from '@aws-sdk/client-athena'; + +import isNil from 'lodash/isNil'; +import Logger from '@cumulus/logger'; + +const log = new Logger({ sender: 'aws-client/AthenaQueryClient' }); + +interface ResultReuseConfiguration { + ResultReuseByAgeConfiguration: { + Enabled: boolean, + MaxAgeInMinutes?: number, + } +} +interface ResultConfiguration { + OutputLocation: string; + EncryptionConfiguration?: { // EncryptionConfiguration + EncryptionOption: 'SSE_S3' | 'SSE_KMS' | 'CSE_KMS'; // required + KmsKey?: string; + }; + ExpectedBucketOwner?: string; + AclConfiguration?: { // AclConfiguration + S3AclOption: 'BUCKET_OWNER_FULL_CONTROL'; // required + }; +} + +interface AthenaQueryClientConfig { + ClientConfig: AthenaClientConfig; + Database: string; + Catalog: string; + ResultConfiguration: ResultConfiguration; + WorkGroup?: string; + ResultReuseConfiguration?: ResultReuseConfiguration; +} + +type MappedObject = { [index: string]: string }; +type MappedData = Array; + +export class AthenaQueryClient { + public database: string; + private client: AthenaClient; + private catalog: string; + private workGroup: string = 'primary'; + private resultConfiguration: ResultConfiguration | undefined; + private resultReuseConfiguration: ResultReuseConfiguration = { + ResultReuseByAgeConfiguration: { + Enabled: true, + MaxAgeInMinutes: 60, + }, + }; + + constructor(config: AthenaQueryClientConfig) { + this.client = new AthenaClient(config.ClientConfig); + this.database = config.Database; + this.catalog = config.Catalog; + + if (config.WorkGroup) this.workGroup = config.WorkGroup; + if (config.ResultConfiguration) { + this.resultConfiguration = config.ResultConfiguration; + } + if (config.ResultReuseConfiguration) { + this.resultReuseConfiguration = config.ResultReuseConfiguration; + } + } + + /** + * Get data from Athena and rerutn it as proper formatted Array of objects + * + * @param {string} sqlQuery - The SQL query string + * @returns {Array} Array of Objects + */ + async query(sqlQuery: string): Promise { + const queryExecutionId = await this.startQueryExecution(sqlQuery); + + const response = await this.checkQueryExecutionStateAndGetData(queryExecutionId); + log.info(`response (${typeof response}) from checkQueryExecutionStateAndGetData: ${JSON.stringify(response)}`); + return response; + } + + /** + * Start Query Execution + * + * @param {string} sqlQuery - The SQL query string + * @returns {string} QueryExecutionId - unique ID of the query run from request + */ + async startQueryExecution(sqlQuery: string): Promise { + const queryExecutionInput = { + QueryString: sqlQuery, + QueryExecutionContext: { + Database: this.database, + Catalog: this.catalog, + }, + ResultConfiguration: this.resultConfiguration, + WorkGroup: this.workGroup, + ResultReuseConfiguration: this.resultReuseConfiguration, + }; + log.info(`about to run query with ${JSON.stringify(queryExecutionInput)}`); + + const { QueryExecutionId } = await this.client.send( + new StartQueryExecutionCommand(queryExecutionInput) + ); + log.info(`from query execution, got back ${QueryExecutionId}, which is a ${typeof QueryExecutionId}`); + + if (QueryExecutionId === undefined) { + throw new Error('QueryExecutionId was returned by Athena StartQueryExecutionCommand as undefined'); + } + return QueryExecutionId; + } + + /** + * Get query execution status and output + * + * @param {string} QueryExecutionId - Id of a query which we sent to Athena + * @returns {GetQueryExecutionCommandOutput} - output from GetQueryExecutionCommand + */ + private async getQueryExecution( + QueryExecutionId: string + ): Promise { + const command = new GetQueryExecutionCommand({ QueryExecutionId }); + return await this.client.send(command); + } + + /** + * Check query exeqution state + * if it is "QUEUED" or "RUNNING", recursively call to check the state + * with increasing polling delays until the state is "SUCCEEDED" and after it we get the data + * + * @param {string} QueryExecutionId - Id of a query which we sent to Athena + * @param {number} delay - polling interval passed in, in millisecs + * @returns {Array} Array of Objects + */ + private async checkQueryExecutionStateAndGetData( + QueryExecutionId: string, + delay: number = 0 + ): Promise { + const response = await this.getQueryExecution(QueryExecutionId); + const state = response.QueryExecution?.Status?.State; + log.info(`response (${typeof response}) and state (${typeof state}) ${state} from GetQueryExecutionCommand. ${JSON.stringify(response)}`); + + if (state === QueryExecutionState.FAILED) { + throw new Error(`Query failed: ${response.QueryExecution!.Status!.StateChangeReason}`); + } else if (state === QueryExecutionState.CANCELLED) { + throw new Error('Query was cancelled'); + } else if (state === QueryExecutionState.SUCCEEDED) { + return await this.getQueryResults(QueryExecutionId); + } else if (state === QueryExecutionState.QUEUED || state === QueryExecutionState.RUNNING) { + // polling intervals: 1000 (1s), 600000 (10m), 3600000 (60m/1h) + let delayPass = delay; + if (delayPass <= 1000) { + delayPass = 1000; + await this.timeout(delayPass); + delayPass += 4000; + } else if (delayPass <= 600000) { + await this.timeout(delayPass); + delayPass *= 2; + } else if (delayPass <= 3600000) { + await this.timeout(delayPass); + delayPass += 600000; + } else { + log.error(`delays have become ${delayPass}, longer than an hour. time to abort`); + throw new Error(`Query ${QueryExecutionId} was queued or running for too long`); + } + + log.info(`about to rerun checkQueryExecutionStateAndGetData with delay ${delayPass} (also ${delayPass / 1000}s)`); + return await this.checkQueryExecutionStateAndGetData(QueryExecutionId, delayPass); + } + log.error(`end of checkQueryExecutionStateAndGetData reached, state ${state} not processed. response: ${JSON.stringify(response)}`); + return undefined; + } + + /** + * Get query execution result + * + * @param {string} QueryExecutionId - Id of a query which we sent to Athena + * @returns {Array} Array of Objects + */ + private async getQueryResults(QueryExecutionId: string): Promise { + const response = await this.client.send(new GetQueryResultsCommand({ + QueryExecutionId, + })); + log.info(`response (${typeof response}) from GetQueryResults: ${JSON.stringify(response)}`); + return this.mapData(response.ResultSet); + } + + /** + * Map data returned from Athena in rows, with each row an object with columns/keys and values. + * + * @param {ResultSet} data - Data in rows returned from Athena Query, in the ResultSet format + * @returns {MappedData} Array of rows of data as MappedObjects + */ + mapData(data: ResultSet | undefined): MappedData { + const mappedData: MappedData = []; + if (data === undefined || data.Rows === undefined || data.Rows.length === 0) return mappedData; + + const columns: string[] = data.Rows[0].Data!.map((column) => column.VarCharValue as string); + + data.Rows.forEach((item, i) => { + if (i === 0) return; + if (item.Data === undefined) return; + + const mappedObject: MappedObject = {}; + item.Data.forEach((datum, j) => { + if (isNil(datum.VarCharValue)) { + mappedObject[columns[j]] = ''; + } else { + mappedObject[columns[j]] = datum.VarCharValue; + } + }); + + mappedData.push(mappedObject); + }); + + return mappedData; + } + + /** + * Simple helper timeout function uses in checkQueryExecutionStateAndGetData function + * + * @param {number} msTime - Time in miliseconds + * @returns {Promise} Promise + */ + private timeout(msTime: number) { + return new Promise((resolve) => setTimeout(resolve, msTime)); + } +} diff --git a/packages/aws-client/src/services.ts b/packages/aws-client/src/services.ts index 04f97ee8ea2..01965f78809 100644 --- a/packages/aws-client/src/services.ts +++ b/packages/aws-client/src/services.ts @@ -1,4 +1,5 @@ import { APIGatewayClient } from '@aws-sdk/client-api-gateway'; +import { AthenaClient } from '@aws-sdk/client-athena'; import { CloudFormation } from '@aws-sdk/client-cloudformation'; import { DynamoDB } from '@aws-sdk/client-dynamodb'; import { DynamoDBDocument, TranslateConfig } from '@aws-sdk/lib-dynamodb'; @@ -19,6 +20,7 @@ import { EC2 } from '@aws-sdk/client-ec2'; import awsClient from './client'; export const apigateway = awsClient(APIGatewayClient, '2015-07-09'); +export const athena = awsClient(AthenaClient, '2017-05-18'); export const ecs = awsClient(ECS, '2014-11-13'); export const ec2 = awsClient(EC2, '2016-11-15'); export const cloudwatchevents = awsClient(CloudWatchEvents, '2015-10-07'); diff --git a/packages/aws-client/src/test-utils.ts b/packages/aws-client/src/test-utils.ts index ce512974996..ceedd702a0d 100644 --- a/packages/aws-client/src/test-utils.ts +++ b/packages/aws-client/src/test-utils.ts @@ -9,6 +9,7 @@ export const inTestMode = () => process.env.NODE_ENV === 'test'; // From https://github.com/localstack/localstack/blob/master/README.md const localStackPorts = { APIGatewayClient: 4566, + AthenaClient: 4566, CloudFormation: 4566, CloudWatchEvents: 4566, DynamoDB: 4566, diff --git a/packages/aws-client/src/types.ts b/packages/aws-client/src/types.ts index a73d376197e..72229eaf93e 100644 --- a/packages/aws-client/src/types.ts +++ b/packages/aws-client/src/types.ts @@ -1,4 +1,5 @@ import { APIGatewayClient } from '@aws-sdk/client-api-gateway'; +import { AthenaClient } from '@aws-sdk/client-athena'; import { CloudWatchEvents } from '@aws-sdk/client-cloudwatch-events'; import { CloudFormation } from '@aws-sdk/client-cloudformation'; import { DynamoDBStreamsClient } from '@aws-sdk/client-dynamodb-streams'; @@ -18,6 +19,7 @@ import { STS } from '@aws-sdk/client-sts'; export type AWSClientTypes = APIGatewayClient | + AthenaClient | DynamoDB | DynamoDBClient | DynamoDBStreamsClient | diff --git a/packages/aws-client/tests/test-AthenaQueryClient.js b/packages/aws-client/tests/test-AthenaQueryClient.js new file mode 100644 index 00000000000..79d1e110194 --- /dev/null +++ b/packages/aws-client/tests/test-AthenaQueryClient.js @@ -0,0 +1,235 @@ +'use strict'; + +// TODO: Remove this comment when localstack is replaced: +// Athena client tests are unable to run in localstack local environment. +// Once localstack is replaced these tests will need to be updated to use +// local AWS instances. This work is outside the scope of CUMULUS-4954. + +const test = require('ava'); +const cryptoRandomString = require('crypto-random-string'); +const sinon = require('sinon'); +const { AthenaClient, StartQueryExecutionCommand } = require('@aws-sdk/client-athena'); + +// TODO: remove mock +const { mockClient } = require('aws-sdk-client-mock'); +const athenaClientMock = mockClient(AthenaClient); + +const { AthenaQueryClient } = require('../AthenaQueryClient'); + +const { + createBucket, + recursivelyDeleteS3Bucket, +} = require('../S3'); + +const randomString = () => cryptoRandomString({ + length: 10, + characters: 'abcdefghijklmnopqrstuvwxyz', // https://docs.aws.amazon.com/athena/latest/ug/tables-databases-columns-names.html +}); + +test.before(async (t) => { + t.context.Bucket = randomString(); + await createBucket(t.context.Bucket); + + t.context.db = `${randomString()}_testdb`; + + // TODO: update clint once localstack is replaced + t.context.client = new AthenaQueryClient({ + ClientConfig: { + region: 'us-east-1', + endpoint: 'http://localhost:4566', + credentials: { + accessKeyId: 'test', + secretAccessKey: 'test', + }, + }, + Database: t.context.db, + ResultConfiguration: { OutputLocation: `s3://${t.context.Bucket}/` }, + }); +}); + +test.afterEach.always(() => { + sinon.restore(); +}); + +test.after.always(async (t) => { + await recursivelyDeleteS3Bucket(t.context.Bucket); +}); + +// TODO: update test once localstack is replced +test('startQueryExecution() initiates a query and receives a QueryExecutionId response', async (t) => { + athenaClientMock.on(StartQueryExecutionCommand).resolves({ + queryId: '12345-abcde-67890', + }); + const tableName = `${randomString()}_table`; + const client = new AthenaClient({ region: 'us-east-1' }); + const command = new StartQueryExecutionCommand({ + QueryString: `CREATE TABLE IF NOT EXISTS ${tableName} +( bucket string, key string, version_id string, is_latest boolean, is_delete_marker boolean);`, + }); + + const response = await client.send(command); + t.is(response.queryId, '12345-abcde-67890'); +}); + +// TODO: fix this test once localstack is replaced +test.skip('mapData returns data in the expected format', (t) => { + const testBucket = 'daac-public-bucket'; + const testKey = `${randomString()}`; + + const expected = [ + { bucket: testBucket, key: testKey, version_id: '', is_latest: true, is_delete_marker: false }, + ]; + + // response is in the shape of GetQueryResultsCommand Output + const response = { + UpdateCount: 0, + ResultSet: { + Rows: [ + { Data: [ + { VarCharValue: 'bucket' }, + { VarCharValue: 'key' }, + { VarCharValue: 'version_id' }, + { VarCharValue: 'is_latest' }, + { VarCharValue: 'is_delete_marker' }, + ] }, + { Data: [ + { VarCharValue: testBucket }, + { VarCharValue: testKey }, + {}, + { VarCharValue: true }, + { VarCharValue: false }, + ] }, + ], + }, + }; + const mappedResult = t.context.client.mapData(response.ResultSet); + + t.deepEqual(expected, mappedResult); +}); + +// TODO: fix this test once localstack is replaced +test.skip('mapData() returns expected result when ResultSet is empty', (t) => { + // responses have emtpy ResultSet.Rows from queries like create tables or views + const response = { + UpdateCount: 0, + ResultSet: { Rows: [], ResultSetMetadata: { ColumnInfo: [] } }, + }; + const mappedResult = t.context.client.mapData(response.ResultSet); + + t.deepEqual([], mappedResult); +}); + +// TODO: update test once localstack is replced +test.skip('query() initiates a query, waits for it to finish, and returns the mapped response', async (t) => { + // could not get ministack duckdb to find a table to perform operations on it, + // even after verifying a create table query succeeeded + // so using the mocked db version of Athena in ministack, which returns mock data + const expected = [{ result: 'mock_value' }]; + + const dbQuery = `CREATE DATABASE IF NOT EXISTS ${t.context.db}`; + const dbResponse = await t.context.client.query(dbQuery); + console.log(`data after createDb: ${JSON.stringify(dbResponse)}`); + + const tableName = `${randomString()}_table`; + // create table + const tableQuery = `CREATE TABLE IF NOT EXISTS ${tableName} + ( bucket string, key string, version_id string, is_latest boolean, is_delete_marker boolean);`; + + await t.context.client.query(tableQuery); + + const testBucket = 'daac-public-bucket'; + const testKey = `${randomString()}`; + // populate table + const addDataQuery = `INSERT INTO ${tableName} VALUES ('${testBucket}', '${testKey}', '', true, false);`; + await t.context.client.query(addDataQuery); + + // get data + const getDataQuery = `SELECT * FROM ${tableName};`; + const results = await t.context.client.query(getDataQuery); + + t.deepEqual(results, expected); +}); + +// TODO: fix this test once localstack is replaced +test.skip('checkQueryExecutionStateAndGetData throws when getQueryExecution returns with a CANCELLED state', async (t) => { + const tableName = `${randomString()}_table`; + const tableQuery = `CREATE TABLE IF NOT EXISTS ${tableName} + ( bucket string, key string, version_id string, is_latest boolean, is_delete_marker boolean);`; + + await t.context.client.query(tableQuery); + + const testBucket = 'daac-public-bucket'; + const testKey = `${randomString()}`; + const addDataQuery = `INSERT INTO ${tableName} VALUES ('${testBucket}', '${testKey}', '', true, false);`; + + await t.context.client.query(addDataQuery); + + const abridgedResponse = { + QueryExecution: { + QueryExecutionId: '1234-abcd-5678-efgh', + Query: '', + ResultConfiguration: { + OutputLocation: `s3://${t.context.Bucket}/`, + }, + QueryExecutionContext: { + Database: t.context.db, + }, + Status: { + State: 'CANCELLED', + SubmissionDateTime: new Date().toISOString(), + }, + }, + }; + + sinon.stub(t.context.client, 'getQueryExecution') + .callsFake(() => Promise.resolve(abridgedResponse)); + + const getDataQuery = `SELECT * FROM ${tableName};`; + await t.throwsAsync( + t.context.client.query(getDataQuery), + { message: 'Query was cancelled' } + ); +}); + +// TODO: fix this test once localstack is replaced +test.skip('checkQueryExecutionStateAndGetData throws when getQueryExecution returns with a FAILED state', async (t) => { + const tableName = `${randomString()}_table`; + const tableQuery = `CREATE TABLE IF NOT EXISTS ${tableName} +( bucket string, key string, version_id string, is_latest boolean, is_delete_marker boolean);`; + + await t.context.client.query(tableQuery); + + const testBucket = 'daac-public-bucket'; + const testKey = `${randomString()}`; + const addDataQuery = `INSERT INTO ${tableName} VALUES ('${testBucket}', '${testKey}', '', true, false);`; + + await t.context.client.query(addDataQuery); + + const abridgedResponse = { + QueryExecution: { + QueryExecutionId: '1234-abcd-5678-efgh', + Query: '', + ResultConfiguration: { + OutputLocation: `s3://${t.context.Bucket}/`, + }, + QueryExecutionContext: { + Database: t.context.db, + }, + Status: { + State: 'FAILED', + StateChangeReason: 'some failure reason', + SubmissionDateTime: new Date().toISOString(), + }, + }, + }; + + sinon.stub(t.context.client, 'getQueryExecution') + .callsFake(() => Promise.resolve(abridgedResponse)); + + const getDataQuery = `SELECT * FROM ${tableName};`; + + await t.throwsAsync( + t.context.client.query(getDataQuery), + { message: 'Query failed: some failure reason' } + ); +});