[RAPTOR-15441] remove cgi package usage (#1806)

Author: yakov-g

custom_model_runner/datarobot_drum/drum/root_predictors/transform_helpers.py

@@ -7,12 +7,12 @@
 import pandas as pd
 import logging
 
-from cgi import FieldStorage
 from io import BytesIO, StringIO
 
 from scipy.io import mmwrite, mmread
 from scipy.sparse import issparse
 from scipy.sparse import csr_matrix
+from werkzeug.formparser import parse_form_data
 
 from datarobot_drum.drum.enum import X_FORMAT_KEY, X_TRANSFORM_KEY
 
@@ -102,20 +102,18 @@ def read_mtx_payload(response_dict, transform_key):
 
 
 def parse_multi_part_response(response):
+    environ = {
+        "wsgi.input": BytesIO(response.content),
+        "REQUEST_METHOD": "POST",
+        "CONTENT_TYPE": response.headers["Content-Type"],
+        "CONTENT_LENGTH": str(len(response.content)),
+    }
+    _, form, files = parse_form_data(environ)
     parsed_response = {}
-    fs = FieldStorage(
-        fp=BytesIO(response.content),
-        headers=response.headers,
-        environ={
-            "REQUEST_METHOD": "POST",
-            "CONTENT_TYPE": response.headers["Content-Type"],
-        },
-    )
-    for child in fs.list:
-        key = child.name
-        value = child.file.read()
-        parsed_response.update({key: value})
-
+    for key, value in form.items():
+        parsed_response[key] = value
+    for key, file_storage in files.items():
+        parsed_response[key] = file_storage.read()
     return parsed_response
 
 

tests/unit/datarobot_drum/drum/root_predictors/test_transform_helpers.py

@@ -7,10 +7,12 @@
 import pandas as pd
 import numpy as np
 import pytest
+import types
 from scipy.sparse import coo_matrix
 
 from datarobot_drum.drum.root_predictors.transform_helpers import (
     validate_and_convert_column_names_for_serialization,
+    parse_multi_part_response,
 )
 
 
@@ -44,3 +46,36 @@ def test_validate_and_convert_column_names_for_serialization(columns, expected_c
     else:
         output_df = validate_and_convert_column_names_for_serialization(df)
         assert list(output_df.columns.values) == expected_columns
+
+
+def test_parse_multi_part_response():
+    file_content1 = b"value1"
+    file_content2 = b"value2"
+    boundary = "boundary123"
+    content = (
+        (
+            f"--{boundary}\r\n"
+            'Content-Disposition: form-data; name="X.format"\r\n\r\n'
+            "csv\r\n"
+            f"--{boundary}\r\n"
+            'Content-Disposition: form-data; name="key1"; filename="file1.txt"\r\n'
+            "Content-Type: text/plain\r\n\r\n"
+        ).encode("utf-8")
+        + file_content1
+        + (
+            f"\r\n--{boundary}\r\n"
+            'Content-Disposition: form-data; name="key2"; filename="file2.txt"\r\n'
+            "Content-Type: text/plain\r\n\r\n"
+        ).encode("utf-8")
+        + file_content2
+        + f"\r\n--{boundary}--\r\n".encode("utf-8")
+    )
+
+    response = types.SimpleNamespace()
+    response.content = content
+    response.headers = {"Content-Type": f"multipart/form-data; boundary={boundary}"}
+
+    result = parse_multi_part_response(response)
+    assert result["X.format"] == "csv"
+    assert result["key1"] == file_content1
+    assert result["key2"] == file_content2