CUMULUS-4954: add AthenaClient by charleshuang80 · Pull Request #4567 · nasa/cumulus

charleshuang80 · 2026-08-14T20:20:47Z

Summary: Summary of changes

Addresses [CUMULUS-4954: Implement Athena Client](https://bugs.earthdata.nasa.gov/browse/CUMULUS-4954)

Changes

Detailed list or prose of changes
...

PR Checklist

Update CHANGELOG
Unit tests
Ad-hoc testing - Deploy changes and test manually
Integration tests

📝 Note:
For most pull requests, please Squash and merge to maintain a clean and readable commit history.

for more information, see https://pre-commit.ci

Updated getQueryResults get retrieve all results

for more information, see https://pre-commit.ci

charleshuang80 · 2026-09-16T17:20:42Z

+// TODO: Remove this comment when localstack is replaced:
+// Athena client tests are unable to run using localstack local environment.
+// Once localstack is replaced, these tests will need to be updated to use
+// local AWS instances. This work should be completed in CUMULUS-5307.


This is probably my mistake, but I think 'instances' should actually be 'services' here.

charleshuang80 · 2026-09-16T17:21:46Z

+});
+
+test.serial('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,


I think this comment is OBE

charleshuang80 · 2026-09-16T17:23:32Z

+  sinon.stub(t.context.client, 'getQueryExecution')
+    .callsFake(() => Promise.resolve(abridgedResponse));
+
+  const expected = [{ result: 'mock_value' }];


Since that comment is OBE, and we're not using an endpoint that returns this, this expected should probably match what we are inserting into the table.

charleshuang80 · 2026-09-16T17:25:03Z

+
+  const abridgedResponse = {
+    QueryExecution: {
+      QueryExecutionId: '1234-abcd-5678-efgh',


Really minor nit, which is kinda not important since we are mocking everything. But, it would probably be better if these Ids were not all the same. We could maybe randomize them or generate a uuid.

charleshuang80 · 2026-09-16T17:25:47Z

+  const abridgedResponse = {
+    QueryExecution: {
+      QueryExecutionId: '1234-abcd-5678-efgh',
+      Query: '',


Also, the query should probably actually reflect the query we are sending for this and the other responses.

charleshuang80 · 2026-09-16T17:29:37Z

+  t.deepEqual([], mappedResult);
+});
+
+test.serial('query() initiates a query, waits for it to finish, and returns the mapped response', async (t) => {


getQueryResults is private so we probably can't call that directly to test and may need to use query, but we should probably test that the NextToken 'pagination' actually works.

Added NextToken test

for more information, see https://pre-commit.ci

charleshuang80 · 2026-10-01T18:08:58Z

+    let nextToken = response?.NextToken?.toString() || '';
+    data = data.concat(await this.mapData(response.ResultSet));
+
+    while (response?.NextToken?.toString()) {


shouldn't the while loop be checking against nextToken?

charleshuang80 · 2026-10-01T18:09:52Z

+
+    while (response?.NextToken?.toString()) {
+      // eslint-disable-next-line no-await-in-loop
+      const response = await this.client.send(new GetQueryResultsCommand({


as the lint checker is pointing out, response is declared above, so that should probably use let and the const can be dropped here

charleshuang80 · 2026-10-01T18:21:44Z

+  };
+
+  athenaClientMock.on(GetQueryResultsCommand).resolvesOnce(
+    getDataResultsResponse


I'm a little confused by the chained and ordered responses here. getDataResultsWithNextToken actually includes a nextToken, but getDataResultsResponse does not. So I would think that getDataResultsWithNextToken should come first and would make the loop want to make another request, and then when getDataResultsResponse is the response that doesn't include a nextToken, then the loop would stop.
Also it looks like there are 5 resolvesOnce and 1 resolves (which would be the default), so I would think there should be 6 rows of data. But the test is expecting 3 rows. Am I missing something?

Fixed next token loop

charleshuang80 · 2026-10-29T14:27:32Z

We need some sort of testing to make sure the functions work. You can either use a test script to call the functions against a real AWS environment, or wait until LocalStack is replaced.

Using mocked tests because we were unable to find an easy replacement for Localstack was discussed in at least 3 or 4 different Arch meetings. It was determined and stated by Nathan that we would proceed with using mocked tests to get this functionality in, and we would carry some tech debt to replace Localstack and have tests written against that replacement. There were no objections raised. So I wrote some tickets to that effect, which I think you are looking at.

Added pRetry to athena client Updated athena client format Removed debug logging

for more information, see https://pre-commit.ci

sylvesterkaczmarek

Thanks for the latest cleanup and test updates. Re-reviewed head 88e90d7 and found a new result-mapping regression that needs correcting before merge.

Athena's actual ResultSetMetadata.ColumnInfo is an array of column descriptors, one per selected column, with fields such as Name, Label, and Type. Both getQueryResults() and mapData() now use Object.keys(ColumnInfo[0]). That returns descriptor property names (Name, Type, etc.), not the SQL result column names. For example, with real metadata [{Name: 'bucket', Type: 'varchar'}, {Name: 'key', Type: 'varchar'}], the mapped result is keyed as {Name: ..., Type: ...} rather than {bucket: ..., key: ...}; additional values may be assigned under an undefined key.

The current mocks hide this because they use ColumnInfo: [{bucket: 'string', key: 'string', ...}], which isn't the Athena API's response shape. AWS reference: https://docs.aws.amazon.com/athena/latest/APIReference/API_GetQueryResults.html

Please derive the columns from ColumnInfo.map(column => column.Name) (with appropriate absent-metadata handling), or use the first-page header row as before, and carry those names into continuation pages. A regression with actual descriptor-shaped metadata, a header on page 1, and data-only rows on page 2 should assert the exact object keys and values. The earlier pagination issue remains fixed, but this metadata change reintroduces incorrect mappings, so I'm requesting changes on the current revision.

rswanso5-nasa · 2026-11-02T16:33:06Z

Updated the PR with the requested changes. I believe that this will resolve the issues addressed above. The only outstanding issue is the that the tests are currently being mocked. Changing those to real tests is blocked by another ticket currently in PR and is outside the scope of this ticket.

sylvesterkaczmarek · 2026-11-02T18:02:37Z

Thanks for the update, @rswanso5-nasa. I checked head 423078b, but the mapping issue from my 30 October review is still present: both getQueryResults() and mapData() use Object.keys(ResultSetMetadata.ColumnInfo[0]). Athena returns ColumnInfo as one descriptor per result column (with properties like Name and Type), so this extracts descriptor field names rather than SQL column names. The mocked ColumnInfo: [{bucket: 'string', key: 'string', ...}] still masks that mismatch. Could the columns come from ColumnInfo.map(column => column.Name) (with suitable missing-metadata handling), with a regression using real descriptor-shaped metadata and a data-only continuation page? I agree replacing the mocked suite with live integration tests can be handled in the separate ticket. Happy to re-review once this mapping is corrected.

Fix format Updated constructor with init values Refactor Athena client

Updated changelog Format change to athena client

sylvesterkaczmarek · 2026-11-05T09:39:23Z

Thanks @rswanso5-nasa for updating the Athena result mapping and descriptor-shaped fixtures. I rechecked head 1d761c6: the original Object.keys(ColumnInfo[0]) problem is fixed, and the new tests use the correct Name fields. One pagination edge remains: formatResultSet() reads ColumnInfo anew on each page and dereferences columnInfo[j].Name. If a continuation response contains data-only rows but omits optional ResultSetMetadata.ColumnInfo, this throws instead of using the first page’s known column names. The pagination mocks currently include full metadata on pages 2 and 3, so they cannot catch this. Could the mapper carry the first page’s column names forward as a fallback and include a data-only continuation page without metadata in the regression? The previously reported descriptor-key mismatch is resolved; this is the remaining resilience concern.

Updated function docs

jennyhliu

Approving the PR with the expectation that the following items will be addressed in follow-up tickets/PRs.

The current getQueryResults function fetches all results at once. Consider providing a separate asynchronous generator function to retrieve the results in batches.
#4567 (comment)
Add comprehensive unit and/or integration tests to ensure code reliability.
Currently, getQueryResults returns an array of objects where all field values are stored as strings, regardless of their actual database schema type. Depending on how this function is consumed, we should consider converting these fields into their native types (e.g., numbers, booleans) using the result set metadata.

Updated tests

charleshuang80 and others added 5 commits August 14, 2026 16:19

CUMULUS-4954: add AthenaClient, no tests yet, no changes to localstack

9226d08

CUMULUS-4954: mock athena tests

5fee5d9

[pre-commit.ci] auto fixes from pre-commit.com hooks

06183f0

for more information, see https://pre-commit.ci

CUMULUS-4954: lint fix

60f3f70

CUMULUS-4954: lint fix

126524f