[RAPTOR-15440] consolidate usage of werkzeug.parse_options_header in one api (#1805)

Author: yakov-g

custom_model_runner/datarobot_drum/drum/root_predictors/generic_predictor.py

@@ -6,7 +6,6 @@
 """
 import urllib
 
-import werkzeug
 from datarobot_drum.drum.adapters.cli.drum_score_adapter import DrumScoreAdapter
 from datarobot_drum.drum.enum import GPU_PREDICTORS
 from datarobot_drum.drum.enum import TARGET_TYPE_ARG_KEYWORD
@@ -21,6 +20,8 @@
     _resolve_outgoing_unstructured_data,
 )
 
+from datarobot_drum.drum.root_predictors.utils import get_mimetype_charset_from_content_type_header
+
 
 class GenericPredictorComponent:
     def __init__(self, params: dict):
@@ -122,10 +123,9 @@ def materialize(self):
     def _materialize_unstructured(self, input_filename, output_filename):
         kwargs_params = {}
         query_params = dict(urllib.parse.parse_qsl(self._params.get("query_params")))
-        mimetype, content_type_params_dict = werkzeug.http.parse_options_header(
+        mimetype, charset = get_mimetype_charset_from_content_type_header(
             self._params.get("content_type")
         )
-        charset = content_type_params_dict.get("charset")
 
         with open(input_filename, "rb") as f:
             data_binary = f.read()

custom_model_runner/datarobot_drum/drum/root_predictors/predict_mixin.py

@@ -6,7 +6,6 @@
 """
 import logging
 
-import werkzeug
 from flask import request, Response, stream_with_context
 from requests_toolbelt import MultipartEncoder
 
@@ -40,6 +39,8 @@
     _resolve_outgoing_unstructured_data,
 )
 
+from datarobot_drum.drum.root_predictors.utils import get_mimetype_charset_from_content_type_header
+
 
 class PredictMixin:
     """
@@ -55,9 +56,7 @@ def _log_if_possible(logger, log_level, message):
 
     @staticmethod
     def _validate_content_type_header(header):
-        ret_mimetype, content_type_params_dict = werkzeug.http.parse_options_header(header)
-        ret_charset = content_type_params_dict.get("charset")
-        return ret_mimetype, ret_charset
+        return get_mimetype_charset_from_content_type_header(header)
 
     @staticmethod
     def _fetch_data_from_request(file_key, logger=None):

custom_model_runner/datarobot_drum/drum/root_predictors/utils.py

@@ -10,9 +10,12 @@
 import shlex
 import subprocess
 import time
+import werkzeug
+
 from queue import Queue, Empty
 from threading import Thread
 
+
 from datarobot_drum.drum.common import get_drum_logger
 from datarobot_drum.drum.enum import (
     ArgumentOptionsEnvVars,
@@ -235,3 +238,9 @@ def _cmd_add_class_labels(
                 labels_str = " ".join(['"{}"'.format(label) for label in labels])
                 cmd += " --class-labels {}".format(labels_str)
     return cmd
+
+
+def get_mimetype_charset_from_content_type_header(header):
+    mimetype, content_type_params_dict = werkzeug.http.parse_options_header(header)
+    charset = content_type_params_dict.get("charset")
+    return mimetype, charset

tests/functional/test_unstructured_mode_per_framework.py

@@ -6,7 +6,6 @@
 """
 import pytest
 import requests
-import werkzeug
 
 from datarobot_drum.drum.enum import ArgumentsOptions
 from datarobot_drum.drum.server import HTTP_422_UNPROCESSABLE_ENTITY
@@ -18,8 +17,23 @@
     _create_custom_model_dir,
 )
 
+# This is only for Java environment case.
+# For the java env we don't install from source,
+# so this func will be available when Java env uses released DRUM with this func.
+# Anyhow, unstructured is not supported by java, so these tests are skipped for Java env.
+try:
+    from datarobot_drum.drum.root_predictors.utils import (
+        get_mimetype_charset_from_content_type_header,
+    )
+except ImportError:
+
+    def get_mimetype_charset_from_content_type_header(header):
+        return None, None
+
+
 from requests_toolbelt import MultipartEncoder
 
+
 from tests.constants import (
     R_NO_ARTIFACTS,
     SKLEARN_NO_ARTIFACTS,
@@ -338,11 +352,11 @@ def test_response_one_var_return(
                             )
                             assert response.ok
                             content_type_header = response.headers["Content-Type"]
-                            mimetype, content_type_params_dict = werkzeug.http.parse_options_header(
+                            mimetype, charset = get_mimetype_charset_from_content_type_header(
                                 content_type_header
                             )
                             assert mimetype == "text/plain"
-                            assert content_type_params_dict["charset"] == UTF8
+                            assert charset == UTF8
                             if data is None:
                                 assert len(response.content) == 0
                             else:
@@ -363,11 +377,11 @@ def test_response_one_var_return(
                             )
                             assert response.ok
                             content_type_header = response.headers["Content-Type"]
-                            mimetype, content_type_params_dict = werkzeug.http.parse_options_header(
+                            mimetype, charset = get_mimetype_charset_from_content_type_header(
                                 content_type_header
                             )
                             assert mimetype == "text/plain"
-                            assert content_type_params_dict["charset"] == UTF8
+                            assert charset == UTF8
                             if data is None:
                                 assert len(response.content) == 0
                             else:
@@ -381,10 +395,9 @@ def test_response_one_var_return(
                     )
                     assert response.ok
                     content_type_header = response.headers["Content-Type"]
-                    mimetype, content_type_params_dict = werkzeug.http.parse_options_header(
+                    mimetype, charset = get_mimetype_charset_from_content_type_header(
                         content_type_header
                     )
                     assert "application/octet-stream" == mimetype
-                    # check params dict is empty
-                    assert not any(content_type_params_dict)
+                    assert charset is None
                     assert response.content == data_bytes

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

@@ -310,3 +310,58 @@ def test_stream_p_open_ignores_empty_lines(self, mock_logger):
         mock_logger.assert_any_call(b"regular line")
         mock_logger.assert_any_call(b"error line")
         assert mock_logger.call_count == 2  # Only two lines should be logged
+
+
+class TestGetMimetypeCharsetFromContentTypeHeader:
+    def test_typical_header(self):
+        from custom_model_runner.datarobot_drum.drum.root_predictors.utils import (
+            get_mimetype_charset_from_content_type_header,
+        )
+
+        mimetype, charset = get_mimetype_charset_from_content_type_header(
+            "text/html; charset=utf-8"
+        )
+
+        assert mimetype == "text/html"
+        assert charset == "utf-8"
+
+    def test_header_without_charset(self):
+        from custom_model_runner.datarobot_drum.drum.root_predictors.utils import (
+            get_mimetype_charset_from_content_type_header,
+        )
+
+        mimetype, charset = get_mimetype_charset_from_content_type_header("application/json")
+
+        assert mimetype == "application/json"
+        assert charset is None
+
+    def test_header_with_additional_params(self):
+        from custom_model_runner.datarobot_drum.drum.root_predictors.utils import (
+            get_mimetype_charset_from_content_type_header,
+        )
+
+        mimetype, charset = get_mimetype_charset_from_content_type_header(
+            "text/plain; charset=iso-8859-1; format=flowed"
+        )
+
+        assert mimetype == "text/plain"
+        assert charset == "iso-8859-1"
+
+    def test_empty_header(self):
+        from custom_model_runner.datarobot_drum.drum.root_predictors.utils import (
+            get_mimetype_charset_from_content_type_header,
+        )
+
+        mimetype, charset = get_mimetype_charset_from_content_type_header("")
+
+        assert mimetype == ""
+        assert charset is None
+
+    def test_none_header(self):
+        from custom_model_runner.datarobot_drum.drum.root_predictors.utils import (
+            get_mimetype_charset_from_content_type_header,
+        )
+
+        mimetype, charset = get_mimetype_charset_from_content_type_header(None)
+        assert mimetype == ""
+        assert charset is None