Skip to content

Commit

Permalink
CHIA-1697: Add new flag to support recursively scanning and following…
Browse files Browse the repository at this point in the history
… links (#18803)

* Add new flag to support recurvely scanning and following links

* black fix

* hardlink_to only in python 3.10+

* Remove unneeded hardlink checks as hardlinks to directories aren't a thing

* Minor update per code review

* use boolean_datacases

* Add new flag to initial-config
  • Loading branch information
emlowe authored Nov 6, 2024
1 parent 315f4bf commit 61cf266
Show file tree
Hide file tree
Showing 3 changed files with 59 additions and 4 deletions.
42 changes: 42 additions & 0 deletions chia/_tests/plotting/test_plot_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand Down Expand Up @@ -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)
20 changes: 16 additions & 4 deletions chia/plotting/util.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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.")
Expand All @@ -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}")
Expand Down
1 change: 1 addition & 0 deletions chia/util/initial-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,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"
Expand Down

0 comments on commit 61cf266

Please sign in to comment.