feat: support --async-builds flag for generate_dag.py (#424)

adlersantos · web-flow · commit 7536df9aeb3d · 2022-07-26T12:06:55.000-04:00
diff --git a/scripts/generate_dag.py b/scripts/generate_dag.py
@@ -51,10 +51,11 @@ def main(
     env: str,
     all_pipelines: bool = False,
     skip_builds: bool = False,
+    async_builds: bool = False,
     format_code: bool = True,
 ):
     if not skip_builds:
-        build_images(dataset_id, env)
+        build_images(dataset_id, env, async_builds)
 
     if all_pipelines:
         for pipeline_dir in list_subdirs(DATASETS_PATH / dataset_id / "pipelines"):
@@ -229,7 +230,7 @@ def copy_files_to_dot_dir(dataset_id: str, pipeline_id: str, env_dir: pathlib.Pa
     )
 
 
-def build_images(dataset_id: str, env: str):
+def build_images(dataset_id: str, env: str, async_builds: bool):
     parent_dir = DATASETS_PATH / dataset_id / "pipelines" / "_images"
     if not parent_dir.exists():
         return
@@ -238,7 +239,7 @@ def build_images(dataset_id: str, env: str):
         dataset_id, parent_dir, PROJECT_ROOT / f".{env}"
     )
     for image_dir in image_dirs:
-        build_and_push_image(dataset_id, image_dir)
+        build_and_push_image(dataset_id, image_dir, async_builds)
 
 
 def copy_image_files_to_dot_dir(
@@ -253,24 +254,27 @@ def copy_image_files_to_dot_dir(
     return list_subdirs(target_dir / "_images")
 
 
-def build_and_push_image(dataset_id: str, image_dir: pathlib.Path):
+def build_and_push_image(
+    dataset_id: str, image_dir: pathlib.Path, async_builds: bool = False
+):
     image_name = f"{dataset_id}__{image_dir.name}"
-    tag = f"gcr.io/{gcp_project_id()}/{image_name}"
+    command = [
+        "gcloud",
+        "builds",
+        "submit",
+        "--async",
+        "--tag",
+        f"gcr.io/{gcp_project_id()}/{image_name}",
+    ]
+
+    if not async_builds:
+        command.remove("--async")
 
     # gcloud builds submit --tag gcr.io/PROJECT_ID/IMAGE_NAME
-    subprocess.check_call(
-        [
-            "gcloud",
-            "builds",
-            "submit",
-            "--tag",
-            str(tag),
-        ],
-        cwd=image_dir,
-    )
+    subprocess.check_call(command, cwd=image_dir)
 
 
-def gcp_project_id(project_id: str = None) -> str:
+def gcp_project_id() -> str:
     _, project_id = google.auth.default()
     return project_id
 
@@ -308,6 +312,16 @@ def gcp_project_id(project_id: str = None) -> str:
     parser.add_argument(
         "--skip-builds", required=False, dest="skip_builds", action="store_true"
     )
+    parser.add_argument(
+        "--async-builds", required=False, dest="async_builds", action="store_false"
+    )
 
     args = parser.parse_args()
-    main(args.dataset, args.pipeline, args.env, args.all_pipelines, args.skip_builds)
+    main(
+        args.dataset,
+        args.pipeline,
+        args.env,
+        args.all_pipelines,
+        args.skip_builds,
+        args.async_builds,
+    )
diff --git a/tests/scripts/test_generate_dag.py b/tests/scripts/test_generate_dag.py
@@ -372,13 +372,58 @@ def test_build_images_copies_image_files_to_env_dir(
 
 def test_build_images_called_when_dataset_has_images_dir(
     dataset_path: pathlib.Path, pipeline_path: pathlib.Path, env: str, mocker
+):
+    copy_config_files_and_set_tmp_folder_names_as_ids(dataset_path, pipeline_path)
+    num_containers = random.randint(1, 3)
+    generate_image_files(dataset_path, num_containers=num_containers)
+    assert (
+        sum([f.is_dir() for f in (dataset_path / "pipelines" / "_images").iterdir()])
+        == num_containers
+    )
+
+    mocker.patch("scripts.generate_dag.build_images")
+    generate_dag.main(dataset_path.name, pipeline_path.name, env, format_code=False)
+
+    async_builds_default = False
+    generate_dag.build_images.assert_called_once_with(
+        dataset_path.name, env, async_builds_default
+    )
+
+
+def test_build_images_called_with_async_false_by_default(
+    dataset_path: pathlib.Path, pipeline_path: pathlib.Path, env: str, mocker
 ):
     copy_config_files_and_set_tmp_folder_names_as_ids(dataset_path, pipeline_path)
     generate_image_files(dataset_path, num_containers=random.randint(1, 3))
 
     mocker.patch("scripts.generate_dag.build_images")
     generate_dag.main(dataset_path.name, pipeline_path.name, env, format_code=False)
-    generate_dag.build_images.assert_called_once_with(dataset_path.name, env)
+
+    async_builds_default = False
+    generate_dag.build_images.assert_called_once_with(
+        dataset_path.name, env, async_builds_default
+    )
+
+
+def test_build_images_called_with_async_builds(
+    dataset_path: pathlib.Path, pipeline_path: pathlib.Path, env: str, mocker
+):
+    copy_config_files_and_set_tmp_folder_names_as_ids(dataset_path, pipeline_path)
+    generate_image_files(dataset_path, num_containers=random.randint(1, 3))
+
+    mocker.patch("scripts.generate_dag.build_images")
+    async_builds = True
+    generate_dag.main(
+        dataset_path.name,
+        pipeline_path.name,
+        env,
+        async_builds=async_builds,
+        format_code=False,
+    )
+
+    generate_dag.build_images.assert_called_once_with(
+        dataset_path.name, env, async_builds
+    )
 
 
 def test_build_images_not_called_given_skip_builds_argument(
