Kleaf: init_ddk: Support for remote prebuilts

* This works by inferring the list of files to download using the
  information from downloads_config.json file distrubuted with the
  builds.

* Download method is expanded to support non-mandotory files, this is to
  not cause errors when files are not found in the build distribution.

Tested: python3 init.py --prebuilts_dir=/abs/path/to/gki_prebuilts \
          --build_id=<placeholder> --kleaf_repo=/abs/path/to/ACK \
	  --url_fmt="http://0.0.0.0:8000/{filename}"
Bug: 328770706
Change-Id: Ib198eadd8211e74d0f2f83b53f25c53b9ee5d209
Signed-off-by: Ulises Mendez Martinez <umendez@google.com>

2 files changed, 87 insertions, 19 deletions

diff --git a/init/init_ddk.py b/init/init_ddk.py
index 3a51dba..2798924 100644
--- a/init/init_ddk.py
+++ b/init/init_ddk.py
@@ -17,6 +17,7 @@
 """Configures the project layout to build DDK modules."""
 
 import argparse
+import concurrent.futures
 import dataclasses
 import json
 import logging
@@ -74,6 +75,7 @@ class KleafProjectSetter:
     url_fmt: str | None
 
     def _symlink_tools_bazel(self):
+        """Creates the symlink tools/bazel."""
         if not self.ddk_workspace or not self.kleaf_repo:
             return
         # Calculate the paths.
@@ -133,6 +135,7 @@ class KleafProjectSetter:
             return path
 
     def _read_download_configs(self) -> str:
+        """Reads the previously downloaded download_configs.json file."""
         download_configs = self.prebuilts_dir / "download_configs.json"
         with open(download_configs, "r", encoding="utf-8") as config:
             # Compress the representation by removing empty spaces to save some space.
@@ -164,6 +167,7 @@ class KleafProjectSetter:
             logging.info("Nothing to update in %s", module_bazel)
 
     def _generate_bazelrc(self):
+        """Creates a Bazel configuration file with the minimum setup required."""
         if not self.ddk_workspace or not self.kleaf_repo:
             return
         bazelrc = self.ddk_workspace / _DEVICE_BAZELRC
@@ -196,55 +200,92 @@ class KleafProjectSetter:
             return None
         return url
 
-    def _download(self, remote_filename: str, out_file_name: str):
+    def _can_download_artifacts(self):
+        """Validates that download are possible within the current context."""
         if not self.url_fmt:
-            logging.error(
-                "Unable to download file %s because --url_fmt was not set.",
-                remote_filename,
-            )
-            return
+            return False
+        if self._get_url("") is None:
+            return False
+        return True
+
+    def _download(
+        self,
+        remote_filename: str,
+        out_file_name: pathlib.Path,
+        mandatory: bool = True,
+    ) -> None:
+        """Given the url_fmt, build_id and build_target downloads a remote file.
+
+        Args:
+            remote_filename: File name for the file, it can contain anchors for,
+              build_id, target and filename.
+            out_file_name: Destination place for the download.
+            mandatory: When set to true, the download fails when the file could
+              not be downloaded.
+        """
         url = self._get_url(remote_filename)
-        if not url:
-            logging.error(
-                "Unable to download %s file because --build_id is missing.",
-                remote_filename,
-            )
-            return
         # Workaround: Rely on host keychain to download files.
         # This is needed otheriwese downloads fail when running this script
         #   using the hermetic Python toolchain.
-        subprocess.check_call(
+        subprocess.run(
             [
                 "python3",
                 pathlib.Path(__file__).parent / "init_download.py",
                 url,
                 out_file_name,
             ],
+            stderr=subprocess.STDOUT if mandatory else subprocess.DEVNULL,
+            check=mandatory,
         )
 
-    def _handle_ddk_workspace(self):
+    def _infer_download_list(self) -> dict[str, dict]:
+        """Infers the list of files to be downloaded using download_configs.json."""
+        download_configs = self.prebuilts_dir / "download_configs.json"
+        with open(download_configs, "w+", encoding="utf-8") as config:
+            self._download("download_configs.json", config.name)
+            return json.load(config)
+
+    def _download_prebuilts(self) -> None:
+        """Downloads prebuilts from a given build_id when provided."""
+        logging.info("Downloading prebuilts into %s", self.prebuilts_dir)
+        files_dict = self._infer_download_list()
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            futures = []
+            for file, config in files_dict.items():
+                dst = self.prebuilts_dir / file
+                dst.parent.mkdir(parents=True, exist_ok=True)
+                futures.append(
+                    executor.submit(
+                        self._download, file, dst, config["mandatory"]
+                    )
+                )
+            for complete_ret in concurrent.futures.as_completed(futures):
+                complete_ret.result()  # Raise exception if any
+
+    def _handle_ddk_workspace(self) -> None:
         if not self.ddk_workspace:
             return
         self.ddk_workspace.mkdir(parents=True, exist_ok=True)
 
-    def _handle_kleaf_repo(self):
+    def _handle_kleaf_repo(self) -> None:
         if not self.kleaf_repo:
             return
         self.kleaf_repo.mkdir(parents=True, exist_ok=True)
         # TODO: b/328770706 - According to the needs, syncing git repos logic should go here.
 
-    def _handle_prebuilts(self):
+    def _handle_prebuilts(self) -> None:
         if not self.ddk_workspace or not self.prebuilts_dir:
             return
         self.prebuilts_dir.mkdir(parents=True, exist_ok=True)
-        # TODO: b/328770706 - When build_id is given dowloand artifacts here.
+        if self._can_download_artifacts():
+            self._download_prebuilts()
 
-    def _run(self):
+    def _run(self) -> None:
         self._symlink_tools_bazel()
         self._generate_module_bazel()
         self._generate_bazelrc()
 
-    def run(self):
+    def run(self) -> None:
         self._handle_ddk_workspace()
         self._handle_kleaf_repo()
         self._handle_prebuilts()
diff --git a/init/init_ddk_test.py b/init/init_ddk_test.py
index fe54e96..db6773f 100644
--- a/init/init_ddk_test.py
+++ b/init/init_ddk_test.py
@@ -14,6 +14,7 @@
 
 """Tests for init_ddk.py"""
 
+import json
 import logging
 import pathlib
 import tempfile
@@ -218,6 +219,32 @@ class KleafProjectSetterTest(parameterized.TestCase):
             self.assertTrue(out_file.exists())
             self.assertEqual(out_file.read_text(), "Hello World!")
 
+    def test_non_mandatory_doesnt_fail(self):
+        """Tests that optional files don't produce errors."""
+        with tempfile.TemporaryDirectory() as tmp:
+            ddk_workspace = pathlib.Path(tmp) / "ddk_workspace"
+            prebuilts_dir = ddk_workspace / "prebuilts_dir"
+            download_configs = ddk_workspace / "download_configs.json"
+            download_configs.parent.mkdir(parents=True, exist_ok=True)
+            download_configs.write_text(json.dumps({
+                "non-existent-file": {
+                    "target_suffix": "non-existent-file",
+                    "mandatory": False,
+                    "remote_filename_fmt": "non-existent-file",
+                }
+            }))
+            with open(download_configs, "r", encoding="utf-8"):
+                url_fmt = f"file://{str(download_configs.parent)}/{{filename}}"
+                init_ddk.KleafProjectSetter(
+                    build_id="12345",
+                    build_target=None,
+                    ddk_workspace=ddk_workspace,
+                    kleaf_repo=None,
+                    local=None,
+                    prebuilts_dir=prebuilts_dir,
+                    url_fmt=url_fmt,
+                ).run()
+
 
 # This could be run as: tools/bazel test //build/kernel:init_ddk_test --test_output=all
 if __name__ == "__main__":