Repository: apache/airflow

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)



Discussion:
Some questions I have I'd like to get people's thoughts on:

1. Should it be a single dag_id or an array of multiple dag_ids?
2. Should we specify only consuming or only producing or both?
3. Maybe we need a separate, perhaps not in the public REST API, endpoint to return the entire dataset-dag graph based on either a dataset_uri or a dag_id?
Actually this could just be in `datasets_summary` endpoint in the webserver.
I can take on this task.

Resolving patch:
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 @@ def get_datasets(
     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 @@ def get_datasets(
 
     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)


Test patch:
diff --git a/tests/api_connexion/endpoints/test_dataset_endpoint.py b/tests/api_connexion/endpoints/test_dataset_endpoint.py
--- a/tests/api_connexion/endpoints/test_dataset_endpoint.py
+++ b/tests/api_connexion/endpoints/test_dataset_endpoint.py
@@ -23,8 +23,15 @@
 import time_machine
 
 from airflow.api_connexion.exceptions import EXCEPTIONS_LINK_MAP
+from airflow.models import DagModel
 from airflow.models.dagrun import DagRun
-from airflow.models.dataset import DatasetDagRunQueue, DatasetEvent, DatasetModel
+from airflow.models.dataset import (
+    DagScheduleDatasetReference,
+    DatasetDagRunQueue,
+    DatasetEvent,
+    DatasetModel,
+    TaskOutletDatasetReference,
+)
 from airflow.security import permissions
 from airflow.utils import timezone
 from airflow.utils.session import provide_session
@@ -252,6 +259,57 @@ def test_filter_datasets_by_uri_pattern_works(self, url, expected_datasets, sess
         dataset_urls = {dataset["uri"] for dataset in response.json["datasets"]}
         assert expected_datasets == dataset_urls
 
+    @pytest.mark.parametrize("dag_ids, expected_num", [("dag1,dag2", 2), ("dag3", 1), ("dag2,dag3", 2)])
+    @provide_session
+    def test_filter_datasets_by_dag_ids_works(self, dag_ids, expected_num, session):
+        session.query(DagModel).delete()
+        session.commit()
+        dag1 = DagModel(dag_id="dag1")
+        dag2 = DagModel(dag_id="dag2")
+        dag3 = DagModel(dag_id="dag3")
+        dataset1 = DatasetModel("s3://folder/key")
+        dataset2 = DatasetModel("gcp://bucket/key")
+        dataset3 = DatasetModel("somescheme://dataset/key")
+        dag_ref1 = DagScheduleDatasetReference(dag_id="dag1", dataset=dataset1)
+        dag_ref2 = DagScheduleDatasetReference(dag_id="dag2", dataset=dataset2)
+        task_ref1 = TaskOutletDatasetReference(dag_id="dag3", task_id="task1", dataset=dataset3)
+        session.add_all([dataset1, dataset2, dataset3, dag1, dag2, dag3, dag_ref1, dag_ref2, task_ref1])
+        session.commit()
+        response = self.client.get(
+            f"/api/v1/datasets?dag_ids={dag_ids}", environ_overrides={"REMOTE_USER": "test"}
+        )
+        assert response.status_code == 200
+        response_data = response.json
+        assert len(response_data["datasets"]) == expected_num
+
+    @pytest.mark.parametrize(
+        "dag_ids, uri_pattern,expected_num",
+        [("dag1,dag2", "folder", 1), ("dag3", "nothing", 0), ("dag2,dag3", "key", 2)],
+    )
+    def test_filter_datasets_by_dag_ids_and_uri_pattern_works(
+        self, dag_ids, uri_pattern, expected_num, session
+    ):
+        session.query(DagModel).delete()
+        session.commit()
+        dag1 = DagModel(dag_id="dag1")
+        dag2 = DagModel(dag_id="dag2")
+        dag3 = DagModel(dag_id="dag3")
+        dataset1 = DatasetModel("s3://folder/key")
+        dataset2 = DatasetModel("gcp://bucket/key")
+        dataset3 = DatasetModel("somescheme://dataset/key")
+        dag_ref1 = DagScheduleDatasetReference(dag_id="dag1", dataset=dataset1)
+        dag_ref2 = DagScheduleDatasetReference(dag_id="dag2", dataset=dataset2)
+        task_ref1 = TaskOutletDatasetReference(dag_id="dag3", task_id="task1", dataset=dataset3)
+        session.add_all([dataset1, dataset2, dataset3, dag1, dag2, dag3, dag_ref1, dag_ref2, task_ref1])
+        session.commit()
+        response = self.client.get(
+            f"/api/v1/datasets?dag_ids={dag_ids}&uri_pattern={uri_pattern}",
+            environ_overrides={"REMOTE_USER": "test"},
+        )
+        assert response.status_code == 200
+        response_data = response.json
+        assert len(response_data["datasets"]) == expected_num
+
 
 class TestGetDatasetsEndpointPagination(TestDatasetEndpoint):
     @pytest.mark.parametrize(