From e33734e7160e840cecbd31ee1b5dc53044ae5e08 Mon Sep 17 00:00:00 2001
From: max <42827971+moiseenkov@users.noreply.github.com>
Date: Wed, 29 Jan 2025 16:25:45 +0100
Subject: [PATCH] [v2-10-test] Add Webserver parameters: max_form_parts,
 max_form_memory_size (#45749) (cherry picked from commit
 bb2aaef342ddee8a8c8ad7120ac8484e6c7f10d4)

Co-authored-by: max <42827971+moiseenkov@users.noreply.github.com>
---
 airflow/config_templates/config.yml  | 16 ++++++++++++++++
 airflow/www/app.py                   |  2 ++
 airflow/www/extensions/init_views.py | 17 +++++++++++++++++
 3 files changed, 35 insertions(+)

diff --git a/airflow/config_templates/config.yml b/airflow/config_templates/config.yml
index 613c5e3394a40..6c65ec7d285ac 100644
--- a/airflow/config_templates/config.yml
+++ b/airflow/config_templates/config.yml
@@ -2120,6 +2120,22 @@ webserver:
       type: boolean
       example: ~
       default: "False"
+    max_form_memory_size:
+      description: |
+        The maximum size in bytes any non-file form field may be in a multipart/form-data body.
+        If this limit is exceeded, a 413 RequestEntityTooLarge error is raised by webserver.
+      version_added: 2.10.5
+      type: integer
+      example: ~
+      default: "500000"
+    max_form_parts:
+      description: |
+        The maximum number of fields that may be present in a multipart/form-data body.
+        If this limit is exceeded, a 413 RequestEntityTooLarge error is raised by webserver.
+      version_added: 2.10.5
+      type: integer
+      example: ~
+      default: "1000"
 email:
   description: |
     Configuration email backend and whether to
diff --git a/airflow/www/app.py b/airflow/www/app.py
index e093e66cfd881..23d79b0186138 100644
--- a/airflow/www/app.py
+++ b/airflow/www/app.py
@@ -78,6 +78,8 @@ def create_app(config=None, testing=False):
     flask_app.config["PERMANENT_SESSION_LIFETIME"] = timedelta(minutes=settings.get_session_lifetime_config())
 
     flask_app.config["MAX_CONTENT_LENGTH"] = conf.getfloat("webserver", "allowed_payload_size") * 1024 * 1024
+    flask_app.config["MAX_FORM_PARTS"] = conf.getint("webserver", "max_form_parts")
+    flask_app.config["MAX_FORM_MEMORY_SIZE"] = conf.getint("webserver", "max_form_memory_size")
 
     webserver_config = conf.get_mandatory_value("webserver", "config_file")
     # Enable customizations in webserver_config.py to be applied via Flask.current_app.
diff --git a/airflow/www/extensions/init_views.py b/airflow/www/extensions/init_views.py
index cc4e1141be707..16ccea91ac137 100644
--- a/airflow/www/extensions/init_views.py
+++ b/airflow/www/extensions/init_views.py
@@ -26,6 +26,7 @@
 from connexion.decorators.validation import RequestBodyValidator
 from connexion.exceptions import BadRequestProblem
 from flask import request
+from werkzeug import Request
 
 from airflow.api_connexion.exceptions import common_error_handler
 from airflow.configuration import conf
@@ -194,6 +195,21 @@ def set_cors_headers_on_response(response):
     return response
 
 
+def init_data_form_parameters():
+    """
+    Initialize custom values for data form parameters.
+
+    This is a workaround for Flask versions prior to 3.1.0.
+    In order to allow users customizing form data parameters, we need these two fields to be configurable.
+    Starting from Flask 3.1.0 these two parameters can be configured through Flask config, but unfortunately,
+    current version of flask supported in Airflow is way older. That's why this workaround was introduced.
+    See https://flask.palletsprojects.com/en/stable/api/#flask.Request.max_form_memory_size
+    # TODO: remove it when Flask upgraded to version 3.1.0 or higher.
+    """
+    Request.max_form_parts = conf.getint("webserver", "max_form_parts")
+    Request.max_form_memory_size = conf.getint("webserver", "max_form_memory_size")
+
+
 class _LazyResolution:
     """
     OpenAPI endpoint that lazily resolves the function on first use.
@@ -286,6 +302,7 @@ def init_api_connexion(app: Flask) -> None:
         validate_responses=True,
         validator_map={"body": _CustomErrorRequestBodyValidator},
     ).blueprint
+    api_bp.before_app_request(init_data_form_parameters)
     api_bp.after_request(set_cors_headers_on_response)
 
     app.register_blueprint(api_bp)