servicex codegen to return transformer image and tag

Author: shriram192

poetry.lock

@@ -14,6 +14,7 @@ Flask-WTF = "^1.0.1"
 itsdangerous = "^2.1.2"
 Werkzeug = "^2.2.2"
 Jinja2 = "^3.1.2"
+requests-toolbelt = "^0.10.1"
 
 [tool.poetry.group.test]
 optional = true

pyproject.toml

@@ -33,13 +33,13 @@
 #
 import os
 
-from flask import Flask
+from flask import Flask, jsonify
 from flask_restful import Api
+
 from servicex_codegen.post_operation import GeneratedCode
 
 
 def handle_invalid_usage(error: BaseException):
-    from flask import jsonify
     response = jsonify({"message": str(error)})
     response.status_code = 400
     return response

servicex_codegen/__init__.py

@@ -35,8 +35,10 @@
 import zipfile
 from tempfile import TemporaryDirectory
 
-from flask import request, Response
+from flask import Response, current_app, request
 from flask_restful import Resource
+from requests_toolbelt import MultipartEncoder
+
 from servicex_codegen.code_generator import CodeGenerator, GeneratedFileResult
 
 
@@ -83,18 +85,29 @@ def stream_generated_code(self, generated_code_result: GeneratedFileResult) -> b
     def post(self):
         try:
             with TemporaryDirectory() as tempdir:
-                code = request.data.decode('utf8')
-                generated_code_result = self.code_generator.generate_code(code, cache_path=tempdir)
+                body = request.get_json()
+                generated_code_result = self.code_generator.generate_code(
+                    body["code"], cache_path=tempdir)
 
                 zip_data = self.stream_generated_code(generated_code_result)
+                transformer_image = body["transformer_image"] if "transformer_image" in body else current_app.config.get(
+                    "DEFAULT_TRANSFORMER_IMAGE")
+                transformer_tag = body["transformer_tag"] if "transformer_tag" in body else current_app.config.get(
+                    "DEFAULT_TRANSFORMER_TAG")
                 # Send the response back to you-know-what.
+
+                m = MultipartEncoder(
+                    fields={'transformer_image': transformer_image, 'transformer_tag': transformer_tag,
+                            'zip_data': zip_data}
+                )
+
                 response = Response(
-                    response=zip_data,
-                    status=200, mimetype='application/octet-stream')
+                    response=m.to_string(),
+                    status=200, mimetype=m.content_type)
                 return response
         except BaseException as e:
             print(str(e))
-            import traceback
             import sys
+            import traceback
             traceback.print_exc(file=sys.stdout)
             return {'Message': str(e)}, 500

servicex_codegen/post_operation.py

@@ -29,7 +29,9 @@
 import os
 import zipfile
 from tempfile import TemporaryDirectory
+from unittest.mock import patch
 
+from requests_toolbelt.multipart import decoder
 from servicex_codegen import create_app
 from servicex_codegen.code_generator import (GenerateCodeException,
                                              GeneratedFileResult)
@@ -51,7 +53,7 @@ def check_zip_file(zip_data: bytes, expected_file_count):
 
 
 class TestPostOperation:
-    def test_post_good_query(self, mocker):
+    def test_post_good_query_with_params(self, mocker):
         """Produce code for a simple good query"""
 
         with TemporaryDirectory() as tempdir, \
@@ -68,11 +70,79 @@ def test_post_good_query(self, mocker):
             }
             app = create_app(config, provided_translator=mock_ast_translator)
             client = app.test_client()
+
+            print(app.config)
+
+            select_stmt = "(call ResultTTree (call Select (call SelectMany (call EventDataset (list 'localds://did_01')"  # noqa: E501
+
+            response = client.post("/servicex/generated-code", json={
+                "transformer_image": "sslhep/servicex_func_adl_xaod_transformer",
+                "transformer_tag": "develop",
+                "code": select_stmt
+            })
+
+        boundary = str(response.data[2:34], 'utf-8')
+        content_type = f"multipart/form-data; boundary={boundary}"
+        decoder_parts = decoder.MultipartDecoder(response.data, content_type)
+
+        transformer_image = str(decoder_parts.parts[0].content, 'utf-8')
+        transformer_tag = str(decoder_parts.parts[1].content, 'utf-8')
+        zip_file = decoder_parts.parts[2].content
+
+        print("Transformer Image: ", transformer_image)
+        print("Transformer Tag: ", transformer_tag)
+        print("Zip File: ", zip_file)
+
+        assert response.status_code == 200
+        check_zip_file(zip_file, 2)
+        # Capture the temporary directory that was generated
+        cache_dir = mock_ast_translator.generate_code.call_args[1]['cache_path']
+        mock_ast_translator.generate_code.assert_called_with(select_stmt,
+                                                             cache_path=cache_dir)
+
+    def test_post_good_query_without_params(self, mocker):
+        """Produce code for a simple good query"""
+
+        with TemporaryDirectory() as tempdir, \
+                open(os.path.join(tempdir, "baz.txt"), 'w'),\
+                open(os.path.join(tempdir, "foo.txt"), 'w'):
+
+            mock_ast_translator = mocker.Mock()
+            mock_ast_translator.generate_code = mocker.Mock(
+                return_value=GeneratedFileResult(hash="1234", output_dir=tempdir)
+            )
+
+            config = {
+                'TARGET_BACKEND': 'uproot',
+                'DEFAULT_TRANSFORMER_IMAGE': 'sslhep/servicex_func_adl_xaod_transformer',
+                'DEFAULT_TRANSFORMER_TAG': 'develop'
+            }
+
+            app = create_app(config, provided_translator=mock_ast_translator)
+            client = app.test_client()
+
+            print(app.config)
+
             select_stmt = "(call ResultTTree (call Select (call SelectMany (call EventDataset (list 'localds://did_01')"  # noqa: E501
 
-            response = client.post("/servicex/generated-code", data=select_stmt)
+            response = client.post("/servicex/generated-code", json={
+                "code": select_stmt
+            })
+
+            boundary = str(response.data[2:34], 'utf-8')
+            content_type = f"multipart/form-data; boundary={boundary}"
+            decoder_parts = decoder.MultipartDecoder(response.data, content_type)
+
+            transformer_image = str(decoder_parts.parts[0].content, 'utf-8')
+            transformer_tag = str(decoder_parts.parts[1].content, 'utf-8')
+            zip_file = decoder_parts.parts[2].content
+
+            print("Transformer Image: ", transformer_image)
+            print("Transformer Tag: ", transformer_tag)
+            print("Zip File: ", zip_file)
+
         assert response.status_code == 200
-        check_zip_file(response.data, 2)
+        check_zip_file(zip_file, 2)
         # Capture the temporary directory that was generated
         cache_dir = mock_ast_translator.generate_code.call_args[1]['cache_path']
         mock_ast_translator.generate_code.assert_called_with(select_stmt,