Yasuke/moto
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Athena: make statement type dependent on query command (#6786)

Browse Source
This commit is contained in:
Michael French 2023-09-08 08:13:26 -05:00 committed by GitHub
parent 056b69bee7
commit 1d3ddc9a19
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
moto/athena/responses.py
View File

@ -55,11 +55,15 @@ class AthenaResponse(BaseResponse):
def get_query_execution(self) -> str: def get_query_execution(self) -> str:
exec_id = self._get_param("QueryExecutionId") exec_id = self._get_param("QueryExecutionId")
execution = self.athena_backend.get_query_execution(exec_id) execution = self.athena_backend.get_query_execution(exec_id)
ddl_commands = ("ALTER", "CREATE", "DESCRIBE", "DROP", "MSCK", "SHOW")
statement_type = "DML"
if execution.query.upper().startswith(ddl_commands):
statement_type = "DDL"
result = { result = {
"QueryExecution": { "QueryExecution": {
"QueryExecutionId": exec_id, "QueryExecutionId": exec_id,
"Query": execution.query, "Query": execution.query,
"StatementType": "DDL", "StatementType": statement_type,
"ResultConfiguration": execution.config, "ResultConfiguration": execution.config,
"QueryExecutionContext": execution.context, "QueryExecutionContext": execution.context,
"Status": { "Status": {

2
tests/test_athena/test_athena.py
View File

@ -140,7 +140,7 @@ def test_get_query_execution():
# #
assert details["QueryExecutionId"] == exex_id assert details["QueryExecutionId"] == exex_id
assert details["Query"] == query assert details["Query"] == query
assert details["StatementType"] == "DDL" assert details["StatementType"] == "DML"
assert details["ResultConfiguration"]["OutputLocation"] == location assert details["ResultConfiguration"]["OutputLocation"] == location
assert details["QueryExecutionContext"]["Database"] == database assert details["QueryExecutionContext"]["Database"] == database
assert details["Status"]["State"] == "SUCCEEDED" assert details["Status"]["State"] == "SUCCEEDED"