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

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

View File

@ -55,11 +55,15 @@ class AthenaResponse(BaseResponse):
def get_query_execution(self) -> str:
exec_id = self._get_param("QueryExecutionId")
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 = {
"QueryExecution": {
"QueryExecutionId": exec_id,
"Query": execution.query,
"StatementType": "DDL",
"StatementType": statement_type,
"ResultConfiguration": execution.config,
"QueryExecutionContext": execution.context,
"Status": {

View File

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