fix: annotation layer json (#9915)

* fix: annotation layer json

* attempt to add a test

* [tests] Fixing test

Co-authored-by: John Bodley <john.bodley@airbnb.com>

superset/views/core.py

@@ -685,6 +685,21 @@ class Superset(BaseSupersetView):
         form_data = get_form_data()[0]
         form_data["layer_id"] = layer_id
         form_data["filters"] = [{"col": "layer_id", "op": "==", "val": layer_id}]
+        # Set all_columns to ensure the TableViz returns the necessary columns to the
+        # frontend.
+        form_data["all_columns"] = [
+            "created_on",
+            "changed_on",
+            "id",
+            "start_dttm",
+            "end_dttm",
+            "layer_id",
+            "short_descr",
+            "long_descr",
+            "json_metadata",
+            "created_by_fk",
+            "changed_by_fk",
+        ]
         datasource = AnnotationDatasource()
         viz_obj = viz.viz_types["table"](datasource, form_data=form_data, force=False)
         payload = viz_obj.get_payload()

tests/core_tests.py

@@ -25,6 +25,7 @@ import json
 import logging
 import os
 from typing import Dict, List, Optional
+from urllib.parse import quote
 
 import pytz
 import random
@@ -50,6 +51,7 @@ from superset.datasets.dao import DatasetDAO
 from superset.db_engine_specs.base import BaseEngineSpec
 from superset.db_engine_specs.mssql import MssqlEngineSpec
 from superset.models import core as models
+from superset.models.annotations import Annotation, AnnotationLayer
 from superset.models.dashboard import Dashboard
 from superset.models.datasource_access_request import DatasourceAccessRequest
 from superset.models.slice import Slice
@@ -177,6 +179,29 @@ class CoreTests(SupersetTestCase):
         resp = self.get_resp(slc.explore_json_url)
         assert '"Jennifer"' in resp
 
+    def test_annotation_json_endpoint(self):
+        # Set up an annotation layer and annotation
+        layer = AnnotationLayer(name="foo", descr="bar")
+        db.session.add(layer)
+        db.session.commit()
+
+        annotation = Annotation(
+            layer_id=layer.id,
+            short_descr="my_annotation",
+            start_dttm=datetime.datetime(2020, 5, 20, 18, 21, 51),
+            end_dttm=datetime.datetime(2020, 5, 20, 18, 31, 51),
+        )
+
+        db.session.add(annotation)
+        db.session.commit()
+
+        resp = self.get_resp(
+            f"/superset/annotation_json/{layer.id}?form_data="
+            + quote(json.dumps({"time_range": "100 years ago : now"}))
+        )
+
+        assert "my_annotation" in resp
+
     def test_old_slice_csv_endpoint(self):
         self.login(username="admin")
         slc = self.get_slice("Girls", db.session)

tests/superset_test_config.py

@@ -76,3 +76,5 @@ CELERY_CONFIG = CeleryConfig
 CUSTOM_TEMPLATE_PROCESSORS = {
     CustomPrestoTemplateProcessor.engine: CustomPrestoTemplateProcessor
 }
+
+PRESERVE_CONTEXT_ON_EXCEPTION = False