fix: datasets related objects, apply filter and openapi spec

superset/datasets/api.py

@@ -42,6 +42,7 @@
 from superset.datasets.schemas import (
     DatasetPostSchema,
     DatasetPutSchema,
+    DatasetRelatedObjectsResponse,
     get_export_ids_schema,
 )
 from superset.views.base import DatasourceFilter, generate_download_headers
@@ -139,6 +140,8 @@ class DatasetRestApi(BaseSupersetModelRestApi):
     filter_rel_fields = {"database": [["id", DatabaseFilter, lambda: []]]}
     allowed_rel_fields = {"database", "owners"}
 
+    openapi_spec_component_schemas = (DatasetRelatedObjectsResponse,)
+
     @expose("/", methods=["POST"])
     @protect()
     @safe
@@ -434,80 +437,42 @@ def related_objects(self, pk: int) -> Response:
               type: integer
           responses:
             200:
-              description: chart and dashboard counts
+            200:
+              description: Query result
               content:
                 application/json:
                   schema:
-                    type: object
-                    properties:
-                      charts:
-                        type: object
-                        properties:
-                          count:
-                            type: integer
-                          result:
-                            type: array
-                            items:
-                              type: object
-                              properties:
-                                id:
-                                  type: integer
-                                slice_name:
-                                  type: string
-                                viz_type:
-                                  type: string
-                      dashboards:
-                        type: object
-                        properties:
-                          count:
-                            type: integer
-                          result:
-                            type: array
-                            items:
-                              type: object
-                              properties:
-                                id:
-                                  type: integer
-                                json_metadata:
-                                  type: object
-                                slug:
-                                  type: string
-                                title:
-                                  type: string
-            400:
-              $ref: '#/components/responses/400'
+                    $ref: "#/components/schemas/DatasetRelatedObjectsResponse"
             401:
               $ref: '#/components/responses/401'
             404:
               $ref: '#/components/responses/404'
             500:
               $ref: '#/components/responses/500'
         """
-        try:
-            data = DatasetDAO.get_related_objects(pk)
-            charts = [
-                {
-                    "id": chart.id,
-                    "slice_name": chart.slice_name,
-                    "viz_type": chart.viz_type,
-                }
-                for chart in data["charts"]
-            ]
-            dashboards = [
-                {
-                    "id": dashboard.id,
-                    "json_metadata": dashboard.json_metadata,
-                    "slug": dashboard.slug,
-                    "title": dashboard.dashboard_title,
-                }
-                for dashboard in data["dashboards"]
-            ]
-            return self.response(
-                200,
-                charts={"count": len(charts), "result": charts},
-                dashboards={"count": len(dashboards), "result": dashboards},
-            )
-        except DatasetNotFoundError:
+        dataset = DatasetDAO.find_by_id(pk)
+        if not dataset:
             return self.response_404()
-        except DatasetForbiddenError:
-            return self.response_403()
+        data = DatasetDAO.get_related_objects(pk)
+        charts = [
+            {
+                "id": chart.id,
+                "slice_name": chart.slice_name,
+                "viz_type": chart.viz_type,
+            }
+            for chart in data["charts"]
+        ]
+        dashboards = [
+            {
+                "id": dashboard.id,
+                "json_metadata": dashboard.json_metadata,
+                "slug": dashboard.slug,
+                "title": dashboard.dashboard_title,
+            }
+            for dashboard in data["dashboards"]
+        ]
+        return self.response(
+            200,
+            charts={"count": len(charts), "result": charts},
+            dashboards={"count": len(dashboards), "result": dashboards},
+        )

superset/datasets/schemas.py

@@ -87,3 +87,35 @@ class DatasetPutSchema(Schema):
     owners = fields.List(fields.Integer())
     columns = fields.List(fields.Nested(DatasetColumnsPutSchema))
     metrics = fields.List(fields.Nested(DatasetMetricsPutSchema))
+
+
+class DatasetRelatedChart(Schema):
+    id = fields.Integer()
+    slice_name = fields.String()
+    viz_type = fields.String()
+
+
+class DatasetRelatedDashboard(Schema):
+    id = fields.Integer()
+    json_metadata = fields.Dict()
+    slug = fields.String()
+    title = fields.String()
+
+
+class DatasetRelatedCharts(Schema):
+    count = fields.Integer(description="Chart count")
+    result = fields.List(
+        fields.Nested(DatasetRelatedChart), description="A list of dashboards"
+    )
+
+
+class DatasetRelatedDashboards(Schema):
+    count = fields.Integer(description="Dashboard count")
+    result = fields.List(
+        fields.Nested(DatasetRelatedDashboard), description="A list of dashboards"
+    )
+
+
+class DatasetRelatedObjectsResponse(Schema):
+    charts = fields.Nested(DatasetRelatedCharts)
+    dashboards = fields.Nested(DatasetRelatedDashboards)

tests/datasets/api_tests.py

@@ -703,7 +703,6 @@ def test_dataset_item_refresh_not_owned(self):
     def test_export_dataset(self):
         """
         Dataset API: Test export dataset
-        :return:
         """
         birth_names_dataset = self.get_birth_names_dataset()
 
@@ -736,7 +735,6 @@ def test_export_dataset(self):
     def test_export_dataset_not_found(self):
         """
         Dataset API: Test export dataset not found
-        :return:
         """
         max_id = db.session.query(func.max(SqlaTable.id)).scalar()
         # Just one does not exist and we get 404
@@ -749,7 +747,6 @@ def test_export_dataset_not_found(self):
     def test_export_dataset_gamma(self):
         """
         Dataset API: Test export dataset has gamma
-        :return:
         """
         birth_names_dataset = self.get_birth_names_dataset()
 
@@ -773,3 +770,21 @@ def test_get_dataset_related_objects(self):
         response = json.loads(rv.data.decode("utf-8"))
         self.assertEqual(response["charts"]["count"], 18)
         self.assertEqual(response["dashboards"]["count"], 2)
+
+    def test_get_dataset_related_objects_not_found(self):
+        """
+        Dataset API: Test related objects not found
+        """
+        max_id = db.session.query(func.max(SqlaTable.id)).scalar()
+        # id does not exist and we get 404
+        invalid_id = max_id + 1
+        uri = f"api/v1/dataset/{invalid_id}/related_objects/"
+        self.login(username="admin")
+        rv = self.client.get(uri)
+        self.assertEqual(rv.status_code, 404)
+        self.logout()
+        self.login(username="gamma")
+        table = self.get_birth_names_dataset()
+        uri = f"api/v1/dataset/{table.id}/related_objects"
+        rv = self.client.get(uri)
+        self.assertEqual(rv.status_code, 404)