Add --remove-invalid option for validation (#19)

Author: piotrekkr

README.md

@@ -37,7 +37,8 @@ but has few shortcomings that can appear quite fast after using it with existing
 
 ## Tech stack
 
-Project is implemented in `Python 3.11` and is using `poetry` as package manager. Main libraries used:
+Project is implemented in `Python 3.11` and is using [`uv`](https://github.com/astral-sh/uv) as package and project manager.
+Main libraries used:
 - [`simple-salesforce`](https://github.com/simple-salesforce/simple-salesforce) - handling Salesforce API
 - [`click`](https://github.com/pallets/click) - working with CLI
 - [`PyYaml`](https://github.com/yaml/pyyaml/) - config parsing
@@ -224,9 +225,19 @@ When validation is complete, show statistics.
 
 ## HOWTOs
 
+### Download and validate with one command
+
+```shell
+docker run --interactive --tty \
+  --volume /path/to/data/directory:/archivist/data \
+  --volume /path/to/your.config.yaml:/archivist/config.yaml \
+  ghcr.io/piotrekkr/salesforce-archivist:latest \
+  download --validate
+```
+
 ### Re-download content version list and document link list
 
-You can remove CSV files from disk, and next download will download full lists again from Salesforce.
+You can remove CSV files from disk, and the next download will fetch full lists again from Salesforce.
 ```shell
 # for chosen type
 rm -rf {data_dir}/{object_type}/*.csv
@@ -240,6 +251,28 @@ rm -rf {data_dir}/*/*.csv
 Already calculated checksums for downloaded files are kept in `{data_dir}/validated_versions.csv`.
 You can remove this file or selected lines from inside this file. This will trigger full validation again.
 
+### How to remove invalid files and redownload them
+
+Validation can show that some checksum or size of downloaded files do not match with values from Salesforce.
+To remove them and download again, you can use `--remove-invalid` flag. This flag will remove invalid files from
+the validated files list and from disk.
+
+```shell
+# download validate and remove invalid in one command
+docker run --interactive --tty \
+  --volume /path/to/data/directory:/archivist/data \
+  --volume /path/to/your.config.yaml:/archivist/config.yaml \
+  ghcr.io/piotrekkr/salesforce-archivist:latest \
+  download --validate --remove-invalid
+
+# remove after validation
+docker run --interactive --tty \
+  --volume /path/to/data/directory:/archivist/data \
+  --volume /path/to/your.config.yaml:/archivist/config.yaml \
+  ghcr.io/piotrekkr/salesforce-archivist:latest \
+  validate --remove-invalid
+```
+
 ## Contributing
 
 // TODO

src/salesforce_archivist/archivist.py

@@ -179,7 +179,11 @@ def _download_attachments(
         global_stats.combine(stats)
 
     def _validate_content_versions_download(
-        self, archivist_obj: ArchivistObject, validated_list: ValidatedList, global_stats: ValidationStats
+        self,
+        archivist_obj: ArchivistObject,
+        validated_list: ValidatedList,
+        global_stats: ValidationStats,
+        remove_invalid: bool = False,
     ) -> bool:
         salesforce = Salesforce(
             archivist_obj=archivist_obj,
@@ -199,12 +203,17 @@ def _validate_content_versions_download(
             download_list=download_list,
             validated_list=validated_list,
             max_workers=self._max_workers,
+            remove_invalid=remove_invalid,
         )
         global_stats.combine(stats)
         return stats.invalid == 0
 
     def _validate_attachments_download(
-        self, archivist_obj: ArchivistObject, validated_list: ValidatedList, global_stats: ValidationStats
+        self,
+        archivist_obj: ArchivistObject,
+        validated_list: ValidatedList,
+        global_stats: ValidationStats,
+        remove_invalid: bool = False,
     ) -> bool:
         salesforce = Salesforce(
             archivist_obj=archivist_obj,
@@ -220,20 +229,21 @@ def _validate_attachments_download(
             download_list=download_list,
             validated_list=validated_list,
             max_workers=self._max_workers,
+            remove_invalid=remove_invalid,
         )
         global_stats.combine(stats)
         return stats.invalid == 0
 
-    def validate(self) -> bool:
+    def validate(self, remove_invalid: bool = False) -> bool:
         validated_list = ValidatedList(self._data_dir)
         if validated_list.data_file_exist():
             validated_list.load_data_from_file()
         global_stats = ValidationStats()
         for archivist_obj in self._objects.values():
             if archivist_obj.obj_type == "Attachment":
-                self._validate_attachments_download(archivist_obj, validated_list, global_stats)
+                self._validate_attachments_download(archivist_obj, validated_list, global_stats, remove_invalid)
             else:
-                self._validate_content_versions_download(archivist_obj, validated_list, global_stats)
+                self._validate_content_versions_download(archivist_obj, validated_list, global_stats, remove_invalid)
         status = "SUCCESS" if global_stats.invalid == 0 else "FAILED"
         color = "green" if global_stats.invalid == 0 else "red"
         click.secho(

src/salesforce_archivist/cli.py

@@ -28,8 +28,9 @@ def cli(ctx: Context) -> None:
 
 @cli.command()
 @click.option("--validate", is_flag=True, default=False, help="Trigger validation after download.")
+@click.option("--remove-invalid", is_flag=True, default=False, help="Remove invalid files after validation.")
 @click.pass_context
-def download(ctx: Context, validate: bool) -> None:
+def download(ctx: Context, validate: bool, remove_invalid: bool) -> None:
     config: ArchivistConfig = ctx.obj["config"]
     sf_client = SalesforceClient(
         instance_url=config.auth.instance_url,
@@ -45,13 +46,14 @@ def download(ctx: Context, validate: bool) -> None:
         max_api_usage_percent=config.max_api_usage_percent,
         max_workers=config.max_workers,
     )
-    if not archivist.download() or validate and not archivist.validate():
+    if not archivist.download() or validate and not archivist.validate(remove_invalid=remove_invalid):
         ctx.exit(code=1)
 
 
 @cli.command()
+@click.option("--remove-invalid", is_flag=True, default=False, help="Remove invalid files after validation.")
 @click.pass_context
-def validate(ctx: Context) -> None:
+def validate(ctx: Context, remove_invalid: bool) -> None:
     config: ArchivistConfig = ctx.obj["config"]
     sf_client = SalesforceClient(
         instance_url=config.auth.instance_url,
@@ -67,7 +69,7 @@ def validate(ctx: Context) -> None:
         max_api_usage_percent=config.max_api_usage_percent,
         max_workers=config.max_workers,
     )
-    if not archivist.validate():
+    if not archivist.validate(remove_invalid=remove_invalid):
         ctx.exit(code=1)
 
 

src/salesforce_archivist/salesforce/salesforce.py

@@ -250,11 +250,11 @@ def validate_download(
         download_list: Union[DownloadContentVersionList, DownloadAttachmentList],
         validated_list: ValidatedList,
         max_workers: int | None = None,
+        remove_invalid: bool = False,
     ) -> ValidationStats:
         try:
             validator = DownloadValidator(
-                validated_list=validated_list,
-                max_workers=max_workers,
+                validated_list=validated_list, max_workers=max_workers, remove_invalid=remove_invalid
             )
             return validator.validate(download_list=download_list)
         finally:

src/salesforce_archivist/salesforce/validation.py

@@ -78,6 +78,9 @@ def is_validated(self, path: str) -> bool:
     def get(self, path: str) -> ValidatedFile | None:
         return self._data.get(path)
 
+    def remove(self, path: str) -> bool:
+        return self._data.pop(path, None) is not None
+
     @property
     def path(self) -> str:
         return self._path
@@ -122,11 +125,12 @@ def combine(self, other: Self) -> None:
 
 
 class DownloadValidator:
-    def __init__(self, validated_list: ValidatedList, max_workers: int | None = None):
+    def __init__(self, validated_list: ValidatedList, max_workers: int | None = None, remove_invalid: bool = False):
         self._validated_list = validated_list
         self._stats = ValidationStats()
         self._lock = threading.Lock()
         self._max_workers = max_workers
+        self._remove_invalid = remove_invalid
 
     def _print_validated_msg(self, msg: str, invalid: bool = False) -> None:
         percent = self._stats.processed / self._stats.total * 100 if self._stats.total > 0 else 0.0
@@ -173,6 +177,9 @@ def _validate_version(self, version: ContentVersion, download_path: str) -> bool
                 self._validated_list.add(
                     ValidatedFile(path=download_path, checksum=checksum, content_size=version.content_size)
                 )
+            if not valid and self._remove_invalid:
+                self._validated_list.remove(download_path)
+                os.remove(download_path)
         except Exception as e:
             msg = "[ KO ] {id} => Exception: {e}".format(id=version.id, e=e)
             valid = False

test/salesforce/test_validation.py

@@ -99,7 +99,7 @@ def test_validated_list_save():
             assert validated_file == loaded_list.get(path=validated_file.path)
 
 
-def test_validated_list_add_get_version():
+def test_validated_list_add_get():
     validated_list = ValidatedList(data_dir="/fake/dir")
     file_1 = ValidatedFile(checksum="checksum1", path="data/path/file_1.txt", content_size=None)
     file_2 = ValidatedFile(checksum=None, path="data/path/file_2.txt", content_size=10)
@@ -108,6 +108,15 @@ def test_validated_list_add_get_version():
     assert validated_list.get(path=file_2.path) is None
 
 
+def test_validated_list_remove():
+    validated_list = ValidatedList(data_dir="/fake/dir")
+    file_1 = ValidatedFile(checksum="checksum1", path="data/path/file_1.txt", content_size=None)
+    validated_list.add(validated_file=file_1)
+    assert validated_list.remove(path=file_1.path)
+    assert not validated_list.remove(path="fake/path")
+    assert len(validated_list) == 0
+
+
 def test_validated_list_is_downloaded():
     validated_list = ValidatedList(data_dir="/fake/dir")
     validated_file_1 = ValidatedFile(path="path/file.txt", checksum="checksum", content_size=None)
@@ -339,6 +348,38 @@ def test_download_validator_validate_object_will_calculate_checksum_and_check_ve
         assert len(validated_list) == 1
 
 
+@pytest.mark.parametrize(
+    "file_data, checksum, should_remove",
+    [
+        ("test", hashlib.md5("test".encode("utf-8")).hexdigest(), False),
+        ("test1", hashlib.md5("test".encode("utf-8")).hexdigest(), True),
+    ],
+)
+def test_download_validator_validate_object_will_remove_invalid_version_file_from_disk_and_validated_list(
+    file_data: str, checksum: str, should_remove
+):
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        archivist_obj = ArchivistObject(data_dir=tmp_dir, obj_type="User")
+        download_path = os.path.join(tmp_dir, "file.txt")
+        with open(download_path, "wb") as file:
+            file.write(file_data.encode("utf-8"))
+
+        version = ContentVersion(
+            version_id="VID1",
+            document_id="DID",
+            checksum=checksum,
+            extension="ext1",
+            title="version1",
+            version_number=1,
+            content_size=10,
+        )
+        validated_list = ValidatedList(data_dir=archivist_obj.obj_dir)
+        validator = DownloadValidator(validated_list=validated_list, remove_invalid=True)
+        validator.validate_object(obj=version, download_path=download_path)
+        assert os.path.exists(download_path) == (not should_remove)
+        assert len(validated_list) == (0 if should_remove else 1)
+
+
 @pytest.mark.parametrize(
     "file_data, size, should_match",
     [

test/test_archivist.py

@@ -346,7 +346,7 @@ def test_archivist_validate_will_load_lists_and_call_validate_method(
     assert load_version_list_mock.call_count == 2
     assert load_attachment_list_mock.call_count == 1
     assert validate_mock.mock_calls == [
-        call(download_list=ANY, validated_list=ANY, max_workers=max_workers),
-        call(download_list=ANY, validated_list=ANY, max_workers=max_workers),
-        call(download_list=ANY, validated_list=ANY, max_workers=max_workers),
+        call(download_list=ANY, validated_list=ANY, max_workers=max_workers, remove_invalid=False),
+        call(download_list=ANY, validated_list=ANY, max_workers=max_workers, remove_invalid=False),
+        call(download_list=ANY, validated_list=ANY, max_workers=max_workers, remove_invalid=False),
     ]