We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Filter datasets by dag_id in rest API
### Description

We have a list datasets endpoint in the API. Right now we can only filter by searching uri string. It would be useful to be able to pass one or multiple `dag_id`s and only return datasets that include that `dag_id` in `consuming_dags` or `producing_tasks`

### Use case/motivation

Make it easy to see what datasets are connected to a dag.

Right now our webserver endpoint to get dataset dependencies returns all datasets+dags and then the UI needs to filter it down.


### Related issues

_No response_

### Are you willing to submit a PR?

- [ ] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `airflow/api_connexion/endpoints/dataset_endpoint.py`
Content:
```
1 # Licensed to the Apache Software Foundation (ASF) under one
2 # or more contributor license agreements.  See the NOTICE file
3 # distributed with this work for additional information
4 # regarding copyright ownership.  The ASF licenses this file
5 # to you under the Apache License, Version 2.0 (the
6 # "License"); you may not use this file except in compliance
7 # with the License.  You may obtain a copy of the License at
8 #
9 #   http://www.apache.org/licenses/LICENSE-2.0
10 #
11 # Unless required by applicable law or agreed to in writing,
12 # software distributed under the License is distributed on an
13 # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14 # KIND, either express or implied.  See the License for the
15 # specific language governing permissions and limitations
16 # under the License.
17 from __future__ import annotations
18 
19 from http import HTTPStatus
20 from typing import TYPE_CHECKING
21 
22 from connexion import NoContent
23 from sqlalchemy import delete, func, select
24 from sqlalchemy.orm import joinedload, subqueryload
25 
26 from airflow.api_connexion import security
27 from airflow.api_connexion.exceptions import NotFound
28 from airflow.api_connexion.parameters import apply_sorting, check_limit, format_datetime, format_parameters
29 from airflow.api_connexion.schemas.dataset_schema import (
30     DatasetCollection,
31     DatasetEventCollection,
32     QueuedEvent,
33     QueuedEventCollection,
34     dataset_collection_schema,
35     dataset_event_collection_schema,
36     dataset_schema,
37     queued_event_collection_schema,
38     queued_event_schema,
39 )
40 from airflow.models.dataset import DatasetDagRunQueue, DatasetEvent, DatasetModel
41 from airflow.utils.db import get_query_count
42 from airflow.utils.session import NEW_SESSION, provide_session
43 from airflow.www.extensions.init_auth_manager import get_auth_manager
44 
45 if TYPE_CHECKING:
46     from sqlalchemy.orm import Session
47 
48     from airflow.api_connexion.types import APIResponse
49 
50 
51 @security.requires_access_dataset("GET")
52 @provide_session
53 def get_dataset(*, uri: str, session: Session = NEW_SESSION) -> APIResponse:
54     """Get a Dataset."""
55     dataset = session.scalar(
56         select(DatasetModel)
57         .where(DatasetModel.uri == uri)
58         .options(joinedload(DatasetModel.consuming_dags), joinedload(DatasetModel.producing_tasks))
59     )
60     if not dataset:
61         raise NotFound(
62             "Dataset not found",
63             detail=f"The Dataset with uri: `{uri}` was not found",
64         )
65     return dataset_schema.dump(dataset)
66 
67 
68 @security.requires_access_dataset("GET")
69 @format_parameters({"limit": check_limit})
70 @provide_session
71 def get_datasets(
72     *,
73     limit: int,
74     offset: int = 0,
75     uri_pattern: str | None = None,
76     order_by: str = "id",
77     session: Session = NEW_SESSION,
78 ) -> APIResponse:
79     """Get datasets."""
80     allowed_attrs = ["id", "uri", "created_at", "updated_at"]
81 
82     total_entries = session.scalars(select(func.count(DatasetModel.id))).one()
83     query = select(DatasetModel)
84     if uri_pattern:
85         query = query.where(DatasetModel.uri.ilike(f"%{uri_pattern}%"))
86     query = apply_sorting(query, order_by, {}, allowed_attrs)
87     datasets = session.scalars(
88         query.options(subqueryload(DatasetModel.consuming_dags), subqueryload(DatasetModel.producing_tasks))
89         .offset(offset)
90         .limit(limit)
91     ).all()
92     return dataset_collection_schema.dump(DatasetCollection(datasets=datasets, total_entries=total_entries))
93 
94 
95 @security.requires_access_dataset("GET")
96 @provide_session
97 @format_parameters({"limit": check_limit})
98 def get_dataset_events(
99     *,
100     limit: int,
101     offset: int = 0,
102     order_by: str = "timestamp",
103     dataset_id: int | None = None,
104     source_dag_id: str | None = None,
105     source_task_id: str | None = None,
106     source_run_id: str | None = None,
107     source_map_index: int | None = None,
108     session: Session = NEW_SESSION,
109 ) -> APIResponse:
110     """Get dataset events."""
111     allowed_attrs = ["source_dag_id", "source_task_id", "source_run_id", "source_map_index", "timestamp"]
112 
113     query = select(DatasetEvent)
114 
115     if dataset_id:
116         query = query.where(DatasetEvent.dataset_id == dataset_id)
117     if source_dag_id:
118         query = query.where(DatasetEvent.source_dag_id == source_dag_id)
119     if source_task_id:
120         query = query.where(DatasetEvent.source_task_id == source_task_id)
121     if source_run_id:
122         query = query.where(DatasetEvent.source_run_id == source_run_id)
123     if source_map_index:
124         query = query.where(DatasetEvent.source_map_index == source_map_index)
125 
126     query = query.options(subqueryload(DatasetEvent.created_dagruns))
127 
128     total_entries = get_query_count(query, session=session)
129     query = apply_sorting(query, order_by, {}, allowed_attrs)
130     events = session.scalars(query.offset(offset).limit(limit)).all()
131     return dataset_event_collection_schema.dump(
132         DatasetEventCollection(dataset_events=events, total_entries=total_entries)
133     )
134 
135 
136 def _generate_queued_event_where_clause(
137     *,
138     dag_id: str | None = None,
139     dataset_id: int | None = None,
140     uri: str | None = None,
141     before: str | None = None,
142     permitted_dag_ids: set[str] | None = None,
143 ) -> list:
144     """Get DatasetDagRunQueue where clause."""
145     where_clause = []
146     if dag_id is not None:
147         where_clause.append(DatasetDagRunQueue.target_dag_id == dag_id)
148     if dataset_id is not None:
149         where_clause.append(DatasetDagRunQueue.dataset_id == dataset_id)
150     if uri is not None:
151         where_clause.append(
152             DatasetDagRunQueue.dataset_id.in_(
153                 select(DatasetModel.id).where(DatasetModel.uri == uri),
154             ),
155         )
156     if before is not None:
157         where_clause.append(DatasetDagRunQueue.created_at < format_datetime(before))
158     if permitted_dag_ids is not None:
159         where_clause.append(DatasetDagRunQueue.target_dag_id.in_(permitted_dag_ids))
160     return where_clause
161 
162 
163 @security.requires_access_dataset("GET")
164 @security.requires_access_dag("GET")
165 @provide_session
166 def get_dag_dataset_queued_event(
167     *, dag_id: str, uri: str, before: str | None = None, session: Session = NEW_SESSION
168 ) -> APIResponse:
169     """Get a queued Dataset event for a DAG."""
170     where_clause = _generate_queued_event_where_clause(dag_id=dag_id, uri=uri, before=before)
171     ddrq = session.scalar(
172         select(DatasetDagRunQueue)
173         .join(DatasetModel, DatasetDagRunQueue.dataset_id == DatasetModel.id)
174         .where(*where_clause)
175     )
176     if ddrq is None:
177         raise NotFound(
178             "Queue event not found",
179             detail=f"Queue event with dag_id: `{dag_id}` and dataset uri: `{uri}` was not found",
180         )
181     queued_event = {"created_at": ddrq.created_at, "dag_id": dag_id, "uri": uri}
182     return queued_event_schema.dump(queued_event)
183 
184 
185 @security.requires_access_dataset("DELETE")
186 @security.requires_access_dag("GET")
187 @provide_session
188 def delete_dag_dataset_queued_event(
189     *, dag_id: str, uri: str, before: str | None = None, session: Session = NEW_SESSION
190 ) -> APIResponse:
191     """Delete a queued Dataset event for a DAG."""
192     where_clause = _generate_queued_event_where_clause(dag_id=dag_id, uri=uri, before=before)
193     delete_stmt = (
194         delete(DatasetDagRunQueue).where(*where_clause).execution_options(synchronize_session="fetch")
195     )
196     result = session.execute(delete_stmt)
197     if result.rowcount > 0:
198         return NoContent, HTTPStatus.NO_CONTENT
199     raise NotFound(
200         "Queue event not found",
201         detail=f"Queue event with dag_id: `{dag_id}` and dataset uri: `{uri}` was not found",
202     )
203 
204 
205 @security.requires_access_dataset("GET")
206 @security.requires_access_dag("GET")
207 @provide_session
208 def get_dag_dataset_queued_events(
209     *, dag_id: str, before: str | None = None, session: Session = NEW_SESSION
210 ) -> APIResponse:
211     """Get queued Dataset events for a DAG."""
212     where_clause = _generate_queued_event_where_clause(dag_id=dag_id, before=before)
213     query = (
214         select(DatasetDagRunQueue, DatasetModel.uri)
215         .join(DatasetModel, DatasetDagRunQueue.dataset_id == DatasetModel.id)
216         .where(*where_clause)
217     )
218     result = session.execute(query).all()
219     total_entries = get_query_count(query, session=session)
220     if not result:
221         raise NotFound(
222             "Queue event not found",
223             detail=f"Queue event with dag_id: `{dag_id}` was not found",
224         )
225     queued_events = [
226         QueuedEvent(created_at=ddrq.created_at, dag_id=ddrq.target_dag_id, uri=uri) for ddrq, uri in result
227     ]
228     return queued_event_collection_schema.dump(
229         QueuedEventCollection(queued_events=queued_events, total_entries=total_entries)
230     )
231 
232 
233 @security.requires_access_dataset("DELETE")
234 @security.requires_access_dag("GET")
235 @provide_session
236 def delete_dag_dataset_queued_events(
237     *, dag_id: str, before: str | None = None, session: Session = NEW_SESSION
238 ) -> APIResponse:
239     """Delete queued Dataset events for a DAG."""
240     where_clause = _generate_queued_event_where_clause(dag_id=dag_id, before=before)
241     delete_stmt = delete(DatasetDagRunQueue).where(*where_clause)
242     result = session.execute(delete_stmt)
243     if result.rowcount > 0:
244         return NoContent, HTTPStatus.NO_CONTENT
245 
246     raise NotFound(
247         "Queue event not found",
248         detail=f"Queue event with dag_id: `{dag_id}` was not found",
249     )
250 
251 
252 @security.requires_access_dataset("GET")
253 @provide_session
254 def get_dataset_queued_events(
255     *, uri: str, before: str | None = None, session: Session = NEW_SESSION
256 ) -> APIResponse:
257     """Get queued Dataset events for a Dataset."""
258     permitted_dag_ids = get_auth_manager().get_permitted_dag_ids(methods=["GET"])
259     where_clause = _generate_queued_event_where_clause(
260         uri=uri, before=before, permitted_dag_ids=permitted_dag_ids
261     )
262     query = (
263         select(DatasetDagRunQueue, DatasetModel.uri)
264         .join(DatasetModel, DatasetDagRunQueue.dataset_id == DatasetModel.id)
265         .where(*where_clause)
266     )
267     total_entries = get_query_count(query, session=session)
268     result = session.execute(query).all()
269     if total_entries > 0:
270         queued_events = [
271             QueuedEvent(created_at=ddrq.created_at, dag_id=ddrq.target_dag_id, uri=uri)
272             for ddrq, uri in result
273         ]
274         return queued_event_collection_schema.dump(
275             QueuedEventCollection(queued_events=queued_events, total_entries=total_entries)
276         )
277     raise NotFound(
278         "Queue event not found",
279         detail=f"Queue event with dataset uri: `{uri}` was not found",
280     )
281 
282 
283 @security.requires_access_dataset("DELETE")
284 @provide_session
285 def delete_dataset_queued_events(
286     *, uri: str, before: str | None = None, session: Session = NEW_SESSION
287 ) -> APIResponse:
288     """Delete queued Dataset events for a Dataset."""
289     permitted_dag_ids = get_auth_manager().get_permitted_dag_ids(methods=["GET"])
290     where_clause = _generate_queued_event_where_clause(
291         uri=uri, before=before, permitted_dag_ids=permitted_dag_ids
292     )
293     delete_stmt = (
294         delete(DatasetDagRunQueue).where(*where_clause).execution_options(synchronize_session="fetch")
295     )
296 
297     result = session.execute(delete_stmt)
298     if result.rowcount > 0:
299         return NoContent, HTTPStatus.NO_CONTENT
300     raise NotFound(
301         "Queue event not found",
302         detail=f"Queue event with dataset uri: `{uri}` was not found",
303     )
304
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/airflow/api_connexion/endpoints/dataset_endpoint.py b/airflow/api_connexion/endpoints/dataset_endpoint.py
--- a/airflow/api_connexion/endpoints/dataset_endpoint.py
+++ b/airflow/api_connexion/endpoints/dataset_endpoint.py
@@ -27,10 +27,12 @@
 from airflow.api_connexion.exceptions import NotFound
 from airflow.api_connexion.parameters import apply_sorting, check_limit, format_datetime, format_parameters
 from airflow.api_connexion.schemas.dataset_schema import (
+    DagScheduleDatasetReference,
     DatasetCollection,
     DatasetEventCollection,
     QueuedEvent,
     QueuedEventCollection,
+    TaskOutletDatasetReference,
     dataset_collection_schema,
     dataset_event_collection_schema,
     dataset_schema,
@@ -73,6 +75,7 @@
     limit: int,
     offset: int = 0,
     uri_pattern: str | None = None,
+    dag_ids: str | None = None,
     order_by: str = "id",
     session: Session = NEW_SESSION,
 ) -> APIResponse:
@@ -81,6 +84,13 @@
 
     total_entries = session.scalars(select(func.count(DatasetModel.id))).one()
     query = select(DatasetModel)
+
+    if dag_ids:
+        dags_list = dag_ids.split(",")
+        query = query.filter(
+            (DatasetModel.consuming_dags.any(DagScheduleDatasetReference.dag_id.in_(dags_list)))
+            | (DatasetModel.producing_tasks.any(TaskOutletDatasetReference.dag_id.in_(dags_list)))
+        )
     if uri_pattern:
         query = query.where(DatasetModel.uri.ilike(f"%{uri_pattern}%"))
     query = apply_sorting(query, order_by, {}, allowed_attrs)

{"golden_diff": "diff --git a/airflow/api_connexion/endpoints/dataset_endpoint.py b/airflow/api_connexion/endpoints/dataset_endpoint.py\n--- a/airflow/api_connexion/endpoints/dataset_endpoint.py\n+++ b/airflow/api_connexion/endpoints/dataset_endpoint.py\n@@ -27,10 +27,12 @@\n from airflow.api_connexion.exceptions import NotFound\n from airflow.api_connexion.parameters import apply_sorting, check_limit, format_datetime, format_parameters\n from airflow.api_connexion.schemas.dataset_schema import (\n+    DagScheduleDatasetReference,\n     DatasetCollection,\n     DatasetEventCollection,\n     QueuedEvent,\n     QueuedEventCollection,\n+    TaskOutletDatasetReference,\n     dataset_collection_schema,\n     dataset_event_collection_schema,\n     dataset_schema,\n@@ -73,6 +75,7 @@\n     limit: int,\n     offset: int = 0,\n     uri_pattern: str | None = None,\n+    dag_ids: str | None = None,\n     order_by: str = \"id\",\n     session: Session = NEW_SESSION,\n ) -> APIResponse:\n@@ -81,6 +84,13 @@\n \n     total_entries = session.scalars(select(func.count(DatasetModel.id))).one()\n     query = select(DatasetModel)\n+\n+    if dag_ids:\n+        dags_list = dag_ids.split(\",\")\n+        query = query.filter(\n+            (DatasetModel.consuming_dags.any(DagScheduleDatasetReference.dag_id.in_(dags_list)))\n+            | (DatasetModel.producing_tasks.any(TaskOutletDatasetReference.dag_id.in_(dags_list)))\n+        )\n     if uri_pattern:\n         query = query.where(DatasetModel.uri.ilike(f\"%{uri_pattern}%\"))\n     query = apply_sorting(query, order_by, {}, allowed_attrs)\n", "issue": "Filter datasets by dag_id in rest API\n### Description\r\n\r\nWe have a list datasets endpoint in the API. Right now we can only filter by searching uri string. It would be useful to be able to pass one or multiple `dag_id`s and only return datasets that include that `dag_id` in `consuming_dags` or `producing_tasks`\r\n\r\n### Use case/motivation\r\n\r\nMake it easy to see what datasets are connected to a dag.\r\n\r\nRight now our webserver endpoint to get dataset dependencies returns all datasets+dags and then the UI needs to filter it down.\r\n\r\n\r\n### Related issues\r\n\r\n_No response_\r\n\r\n### Are you willing to submit a PR?\r\n\r\n- [ ] Yes I am willing to submit a PR!\r\n\r\n### Code of Conduct\r\n\r\n- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)\r\n\n", "before_files": [{"content": "# Licensed to the Apache Software Foundation (ASF) under one\n# or more contributor license agreements.  See the NOTICE file\n# distributed with this work for additional information\n# regarding copyright ownership.  The ASF licenses this file\n# to you under the Apache License, Version 2.0 (the\n# \"License\"); you may not use this file except in compliance\n# with the License.  You may obtain a copy of the License at\n#\n#   http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing,\n# software distributed under the License is distributed on an\n# \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY\n# KIND, either express or implied.  See the License for the\n# specific language governing permissions and limitations\n# under the License.\nfrom __future__ import annotations\n\nfrom http import HTTPStatus\nfrom typing import TYPE_CHECKING\n\nfrom connexion import NoContent\nfrom sqlalchemy import delete, func, select\nfrom sqlalchemy.orm import joinedload, subqueryload\n\nfrom airflow.api_connexion import security\nfrom airflow.api_connexion.exceptions import NotFound\nfrom airflow.api_connexion.parameters import apply_sorting, check_limit, format_datetime, format_parameters\nfrom airflow.api_connexion.schemas.dataset_schema import (\n    DatasetCollection,\n    DatasetEventCollection,\n    QueuedEvent,\n    QueuedEventCollection,\n    dataset_collection_schema,\n    dataset_event_collection_schema,\n    dataset_schema,\n    queued_event_collection_schema,\n    queued_event_schema,\n)\nfrom airflow.models.dataset import DatasetDagRunQueue, DatasetEvent, DatasetModel\nfrom airflow.utils.db import get_query_count\nfrom airflow.utils.session import NEW_SESSION, provide_session\nfrom airflow.www.extensions.init_auth_manager import get_auth_manager\n\nif TYPE_CHECKING:\n    from sqlalchemy.orm import Session\n\n    from airflow.api_connexion.types import APIResponse\n\n\n@security.requires_access_dataset(\"GET\")\n@provide_session\ndef get_dataset(*, uri: str, session: Session = NEW_SESSION) -> APIResponse:\n    \"\"\"Get a Dataset.\"\"\"\n    dataset = session.scalar(\n        select(DatasetModel)\n        .where(DatasetModel.uri == uri)\n        .options(joinedload(DatasetModel.consuming_dags), joinedload(DatasetModel.producing_tasks))\n    )\n    if not dataset:\n        raise NotFound(\n            \"Dataset not found\",\n            detail=f\"The Dataset with uri: `{uri}` was not found\",\n        )\n    return dataset_schema.dump(dataset)\n\n\n@security.requires_access_dataset(\"GET\")\n@format_parameters({\"limit\": check_limit})\n@provide_session\ndef get_datasets(\n    *,\n    limit: int,\n    offset: int = 0,\n    uri_pattern: str | None = None,\n    order_by: str = \"id\",\n    session: Session = NEW_SESSION,\n) -> APIResponse:\n    \"\"\"Get datasets.\"\"\"\n    allowed_attrs = [\"id\", \"uri\", \"created_at\", \"updated_at\"]\n\n    total_entries = session.scalars(select(func.count(DatasetModel.id))).one()\n    query = select(DatasetModel)\n    if uri_pattern:\n        query = query.where(DatasetModel.uri.ilike(f\"%{uri_pattern}%\"))\n    query = apply_sorting(query, order_by, {}, allowed_attrs)\n    datasets = session.scalars(\n        query.options(subqueryload(DatasetModel.consuming_dags), subqueryload(DatasetModel.producing_tasks))\n        .offset(offset)\n        .limit(limit)\n    ).all()\n    return dataset_collection_schema.dump(DatasetCollection(datasets=datasets, total_entries=total_entries))\n\n\n@security.requires_access_dataset(\"GET\")\n@provide_session\n@format_parameters({\"limit\": check_limit})\ndef get_dataset_events(\n    *,\n    limit: int,\n    offset: int = 0,\n    order_by: str = \"timestamp\",\n    dataset_id: int | None = None,\n    source_dag_id: str | None = None,\n    source_task_id: str | None = None,\n    source_run_id: str | None = None,\n    source_map_index: int | None = None,\n    session: Session = NEW_SESSION,\n) -> APIResponse:\n    \"\"\"Get dataset events.\"\"\"\n    allowed_attrs = [\"source_dag_id\", \"source_task_id\", \"source_run_id\", \"source_map_index\", \"timestamp\"]\n\n    query = select(DatasetEvent)\n\n    if dataset_id:\n        query = query.where(DatasetEvent.dataset_id == dataset_id)\n    if source_dag_id:\n        query = query.where(DatasetEvent.source_dag_id == source_dag_id)\n    if source_task_id:\n        query = query.where(DatasetEvent.source_task_id == source_task_id)\n    if source_run_id:\n        query = query.where(DatasetEvent.source_run_id == source_run_id)\n    if source_map_index:\n        query = query.where(DatasetEvent.source_map_index == source_map_index)\n\n    query = query.options(subqueryload(DatasetEvent.created_dagruns))\n\n    total_entries = get_query_count(query, session=session)\n    query = apply_sorting(query, order_by, {}, allowed_attrs)\n    events = session.scalars(query.offset(offset).limit(limit)).all()\n    return dataset_event_collection_schema.dump(\n        DatasetEventCollection(dataset_events=events, total_entries=total_entries)\n    )\n\n\ndef _generate_queued_event_where_clause(\n    *,\n    dag_id: str | None = None,\n    dataset_id: int | None = None,\n    uri: str | None = None,\n    before: str | None = None,\n    permitted_dag_ids: set[str] | None = None,\n) -> list:\n    \"\"\"Get DatasetDagRunQueue where clause.\"\"\"\n    where_clause = []\n    if dag_id is not None:\n        where_clause.append(DatasetDagRunQueue.target_dag_id == dag_id)\n    if dataset_id is not None:\n        where_clause.append(DatasetDagRunQueue.dataset_id == dataset_id)\n    if uri is not None:\n        where_clause.append(\n            DatasetDagRunQueue.dataset_id.in_(\n                select(DatasetModel.id).where(DatasetModel.uri == uri),\n            ),\n        )\n    if before is not None:\n        where_clause.append(DatasetDagRunQueue.created_at < format_datetime(before))\n    if permitted_dag_ids is not None:\n        where_clause.append(DatasetDagRunQueue.target_dag_id.in_(permitted_dag_ids))\n    return where_clause\n\n\n@security.requires_access_dataset(\"GET\")\n@security.requires_access_dag(\"GET\")\n@provide_session\ndef get_dag_dataset_queued_event(\n    *, dag_id: str, uri: str, before: str | None = None, session: Session = NEW_SESSION\n) -> APIResponse:\n    \"\"\"Get a queued Dataset event for a DAG.\"\"\"\n    where_clause = _generate_queued_event_where_clause(dag_id=dag_id, uri=uri, before=before)\n    ddrq = session.scalar(\n        select(DatasetDagRunQueue)\n        .join(DatasetModel, DatasetDagRunQueue.dataset_id == DatasetModel.id)\n        .where(*where_clause)\n    )\n    if ddrq is None:\n        raise NotFound(\n            \"Queue event not found\",\n            detail=f\"Queue event with dag_id: `{dag_id}` and dataset uri: `{uri}` was not found\",\n        )\n    queued_event = {\"created_at\": ddrq.created_at, \"dag_id\": dag_id, \"uri\": uri}\n    return queued_event_schema.dump(queued_event)\n\n\n@security.requires_access_dataset(\"DELETE\")\n@security.requires_access_dag(\"GET\")\n@provide_session\ndef delete_dag_dataset_queued_event(\n    *, dag_id: str, uri: str, before: str | None = None, session: Session = NEW_SESSION\n) -> APIResponse:\n    \"\"\"Delete a queued Dataset event for a DAG.\"\"\"\n    where_clause = _generate_queued_event_where_clause(dag_id=dag_id, uri=uri, before=before)\n    delete_stmt = (\n        delete(DatasetDagRunQueue).where(*where_clause).execution_options(synchronize_session=\"fetch\")\n    )\n    result = session.execute(delete_stmt)\n    if result.rowcount > 0:\n        return NoContent, HTTPStatus.NO_CONTENT\n    raise NotFound(\n        \"Queue event not found\",\n        detail=f\"Queue event with dag_id: `{dag_id}` and dataset uri: `{uri}` was not found\",\n    )\n\n\n@security.requires_access_dataset(\"GET\")\n@security.requires_access_dag(\"GET\")\n@provide_session\ndef get_dag_dataset_queued_events(\n    *, dag_id: str, before: str | None = None, session: Session = NEW_SESSION\n) -> APIResponse:\n    \"\"\"Get queued Dataset events for a DAG.\"\"\"\n    where_clause = _generate_queued_event_where_clause(dag_id=dag_id, before=before)\n    query = (\n        select(DatasetDagRunQueue, DatasetModel.uri)\n        .join(DatasetModel, DatasetDagRunQueue.dataset_id == DatasetModel.id)\n        .where(*where_clause)\n    )\n    result = session.execute(query).all()\n    total_entries = get_query_count(query, session=session)\n    if not result:\n        raise NotFound(\n            \"Queue event not found\",\n            detail=f\"Queue event with dag_id: `{dag_id}` was not found\",\n        )\n    queued_events = [\n        QueuedEvent(created_at=ddrq.created_at, dag_id=ddrq.target_dag_id, uri=uri) for ddrq, uri in result\n    ]\n    return queued_event_collection_schema.dump(\n        QueuedEventCollection(queued_events=queued_events, total_entries=total_entries)\n    )\n\n\n@security.requires_access_dataset(\"DELETE\")\n@security.requires_access_dag(\"GET\")\n@provide_session\ndef delete_dag_dataset_queued_events(\n    *, dag_id: str, before: str | None = None, session: Session = NEW_SESSION\n) -> APIResponse:\n    \"\"\"Delete queued Dataset events for a DAG.\"\"\"\n    where_clause = _generate_queued_event_where_clause(dag_id=dag_id, before=before)\n    delete_stmt = delete(DatasetDagRunQueue).where(*where_clause)\n    result = session.execute(delete_stmt)\n    if result.rowcount > 0:\n        return NoContent, HTTPStatus.NO_CONTENT\n\n    raise NotFound(\n        \"Queue event not found\",\n        detail=f\"Queue event with dag_id: `{dag_id}` was not found\",\n    )\n\n\n@security.requires_access_dataset(\"GET\")\n@provide_session\ndef get_dataset_queued_events(\n    *, uri: str, before: str | None = None, session: Session = NEW_SESSION\n) -> APIResponse:\n    \"\"\"Get queued Dataset events for a Dataset.\"\"\"\n    permitted_dag_ids = get_auth_manager().get_permitted_dag_ids(methods=[\"GET\"])\n    where_clause = _generate_queued_event_where_clause(\n        uri=uri, before=before, permitted_dag_ids=permitted_dag_ids\n    )\n    query = (\n        select(DatasetDagRunQueue, DatasetModel.uri)\n        .join(DatasetModel, DatasetDagRunQueue.dataset_id == DatasetModel.id)\n        .where(*where_clause)\n    )\n    total_entries = get_query_count(query, session=session)\n    result = session.execute(query).all()\n    if total_entries > 0:\n        queued_events = [\n            QueuedEvent(created_at=ddrq.created_at, dag_id=ddrq.target_dag_id, uri=uri)\n            for ddrq, uri in result\n        ]\n        return queued_event_collection_schema.dump(\n            QueuedEventCollection(queued_events=queued_events, total_entries=total_entries)\n        )\n    raise NotFound(\n        \"Queue event not found\",\n        detail=f\"Queue event with dataset uri: `{uri}` was not found\",\n    )\n\n\n@security.requires_access_dataset(\"DELETE\")\n@provide_session\ndef delete_dataset_queued_events(\n    *, uri: str, before: str | None = None, session: Session = NEW_SESSION\n) -> APIResponse:\n    \"\"\"Delete queued Dataset events for a Dataset.\"\"\"\n    permitted_dag_ids = get_auth_manager().get_permitted_dag_ids(methods=[\"GET\"])\n    where_clause = _generate_queued_event_where_clause(\n        uri=uri, before=before, permitted_dag_ids=permitted_dag_ids\n    )\n    delete_stmt = (\n        delete(DatasetDagRunQueue).where(*where_clause).execution_options(synchronize_session=\"fetch\")\n    )\n\n    result = session.execute(delete_stmt)\n    if result.rowcount > 0:\n        return NoContent, HTTPStatus.NO_CONTENT\n    raise NotFound(\n        \"Queue event not found\",\n        detail=f\"Queue event with dataset uri: `{uri}` was not found\",\n    )\n", "path": "airflow/api_connexion/endpoints/dataset_endpoint.py"}], "after_files": [{"content": "# Licensed to the Apache Software Foundation (ASF) under one\n# or more contributor license agreements.  See the NOTICE file\n# distributed with this work for additional information\n# regarding copyright ownership.  The ASF licenses this file\n# to you under the Apache License, Version 2.0 (the\n# \"License\"); you may not use this file except in compliance\n# with the License.  You may obtain a copy of the License at\n#\n#   http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing,\n# software distributed under the License is distributed on an\n# \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY\n# KIND, either express or implied.  See the License for the\n# specific language governing permissions and limitations\n# under the License.\nfrom __future__ import annotations\n\nfrom http import HTTPStatus\nfrom typing import TYPE_CHECKING\n\nfrom connexion import NoContent\nfrom sqlalchemy import delete, func, select\nfrom sqlalchemy.orm import joinedload, subqueryload\n\nfrom airflow.api_connexion import security\nfrom airflow.api_connexion.exceptions import NotFound\nfrom airflow.api_connexion.parameters import apply_sorting, check_limit, format_datetime, format_parameters\nfrom airflow.api_connexion.schemas.dataset_schema import (\n    DagScheduleDatasetReference,\n    DatasetCollection,\n    DatasetEventCollection,\n    QueuedEvent,\n    QueuedEventCollection,\n    TaskOutletDatasetReference,\n    dataset_collection_schema,\n    dataset_event_collection_schema,\n    dataset_schema,\n    queued_event_collection_schema,\n    queued_event_schema,\n)\nfrom airflow.models.dataset import DatasetDagRunQueue, DatasetEvent, DatasetModel\nfrom airflow.utils.db import get_query_count\nfrom airflow.utils.session import NEW_SESSION, provide_session\nfrom airflow.www.extensions.init_auth_manager import get_auth_manager\n\nif TYPE_CHECKING:\n    from sqlalchemy.orm import Session\n\n    from airflow.api_connexion.types import APIResponse\n\n\n@security.requires_access_dataset(\"GET\")\n@provide_session\ndef get_dataset(*, uri: str, session: Session = NEW_SESSION) -> APIResponse:\n    \"\"\"Get a Dataset.\"\"\"\n    dataset = session.scalar(\n        select(DatasetModel)\n        .where(DatasetModel.uri == uri)\n        .options(joinedload(DatasetModel.consuming_dags), joinedload(DatasetModel.producing_tasks))\n    )\n    if not dataset:\n        raise NotFound(\n            \"Dataset not found\",\n            detail=f\"The Dataset with uri: `{uri}` was not found\",\n        )\n    return dataset_schema.dump(dataset)\n\n\n@security.requires_access_dataset(\"GET\")\n@format_parameters({\"limit\": check_limit})\n@provide_session\ndef get_datasets(\n    *,\n    limit: int,\n    offset: int = 0,\n    uri_pattern: str | None = None,\n    dag_ids: str | None = None,\n    order_by: str = \"id\",\n    session: Session = NEW_SESSION,\n) -> APIResponse:\n    \"\"\"Get datasets.\"\"\"\n    allowed_attrs = [\"id\", \"uri\", \"created_at\", \"updated_at\"]\n\n    total_entries = session.scalars(select(func.count(DatasetModel.id))).one()\n    query = select(DatasetModel)\n\n    if dag_ids:\n        dags_list = dag_ids.split(\",\")\n        query = query.filter(\n            (DatasetModel.consuming_dags.any(DagScheduleDatasetReference.dag_id.in_(dags_list)))\n            | (DatasetModel.producing_tasks.any(TaskOutletDatasetReference.dag_id.in_(dags_list)))\n        )\n    if uri_pattern:\n        query = query.where(DatasetModel.uri.ilike(f\"%{uri_pattern}%\"))\n    query = apply_sorting(query, order_by, {}, allowed_attrs)\n    datasets = session.scalars(\n        query.options(subqueryload(DatasetModel.consuming_dags), subqueryload(DatasetModel.producing_tasks))\n        .offset(offset)\n        .limit(limit)\n    ).all()\n    return dataset_collection_schema.dump(DatasetCollection(datasets=datasets, total_entries=total_entries))\n\n\n@security.requires_access_dataset(\"GET\")\n@provide_session\n@format_parameters({\"limit\": check_limit})\ndef get_dataset_events(\n    *,\n    limit: int,\n    offset: int = 0,\n    order_by: str = \"timestamp\",\n    dataset_id: int | None = None,\n    source_dag_id: str | None = None,\n    source_task_id: str | None = None,\n    source_run_id: str | None = None,\n    source_map_index: int | None = None,\n    session: Session = NEW_SESSION,\n) -> APIResponse:\n    \"\"\"Get dataset events.\"\"\"\n    allowed_attrs = [\"source_dag_id\", \"source_task_id\", \"source_run_id\", \"source_map_index\", \"timestamp\"]\n\n    query = select(DatasetEvent)\n\n    if dataset_id:\n        query = query.where(DatasetEvent.dataset_id == dataset_id)\n    if source_dag_id:\n        query = query.where(DatasetEvent.source_dag_id == source_dag_id)\n    if source_task_id:\n        query = query.where(DatasetEvent.source_task_id == source_task_id)\n    if source_run_id:\n        query = query.where(DatasetEvent.source_run_id == source_run_id)\n    if source_map_index:\n        query = query.where(DatasetEvent.source_map_index == source_map_index)\n\n    query = query.options(subqueryload(DatasetEvent.created_dagruns))\n\n    total_entries = get_query_count(query, session=session)\n    query = apply_sorting(query, order_by, {}, allowed_attrs)\n    events = session.scalars(query.offset(offset).limit(limit)).all()\n    return dataset_event_collection_schema.dump(\n        DatasetEventCollection(dataset_events=events, total_entries=total_entries)\n    )\n\n\ndef _generate_queued_event_where_clause(\n    *,\n    dag_id: str | None = None,\n    dataset_id: int | None = None,\n    uri: str | None = None,\n    before: str | None = None,\n    permitted_dag_ids: set[str] | None = None,\n) -> list:\n    \"\"\"Get DatasetDagRunQueue where clause.\"\"\"\n    where_clause = []\n    if dag_id is not None:\n        where_clause.append(DatasetDagRunQueue.target_dag_id == dag_id)\n    if dataset_id is not None:\n        where_clause.append(DatasetDagRunQueue.dataset_id == dataset_id)\n    if uri is not None:\n        where_clause.append(\n            DatasetDagRunQueue.dataset_id.in_(\n                select(DatasetModel.id).where(DatasetModel.uri == uri),\n            ),\n        )\n    if before is not None:\n        where_clause.append(DatasetDagRunQueue.created_at < format_datetime(before))\n    if permitted_dag_ids is not None:\n        where_clause.append(DatasetDagRunQueue.target_dag_id.in_(permitted_dag_ids))\n    return where_clause\n\n\n@security.requires_access_dataset(\"GET\")\n@security.requires_access_dag(\"GET\")\n@provide_session\ndef get_dag_dataset_queued_event(\n    *, dag_id: str, uri: str, before: str | None = None, session: Session = NEW_SESSION\n) -> APIResponse:\n    \"\"\"Get a queued Dataset event for a DAG.\"\"\"\n    where_clause = _generate_queued_event_where_clause(dag_id=dag_id, uri=uri, before=before)\n    ddrq = session.scalar(\n        select(DatasetDagRunQueue)\n        .join(DatasetModel, DatasetDagRunQueue.dataset_id == DatasetModel.id)\n        .where(*where_clause)\n    )\n    if ddrq is None:\n        raise NotFound(\n            \"Queue event not found\",\n            detail=f\"Queue event with dag_id: `{dag_id}` and dataset uri: `{uri}` was not found\",\n        )\n    queued_event = {\"created_at\": ddrq.created_at, \"dag_id\": dag_id, \"uri\": uri}\n    return queued_event_schema.dump(queued_event)\n\n\n@security.requires_access_dataset(\"DELETE\")\n@security.requires_access_dag(\"GET\")\n@provide_session\ndef delete_dag_dataset_queued_event(\n    *, dag_id: str, uri: str, before: str | None = None, session: Session = NEW_SESSION\n) -> APIResponse:\n    \"\"\"Delete a queued Dataset event for a DAG.\"\"\"\n    where_clause = _generate_queued_event_where_clause(dag_id=dag_id, uri=uri, before=before)\n    delete_stmt = (\n        delete(DatasetDagRunQueue).where(*where_clause).execution_options(synchronize_session=\"fetch\")\n    )\n    result = session.execute(delete_stmt)\n    if result.rowcount > 0:\n        return NoContent, HTTPStatus.NO_CONTENT\n    raise NotFound(\n        \"Queue event not found\",\n        detail=f\"Queue event with dag_id: `{dag_id}` and dataset uri: `{uri}` was not found\",\n    )\n\n\n@security.requires_access_dataset(\"GET\")\n@security.requires_access_dag(\"GET\")\n@provide_session\ndef get_dag_dataset_queued_events(\n    *, dag_id: str, before: str | None = None, session: Session = NEW_SESSION\n) -> APIResponse:\n    \"\"\"Get queued Dataset events for a DAG.\"\"\"\n    where_clause = _generate_queued_event_where_clause(dag_id=dag_id, before=before)\n    query = (\n        select(DatasetDagRunQueue, DatasetModel.uri)\n        .join(DatasetModel, DatasetDagRunQueue.dataset_id == DatasetModel.id)\n        .where(*where_clause)\n    )\n    result = session.execute(query).all()\n    total_entries = get_query_count(query, session=session)\n    if not result:\n        raise NotFound(\n            \"Queue event not found\",\n            detail=f\"Queue event with dag_id: `{dag_id}` was not found\",\n        )\n    queued_events = [\n        QueuedEvent(created_at=ddrq.created_at, dag_id=ddrq.target_dag_id, uri=uri) for ddrq, uri in result\n    ]\n    return queued_event_collection_schema.dump(\n        QueuedEventCollection(queued_events=queued_events, total_entries=total_entries)\n    )\n\n\n@security.requires_access_dataset(\"DELETE\")\n@security.requires_access_dag(\"GET\")\n@provide_session\ndef delete_dag_dataset_queued_events(\n    *, dag_id: str, before: str | None = None, session: Session = NEW_SESSION\n) -> APIResponse:\n    \"\"\"Delete queued Dataset events for a DAG.\"\"\"\n    where_clause = _generate_queued_event_where_clause(dag_id=dag_id, before=before)\n    delete_stmt = delete(DatasetDagRunQueue).where(*where_clause)\n    result = session.execute(delete_stmt)\n    if result.rowcount > 0:\n        return NoContent, HTTPStatus.NO_CONTENT\n\n    raise NotFound(\n        \"Queue event not found\",\n        detail=f\"Queue event with dag_id: `{dag_id}` was not found\",\n    )\n\n\n@security.requires_access_dataset(\"GET\")\n@provide_session\ndef get_dataset_queued_events(\n    *, uri: str, before: str | None = None, session: Session = NEW_SESSION\n) -> APIResponse:\n    \"\"\"Get queued Dataset events for a Dataset.\"\"\"\n    permitted_dag_ids = get_auth_manager().get_permitted_dag_ids(methods=[\"GET\"])\n    where_clause = _generate_queued_event_where_clause(\n        uri=uri, before=before, permitted_dag_ids=permitted_dag_ids\n    )\n    query = (\n        select(DatasetDagRunQueue, DatasetModel.uri)\n        .join(DatasetModel, DatasetDagRunQueue.dataset_id == DatasetModel.id)\n        .where(*where_clause)\n    )\n    total_entries = get_query_count(query, session=session)\n    result = session.execute(query).all()\n    if total_entries > 0:\n        queued_events = [\n            QueuedEvent(created_at=ddrq.created_at, dag_id=ddrq.target_dag_id, uri=uri)\n            for ddrq, uri in result\n        ]\n        return queued_event_collection_schema.dump(\n            QueuedEventCollection(queued_events=queued_events, total_entries=total_entries)\n        )\n    raise NotFound(\n        \"Queue event not found\",\n        detail=f\"Queue event with dataset uri: `{uri}` was not found\",\n    )\n\n\n@security.requires_access_dataset(\"DELETE\")\n@provide_session\ndef delete_dataset_queued_events(\n    *, uri: str, before: str | None = None, session: Session = NEW_SESSION\n) -> APIResponse:\n    \"\"\"Delete queued Dataset events for a Dataset.\"\"\"\n    permitted_dag_ids = get_auth_manager().get_permitted_dag_ids(methods=[\"GET\"])\n    where_clause = _generate_queued_event_where_clause(\n        uri=uri, before=before, permitted_dag_ids=permitted_dag_ids\n    )\n    delete_stmt = (\n        delete(DatasetDagRunQueue).where(*where_clause).execution_options(synchronize_session=\"fetch\")\n    )\n\n    result = session.execute(delete_stmt)\n    if result.rowcount > 0:\n        return NoContent, HTTPStatus.NO_CONTENT\n    raise NotFound(\n        \"Queue event not found\",\n        detail=f\"Queue event with dataset uri: `{uri}` was not found\",\n    )\n", "path": "airflow/api_connexion/endpoints/dataset_endpoint.py"}]}