@@ -35,14 +35,16 @@ def for_exec_def(self, workspace_id: str, exec_def: ExecutionDefinition) -> Exec
3535 exec_def: execution definition - this prescribes what to calculate, how to place labels and metric values
3636 into dimensions
3737 """
38- response = self ._actions_api .compute_report (workspace_id , exec_def .as_api_model (), _check_return_type = False )
38+ response , _ , headers = self ._actions_api .compute_report (
39+ workspace_id , exec_def .as_api_model (), _check_return_type = False , _return_http_data_only = False
40+ )
3941
4042 return Execution (
4143 api_client = self ._api_client ,
4244 workspace_id = workspace_id ,
4345 exec_def = exec_def ,
4446 response = response ,
45- cancel_token = response . headers .get ("X-GDC-CANCEL-TOKEN " )
47+ cancel_token = headers .get ("X-Gdc-Cancel-Token " )
4648 if exec_def .is_cancellable or self ._api_client .executions_cancellable
4749 else None ,
4850 )
@@ -112,23 +114,25 @@ def ai_chat_history_reset(self, workspace_id: str) -> None:
112114 chat_history_request = ChatHistoryRequest (reset = True )
113115 self ._actions_api .ai_chat_history (workspace_id , chat_history_request , _check_return_type = False )
114116
115- def cancel_executions (self , executions : list [Execution ]) -> None :
117+ def cancel_executions (self , executions : list [tuple [ str , str ] ]) -> None :
116118 """
117119 Try to cancel given executions using the cancel api endpoint.
120+ Order of token applications is not guaranteed.
118121
119122 *Note that this is currently a noop, we will be enabling this functionality soon.*
120123
121124 Args:
122- executions: list of executions to send for cancellation
125+ executions: list of tuples [workspace_id, cancel_token] to send for cancellation
123126 """
124127 workspace_to_tokens : dict [str , set [str ]] = {}
125128
126129 for execution in executions :
127- if not workspace_to_tokens [execution .workspace_id ]:
128- workspace_to_tokens [execution .workspace_id ] = set ()
130+ workspace_id , cancel_token = execution
131+
132+ if workspace_id not in workspace_to_tokens :
133+ workspace_to_tokens [workspace_id ] = set ()
129134
130- if execution .cancel_token :
131- workspace_to_tokens [execution .workspace_id ].add (execution .cancel_token )
135+ workspace_to_tokens [workspace_id ].add (cancel_token )
132136
133137 for workspace_id , token_ids in workspace_to_tokens .items ():
134138 self ._actions_api .cancel_executions (workspace_id , AfmCancelTokens (list (token_ids )))
0 commit comments