chore: polish post processing step used for librarian client library generation (googleapis#14206)

parthea · web-flow · commit ebede14f5504 · 2025-08-12T19:53:47.000-04:00
This PR adds 2 new functions `_copy_files_needed_for_post_processing`
and `_clean_up_files_after_post_processing`.

`_copy_files_needed_for_post_processing` - This function is used to copy
files from the `input` directory . Specifically, `.repo-metadata.json`
and `scripts/client-post-processing` are needed as part of client
library generation.

`_clean_up_files_after_post_processing` - This function removes files
which should not be part of the generated code because it would create
an unnecessary diff. The files omitted are:

- `CHANGELOG.md`
- `docs/CHANGELOG.md`
- `docs/README.rst`
- `scripts/client-post-processing/*.yaml`
- `**/gapic_version.py`
- `samples/generated_samples/snippet_metadata*.json`
diff --git a/.generator/cli.py b/.generator/cli.py
@@ -13,22 +13,23 @@
 # limitations under the License.
 
 import argparse
+import glob
 import json
 import logging
 import os
 import re
+import shutil
 import subprocess
 import sys
-import subprocess
 from typing import Dict, List
 
 try:
     import synthtool
-    from synthtool import gcp
+    from synthtool.languages import python_mono_repo
 
     SYNTHTOOL_INSTALLED = True
     SYNTHTOOL_IMPORT_ERROR = None
-except ImportError as e:
+except ImportError as e: # pragma: NO COVER
     SYNTHTOOL_IMPORT_ERROR = e
     SYNTHTOOL_INSTALLED = False
 
@@ -90,7 +91,7 @@ def _determine_bazel_rule(api_path: str, source: str) -> str:
 
         # This check is for a logical failure (no match), not a runtime exception.
         # It's good to keep it for clear error messaging.
-        if not match:
+        if not match: # pragma: NO COVER
             raise ValueError(
                 f"No Bazel rule with a name ending in '-py' found in {build_file_path}"
             )
@@ -136,7 +137,19 @@ def _build_bazel_target(bazel_rule: str, source: str):
     """
     logger.info(f"Executing build for rule: {bazel_rule}")
     try:
-        command = ["bazelisk", "build", bazel_rule]
+        # We're using the prewarmed bazel cache from the docker image to speed up the bazelisk commands.
+        # Previously built artifacts are stored in `/bazel_cache/_bazel_ubuntu/output_base` and will be
+        # used to speed up the build. `disk_cache` is used as the 'remote cache' and is also prewarmed as part of
+        # the docker image.
+        # See https://bazel.build/remote/caching#disk-cache which explains using a file system as a 'remote cache'.
+        command = [
+            "bazelisk",
+            "--output_base=/bazel_cache/_bazel_ubuntu/output_base",
+            "build",
+            "--disk_cache=/bazel_cache/_bazel_ubuntu/cache/repos",
+            "--incompatible_strict_action_env",
+            bazel_rule,
+        ]
         subprocess.run(
             command,
             cwd=source,
@@ -171,7 +184,14 @@ def _locate_and_extract_artifact(
     try:
         # 1. Find the bazel-bin output directory.
         logger.info("Locating Bazel output directory...")
-        info_command = ["bazelisk", "info", "bazel-bin"]
+        # Previously built artifacts are stored in `/bazel_cache/_bazel_ubuntu/output_base`.
+        # See `--output_base` in `_build_bazel_target`
+        info_command = [
+            "bazelisk",
+            "--output_base=/bazel_cache/_bazel_ubuntu/output_base",
+            "info",
+            "bazel-bin",
+        ]
         result = subprocess.run(
             info_command,
             cwd=source,
@@ -206,17 +226,87 @@ def _locate_and_extract_artifact(
         ) from e
 
 
-def _run_post_processor():
-    """Runs the synthtool post-processor on the output directory."""
+def _run_post_processor(output: str, library_id: str):
+    """Runs the synthtool post-processor on the output directory.
+
+    Args:
+        output(str): Path to the directory in the container where code
+            should be generated.
+        library_id(str): The library id to be used for post processing.
+
+    """
+    os.chdir(output)
+    path_to_library = f"packages/{library_id}"
     logger.info("Running Python post-processor...")
     if SYNTHTOOL_INSTALLED:
-        command = ["python3", "-m", "synthtool.languages.python_mono_repo"]
-        subprocess.run(command, cwd=OUTPUT_DIR, text=True, check=True)
+        python_mono_repo.owlbot_main(path_to_library)
     else:
-        raise SYNTHTOOL_IMPORT_ERROR
+        raise SYNTHTOOL_IMPORT_ERROR # pragma: NO COVER
     logger.info("Python post-processor ran successfully.")
 
 
+def _copy_files_needed_for_post_processing(output: str, input: str, library_id: str):
+    """Copy files to the output directory whcih are needed during the post processing
+    step, such as .repo-metadata.json and script/client-post-processing, using
+    the input directory as the source.
+
+    Args:
+        output(str): Path to the directory in the container where code
+            should be generated.
+        input(str): The path to the directory in the container
+            which contains additional generator input.
+        library_id(str): The library id to be used for post processing.
+    """
+
+    path_to_library = f"packages/{library_id}"
+
+    # We need to create these directories so that we can copy files necessary for post-processing.
+    os.makedirs(f"{output}/{path_to_library}")
+    os.makedirs(f"{output}/{path_to_library}/scripts/client-post-processing")
+    shutil.copy(
+        f"{input}/{path_to_library}/.repo-metadata.json",
+        f"{output}/{path_to_library}/.repo-metadata.json",
+    )
+
+    # copy post-procesing files
+    for post_processing_file in glob.glob(f"{input}/client-post-processing/*.yaml"): # pragma: NO COVER
+        with open(post_processing_file, "r") as post_processing:
+            if f"{path_to_library}/" in post_processing.read():
+                shutil.copy(
+                    post_processing_file,
+                    f"{output}/{path_to_library}/scripts/client-post-processing",
+                )
+
+
+def _clean_up_files_after_post_processing(output: str, library_id: str):
+    """
+    Clean up files which should not be included in the generated client
+
+    Args:
+        output(str): Path to the directory in the container where code
+            should be generated.
+        library_id(str): The library id to be used for post processing.
+    """
+    path_to_library = f"packages/{library_id}"
+    shutil.rmtree(f"{output}/{path_to_library}/.nox")
+    os.remove(f"{output}/{path_to_library}/CHANGELOG.md")
+    os.remove(f"{output}/{path_to_library}/docs/CHANGELOG.md")
+    os.remove(f"{output}/{path_to_library}/docs/README.rst")
+    for post_processing_file in glob.glob(
+        f"{output}/{path_to_library}/scripts/client-post-processing/*.yaml"
+    ): # pragma: NO COVER
+        os.remove(post_processing_file)
+    for gapic_version_file in glob.glob(
+        f"{output}/{path_to_library}/**/gapic_version.py", recursive=True
+    ): # pragma: NO COVER
+        os.remove(gapic_version_file)
+    for snippet_metadata_file in glob.glob(
+        f"{output}/{path_to_library}/samples/generated_samples/snippet_metadata*.json"
+    ): # pragma: NO COVER
+        os.remove(snippet_metadata_file)
+    shutil.rmtree(f"{output}/owl-bot-staging")
+
+
 def handle_generate(
     librarian: str = LIBRARIAN_DIR,
     source: str = SOURCE_DIR,
@@ -238,7 +328,7 @@ def handle_generate(
             API protos.
         output(str): Path to the directory in the container where code
             should be generated.
-        input(str): The path path to the directory in the container
+        input(str): The path to the directory in the container
             which contains additional generator input.
 
     Raises:
@@ -249,7 +339,6 @@ def handle_generate(
         # Read a generate-request.json file
         request_data = _read_json_file(f"{librarian}/{GENERATE_REQUEST_FILE}")
         library_id = _get_library_id(request_data)
-
         for api in request_data.get("apis", []):
             api_path = api.get("path")
             if api_path:
@@ -258,7 +347,15 @@ def handle_generate(
                 _locate_and_extract_artifact(
                     bazel_rule, library_id, source, output, api_path
                 )
-                _run_post_processor(output, f"packages/{library_id}")
+
+        _copy_files_needed_for_post_processing(output, input, library_id)
+        _run_post_processor(output, library_id)
+        _clean_up_files_after_post_processing(output, library_id)
+
+        # Write the `generate-response.json` using `generate-request.json` as the source
+        with open(f"{librarian}/generate-response.json", "w") as f:
+            json.dump(request_data, f, indent=4)
+            f.write("\n")
 
     except Exception as e:
         raise ValueError("Generation failed.") from e
@@ -293,6 +390,7 @@ def _run_individual_session(nox_session: str, library_id: str):
         nox_session(str): The nox session to run
         library_id(str): The library id under test
     """
+
     command = [
         "nox",
         "-s",
@@ -324,7 +422,7 @@ def handle_build(librarian: str = LIBRARIAN_DIR):
     logger.info("'build' command executed.")
 
 
-if __name__ == "__main__":
+if __name__ == "__main__": # pragma: NO COVER
     parser = argparse.ArgumentParser(description="A simple CLI tool.")
     subparsers = parser.add_subparsers(
         dest="command", required=True, help="Available commands"
@@ -373,7 +471,6 @@ def handle_build(librarian: str = LIBRARIAN_DIR):
         sys.exit(1)
 
     args = parser.parse_args()
-    args.func()
 
     # Pass specific arguments to the handler functions for generate/build
     if args.command == "generate":
diff --git a/.generator/requirements-test.in b/.generator/requirements-test.in
@@ -13,4 +13,6 @@
 # limitations under the License.
 
 pytest
+pytest-cov
 pytest-mock
+gcp-synthtool @ git+https://github.com/googleapis/synthtool@5aa438a342707842d11fbbb302c6277fbf9e4655
diff --git a/.generator/test_cli.py b/.generator/test_cli.py
@@ -25,6 +25,8 @@
     LIBRARIAN_DIR,
     REPO_DIR,
     _build_bazel_target,
+    _clean_up_files_after_post_processing,
+    _copy_files_needed_for_post_processing,
     _determine_bazel_rule,
     _get_library_id,
     _locate_and_extract_artifact,
@@ -223,6 +225,9 @@ def test_locate_and_extract_artifact_fails(mocker, caplog):
         _locate_and_extract_artifact(
             "//google/cloud/language/v1:rule-py",
             "google-cloud-language",
+            "source",
+            "output",
+            "google/cloud/language/v1",
         )
 
 
@@ -232,27 +237,18 @@ def test_run_post_processor_success(mocker, caplog):
     """
     caplog.set_level(logging.INFO)
     mocker.patch("cli.SYNTHTOOL_INSTALLED", return_value=True)
-    mock_subprocess = mocker.patch("cli.subprocess.run")
-
-    _run_post_processor()
+    mock_chdir = mocker.patch("cli.os.chdir")
+    mock_owlbot_main = mocker.patch(
+        "cli.synthtool.languages.python_mono_repo.owlbot_main"
+    )
+    _run_post_processor("output", "google-cloud-language")
 
-    mock_subprocess.assert_called_once()
+    mock_chdir.assert_called_once()
 
-    assert mock_subprocess.call_args.kwargs["cwd"] == "output"
+    mock_owlbot_main.assert_called_once_with("packages/google-cloud-language")
     assert "Python post-processor ran successfully." in caplog.text
 
 
-def test_locate_and_extract_artifact_fails(mocker, caplog):
-    """
-    Tests that an exception is raised if the subprocess command fails.
-    """
-    caplog.set_level(logging.INFO)
-    mocker.patch("cli.SYNTHTOOL_INSTALLED", return_value=True)
-
-    with pytest.raises(FileNotFoundError):
-        _run_post_processor()
-
-
 def test_handle_generate_success(caplog, mock_generate_request_file, mocker):
     """
     Tests the successful execution path of handle_generate.
@@ -265,10 +261,23 @@ def test_handle_generate_success(caplog, mock_generate_request_file, mocker):
     mock_build_target = mocker.patch("cli._build_bazel_target")
     mock_locate_and_extract_artifact = mocker.patch("cli._locate_and_extract_artifact")
     mock_run_post_processor = mocker.patch("cli._run_post_processor")
+    mock_copy_files_needed_for_post_processing = mocker.patch(
+        "cli._copy_files_needed_for_post_processing"
+    )
+    mock_clean_up_files_after_post_processing = mocker.patch(
+        "cli._clean_up_files_after_post_processing"
+    )
 
     handle_generate()
 
     mock_determine_rule.assert_called_once_with("google/cloud/language/v1", "source")
+    mock_run_post_processor.assert_called_once_with("output", "google-cloud-language")
+    mock_copy_files_needed_for_post_processing.assert_called_once_with(
+        "output", "input", "google-cloud-language"
+    )
+    mock_clean_up_files_after_post_processing.assert_called_once_with(
+        "output", "google-cloud-language"
+    )
 
 
 def test_handle_generate_fail(caplog):
@@ -406,3 +415,17 @@ def test_invalid_json(mocker):
 
     with pytest.raises(json.JSONDecodeError):
         _read_json_file("fake/path.json")
+
+def test_copy_files_needed_for_post_processing_success(mocker):
+    mock_makedirs = mocker.patch("os.makedirs")
+    mock_shutil_copy = mocker.patch("shutil.copy")
+    _copy_files_needed_for_post_processing('output', 'input', 'library_id')
+    
+    mock_makedirs.assert_called()
+    mock_shutil_copy.assert_called_once()
+
+
+def test_clean_up_files_after_post_processing_success(mocker):
+    mock_shutil_rmtree = mocker.patch("shutil.rmtree")
+    mock_os_remove = mocker.patch("os.remove")
+    _clean_up_files_after_post_processing('output', 'library_id')
diff --git a/.github/workflows/generator.yml b/.github/workflows/generator.yml
@@ -29,3 +29,7 @@ jobs:
     - name: Run generator_cli tests
       run: |
         pytest .generator/test_cli.py
+    - name: Check coverage
+      run: |
+        pytest --cov=. --cov-report=term-missing --cov-fail-under=100
+      working-directory: .generator
