refactor(RHTAPBUGS-919): create components in batches (#92)

* refactor(RHTAPBUGS-919): create components in batches

By creating sbom components in Pyxis in batches
of 5, the overall script runtime is reduced by around
47 % in my testing with an sbom with 113 components.

The batch size is configurable, but currently set to 5.

* fix: flake8 and black conflict

Flake8 was disagreeing with a formatting change
black did. It turns out this is a known issue
and it seems the recommended solution is to ignore
this warning in flake8.

See here for more details:
psf/black#315

Signed-off-by: Martin Malina <mmalina@redhat.com>

.flake8

@@ -1,2 +1,3 @@
 [flake8]
 max-line-length = 95
+extend-ignore = E203

pyxis/pyxis.py

@@ -82,7 +82,7 @@ def post(url: str, body: Dict[str, Any]) -> requests.Response:
  return resp
 
 
-def graphql_query(graphql_api: str, body: Dict[str, Any], query: str) -> Dict[str, Any]:
+def graphql_query(graphql_api: str, body: Dict[str, Any]) -> Dict[str, Any]:
  """Make a request to Pyxis GraphQL API
 
  This will make a POST request and then check the result
@@ -91,26 +91,22 @@ def graphql_query(graphql_api: str, body: Dict[str, Any], query: str) -> Dict[st
  Args:
  graphql_api (str): Pyxis GraphQL API URL
  body (Dict[str, Any]): Request payload
- query (str): Name of the Pyxis GraphQL query/mutation used
 
  :return: Pyxis response
  """
  resp = post(graphql_api, body)
  resp_json = resp.json()
 
- error_msg = f"Pyxis GraphQL query '{query}' failed"
- if resp_json.get("data") is None or query not in resp_json["data"]:
+ error_msg = "Pyxis GraphQL query failed"
+ if resp_json.get("data") is None or any(
+ [query["error"] is not None for query in resp_json["data"].values()]
+ ):
  LOGGER.error(error_msg)
  LOGGER.error(f"Pyxis response: {resp_json}")
  LOGGER.error(f"Pyxis trace_id: {resp.headers.get('trace_id')}")
  raise RuntimeError(error_msg)
- elif resp_json["data"][query]["error"] is not None:
- error_msg = f"{error_msg}: {resp_json['data'][query]['error']['detail']}"
- LOGGER.error(error_msg)
- LOGGER.error(f"Pyxis trace_id: {resp.headers.get('trace_id')}")
- raise RuntimeError(error_msg)
 
- return resp_json["data"][query]["data"]
+ return resp_json["data"]
 
 
 def put(url: str, body: Dict[str, Any]) -> Dict[str, Any]:

pyxis/test_pyxis.py

@@ -86,9 +86,9 @@ def test_graphql_query__success(mock_post: MagicMock):
  }
  }
 
- data = pyxis.graphql_query(API_URL, REQUEST_BODY, QUERY)
+ data = pyxis.graphql_query(API_URL, REQUEST_BODY)
 
- assert data == mock_data
+ assert data[QUERY]["data"] == mock_data
  mock_post.assert_called_once_with(API_URL, REQUEST_BODY)
 
 
@@ -102,7 +102,7 @@ def test_graphql_query__general_graphql_error(mock_post: MagicMock):
  }
 
  with pytest.raises(RuntimeError):
- pyxis.graphql_query(API_URL, REQUEST_BODY, QUERY)
+ pyxis.graphql_query(API_URL, REQUEST_BODY)
 
  mock_post.assert_called_once_with(API_URL, REQUEST_BODY)
 
@@ -121,7 +121,7 @@ def test_graphql_query__pyxis_error(mock_post: MagicMock):
  }
 
  with pytest.raises(RuntimeError):
- pyxis.graphql_query(API_URL, REQUEST_BODY, QUERY)
+ pyxis.graphql_query(API_URL, REQUEST_BODY)
 
  mock_post.assert_called_once_with(API_URL, REQUEST_BODY)
 

pyxis/test_upload_sbom.py

@@ -7,7 +7,8 @@
  get_image,
  create_content_manifest,
  get_existing_bom_refs,
- create_content_manifest_component,
+ create_content_manifest_components,
+ get_template,
  load_sbom_components,
  check_bom_ref_duplicates,
  convert_keys,
@@ -57,7 +58,7 @@ def test_upload_sbom_with_retry__fails(mock_upload_sbom):
  assert mock_upload_sbom.call_count == 2
 
 
-@patch("upload_sbom.create_content_manifest_component")
+@patch("upload_sbom.create_content_manifest_components")
 @patch("upload_sbom.get_existing_bom_refs")
 @patch("upload_sbom.load_sbom_components")
 @patch("upload_sbom.create_content_manifest")
@@ -67,7 +68,7 @@ def test_upload_sbom__success(
  mock_create_content_manifest,
  mock_load_sbom_components,
  mock_get_existing_bom_refs,
- mock_create_content_manifest_component,
+ mock_create_content_manifest_components,
 ):
  """
  Basic use case - nothing exists in Pyxis yet and all components are successfully created
@@ -86,22 +87,26 @@ def test_upload_sbom__success(
 
  mock_create_content_manifest.assert_called_once_with(GRAPHQL_API, IMAGE_ID)
  mock_get_existing_bom_refs.assert_not_called()
- assert mock_create_content_manifest_component.call_args_list == [
- call(GRAPHQL_API, MANIFEST_ID, {"bom_ref": "aaa"}),
- call(GRAPHQL_API, MANIFEST_ID, {"bom_ref": "bbb"}),
- call(GRAPHQL_API, MANIFEST_ID, {"type": "library"}),
- ]
+ mock_create_content_manifest_components.assert_called_once_with(
+ GRAPHQL_API,
+ MANIFEST_ID,
+ [
+ {"bom_ref": "aaa"},
+ {"bom_ref": "bbb"},
+ {"type": "library"},
+ ],
+ )
 
 
-@patch("upload_sbom.create_content_manifest_component")
+@patch("upload_sbom.create_content_manifest_components")
 @patch("upload_sbom.load_sbom_components")
 @patch("upload_sbom.create_content_manifest")
 @patch("upload_sbom.get_image")
 def test_upload_sbom__manifest_and_one_component_exist(
  mock_get_image,
  mock_create_content_manifest,
  mock_load_sbom_components,
- mock_create_content_manifest_component,
+ mock_create_content_manifest_components,
 ):
  """Creation of the manifest and the first component is skipped"""
  mock_get_image.return_value = {
@@ -119,20 +124,20 @@ def test_upload_sbom__manifest_and_one_component_exist(
  upload_sbom(GRAPHQL_API, IMAGE_ID, SBOM_PATH)
 
  mock_create_content_manifest.assert_not_called()
- mock_create_content_manifest_component.assert_called_once_with(
- GRAPHQL_API, MANIFEST_ID, {"bom_ref": "bbb"}
+ mock_create_content_manifest_components.assert_called_once_with(
+ GRAPHQL_API, MANIFEST_ID, [{"bom_ref": "bbb"}]
  )
 
 
-@patch("upload_sbom.create_content_manifest_component")
+@patch("upload_sbom.create_content_manifest_components")
 @patch("upload_sbom.load_sbom_components")
 @patch("upload_sbom.create_content_manifest")
 @patch("upload_sbom.get_image")
 def test_upload_sbom__all_components_exist(
  mock_get_image,
  mock_create_content_manifest,
  mock_load_sbom_components,
- mock_create_content_manifest_component,
+ mock_create_content_manifest_components,
 ):
  """Creation of manifest and all components is skipped"""
  mock_get_image.return_value = {
@@ -150,7 +155,7 @@ def test_upload_sbom__all_components_exist(
  upload_sbom(GRAPHQL_API, IMAGE_ID, SBOM_PATH)
 
  mock_create_content_manifest.assert_not_called()
- mock_create_content_manifest_component.assert_not_called()
+ mock_create_content_manifest_components.assert_not_called()
 
 
 def generate_pyxis_response(query_name, data=None, error=False):
@@ -170,15 +175,15 @@ def generate_pyxis_response(query_name, data=None, error=False):
  return response
 
 
-@patch("upload_sbom.create_content_manifest_component")
+@patch("upload_sbom.create_content_manifest_components")
 @patch("upload_sbom.load_sbom_components")
 @patch("upload_sbom.create_content_manifest")
 @patch("upload_sbom.get_image")
 def test_upload_sbom__existing_bom_ref_is_skipped(
  mock_get_image,
  mock_create_content_manifest,
  mock_load_sbom_components,
- mock_create_content_manifest_component,
+ mock_create_content_manifest_components,
 ):
  """One component already exists in Pyxis. Our sbom contains two
  components. So we want to upload only the second one to Pyxis,
@@ -200,7 +205,7 @@ def test_upload_sbom__existing_bom_ref_is_skipped(
  upload_sbom(GRAPHQL_API, IMAGE_ID, SBOM_PATH)
 
  mock_create_content_manifest.assert_not_called()
- mock_create_content_manifest_component.assert_not_called()
+ mock_create_content_manifest_components.assert_called_with(GRAPHQL_API, MANIFEST_ID, [])
 
 
 @patch("pyxis.post")
@@ -316,28 +321,83 @@ def test_get_existing_bom_refs__no_components_result_in_empty_set():
  assert bom_refs == set()
 
 
-@patch("pyxis.post")
-def test_create_content_manifest_component__success(mock_post):
- mock_post.return_value = generate_pyxis_response(
- "create_content_manifest_component_for_manifest", {"_id": COMPONENT_ID}
+@patch("pyxis.graphql_query")
+@patch("upload_sbom.get_template")
+def test_create_content_manifest_components__success(mock_get_template, mock_graphql_query):
+ create_content_manifest_components(
+ GRAPHQL_API, MANIFEST_ID, [COMPONENT_DICT], batch_size=2
  )
 
- id = create_content_manifest_component(GRAPHQL_API, MANIFEST_ID, COMPONENT_DICT)
+ mock_get_template.assert_called_once_with()
+ mock_get_template.return_value.render.assert_called_once_with(components=[COMPONENT_DICT])
+ mock_graphql_query.assert_called_once()
 
- assert id == COMPONENT_ID
- mock_post.assert_called_once()
 
+@patch("pyxis.graphql_query")
+@patch("upload_sbom.get_template")
+def test_create_content_manifest_component__multiple_batches(
+ mock_get_template, mock_graphql_query
+):
+ comp1 = {"bom_ref": "aaa"}
+ comp2 = {"bom_ref": "bbb"}
+ comp3 = {"type": "library"}
+ components = [comp1, comp2, comp3]
 
-@patch("pyxis.post")
-def test_create_content_manifest_component__error(mock_post):
- mock_post.return_value = generate_pyxis_response(
- "create_content_manifest_component_for_manifest", error=True
- )
+ create_content_manifest_components(GRAPHQL_API, MANIFEST_ID, components, batch_size=2)
 
- with pytest.raises(RuntimeError):
- create_content_manifest_component(GRAPHQL_API, MANIFEST_ID, COMPONENT_DICT)
+ mock_get_template.assert_called_once_with()
+ assert mock_get_template.return_value.render.call_args_list == [
+ call(components=[comp1, comp2]),
+ call(components=[comp3]),
+ ]
+ assert mock_graphql_query.call_count == 2
+ assert mock_graphql_query.call_args_list == [
+ call(
+ GRAPHQL_API,
+ {
+ "query": mock_get_template.return_value.render.return_value,
+ "variables": {
+ "id": MANIFEST_ID,
+ "input0": comp1,
+ "input1": comp2,
+ },
+ },
+ ),
+ call(
+ GRAPHQL_API,
+ {
+ "query": mock_get_template.return_value.render.return_value,
+ "variables": {
+ "id": MANIFEST_ID,
+ "input0": comp3,
+ },
+ },
+ ),
+ ]
 
- mock_post.assert_called_once()
+
+@patch("pyxis.graphql_query")
+@patch("upload_sbom.get_template")
+def test_create_content_manifest_component__no_components(
+ mock_get_template, mock_graphql_query
+):
+ create_content_manifest_components(GRAPHQL_API, MANIFEST_ID, [])
+
+ mock_get_template.assert_not_called()
+ mock_graphql_query.assert_not_called()
+
+
+@patch("upload_sbom.Template")
+@patch("builtins.open")
+@patch("upload_sbom.os")
+def test_get_template(mock_os, mock_open, mock_template):
+ template_path = mock_os.path.join.return_value
+
+ template = get_template()
+
+ mock_open.assert_called_with(template_path)
+ assert template == mock_template.return_value
+ mock_open.return_value.__enter__.return_value.read.assert_called_once_with()
 
 
 @patch("json.load")