CHIA-1697: Add new flag to support recursively scanning and following links

chia/_tests/plotting/test_plot_manager.py

@@ -14,6 +14,7 @@
 from chia_rs import G1Element
 
 from chia._tests.plotting.util import get_test_plots
+from chia._tests.util.misc import boolean_datacases
 from chia._tests.util.time_out_assert import time_out_assert
 from chia.plotting.cache import CURRENT_VERSION, CacheDataV1
 from chia.plotting.manager import Cache, PlotManager
@@ -33,6 +34,9 @@
 
 log = logging.getLogger(__name__)
 
+# These tests are not dependent on consensus rules
+pytestmark = [pytest.mark.limit_consensus_modes(reason="does not depend on consensus rules")]
+
 
 @dataclass
 class MockDiskProver:
@@ -737,3 +741,41 @@ async def test_recursive_plot_scan(environment: Environment) -> None:
     add_plot_directory(env.root_path, str(sub_dir_1_0_1.path))
     expected_result.loaded = []
     await env.refresh_tester.run(expected_result)
+
+
+@boolean_datacases(name="follow_links", false="no_follow", true="follow")
+@pytest.mark.anyio
+async def test_recursive_plot_scan_symlinks(environment: Environment, follow_links: bool) -> None:
+    env: Environment = environment
+    # Create a directory tree with some subdirectories containing plots, others not.
+    root_plot_dir = env.root_path / "root"
+    sub_dir_0: Directory = Directory(root_plot_dir / "0", env.dir_1.plots[0:2])
+    sub_dir_0_1: Directory = Directory(sub_dir_0.path / "1", env.dir_1.plots[2:3])
+
+    # Create a plot directory outside of the root plot directory
+    other_plot_dir: Directory = Directory(env.root_path / "other", env.dir_1.plots[3:7])
+
+    # Create a symlink to the other plot directory from inside the root plot directory
+    symlink_path = Path(root_plot_dir / "other")
+    symlink_path.symlink_to(env.root_path / "other")
+
+    # Adding the root without `recursive_plot_scan` and running a test should not load any plots (match an empty result)
+    expected_result = PlotRefreshResult()
+    add_plot_directory(env.root_path, str(root_plot_dir))
+    await env.refresh_tester.run(expected_result)
+
+    if follow_links:
+        expected_plot_list = sub_dir_0.plot_info_list() + sub_dir_0_1.plot_info_list() + other_plot_dir.plot_info_list()
+    else:
+        expected_plot_list = sub_dir_0.plot_info_list() + sub_dir_0_1.plot_info_list()
+
+    # Set the recursive scan flag in the config and symlink flag
+    with lock_and_load_config(env.root_path, "config.yaml") as config:
+        config["harvester"]["recursive_plot_scan"] = True
+        config["harvester"]["recursive_follow_links"] = follow_links
+        save_config(env.root_path, "config.yaml", config)
+
+    # With the flag enabled it should load all expected plots
+    expected_result.loaded = expected_plot_list  # type: ignore[assignment]
+    expected_result.processed = len(expected_plot_list)
+    await env.refresh_tester.run(expected_result)

chia/plotting/util.py

@@ -110,13 +110,14 @@ def get_plot_filenames(root_path: Path) -> dict[Path, list[Path]]:
     all_files: dict[Path, list[Path]] = {}
     config = load_config(root_path, "config.yaml")
     recursive_scan: bool = config["harvester"].get("recursive_plot_scan", DEFAULT_RECURSIVE_PLOT_SCAN)
+    recursive_follow_links: bool = config["harvester"].get("recursive_follow_links", False)
     for directory_name in get_plot_directories(root_path, config):
         try:
             directory = Path(directory_name).resolve()
         except (OSError, RuntimeError):
             log.exception(f"Failed to resolve {directory_name}")
             continue
-        all_files[directory] = get_filenames(directory, recursive_scan)
+        all_files[directory] = get_filenames(directory, recursive_scan, recursive_follow_links)
     return all_files
 
 
@@ -219,7 +220,7 @@ def update_harvester_config(
         save_config(root_path, "config.yaml", config)
 
 
-def get_filenames(directory: Path, recursive: bool) -> list[Path]:
+def get_filenames(directory: Path, recursive: bool, follow_links: bool) -> list[Path]:
     try:
         if not directory.exists():
             log.warning(f"Directory: {directory} does not exist.")
@@ -229,8 +230,19 @@ def get_filenames(directory: Path, recursive: bool) -> list[Path]:
         return []
     all_files: list[Path] = []
     try:
-        glob_function = directory.rglob if recursive else directory.glob
-        all_files = [child for child in glob_function("*.plot") if child.is_file() and not child.name.startswith("._")]
+        if follow_links and recursive:
+            import glob
+
+            files = glob.glob(str(directory / "**" / "*.plot"), recursive=True)
+            for file in files:
+                filepath = Path(file).resolve()
+                if filepath.is_file() and not filepath.name.startswith("._"):
+                    all_files.append(filepath)
+        else:
+            glob_function = directory.rglob if recursive else directory.glob
+            all_files = [
+                child for child in glob_function("*.plot") if child.is_file() and not child.name.startswith("._")
+            ]
         log.debug(f"get_filenames: {len(all_files)} files found in {directory}, recursive: {recursive}")
     except Exception as e:
         log.warning(f"Error reading directory {directory} {e}")

chia/util/initial-config.yaml

@@ -160,6 +160,7 @@ harvester:
   # Plots are searched for in the following directories
   plot_directories: []
   recursive_plot_scan: False # If True the harvester scans plots recursively in the provided directories.
+  recursive_follow_links: False # If True the harvester follows symlinks when scanning for plots recursively
 
   ssl:
     private_crt: "config/ssl/harvester/private_harvester.crt"