dpgaspar commented on a change in pull request #11665:
URL: 
https://github.com/apache/incubator-superset/pull/11665#discussion_r535504766



##########
File path: tests/charts/api_tests.py
##########
@@ -1383,3 +1386,44 @@ def test_import_chart_invalid(self):
         assert response == {
             "message": {"metadata.yaml": {"type": ["Must be equal to Slice."]}}
         }
+
+    @pytest.mark.usefixtures("create_annotation_layers")

Review comment:
       We should think about moving all these tests away from 
`unittest.TestCase` if possible, there are some really cool pytest features 
that we are missing out

##########
File path: superset/common/query_context.py
##########
@@ -220,7 +222,80 @@ def cache_key(self, query_obj: QueryObject, **kwargs: Any) 
-> Optional[str]:
         )
         return cache_key
 
-    def get_df_payload(  # pylint: disable=too-many-statements
+    @staticmethod
+    def get_native_annotation_data(query_obj: QueryObject) -> Dict[str, Any]:
+        annotation_data = {}
+        annotation_layers = [
+            layer
+            for layer in query_obj.annotation_layers
+            if layer["sourceType"] == "NATIVE"
+        ]
+        layer_ids = [layer["value"] for layer in annotation_layers]
+        layer_objects = {
+            layer_object.id: layer_object
+            for layer_object in AnnotationLayerDAO.find_by_ids(layer_ids)
+        }
+
+        # annotations
+        for layer in annotation_layers:
+            layer_id = layer["value"]
+            layer_name = layer["name"]
+            columns = [
+                "start_dttm",
+                "end_dttm",
+                "short_descr",
+                "long_descr",
+                "json_metadata",
+            ]
+            layer_object = layer_objects[layer_id]
+            records = [
+                {column: getattr(annotation, column) for column in columns}
+                for annotation in layer_object.annotation
+            ]
+            result = {"columns": columns, "records": records}
+            annotation_data[layer_name] = result
+        return annotation_data
+
+    @staticmethod
+    def get_viz_annotation_data(
+        annotation_layer: Dict[str, Any], force: bool
+    ) -> Dict[str, Any]:
+        slice_id = annotation_layer["value"]
+        slc = db.session.query(Slice).filter_by(id=slice_id).one_or_none()

Review comment:
       nit: could make sense to use `ChartDao.find_by_id` instead, note DAO's 
are starting to support being called from celery workers and the Security chart 
filter is always applied there

##########
File path: superset/charts/api.py
##########
@@ -485,6 +485,7 @@ def data(self) -> Response:
         """
         if request.is_json:
             json_body = request.json
+            print(json_body)

Review comment:
       remove?




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
[email protected]



---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to