refactor(api_connextion): use session.row_count to decide whether to …

…return NotFound for queueEvent batch delete
apache · Feb 7, 2024 · a9ae4c2 · a9ae4c2
1 parent 576f9de
commit a9ae4c2
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 22 deletions.
diff --git a/airflow/api_connexion/endpoints/dataset_endpoint.py b/airflow/api_connexion/endpoints/dataset_endpoint.py
@@ -213,16 +213,15 @@ def delete_dag_dataset_queue_events(
 ) -> APIResponse:
     """Delete queued Dataset events for a DAG."""
     where_clauses = _build_get_ddrqs_where_clause(dag_id=dag_id, before=before)
-    query = select(DatasetDagRunQueue).where(*where_clauses)
-    ddrqs = session.scalars(query).all()
-    if not ddrqs:
-        raise NotFound(
-            "Queue event not found",
-            detail=f"Queue event with dag_id: `{dag_id}` was not found",
-        )
-    delete_stmt = delete(DatasetDagRunQueue).where(*where_clauses)
-    session.execute(delete_stmt)
-    return NoContent, HTTPStatus.NO_CONTENT
+    if where_clauses:
+        delete_stmt = delete(DatasetDagRunQueue).where(*where_clauses)
+        s = session.execute(delete_stmt)
+        if s.rowcount:
+            return NoContent, HTTPStatus.NO_CONTENT
+    raise NotFound(
+        "Queue event not found",
+        detail=f"Queue event with dag_id: `{dag_id}` was not found",
+    )
 
 
 def _build_get_dataset_ddrqs_where_clause(uri: str, session: Session, before: str | None = None):
@@ -266,15 +265,12 @@ def delete_dataset_queue_events(
 ) -> APIResponse:
     """Delete queued Dataset events for a Dataset"""
     where_clauses = _build_get_dataset_ddrqs_where_clause(uri=uri, session=session, before=before)
-    ddrqs = None
     if where_clauses:
-        query = select(DatasetDagRunQueue).where(*where_clauses)
-        ddrqs = session.scalars(query).all()
-    if not where_clauses or not ddrqs:
-        raise NotFound(
-            "Queue event not found",
-            detail=f"Queue event with dataset uri: `{uri}` was not found",
-        )
-    delete_stmt = delete(DatasetDagRunQueue).where(*where_clauses)
-    session.execute(delete_stmt)
-    return NoContent, HTTPStatus.NO_CONTENT
+        delete_stmt = delete(DatasetDagRunQueue).where(*where_clauses)
+        s = session.execute(delete_stmt)
+        if s.rowcount:
+            return NoContent, HTTPStatus.NO_CONTENT
+    raise NotFound(
+        "Queue event not found",
+        detail=f"Queue event with dataset uri: `{uri}` was not found",
+    )
diff --git a/tests/api_connexion/endpoints/test_dataset_endpoint.py b/tests/api_connexion/endpoints/test_dataset_endpoint.py
@@ -820,7 +820,7 @@ def test_should_respond_404(self):
 
         assert response.status_code == 404
         assert {
-            "detail": "Queue event with datast_uri: `not_exists` was not found",
+            "detail": "Queue event with dataset uri: `not_exists` was not found",
             "status": 404,
             "title": "Queue event not found",
             "type": EXCEPTIONS_LINK_MAP[404],