This commit is contained in:
Bert Blommers 2020-05-16 15:03:26 +01:00
commit ffb521f86b
3 changed files with 79 additions and 69 deletions

View file

@ -50,7 +50,6 @@ class WorkGroup(TaggableResourceMixin, BaseModel):
class Execution(BaseModel):
def __init__(self, query, context, config, workgroup):
self.id = str(uuid4())
self.query = query
@ -97,11 +96,13 @@ class AthenaBackend(BaseBackend):
"State": wg.state,
"Configuration": wg.configuration,
"Description": wg.description,
"CreationTime": time.time()
"CreationTime": time.time(),
}
def start_query_execution(self, query, context, config, workgroup):
execution = Execution(query=query, context=context, config=config, workgroup=workgroup)
execution = Execution(
query=query, context=context, config=config, workgroup=workgroup
)
self.executions[execution.id] = execution
return execution.id

View file

@ -43,32 +43,34 @@ class AthenaResponse(BaseResponse):
workgroup = self._get_param("WorkGroup")
if workgroup and not self.athena_backend.get_work_group(workgroup):
return self.error("WorkGroup does not exist", 400)
id = self.athena_backend.start_query_execution(query=query, context=context, config=config, workgroup=workgroup)
id = self.athena_backend.start_query_execution(
query=query, context=context, config=config, workgroup=workgroup
)
return json.dumps({"QueryExecutionId": id})
def get_query_execution(self):
exec_id = self._get_param("QueryExecutionId")
execution = self.athena_backend.get_execution(exec_id)
result = {
'QueryExecution': {
'QueryExecutionId': exec_id,
'Query': execution.query,
'StatementType': 'DDL',
'ResultConfiguration': execution.config,
'QueryExecutionContext': execution.context,
'Status': {
'State': execution.status,
'SubmissionDateTime': execution.start_time
"QueryExecution": {
"QueryExecutionId": exec_id,
"Query": execution.query,
"StatementType": "DDL",
"ResultConfiguration": execution.config,
"QueryExecutionContext": execution.context,
"Status": {
"State": execution.status,
"SubmissionDateTime": execution.start_time,
},
'Statistics': {
'EngineExecutionTimeInMillis': 0,
'DataScannedInBytes': 0,
'TotalExecutionTimeInMillis': 0,
'QueryQueueTimeInMillis': 0,
'QueryPlanningTimeInMillis': 0,
'ServiceProcessingTimeInMillis': 0
"Statistics": {
"EngineExecutionTimeInMillis": 0,
"DataScannedInBytes": 0,
"TotalExecutionTimeInMillis": 0,
"QueryQueueTimeInMillis": 0,
"QueryPlanningTimeInMillis": 0,
"ServiceProcessingTimeInMillis": 0,
},
'WorkGroup': execution.workgroup
"WorkGroup": execution.workgroup,
}
}
return json.dumps(result)
@ -80,11 +82,6 @@ class AthenaResponse(BaseResponse):
def error(self, msg, status):
return (
json.dumps(
{
"__type": "InvalidRequestException",
"Message": msg,
}
),
json.dumps({"__type": "InvalidRequestException", "Message": msg,}),
dict(status=status),
)