From f7167fa3b6dea8d09d26105bfb1b62bd36e42f0d Mon Sep 17 00:00:00 2001
From: Haibo Lin <hlin@redhat.com>
Date: Jul 17 2020 01:08:36 +0000
Subject: Allow including scratch module builds


JIRA: RHELCMP-439
Signed-off-by: Haibo Lin <hlin@redhat.com>

---

diff --git a/pungi/checks.py b/pungi/checks.py
index e98db28..4adf2f6 100644
--- a/pungi/checks.py
+++ b/pungi/checks.py
@@ -786,6 +786,14 @@ def make_schema():
                 "type": "boolean",
                 "default": True,
             },
+            "pkgset_scratch_modules": {
+                "type": "object",
+                "patternProperties": {
+                    "^.+$": {"$ref": "#/definitions/list_of_strings"}
+                },
+                "additionalProperties": False,
+            },
+            "mbs_api_url": {"type": "string"},
             "disc_types": {"type": "object", "default": {}},
             "paths_module": {"type": "string"},
             "skip_phases": {
@@ -1258,6 +1266,7 @@ CONFIG_DEPS = {
         "conflicts": ((lambda x: not x, ["base_product_name", "base_product_short"]),),
     },
     "product_id": {"conflicts": [(lambda x: not x, ["product_id_allow_missing"])]},
+    "pkgset_scratch_modules": {"requires": ((lambda x: x, ["mbs_api_url"]),)},
     "pkgset_source": {
         "requires": [(lambda x: x == "repos", ["pkgset_repos"])],
         "conflicts": [
diff --git a/pungi/paths.py b/pungi/paths.py
index 1531ad0..43b084c 100644
--- a/pungi/paths.py
+++ b/pungi/paths.py
@@ -490,6 +490,8 @@ class WorkPaths(object):
 
     def module_defaults_dir(self, create_dir=True):
         """
+        Example:
+            work/global/module_defaults
         """
         path = os.path.join(self.topdir(create_dir=create_dir), "module_defaults")
         if create_dir:
diff --git a/pungi/phases/pkgset/sources/source_koji.py b/pungi/phases/pkgset/sources/source_koji.py
index f5709c6..309ed65 100644
--- a/pungi/phases/pkgset/sources/source_koji.py
+++ b/pungi/phases/pkgset/sources/source_koji.py
@@ -26,9 +26,10 @@ from kobo.shortcuts import force_list
 
 import pungi.wrappers.kojiwrapper
 from pungi.wrappers.comps import CompsWrapper
+from pungi.wrappers.mbs import MBSWrapper
 import pungi.phases.pkgset.pkgsets
 from pungi.arch import getBaseArch
-from pungi.util import retry, get_arch_variant_data
+from pungi.util import retry, get_arch_variant_data, get_variant_data
 from pungi.module_util import Modulemd
 
 from pungi.phases.pkgset.common import MaterializedPackageSet, get_all_arches
@@ -272,6 +273,51 @@ def _add_module_to_variant(
     return nsvc
 
 
+def _add_scratch_modules_to_variant(
+    compose, variant, scratch_modules, variant_tags, tag_to_mmd
+):
+    if compose.compose_type != "test" and scratch_modules:
+        compose.log_warning("Only test composes could include scratch module builds")
+        return
+
+    mbs = MBSWrapper(compose.conf["mbs_api_url"])
+    for nsvc in scratch_modules:
+        module_build = mbs.get_module_build_by_nsvc(nsvc)
+        if not module_build:
+            continue
+        try:
+            final_modulemd = mbs.final_modulemd(module_build["id"])
+        except Exception:
+            compose.log_error("Unable to get modulemd for build %s" % module_build)
+            raise
+        tag = module_build["koji_tag"]
+        variant_tags[variant].append(tag)
+        tag_to_mmd.setdefault(tag, {})
+        for arch in variant.arches:
+            try:
+                mmd = Modulemd.ModuleStream.read_string(
+                    final_modulemd[arch], strict=True
+                )
+                variant.arch_mmds.setdefault(arch, {})[nsvc] = mmd
+            except KeyError:
+                continue
+            tag_to_mmd[tag].setdefault(arch, set()).add(mmd)
+
+            if tag_to_mmd[tag]:
+                compose.log_info(
+                    "Module '%s' in variant '%s' will use Koji tag '%s' "
+                    "(as a result of querying module '%s')",
+                    nsvc,
+                    variant,
+                    tag,
+                    module_build["name"],
+                )
+
+                # Store mapping NSVC --> koji_tag into variant. This is needed
+                # in createrepo phase where metadata is exposed by productmd
+                variant.module_uid_to_koji_tag[nsvc] = tag
+
+
 def _is_filtered_out(compose, variant, arch, module_name, module_stream):
     """Check if module with given name and stream is filter out from this stream.
     """
@@ -618,6 +664,14 @@ def populate_global_pkgset(compose, koji_wrapper, path_prefix, event):
                 compose, koji_wrapper, event, variant, variant_tags, tag_to_mmd
             )
 
+        variant_scratch_modules = get_variant_data(
+            compose.conf, "pkgset_scratch_modules", variant
+        )
+        if variant_scratch_modules:
+            _add_scratch_modules_to_variant(
+                compose, variant, variant_scratch_modules, variant_tags, tag_to_mmd
+            )
+
         # Ensure that every tag added to `variant_tags` is added also to
         # `compose_tags`.
         for variant_tag in variant_tags[variant]:
diff --git a/pungi/wrappers/mbs.py b/pungi/wrappers/mbs.py
new file mode 100644
index 0000000..730ea0e
--- /dev/null
+++ b/pungi/wrappers/mbs.py
@@ -0,0 +1,43 @@
+import os
+import requests
+
+
+class MBSWrapper(object):
+    def __init__(self, api_url):
+        """
+        :param string api_url: e.g. https://example.com/module-build-service/2
+        """
+        self.api_url = api_url
+
+    def _get(self, resource, params=None):
+        """Get specified resource.
+
+        :param string resource: e.g. module-builds, final-modulemd
+        :param dict data:
+        """
+        url = os.path.join(self.api_url, resource)
+        try:
+            resp = requests.get(url, params=params)
+        except Exception as e:
+            raise Exception(
+                "Failed to query URL %s with params %s - %s" % (url, params, str(e))
+            )
+        resp.raise_for_status()
+        return resp
+
+    def module_builds(self, filters=None):
+        return self._get("module-builds", filters).json()
+
+    def get_module_build_by_nsvc(self, nsvc):
+        nsvc_list = nsvc.split(":")
+        if len(nsvc_list) != 4:
+            raise ValueError("Invalid N:S:V:C - %s" % nsvc)
+        filters = dict(zip(["name", "stream", "version", "context"], nsvc_list))
+        resp = self.module_builds(filters)
+        if resp["items"]:
+            return resp["items"][0]
+        else:
+            return None
+
+    def final_modulemd(self, module_build_id):
+        return self._get("final-modulemd/%s" % module_build_id).json()
