diff --git a/cdk-ops/lambda/libdatawarehouse/src/libdatawarehouse/data_api.py b/cdk-ops/lambda/libdatawarehouse/src/libdatawarehouse/data_api.py
index 095d669..826ea90 100644
--- a/cdk-ops/lambda/libdatawarehouse/src/libdatawarehouse/data_api.py
+++ b/cdk-ops/lambda/libdatawarehouse/src/libdatawarehouse/data_api.py
@@ -14,7 +14,8 @@ def wait_for_results(
     client,
     statement_id: str,
     polling_interval: float = 0.05,
-    polling_timeout: int = round(300 / 0.05),
+    timeout: float = 300.0,
+    cancel_at_timeout: bool = True,
 ) -> Tuple[Optional[str], Dict]:
     """Waits for a given statement to finish.
 
@@ -23,15 +24,20 @@ def wait_for_results(
     :param float polling_interval: interval in seconds between two consecutive
     pollings.
 
-    :param int polling_timeout: timeout represented as the number of pollings.
+    :param float timeout: timeout in seconds.
+
+    :param bool cancel_at_timeout: whether cancels the statement when it times
+    out.
     """
-    polling_counter = 0
+    start_time = time.time()
     while True:
         res = client.describe_statement(Id=statement_id)
         status = res['Status']
         if status not in RUNNING_STATUSES:
             return status, res
-        polling_counter += 1
-        if polling_counter >= polling_timeout:
+        elapsed = time.time() - start_time
+        if elapsed >= timeout:
+            if cancel_at_timeout:
+                client.cancel_statement(Id=statement_id)
             return None, res
         time.sleep(polling_interval)