SQL API: AWS Athena Datasource Failed

I’m encountering an issue with the SQL API integration. Specifically, I’m trying to configure Cognito authentication to access the SQL API. While I’m able to establish a connection and pass the required data for model building via the securityContext in the check_sql_auth function, I'm running into problems with query execution. After connecting to the SQL API (e.g., through psql-cli), I can view all tables associated with the user by running commands like \dt. However, when attempting to execute a SELECT query, I receive the following error:

ERROR: Arrow error: Compute error: Error: The security token included in the request is invalid.

The data source for the models is AWS Athena. Upon inspecting the Cube logs, I noticed errors indicating that the security token is not being recognized during the connection to Athena:

"error":"UnrecognizedClientException: The security token included in the request is invalid.\n at throwDefaultError (/cube/node_modules/@smithy/smithy-client/dist-cjs/index.js:838:20)\n at /cube/node_modules/@smithy/smithy-client/dist-cjs/index.js:847:5\n at de_CommandError (/cube/node_modules/@aws-sdk/client-athena/dist-cjs/index.js:1829:14)\n at processTicksAndRejections (node:internal/process/task_queues:95:5)\n at /cube/node_modules/@smithy/middleware-serde/dist-cjs/index.js:35:20\n at /cube/node_modules/@smithy/core/dist-cjs/index.js:165:18\n at /cube/node_modules/@smithy/middleware-retry/dist-cjs/index.js:320:38\n at /cube/node_modules/@aws-sdk/middleware-logger/dist-cjs/index.js:33:22\n at AthenaDriver.testConnection (/cube/node_modules/@cubejs-backend/athena-driver/src/AthenaDriver.ts:201:5)\n at /cube/node_modules/@cubejs-backend/server-core/src/core/server.ts:591:15\n at Object.query (/cube/node_modules/@cubejs-backend/query-orchestrator/src/orchestrator/QueryCache.ts:593:26)\n at QueryQueue.processQuery (/cube/node_modules/@cubejs-backend/query-orchestrator/src/orchestrator/QueryQueue.js:842:25)"

Here’s an excerpt from my check_sql_auth function:

@config('check_sql_auth')
def check_sql_auth(req: dict, user_name: str, password: str) -> dict:

    token = get_tenant_token(user_name, password)

    if token:
        payload = get_payload({'token': token}, None)
        return {
            "password": password,
            "securityContext": {
                USER_UUID: payload.get("user-uuid"),
                DATA_SOURCE: payload.get('data-source'),
            }
        }

    raise Exception("Incorrect user name or password")

Has anyone experienced a similar issue or could provide guidance on what might be going wrong? I would appreciate any suggestions or troubleshooting steps.

Version: [0.35.29]

cube-js / cube

SQL API: AWS Athena Datasource Failed #8778