Correctly set json_provider_class on Flask app so it uses our encoder

[ashb]

Setting json_provider_class where we did had no effect, as it turns out Flask() sets self.json = self.json_provider_class(self), so we were setting it too late to take any affect.

Closes #26546, #26527

[tirkarthi]

I was working on a PR for this too. The UI shows error with below traceback for this patch and I am not sure why tests didn't catch this. The places need to use app.json.dumps instead of AirflowJsonEncoder like below patch along with other places too.

Edit : Tests fail, sorry I was running tests probably without patch.

[2022-09-21 18:43:28,419] {_internal.py:224} INFO - 127.0.0.1 - - [21/Sep/2022 18:43:28] "POST /last_dagruns HTTP/1.1" 500 -
Traceback (most recent call last):
  File "/home/karthikeyan/stuff/python/airflow/.env/lib/python3.10/site-packages/flask/app.py", line 2548, in __call__
    return self.wsgi_app(environ, start_response)
  File "/home/karthikeyan/stuff/python/airflow/.env/lib/python3.10/site-packages/flask/app.py", line 2528, in wsgi_app
    response = self.handle_exception(e)
  File "/home/karthikeyan/stuff/python/airflow/.env/lib/python3.10/site-packages/flask/app.py", line 2525, in wsgi_app
    response = self.full_dispatch_request()
  File "/home/karthikeyan/stuff/python/airflow/.env/lib/python3.10/site-packages/flask/app.py", line 1822, in full_dispatch_request
    rv = self.handle_user_exception(e)
  File "/home/karthikeyan/stuff/python/airflow/.env/lib/python3.10/site-packages/flask/app.py", line 1820, in full_dispatch_request
    rv = self.dispatch_request()
  File "/home/karthikeyan/stuff/python/airflow/.env/lib/python3.10/site-packages/flask/app.py", line 1796, in dispatch_request
    return self.ensure_sync(self.view_functions[rule.endpoint])(**view_args)
  File "/home/karthikeyan/stuff/python/airflow/airflow/www/auth.py", line 47, in decorated
    return func(*args, **kwargs)
  File "/home/karthikeyan/stuff/python/airflow/airflow/utils/session.py", line 75, in wrapper
    return func(*args, session=session, **kwargs)
  File "/home/karthikeyan/stuff/python/airflow/airflow/www/views.py", line 1049, in last_dagruns
    return wwwutils.json_response(resp)
  File "/home/karthikeyan/stuff/python/airflow/airflow/www/utils.py", line 328, in json_response
    response=json.dumps(obj, indent=4, cls=AirflowJsonEncoder), status=200, mimetype="application/json"
  File "/usr/lib/python3.10/json/__init__.py", line 234, in dumps
    return cls(
  File "/home/karthikeyan/stuff/python/airflow/airflow/utils/json.py", line 47, in __init__
    super().__init__(*args, **kwargs)
TypeError: JSONProvider.__init__() got an unexpected keyword argument 'skipkeys'

diff --git a/airflow/www/utils.py b/airflow/www/utils.py
index d0efa611d5..5e57ade1d8 100644
--- a/airflow/www/utils.py
+++ b/airflow/www/utils.py
@@ -45,6 +45,7 @@ from airflow.models import errors
 from airflow.models.dagwarning import DagWarning
 from airflow.models.taskinstance import TaskInstance
 from airflow.utils import timezone
+from airflow.utils.airflow_flask_app import get_airflow_app
 from airflow.utils.code_utils import get_python_source
 from airflow.utils.helpers import alchemy_to_dict
 from airflow.utils.json import AirflowJsonEncoder
@@ -324,8 +325,9 @@ def epoch(dttm):
 
 def json_response(obj):
     """Returns a json response from a json serializable python object"""
+    app = get_airflow_app()
     return Response(
-        response=json.dumps(obj, indent=4, cls=AirflowJsonEncoder), status=200, mimetype="application/json"
+        response=app.json.dumps(obj, indent=4), status=200, mimetype="application/json"
     )
 
 
diff --git a/airflow/www/views.py b/airflow/www/views.py
index 2c6701505c..ce334bfb65 100644
--- a/airflow/www/views.py
+++ b/airflow/www/views.py
@@ -3397,8 +3397,9 @@ class Airflow(AirflowBaseView):
     @action_logging
     def task_instances(self):
         """Shows task instances."""
+        app = get_airflow_app()
         dag_id = request.args.get('dag_id')
-        dag = get_airflow_app().dag_bag.get_dag(dag_id)
+        dag = app.dag_bag.get_dag(dag_id)
 
         dttm = request.args.get('execution_date')
         if dttm:
@@ -3412,7 +3413,7 @@ class Airflow(AirflowBaseView):
                 for ti in dag.get_task_instances(dttm, dttm)
             }
 
-        return json.dumps(task_instances, cls=utils_json.AirflowJsonEncoder)
+        return json.dumps(task_instances, dumps=app.json.dumps)
 
     @expose('/object/grid_data')
     @auth.has_access(
@@ -3423,8 +3424,9 @@ class Airflow(AirflowBaseView):
     )
     def grid_data(self):
         """Returns grid data"""
+        app = get_airflow_app()
         dag_id = request.args.get('dag_id')
-        dag = get_airflow_app().dag_bag.get_dag(dag_id)
+        dag = app.dag_bag.get_dag(dag_id)
 
         if not dag:
             return {'error': f"can't find dag {dag_id}"}, 404
@@ -3467,7 +3469,7 @@ class Airflow(AirflowBaseView):
             }
         # avoid spaces to reduce payload size
         return (
-            htmlsafe_json_dumps(data, separators=(',', ':'), cls=utils_json.AirflowJsonEncoder),
+            htmlsafe_json_dumps(data, separators=(',', ':'), dumps=app.json.dumps),
             {'Content-Type': 'application/json; charset=utf-8'},
         )
 
@@ -3475,7 +3477,8 @@ class Airflow(AirflowBaseView):
     @auth.has_access([(permissions.ACTION_CAN_READ, permissions.RESOURCE_DAG)])
     def next_run_datasets(self, dag_id):
         """Returns datasets necessary, and their status, for the next dag run"""
-        dag = get_airflow_app().dag_bag.get_dag(dag_id)
+        app = get_airflow_app()
+        dag = app.dag_bag.get_dag(dag_id)
         if not dag:
             return {'error': f"can't find dag {dag_id}"}, 404
 
@@ -3510,7 +3513,7 @@ class Airflow(AirflowBaseView):
                 .all()
             ]
         return (
-            htmlsafe_json_dumps(data, separators=(',', ':'), cls=utils_json.AirflowJsonEncoder),
+            htmlsafe_json_dumps(data, separators=(',', ':'), dumps=app.json.dumps),
             {'Content-Type': 'application/json; charset=utf-8'},
         )
 
@@ -3524,6 +3527,7 @@ class Airflow(AirflowBaseView):
         """Returns dataset dependencies graph."""
         nodes_dict: dict[str, Any] = {}
         edge_tuples: set[dict[str, str]] = set()
+        app = get_airflow_app()
 
         for dag, dependencies in SerializedDagModel.get_dag_dependencies().items():
             dag_node_id = f"dag:{dag}"
@@ -3547,7 +3551,7 @@ class Airflow(AirflowBaseView):
         }
 
         return (
-            htmlsafe_json_dumps(data, separators=(',', ':'), cls=utils_json.AirflowJsonEncoder),
+            htmlsafe_json_dumps(data, separators=(',', ':'), dumps=app.json.dumps),
             {'Content-Type': 'application/json; charset=utf-8'},
         )

[ashb]

@tirkarthi Looks like your fix is better/bigger than mine, so do you want to add a separate PR for those cases? Edit: Oh I see, nm, I'm on it.

I'll remove the "close" marker from this PR as I don't think it fully fixes it? (I admit to not testing in the app, just using the minimal repro case from the closed PR)

Edit 2: a better fix then get_airflow_app() is to use https://flask.palletsprojects.com/en/2.2.x/api/?highlight=json#module-flask.json I think.

[tirkarthi]

@ashb Feel free to use the above patch. I guess there is some change due to inheriting from DefaultJSONProvider where the cls argument used inside htmlsafe_json_dumps fails along with other places using json.dumps from stdlib. Since htmlsafe_json_dumps has dumps argument passing app.json.dumps will have similar effect as earlier where cls was passed. json.dumps can be replaced by app.json.dumps.

[tirkarthi]

yes, flask.json seems better. There is note from 2.2 that current_app.json.dumps is used internally that will make the patch simpler. https://flask.palletsprojects.com/en/2.2.x/api/?highlight=json#flask.json.dumps

[ashb]

Was a little bit more complex to have a provider work correctly, so asking for re-review.

[dstandish]

At some point in the past, if user had simplejson installed, then webserver would blow up, that's why we imported from flask.json. It might be worth doing pip install simplejson just to make sure that it is compatible

[dstandish]

it was this commit: ea3d42a

[ashb]

Tested that, doesn't blow up (by which I mean tests still create the app and pass)

[dstandish]

cool -- i just tried it too, installed it and launched webserver and navigated to a view or two

[ashb]

Failure was fixed on main already